Upload folder using huggingface_hub

decider/engine.py

@@ -19,25 +19,67 @@ def _bucket(x, buckets):
     return None
 
 
+def fused_causal_conv1d_fn(hidden_states, weight, bias=None, activation=None, **kwargs):
+    """Depthwise causal conv (kernel k) as k shifted multiply-adds: fuses under torch.compile,
+    unlike the cuDNN grouped conv fallback (which was ~11% of batched GPU time)."""
+    B, C, T = hidden_states.shape; k = weight.shape[-1]
+    x = F.pad(hidden_states.to(weight.dtype), (k - 1, 0))
+    out = x[:, :, k - 1:k - 1 + T] * weight[:, k - 1][None, :, None]
+    for j in range(k - 1):
+        out = out + x[:, :, j:j + T] * weight[:, j][None, :, None]
+    if bias is not None:
+        out = out + bias[None, :, None]
+    if activation == "silu":
+        out = F.silu(out)
+    elif activation is not None:
+        from transformers.activations import ACT2FN
+        out = ACT2FN[activation](out)
+    return out.to(hidden_states.dtype)
+
+
+def patch_conv():
+    from transformers.models.qwen3_5 import modeling_qwen3_5 as mq
+    mq.causal_conv1d_fn = fused_causal_conv1d_fn
+
+
 class Engine:
-    def __init__(self, path, device="cuda", dtype=torch.bfloat16, use_graphs=True, max_ctx_tokens=1536):
+    """compile: torch.compile the forward (needs use_cache=False; ~1.4x batched, fuses elementwise work).
+    fp8: e4m3 weights + per-token activation scaling on the big linears (Hopper tensor cores).
+    conv_patch: fusable depthwise causal conv instead of the cuDNN fallback."""
+    def __init__(self, path, device="cuda", dtype=torch.bfloat16, use_graphs=True, max_ctx_tokens=1536,
+                 compile=True, fp8=False, conv_patch=True):
+        if conv_patch:
+            patch_conv()
         self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval()
         self.tok = self.m.tok; self.dev = device; self.use_graphs = use_graphs; self.max_ctx = max_ctx_tokens
         self.core, self.W = self.m.lm.model, self.m.lm.lm_head.weight[self.m.letters].detach().clone()
+        self.cfg = dict(compile=compile, fp8=fp8, conv_patch=conv_patch, graphs=use_graphs)
+        if fp8:
+            from .fp8 import convert_to_fp8
+            self.cfg["fp8_layers"] = convert_to_fp8(self.core)
+        if compile:
+            import torch._dynamo
+            torch._dynamo.config.cache_size_limit = 128
+            self._fwd_impl = torch.compile(self._fwd_eager, dynamic=False)
+        else:
+            self._fwd_impl = self._fwd_eager
         self.graphs = {}                       # (B, T) -> (static_ids, static_out, graph)
         self.pool = torch.cuda.graph_pool_handle() if use_graphs else None
         self.stats = dict(graph_captures=0, forwards=0)
 
+    def _fwd_eager(self, ids):
+        h = self.core(input_ids=ids, use_cache=False).last_hidden_state
+        return F.linear(h, self.W).float()                          # [B, T, K]
+
     @torch.no_grad()
     def _fwd(self, ids):
-        h = self.core(input_ids=ids).last_hidden_state
-        return F.linear(h, self.W).float()                          # [B, T, K]
+        return self._fwd_impl(ids)
 
     def _capture(self, B, T):
         s_ids = torch.full((B, T), self.tok.pad_token_id, dtype=torch.long, device=self.dev)
         st = torch.cuda.Stream(); st.wait_stream(torch.cuda.current_stream())
         with torch.cuda.stream(st):
-            for _ in range(2): self._fwd(s_ids)                     # warm-up: triton autotune etc.
+            for _ in range(3): self._fwd(s_ids)                     # warm-up: compile / triton autotune
         torch.cuda.current_stream().wait_stream(st)
         g = torch.cuda.CUDAGraph()
         with torch.cuda.graph(g, pool=self.pool):
@@ -89,8 +131,9 @@ if __name__ == "__main__":
     from . import data as D
     from .infer import Decider
     path = sys.argv[1] if len(sys.argv) > 1 else "runs/r3_v2/model"
+    cfg = dict(compile="nocompile" not in sys.argv[2:], fp8="fp8" in sys.argv[2:], conv_patch="noconv" not in sys.argv[2:])
     _, evals = D.load_cache("data/tasks.pkl")
-    eng = Engine(path)
+    eng = Engine(path, **cfg); print("engine cfg", eng.cfg)
     rng = random.Random(0)
     exs = evals["support_tickets"][:64] + evals["clinc_oos"][:64] + evals["race"][:32]
     items = [build(e, eng.tok, rng, max_ctx_tokens=1536) for e in exs]

decider/fp8.py

@@ -0,0 +1,54 @@
+"""FP8 (e4m3) linear layers for Hopper via torch._scaled_mm.
+Weights: per-output-channel scales, quantised once.  Activations: per-token dynamic scales.
+Under torch.compile the quantisation ops fuse into the surrounding elementwise work."""
+import torch, torch.nn as nn
+
+E4M3_MAX = 448.0
+
+
+def _quant_rowwise(x):
+    s = x.abs().amax(dim=-1, keepdim=True).float().clamp(min=1e-12) / E4M3_MAX
+    return (x.float() / s).clamp(-E4M3_MAX, E4M3_MAX).to(torch.float8_e4m3fn), s
+
+
+class FP8Linear(nn.Module):
+    def __init__(self, lin: nn.Linear):
+        super().__init__()
+        wq, sw = _quant_rowwise(lin.weight.detach())          # [N,K] fp8, [N,1]
+        self.register_buffer("wq", wq.contiguous())               # [N,K]; passed as wq.t() -> [K,N] column-major, as _scaled_mm wants
+        self.register_buffer("sw_t", sw.t().contiguous())        # [1,N]
+        self.bias = None if lin.bias is None else nn.Parameter(lin.bias.detach().clone(), requires_grad=False)
+        self.in_features, self.out_features = lin.in_features, lin.out_features
+        self.out_dtype = lin.weight.dtype
+
+    def forward(self, x):
+        shp = x.shape[:-1]
+        x2 = x.reshape(-1, self.in_features)
+        xq, sx = _quant_rowwise(x2)
+        y = torch._scaled_mm(xq, self.wq.t(), scale_a=sx, scale_b=self.sw_t, bias=self.bias, out_dtype=self.out_dtype)
+        return y.reshape(*shp, self.out_features)
+
+
+def convert_to_fp8(model, skip=("lm_head",), min_dim=1024):
+    """Replace nn.Linear (with in/out >= min_dim) by FP8Linear in place. Returns count."""
+    n = 0
+    for name, mod in list(model.named_modules()):
+        for cname, child in list(mod.named_children()):
+            full = f"{name}.{cname}" if name else cname
+            if isinstance(child, nn.Linear) and not any(s in full for s in skip) and min(child.in_features, child.out_features) >= min_dim:
+                setattr(mod, cname, FP8Linear(child)); n += 1
+    return n
+
+
+if __name__ == "__main__":
+    import time
+    lin = nn.Linear(2048, 6144, bias=False).cuda().to(torch.bfloat16)
+    f8 = FP8Linear(lin)
+    x = torch.randn(8192, 2048, device="cuda", dtype=torch.bfloat16)
+    ref = lin(x); got = f8(x)
+    print("rel err", ((ref.float() - got.float()).abs().mean() / ref.float().abs().mean()).item())
+    for f, name in [(lin, "bf16 linear"), (f8, "fp8 linear (eager)"), (torch.compile(f8), "fp8 linear (compiled)")]:
+        for _ in range(3): f(x)
+        torch.cuda.synchronize(); t = time.time()
+        for _ in range(20): f(x)
+        torch.cuda.synchronize(); print(f"{name:24s} {(time.time()-t)/20*1000:.3f} ms")

decider/infer.py

@@ -24,8 +24,16 @@ class Example:
 
 
 class Decider:
-    def __init__(self, path, device="cuda", dtype=torch.bfloat16, temperature=1.0, abstain_below=0.0):
-        self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval()
+    """use_graphs=True (default on CUDA) routes scoring through decider.engine.Engine: shape-bucketed
+    CUDA graphs, ~7x lower single-request latency than eager. Set False for CPU or debugging."""
+    def __init__(self, path, device="cuda", dtype=torch.bfloat16, temperature=1.0, abstain_below=0.0, use_graphs=None):
+        if use_graphs is None:
+            use_graphs = str(device).startswith("cuda")
+        if use_graphs:
+            from .engine import Engine
+            self.eng = Engine(path, device=device, dtype=dtype); self.m = self.eng.m
+        else:
+            self.eng = None; self.m = DecisionModel(path, dtype=dtype, grad_ckpt=False).to(device).eval()
         self.dev = device; self.T = temperature; self.abstain_below = abstain_below
 
     @torch.no_grad()
@@ -40,10 +48,13 @@ class Decider:
             def shuffle(self, x): pass
             def sample(self, xs, k): return xs[:k]
         items = [build(e, self.m.tok, _NoShuffle(), max_ctx_tokens=max_ctx_tokens) for e in exs]
-        b = collate(items, self.m.tok.pad_token_id)
-        logits = self.m.slot_logits(b["input_ids"].to(self.dev), b["attention_mask"].to(self.dev), b["slot_idx"].to(self.dev),
-                                    b["slot_batch"].to(self.dev), b["nopts"].to(self.dev))
-        probs = torch.softmax(logits / self.T, -1).cpu()
+        if self.eng is not None:
+            probs = torch.cat(self.eng.score_items(items, temperature=self.T))
+        else:
+            b = collate(items, self.m.tok.pad_token_id)
+            logits = self.m.slot_logits(b["input_ids"].to(self.dev), b["attention_mask"].to(self.dev), b["slot_idx"].to(self.dev),
+                                        b["slot_batch"].to(self.dev), b["nopts"].to(self.dev))
+            probs = torch.softmax(logits / self.T, -1).cpu()
         out, k = [], 0
         for context, qs in requests:
             res = []