phanerozoic commited on Jul 14

Commit

4bca4f5

verified ·

1 Parent(s): af194d8

Uploaded using `kernel-builder`.

Browse files

Files changed (45) hide show

build/torch211-cxx11-cu126-x86_64-linux/__init__.py +964 -0
build/torch211-cxx11-cu126-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch211-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu126-x86_64-linux/metadata.json +36 -0
build/torch211-cxx11-cu126-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +964 -0
build/torch211-cxx11-cu128-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +38 -0
build/torch211-cxx11-cu128-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +964 -0
build/torch211-cxx11-cu130-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +38 -0
build/torch211-cxx11-cu130-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch212-cxx11-cu126-x86_64-linux/__init__.py +964 -0
build/torch212-cxx11-cu126-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch212-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu126-x86_64-linux/metadata.json +36 -0
build/torch212-cxx11-cu126-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +964 -0
build/torch212-cxx11-cu130-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +38 -0
build/torch212-cxx11-cu130-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +964 -0
build/torch212-cxx11-cu132-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +38 -0
build/torch212-cxx11-cu132-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch213-cxx11-cu126-x86_64-linux/__init__.py +964 -0
build/torch213-cxx11-cu126-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch213-cxx11-cu126-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu126-x86_64-linux/metadata.json +36 -0
build/torch213-cxx11-cu126-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch213-cxx11-cu130-x86_64-linux/__init__.py +964 -0
build/torch213-cxx11-cu130-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +38 -0
build/torch213-cxx11-cu130-x86_64-linux/model_as_a_kernel/__init__.py +26 -0
build/torch213-cxx11-cu132-x86_64-linux/__init__.py +964 -0
build/torch213-cxx11-cu132-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so +3 -0
build/torch213-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu132-x86_64-linux/metadata.json +38 -0
build/torch213-cxx11-cu132-x86_64-linux/model_as_a_kernel/__init__.py +26 -0

build/torch211-cxx11-cu126-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch211-cxx11-cu126-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:6461efe05f2ae10a20c69a3648f9bb7bb7af9c6914b95b00f19c754e83298201
+size 3839520

build/torch211-cxx11-cu126-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch211-cxx11-cu126-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,36 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "ZGHv4F8q4Qogxpo2SPm7e7evnGkUuVsA8Zx1ToMpggE=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu126-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu128-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch211-cxx11-cu128-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:bc112475d41da2f89920294225377bf527635de3ba393a1a6e9f4106c0e34ebb
+size 6395056

build/torch211-cxx11-cu128-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,38 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "vBEkddQdoviZIClCJTd79SdjXeO6OToabp9BBsDjTrs=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu128-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch211-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch211-cxx11-cu130-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:590bdbf325381981283ce087722012cbe44e97ac57595475fbcb12b043e7a79c
+size 6347568

build/torch211-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,38 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "WQvb8yU4GYEoPOCHciASy+ROl6xXWVR1+8sSsEPnp5w=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch211-cxx11-cu130-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu126-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch212-cxx11-cu126-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e7e3d83357592292c3b6901445991763e8e7c61f66373bb2191da1bfd2d6860d
+size 3845320

build/torch212-cxx11-cu126-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch212-cxx11-cu126-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,36 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "5+PYM1dZIpLDtpAURZkXY+jnxh9mNzuyGR2hv9LWhg0=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu126-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch212-cxx11-cu130-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:2b82256093271ec83b0ab2150b9bdc2aded273ef1abeb215467793cc5658e970
+size 6354080

build/torch212-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,38 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "K4IlYJMnHsg7CrIVC5vcKt7Sc+8avrIVRneTzFZY6XA=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu130-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch212-cxx11-cu132-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:effc2eb17efb6c24253a3cf8f08442c69792e6dbda74ed6909650ea26dddcf60
+size 6374560

build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,38 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "7/wusX77bCQlOjz48IRCxpeS5tvadO1pCWUOom3dz2A=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch213-cxx11-cu126-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch213-cxx11-cu126-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:058f096551b99c0ad1db80dcc08bb6a4a95ec9924515919892a196e8e7ddf3bc
+size 3849856

build/torch213-cxx11-cu126-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch213-cxx11-cu126-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,36 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "BY8JZVG5nArR24DcwIu2pKleyZJFFZGYkqGW6Ofd87w=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu126-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch213-cxx11-cu130-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:c23353414f62f5288193b6608c66a56b3b74ed86a4ab404dbea27dd31cd33f63
+size 6358528

build/torch213-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch213-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,38 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "wjNTQU9i9SiBk7ZgjGalazt07Yakq0BNvqJ90xzTP2M=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu130-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch213-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,964 @@

+"""model-as-a-kernel: an entire llama-family decode step as one kernel launch.
+A phase interpreter runs the whole forward pass (embedding, every layer's
+norms, QKV, rope, attention over the KV cache, SwiGLU MLP, the LM head, and
+the greedy argmax) inside a single persistent kernel, with a software grid
+barrier between phases. The in-kernel step loop makes an entire greedy
+generation, prompt consumption included, one kernel launch: each step's
+argmax feeds the next step's embedding read without leaving the device.
+"""
+import os
+import struct
+from typing import List, Optional, Sequence
+import torch
+try:
+    from ._ops import ops
+except ImportError:  # local development build (torch.utils.cpp_extension)
+    class _LocalOps:
+        def __getattr__(self, name):
+            return getattr(torch.ops.mak_ext, name)
+    ops = _LocalOps()
+__all__ = ["MegaModel"]
+_OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
+_OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
+_OP_KV_APPEND, _OP_NORMRES, _OP_PLEMIX = 6, 7, 8
+_OP_GEMV_PLAIN, _OP_NORMB, _OP_GLUB, _OP_ATTNFINB = 9, 10, 11, 12
+_OP_GEMV_Q4 = 13
+_IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
+_IT_RMSNORM_G, _IT_GELU_GLU = 4, 5
+_EP_GELU_PLE, _EP_F32_AMAX_CAP = 5, 6
+_ITF_REDUCE, _ITF_WRITEBACK = 8, 16
+_EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
+_NSLICE = 4
+# Widest [B][K] input panel (bf16 elements) a batched projection stages in
+# shared before it spills to the global-scratch path. 16384 (32 KB) keeps
+# two blocks resident per SM; the persistent kernel shares one occupancy
+# across all phases, so a wider panel would drop every phase to one block
+# per SM, which measures slower than reading the input from L2.
+_STAGE_BUDGET = 16384
+def _fbits(f: float) -> int:
+    return struct.unpack("<i", struct.pack("<f", float(f)))[0]
+def _row(op, p1=0, p2=0, p3=0, p4=0, p5=0, p6=0, n=0, k=0, it=0, epi=0,
+         hq=0, hkv=0, d=0, f0=0.0, i0=0):
+    return [int(op), int(p1), int(p2), int(p3), int(p4), int(p5), int(p6),
+            int(n), int(k), int(it), int(epi), int(hq), int(hkv), int(d),
+            _fbits(f0), int(i0)]
+class MegaModel:
+    """A packed llama-family decoder whose decode step is one kernel launch.
+    Supported architecture: RMSNorm decoder blocks with rotary attention
+    (GQA, optional Qwen3-style per-head qk RMSNorm), SwiGLU MLP, no attention
+    or MLP biases, plain (unscaled) rope, batch size 1. This covers SmolLM2,
+    TinyLlama, Qwen3 dense, Llama-class checkpoints without rope scaling,
+    and similar models.
+    """
+    def __init__(self, weights: dict, config: dict, device="cuda",
+                 max_seq: int = 4096, max_gen: int = 4096):
+        cfg = dict(config)
+        self.L = int(cfg["num_hidden_layers"])
+        self.hidden = int(cfg["hidden_size"])
+        self.Hq = int(cfg["num_attention_heads"])
+        self.Hkv = int(cfg.get("num_key_value_heads") or self.Hq)
+        self.D = int(cfg.get("head_dim") or self.hidden // self.Hq)
+        self.I = int(cfg["intermediate_size"])
+        self.V = int(cfg["vocab_size"])
+        self.theta = float(cfg.get("rope_theta", 10000.0))
+        self.eps = float(cfg.get("rms_norm_eps", 1e-6))
+        self.qk_norm = bool(cfg.get("qk_norm", False))
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        if self.D % 2 or self.D > 256:
+            raise ValueError("head_dim must be even and <= 256")
+        if self.hidden % 8 or self.I % 8 or (self.Hq * self.D) % 8:
+            raise ValueError(
+                "hidden, intermediate, and Hq*head_dim must be multiples of 8")
+        if self.Hq % self.Hkv:
+            raise ValueError("num_attention_heads must be divisible by num_key_value_heads")
+        dt = torch.bfloat16
+        dev = self.device
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        # Row-tile-of-8 weight layout [N/8][K/8][8][8] (MAK_TILED=1): the
+        # eight warps of a block read one contiguous 4KB stream per step
+        # instead of eight 4KB-strided pieces. Measured slower on every
+        # card (H200 +15 percent, RTX PRO +9, Ada +14 — per-warp stream
+        # locality beats per-block contiguity on all of them), so the
+        # default is the plain layout; the flag rides in the program and
+        # both layouts compute identical bits.
+        env_t = os.environ.get("MAK_TILED", "").strip()
+        self.tiled = env_t == "1"
+        def pack_tiled(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            if not self.tiled:
+                self._keep.append(t)
+                return t
+            n, k = t.shape
+            assert k % 8 == 0
+            n8 = (n + 7) // 8 * 8
+            if n8 != n:
+                t = torch.cat(
+                    [t, torch.zeros(n8 - n, k, dtype=dt, device=dev)], 0)
+            t = t.view(n8 // 8, 8, k // 8, 8).permute(0, 2, 1, 3).contiguous()
+            self._keep.append(t)
+            return t
+        self._keep: List[torch.Tensor] = []
+        embed = pack(weights["embed"])
+        final_norm = pack(weights["norm"])
+        lm_head = weights["lm_head"]  # packed below (dense-tiled or nf4)
+        inv_freq = weights.get("inv_freq")
+        if inv_freq is None:
+            d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
+            inv_freq = 1.0 / (self.theta ** (d_idx / self.D))
+        inv_freq = inv_freq.detach().to(device=dev,
+                                        dtype=torch.float32).contiguous()
+        if inv_freq.numel() != self.D // 2:
+            raise ValueError("inv_freq length must be head_dim / 2")
+        self._keep.append(inv_freq)
+        self._invf = inv_freq
+        qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
+        def _is_q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def pack_w(w):
+            # nf4-packed weights move to the device intact; dense weights take
+            # the usual (optionally row-tiled) bf16 packing.
+            if _is_q4(w):
+                pk = w["packed"].to(dev).contiguous()
+                am = w["absmax"].to(device=dev, dtype=torch.float32).contiguous()
+                self._keep.append(pk)
+                self._keep.append(am)
+                return {"q4": True, "packed": pk, "absmax": am,
+                        "N": int(w["N"]), "K": int(w["K"])}
+            return pack_tiled(w)
+        def _wshape(w, expected):
+            got = (int(w["N"]), int(w["K"])) if _is_q4(w) else tuple(w.shape)
+            assert got == expected, (got, expected)
+        self._has_q4 = False
+        self._layers = []
+        for lw in weights["layers"]:
+            _wshape(lw["wqkv"], (qdim + 2 * kvdim, self.hidden))
+            _wshape(lw["wo"], (self.hidden, qdim))
+            _wshape(lw["wgu"], (2 * self.I, self.hidden))
+            _wshape(lw["wdown"], (self.hidden, self.I))
+            layer = {
+                "wqkv": pack_w(lw["wqkv"]), "wo": pack_w(lw["wo"]),
+                "wgu": pack_w(lw["wgu"]), "wdown": pack_w(lw["wdown"]),
+                "ln1": pack(lw["ln1"]), "ln2": pack(lw["ln2"]),
+                "qn": pack(lw["qn"]) if self.qk_norm else None,
+                "kn": pack(lw["kn"]) if self.qk_norm else None,
+            }
+            self._has_q4 |= any(_is_q4(layer[n]) for n in
+                                ("wqkv", "wo", "wgu", "wdown"))
+            self._layers.append(layer)
+        assert len(self._layers) == self.L
+        assert embed.shape == (self.V, self.hidden)
+        lm_head = pack_w(lm_head)
+        self._has_q4 |= _is_q4(lm_head)
+        if not _is_q4(lm_head):
+            assert lm_head.numel() >= self.V * self.hidden  # row-tiled, padded
+        self._w_embed, self._w_fnorm, self._w_lmhead = (embed, final_norm,
+                                                        lm_head)
+        self._qdim, self._kvdim = qdim, kvdim
+        # Working buffers. Pointers to these are baked into the program, so
+        # they (like the packed weights) live for the model's lifetime.
+        self._maxk = max(self.hidden, qdim, self.I)
+        # prefill chunk size, bounded by the staging shared-memory budget
+        # (16K bf16 elements; the cp.async weight ring rides alongside)
+        self._chunk_m = max(1, min(8, 16384 // self._maxk))
+        cm = 8  # buffers sized for the maximum chunk
+        self._hidden = torch.empty(cm * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(cm * (qdim + 2 * kvdim), dtype=dt, device=dev)
+        self._gu = torch.empty(cm * 2 * self.I, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        # attention chunk length (kernel maximum 128); shorter chunks give
+        # the attention phase more grid parallelism at the cost of more
+        # softmax partials to finalize
+        env_ch = os.environ.get("MAK_CHUNK", "").strip()
+        self._chunk = int(env_ch) if env_ch in ("32", "64", "128") else 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._partials = torch.empty(
+            cm * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._vcache = torch.zeros(self.L, self.Hkv, self.max_seq, self.D,
+                                   dtype=dt, device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64, device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._batch = 0
+        self._pos_b = None
+        self._kv_bstride = 0
+        self._gemma = False
+        # transformed-input scratch: the nf4 and wide-batch paths write the
+        # dense bf16 input here (one row per active token) for the following
+        # plain/quant GEMV; sized for a prefill chunk (up to 8 rows)
+        self._xg = torch.empty(cm * self._maxk, dtype=dt, device=dev)
+        self._build_programs()
+    # ------------------------------------------------------------------
+    def _kvp(self, cache, li: int, b: int) -> int:
+        t = cache[li]
+        return (t[b] if t.dim() == 4 else t).data_ptr()
+    def _build_rows(self, kv_slice: int = 0, batch: bool = False,
+                    big: bool = False):
+        """Program rows for the llama-family path. kv_slice selects the
+        batch slice the KV pointers address (per-sequence prefill);
+        batch=True emits the batched-decode tail (per-row fp32 logits and
+        explicit argmax phases) instead of the fused LM-head amax. In batch
+        mode each projection stays fused (input staged in shared) when its
+        [B][K] panel fits the budget; the projections whose K is too wide
+        transform into a global scratch and read it with a plain GEMV."""
+        _tbit = (1 << 20) if self.tiled else 0  # GEMV weight layout flag
+        qdim, kvdim = self._qdim, self._kvdim
+        embed, lm_head = self._w_embed, self._w_lmhead
+        final_norm = self._w_fnorm
+        xg = self._xg.data_ptr()
+        B = self._batch if batch else 1
+        rows, names = [], []
+        staged_elems = [0]  # widest [B][K] panel any fused projection stages
+        def fits(K):
+            return B * K <= _STAGE_BUDGET
+        def note(K):
+            staged_elems[0] = max(staged_elems[0], B * K)
+        def _q4(w):
+            return isinstance(w, dict) and w.get("q4")
+        def _wgemv(out, w, N, K, epi, resid):
+            # the GEMV after a transform-to-scratch phase: nf4-dequant when
+            # the weight is packed, otherwise a plain bf16 GEMV
+            if _q4(w):
+                rows.append(_row(_OP_GEMV_Q4, p1=xg, p2=w["packed"].data_ptr(),
+                                 p3=out, p5=resid, p6=w["absmax"].data_ptr(),
+                                 n=N, k=K, epi=epi))
+            else:
+                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w.data_ptr(),
+                                 p3=out, p5=resid, n=N, k=K, it=_tbit,
+                                 epi=epi))
+        def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
+                      variant=_IT_RMSNORM):
+            if _q4(w) or (big and not fits(K)):
+                rows.append(_row(_OP_NORMB, p1=inp, p2=gamma, p3=xg, k=K,
+                                 it=variant, f0=self.eps))
+                names.append(name + ".norm")
+                _wgemv(out, w, N, K, epi, resid)
+                names.append(name)
+            else:
+                note(K)
+                rows.append(_row(_OP_GEMV, p1=inp, p2=w.data_ptr(), p3=out,
+                                 p4=gamma, p5=resid, n=N, k=K,
+                                 it=variant | _tbit, epi=epi, f0=self.eps))
+                names.append(name)
+        def glu_gemv(name, gu, w, out, resid, variant=_IT_SWIGLU):
+            if _q4(w) or (big and not fits(self.I)):
+                rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
+                names.append(name + ".glu")
+                _wgemv(out, w, self.hidden, self.I, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(self.I)
+                rows.append(_row(_OP_GEMV, p1=gu, p2=w.data_ptr(), p3=out,
+                                 p5=resid, n=self.hidden, k=self.I,
+                                 it=variant | _tbit, epi=_EP_RESID))
+                names.append(name)
+        def attn_gemv(name, w, out, resid):
+            if _q4(w) or (big and not fits(qdim)):
+                rows.append(_row(_OP_ATTNFINB, p1=self._partials.data_ptr(),
+                                 p3=xg, k=qdim, hq=self.Hq,
+                                 hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name + ".fin")
+                _wgemv(out, w, self.hidden, qdim, _EP_RESID, resid)
+                names.append(name)
+            else:
+                note(qdim)
+                rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                                 p2=w.data_ptr(), p3=out, p5=resid,
+                                 n=self.hidden, k=qdim,
+                                 it=_IT_ATTNFIN | _tbit, epi=_EP_RESID,
+                                 hq=self.Hq, hkv=self._chunk << 16, d=self.D,
+                                 i0=self._maxch))
+                names.append(name)
+        rows.append(_row(_OP_EMBED, p2=embed.data_ptr(),
+                         p3=self._hidden.data_ptr(),
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden))
+        names.append("embed")
+        scale = 1.0 / (self.D ** 0.5)
+        hid = self._hidden.data_ptr()
+        for li, lw in enumerate(self._layers):
+            kc = self._kvp(self._kcache, li, kv_slice)
+            vc = self._kvp(self._vcache, li, kv_slice)
+            norm_gemv(f"L{li}.qkv", lw["ln1"].data_ptr(),
+                      lw["wqkv"], self._qkv.data_ptr(),
+                      qdim + 2 * kvdim, self.hidden, _EP_STORE, hid)
+            rows.append(_row(6, p1=self._qkv.data_ptr(), p2=kc, p3=vc,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq, hkv=self.Hkv, d=self.D))
+            names.append(f"L{li}.kvappend")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(), p2=kc,
+                             p3=vc, p4=self._partials.data_ptr(),
+                             p5=lw["qn"].data_ptr() if self.qk_norm else 0,
+                             p6=lw["kn"].data_ptr() if self.qk_norm else 0,
+                             n=self.max_seq, k=1 if self.qk_norm else 0,
+                             it=_fbits(self.eps),
+                             epi=self._invf.data_ptr(),
+                             hq=self.Hq,
+                             hkv=self.Hkv | (self._chunk << 16),
+                             d=self.D, f0=scale, i0=self._maxch))
+            names.append(f"L{li}.attn")
+            attn_gemv(f"L{li}.o", lw["wo"], hid, hid)
+            norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
+                      lw["wgu"], self._gu.data_ptr(), 2 * self.I,
+                      self.hidden, _EP_STORE, hid)
+            glu_gemv(f"L{li}.down", self._gu.data_ptr(), lw["wdown"], hid, hid)
+        if batch or _q4(lm_head):
+            norm_gemv("lm_head", final_norm.data_ptr(), lm_head,
+                      self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
+                      hid)
+            rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
+                             p3=self._parts.data_ptr(), n=self.V))
+            names.append("argmax.part")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr(),
+                             k=self.max_gen))
+            names.append("argmax.fin")
+        else:
+            rows.append(_row(_OP_GEMV, p1=hid, p2=lm_head.data_ptr(),
+                             p3=self._logits.data_ptr(),
+                             p4=final_norm.data_ptr(),
+                             p6=self._parts.data_ptr(), n=self.V,
+                             k=self.hidden, it=_IT_RMSNORM | _tbit,
+                             epi=_EP_F32_AMAX, hkv=1, f0=self.eps))
+            names.append("lm_head")
+            rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                             p3=self._token.data_ptr(),
+                             p4=self._tokens_out.data_ptr()))
+            names.append("argmax.fin")
+        se = max(staged_elems[0], self._maxk)
+        return rows, names, se
+    def _build_programs(self):
+        dev = self.device
+        rows, names, _ = self._build_rows(kv_slice=0, batch=False)
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        if self._batch:
+            big = self._batch > min(8, 16384 // self._maxk)
+            rows, _, se = self._build_rows(kv_slice=0, batch=True, big=big)
+            self._prog_batch = torch.tensor(rows, dtype=torch.int64,
+                                            device=dev)
+            self._batch_stage_elems = int(se)
+            self._progs_prefill = [self._prog]
+            for b in range(1, self._batch):
+                rows, _, _ = self._build_rows(kv_slice=b, batch=False)
+                self._progs_prefill.append(
+                    torch.tensor(rows, dtype=torch.int64, device=dev))
+    def batch_max(self) -> int:
+        """Largest batch size this model supports. Up to 8 sequences stage
+        in shared; beyond that the transformed input moves to a global
+        scratch, bounded by the kernel's register accumulators."""
+        try:
+            hard = int(ops.mak_batch_maxb())
+        except (AttributeError, RuntimeError):
+            hard = 8
+        return 1 if self._gemma else hard
+    # ------------------------------------------------------------------
+    def enable_batch(self, B: int):
+        """Allocate per-sequence KV caches, token buffers, and programs for
+        batched decode. B up to 8 stages the input in shared; larger B (up
+        to batch_max()) routes the transformed input through a global
+        scratch so the batch width is not bounded by the projection size."""
+        if self._gemma:
+            raise ValueError("batched decode covers the llama-family path")
+        maxb = self.batch_max()
+        if not 1 <= B <= maxb:
+            raise ValueError(f"B must be in [1, {maxb}] for this model")
+        if self._batch == B:
+            return
+        dt, dev = torch.bfloat16, self.device
+        self._kcache = torch.zeros(self.L, B, self.Hkv, self.max_seq,
+                                   self.D, dtype=dt, device=dev)
+        self._vcache = torch.zeros_like(self._kcache)
+        self._token = torch.zeros(B, dtype=torch.int32, device=dev)
+        self._tokens_out = torch.zeros(B * self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._logits = torch.empty(B * self.V, dtype=torch.float32,
+                                   device=dev)
+        self._parts = torch.zeros(B * self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        self._pos_b = torch.zeros(B, dtype=torch.int32, device=dev)
+        # transformed-input scratch: one row per active token, and prefill
+        # (during generate_batch) still uses up to chunk_m rows
+        self._xg = torch.empty(max(8, B) * self._maxk, dtype=dt, device=dev)
+        # working buffers hold one row per active sequence; prefill chunks
+        # still use up to chunk_m rows, so keep at least the fused capacity
+        nb = max(8, B)
+        qkvd = self._qdim + 2 * self._kvdim
+        self._hidden = torch.empty(nb * self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty(nb * qkvd, dtype=dt, device=dev)
+        self._gu = torch.empty(nb * 2 * self.I, dtype=dt, device=dev)
+        self._partials = torch.empty(
+            nb * self.Hq * self._maxch * (self.D + 2), dtype=torch.float32,
+            device=dev)
+        self._kv_bstride = self.Hkv * self.max_seq * self.D
+        self._batch = B
+        self._build_programs()
+    def decode_batch(self, tokens: Sequence[int], positions: Sequence[int],
+                     steps: int = 1) -> torch.Tensor:
+        """Batched greedy decode: sequence b consumes tokens[b] at
+        positions[b]; `steps` in-kernel steps run with per-sequence argmax
+        feedback. Returns the generated tokens [B, steps]; the final
+        step's fp32 logits are live in batch_logits()."""
+        B = self._batch
+        if B < 1:
+            raise ValueError("call enable_batch(B) first")
+        if len(tokens) != B or len(positions) != B:
+            raise ValueError("tokens and positions must have B entries")
+        if steps < 1 or steps > self.max_gen:
+            raise ValueError("steps out of range")
+        if max(positions) + steps > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        self._token.copy_(torch.tensor(list(tokens), dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor(list(positions), dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b, steps,
+                          0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        return self._tokens_out.view(B, self.max_gen)[:, :steps]
+    def batch_logits(self) -> torch.Tensor:
+        """fp32 logits [B, V] of the most recent batched step."""
+        return self._logits.view(max(self._batch, 1), self.V)
+    def generate_batch(self, prompts: Sequence[Sequence[int]],
+                       max_new: int) -> List[List[int]]:
+        """Greedy generation for a batch of prompts: per-sequence prefill,
+        then one batched decode launch. Each sequence's token stream is
+        bit-identical to its own single-sequence generate()."""
+        B = len(prompts)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        self.enable_batch(B)
+        firsts: List[int] = []
+        for b, ids in enumerate(prompts):
+            P = len(ids)
+            if P < 1 or P + max_new > self.max_seq:
+                raise ValueError("sequence exceeds max_seq")
+            self._prompt_in[:P].copy_(
+                torch.tensor(list(ids), dtype=torch.int32))
+            ops.mak_run_seq(self._progs_prefill[b], self._bar, 0, P, 1 - P,
+                            P, self._maxk, self._chunk_m)
+            firsts.append(int(self._token[0].item()))
+        if max_new == 1:
+            return [[t] for t in firsts]
+        self._token.copy_(torch.tensor(firsts, dtype=torch.int32))
+        self._pos_b.copy_(torch.tensor([len(p) for p in prompts],
+                                       dtype=torch.int32))
+        ops.mak_run_batch(self._prog_batch, self._bar, self._pos_b,
+                          max_new - 1, 0, self._maxk, self._kv_bstride,
+                          self._batch_stage_elems)
+        rest = (self._tokens_out.view(B, self.max_gen)[:, :max_new - 1]
+                .cpu().tolist())
+        return [[firsts[b]] + rest[b] for b in range(B)]
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _dense_weight(mod):
+        """The effective bf16 weight of a linear layer, dequantizing packed
+        quantized layers (bitsandbytes 4-bit and 8-bit) so a quantized
+        checkpoint loads directly. Dequantization is exact and does not
+        depend on the card, so inference on the result keeps every
+        guarantee; the weights materialize to bf16, so the quantized memory
+        footprint is not preserved (that needs the in-kernel packed path)."""
+        w = mod.weight
+        qs = getattr(w, "quant_state", None)
+        if qs is not None:  # bitsandbytes 4-bit (nf4 / fp4)
+            import bitsandbytes as bnb
+            return bnb.functional.dequantize_4bit(w.data, qs).to(torch.bfloat16)
+        if getattr(w, "SCB", None) is not None or hasattr(mod, "SCB"):
+            # bitsandbytes 8-bit (LLM.int8): row scales in SCB, int8 in CB
+            scb = w.SCB if getattr(w, "SCB", None) is not None else mod.SCB
+            cb = w.data if w.data.dtype == torch.int8 else mod.CB
+            return (cb.to(torch.float32)
+                    * (scb.to(torch.float32) / 127.0).unsqueeze(1)
+                    ).to(torch.bfloat16)
+        return w
+    @staticmethod
+    def _q4_parts(mod):
+        """(packed [N, K/2] uint8, absmax [N*K/64] fp32) for a bitsandbytes
+        nf4 layer whose weights stay packed, else None. Double-quantized
+        absmax is materialized to fp32 so the kernel needs one scale array."""
+        w = getattr(mod, "weight", None)
+        qs = getattr(w, "quant_state", None)
+        if qs is None or getattr(qs, "quant_type", None) != "nf4":
+            return None
+        import bitsandbytes as bnb
+        n, k = int(qs.shape[0]), int(qs.shape[1])
+        if k % 64:
+            return None  # kernel assumes 64 | K for per-block absmax
+        packed = w.data.reshape(n, k // 2).contiguous()
+        if getattr(qs, "nested", False):
+            absmax = bnb.functional.dequantize_blockwise(
+                qs.absmax, qs.state2) + qs.offset
+        else:
+            absmax = qs.absmax
+        return packed, absmax.float().reshape(-1).contiguous()
+    @classmethod
+    def _proj(cls, mods):
+        """A projection, concatenated across `mods` (q/k/v or gate/up). Stays
+        nf4-packed when every part is nf4, else a dense bf16 tensor."""
+        parts = [cls._q4_parts(m) for m in mods]
+        if all(p is not None for p in parts):
+            packed = torch.cat([p[0] for p in parts], 0)
+            absmax = torch.cat([p[1] for p in parts], 0)
+            return {"q4": True, "packed": packed, "absmax": absmax,
+                    "N": packed.shape[0], "K": packed.shape[1] * 2}
+        return torch.cat([cls._dense_weight(m) for m in mods], 0)
+    @classmethod
+    def from_pretrained(cls, model, device="cuda", max_seq: int = 4096,
+                        max_gen: int = 4096):
+        """Build from a transformers model (object or repo id). A quantized
+        checkpoint (bitsandbytes 4-bit/8-bit) is accepted directly; its
+        weights are dequantized to bf16 at load."""
+        if isinstance(model, str):
+            from transformers import AutoModelForCausalLM
+            model = AutoModelForCausalLM.from_pretrained(
+                model, torch_dtype=torch.bfloat16)
+        if "gemma4" in getattr(model.config, "model_type", ""):
+            return cls._from_gemma4(model, device, max_seq, max_gen)
+        hf = model.config
+        dec = model.model
+        attn0 = dec.layers[0].self_attn
+        qk_norm = hasattr(attn0, "q_norm") and attn0.q_norm is not None
+        if getattr(attn0.q_proj, "bias", None) is not None:
+            raise ValueError("attention biases are not supported")
+        # rope config across transformers API generations: v5 keeps theta and
+        # type in a rope_parameters dict; v4 used a rope_theta attribute and
+        # an optional rope_scaling dict.
+        rp = getattr(hf, "rope_parameters", None)
+        if not isinstance(rp, dict):
+            rs = getattr(hf, "rope_scaling", None)
+            rp = rs if isinstance(rs, dict) else {}
+        theta = rp.get("rope_theta")
+        if theta is None:
+            theta = getattr(hf, "rope_theta", 10000.0)
+        rope_type = rp.get("rope_type", rp.get("type", "default"))
+        if rope_type not in ("default", "llama3"):
+            raise ValueError(f"rope type {rope_type!r} is not supported")
+        # take the frequency table transformers computed (scaling included)
+        rot = dec.rotary_emb
+        att_scale = float(getattr(rot, "attention_scaling", 1.0))
+        if att_scale != 1.0:
+            raise ValueError("rope attention_scaling != 1 is not supported")
+        inv_freq = rot.inv_freq.detach().float()
+        pj = cls._proj
+        layers = []
+        for lyr in dec.layers:
+            a, m = lyr.self_attn, lyr.mlp
+            layers.append({
+                "wqkv": pj([a.q_proj, a.k_proj, a.v_proj]),
+                "wo": pj([a.o_proj]),
+                "wgu": pj([m.gate_proj, m.up_proj]),
+                "wdown": pj([m.down_proj]),
+                "ln1": lyr.input_layernorm.weight,
+                "ln2": lyr.post_attention_layernorm.weight,
+                "qn": a.q_norm.weight if qk_norm else None,
+                "kn": a.k_norm.weight if qk_norm else None,
+            })
+        weights = {"embed": dec.embed_tokens.weight, "norm": dec.norm.weight,
+                   "lm_head": pj([model.lm_head]), "layers": layers,
+                   "inv_freq": inv_freq}
+        config = {
+            "num_hidden_layers": hf.num_hidden_layers,
+            "hidden_size": hf.hidden_size,
+            "num_attention_heads": hf.num_attention_heads,
+            "num_key_value_heads": getattr(hf, "num_key_value_heads", None),
+            "head_dim": getattr(hf, "head_dim", None),
+            "intermediate_size": hf.intermediate_size,
+            "vocab_size": hf.vocab_size,
+            "rope_theta": float(theta),
+            "rms_norm_eps": hf.rms_norm_eps,
+            "qk_norm": qk_norm,
+        }
+        return cls(weights, config, device=device, max_seq=max_seq,
+                   max_gen=max_gen)
+    # ------------------------------------------------------------------
+    @classmethod
+    def _from_gemma4(cls, model, device, max_seq, max_gen):
+        """gemma-4-E2B-it program: gemma norms (fp32, raw weight),
+        gelu-tanh gating, post-sublayer norm phases, per-layer-input
+        pathway, sliding windows, shared KV, dual head dims, softcapped
+        head. Scope is the E2B text model only."""
+        self = cls.__new__(cls)
+        cfg = model.config.text_config if hasattr(model.config, "text_config") \
+            else model.config
+        dec = model.model.language_model if hasattr(model.model, "language_model") \
+            else model.model
+        assert cfg.model_type == "gemma4_text", cfg.model_type
+        assert not cfg.enable_moe_block
+        self.L = int(cfg.num_hidden_layers)
+        self.hidden = int(cfg.hidden_size)
+        self.Hq = int(cfg.num_attention_heads)
+        self.V = int(cfg.vocab_size)
+        self.eps = float(cfg.rms_norm_eps)
+        self.qk_norm = True
+        self.max_seq = int(max_seq)
+        self.max_gen = int(max_gen)
+        self.device = torch.device(device)
+        self.tiled = False
+        win = int(cfg.sliding_window)
+        cap = float(cfg.final_logit_softcapping)
+        PL = int(cfg.hidden_size_per_layer_input)
+        dt = torch.bfloat16
+        dev = self.device
+        self._keep = []
+        def pack(t):
+            t = t.detach().to(device=dev, dtype=dt).contiguous()
+            self._keep.append(t)
+            return t
+        def packf(t):
+            t = t.detach().to(device=dev, dtype=torch.float32).contiguous()
+            self._keep.append(t)
+            return t
+        embed = pack(dec.embed_tokens.weight)
+        ple_tab = pack(dec.embed_tokens_per_layer.weight)
+        plm_w = pack(dec.per_layer_model_projection.weight)
+        pln_w = pack(dec.per_layer_projection_norm.weight)
+        fnorm = pack(dec.norm.weight)
+        rot = dec.rotary_emb
+        invf = {lt: packf(getattr(rot, f"{lt}_inv_freq").float())
+                for lt in set(cfg.layer_types)}
+        ascale = {lt: float(getattr(rot, f"{lt}_attention_scaling"))
+                  for lt in set(cfg.layer_types)}
+        layers = []
+        for li, lyr in enumerate(dec.layers):
+            a = lyr.self_attn
+            D_l = int(a.head_dim)
+            shared = bool(a.is_kv_shared_layer)
+            ent = {
+                "D": D_l, "shared": shared,
+                "lscale": float(lyr.layer_scalar.float().item()),
+                "type": cfg.layer_types[li],
+                "I": int(lyr.mlp.intermediate_size),
+                "qn": pack(a.q_norm.weight),
+                "o": pack(a.o_proj.weight),
+                "ln_in": pack(lyr.input_layernorm.weight),
+                "ln_pa": pack(lyr.post_attention_layernorm.weight),
+                "ln_pf": pack(lyr.pre_feedforward_layernorm.weight),
+                "ln_ff": pack(lyr.post_feedforward_layernorm.weight),
+                "ln_pl": pack(lyr.post_per_layer_input_norm.weight),
+                "plig": pack(lyr.per_layer_input_gate.weight),
+                "plpr": pack(lyr.per_layer_projection.weight),
+                "wgu": pack(torch.cat([lyr.mlp.gate_proj.weight,
+                                       lyr.mlp.up_proj.weight], 0)),
+                "wdn": pack(lyr.mlp.down_proj.weight),
+            }
+            if shared:
+                ent["wqkv"] = pack(a.q_proj.weight)
+                ent["kn"] = None
+            else:
+                ent["wqkv"] = pack(torch.cat(
+                    [a.q_proj.weight, a.k_proj.weight, a.v_proj.weight], 0))
+                ent["kn"] = pack(a.k_norm.weight)
+                ent["kc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                ent["vc"] = torch.zeros(self.max_seq, D_l, dtype=dt,
+                                        device=dev)
+                self._keep += [ent["kc"], ent["vc"]]
+            layers.append(ent)
+        # shared layers read the last non-shared layer of their type
+        src = {}
+        for li, ent in enumerate(layers):
+            if not ent["shared"]:
+                src[ent["type"]] = ent
+        for ent in layers:
+            if ent["shared"]:
+                ent["kc"] = src[ent["type"]]["kc"]
+                ent["vc"] = src[ent["type"]]["vc"]
+        maxD = max(e["D"] for e in layers)
+        maxI = max(e["I"] for e in layers)
+        self._maxk = max(self.hidden, self.Hq * maxD, maxI)
+        self._chunk_m = 1
+        self._chunk = 128
+        self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
+        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._tmp = torch.empty(self.hidden, dtype=dt, device=dev)
+        self._qkv = torch.empty((self.Hq + 2) * maxD, dtype=dt, device=dev)
+        self._gu = torch.empty(2 * maxI, dtype=dt, device=dev)
+        self._t256 = torch.empty(PL, dtype=dt, device=dev)
+        self._ctxraw = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._pletok = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._ple = torch.empty(self.L * PL, dtype=dt, device=dev)
+        self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
+        self._partials = torch.empty(
+            self.Hq * self._maxch * (maxD + 2), dtype=torch.float32,
+            device=dev)
+        self._token = torch.zeros(1, dtype=torch.int32, device=dev)
+        self._prompt_in = torch.zeros(self.max_seq, dtype=torch.int32,
+                                      device=dev)
+        self._tokens_out = torch.zeros(self.max_gen, dtype=torch.int32,
+                                       device=dev)
+        self._bar = torch.zeros(34, dtype=torch.int32, device=dev)
+        probe = torch.zeros(1, 16, dtype=torch.int64, device=dev)
+        self._nblocks = int(ops.mak_num_blocks(probe, self._maxk))
+        self._parts = torch.zeros(self._nblocks, dtype=torch.int64,
+                                  device=dev)
+        try:
+            self._bw_per_sm = float(ops.mak_bw_per_sm(probe))
+        except (AttributeError, RuntimeError):
+            self._bw_per_sm = 0.0
+        self._invf = invf[layers[0]["type"]]
+        self.Hkv, self.D, self.I = 1, maxD, maxI
+        self.theta = 0.0
+        emb_scale = float(torch.tensor(self.hidden ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        ple_scale = float(torch.tensor(PL ** 0.5,
+                                       dtype=torch.float32).to(dt))
+        proj_scale = float(torch.tensor(self.hidden ** -0.5,
+                                        dtype=torch.float32))
+        hid = self._hidden.data_ptr()
+        tmp = self._tmp.data_ptr()
+        rows, names = [], []
+        rows.append(_row(_OP_EMBED, p1=ple_tab.data_ptr(),
+                         p2=embed.data_ptr(), p3=hid,
+                         p4=self._parts.data_ptr(),
+                         p5=self._prompt_in.data_ptr(),
+                         p6=self._token.data_ptr(),
+                         n=self._tokens_out.data_ptr(), k=self.hidden,
+                         it=self._pletok.data_ptr(),
+                         hkv=self.L * PL, f0=emb_scale,
+                         i0=_fbits(ple_scale)))
+        names.append("embed")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=plm_w.data_ptr(),
+                         p3=self._ctxraw.data_ptr(), n=self.L * PL,
+                         k=self.hidden, it=_IT_NONE, epi=_EP_STORE))
+        names.append("plm")
+        rows.append(_row(8, p1=self._ctxraw.data_ptr(),
+                         p2=pln_w.data_ptr(), p3=self._ple.data_ptr(),
+                         p4=self._pletok.data_ptr(), n=self.L, k=PL,
+                         f0=self.eps, i0=_fbits(proj_scale)))
+        names.append("plemix")
+        for li, e in enumerate(layers):
+            D_l, qdim = e["D"], self.Hq * e["D"]
+            lt = e["type"]
+            w_l = win if lt == "sliding_attention" else 0
+            asb = _fbits(ascale[lt]) if ascale[lt] != 1.0 else 0
+            flags = 1 | 2 | (0 if e["shared"] else 4) | \
+                (8 if e["shared"] else 0)
+            n_qkv = qdim if e["shared"] else qdim + 2 * D_l
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wqkv"].data_ptr(),
+                             p3=self._qkv.data_ptr(),
+                             p4=e["ln_in"].data_ptr(), n=n_qkv,
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.qkv")
+            rows.append(_row(_OP_ATTN, p1=self._qkv.data_ptr(),
+                             p2=e["kc"].data_ptr(), p3=e["vc"].data_ptr(),
+                             p4=self._partials.data_ptr(),
+                             p5=e["qn"].data_ptr(),
+                             p6=e["kn"].data_ptr() if e["kn"] is not None
+                             else 0,
+                             n=self.max_seq,
+                             k=flags | (w_l << 16), it=_fbits(self.eps),
+                             epi=invf[lt].data_ptr(), hq=self.Hq,
+                             hkv=1 | (self._chunk << 16), d=D_l, f0=1.0,
+                             i0=self._maxch | (asb << 32)))
+            names.append(f"L{li}.attn")
+            rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
+                             p2=e["o"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_STORE, hq=self.Hq,
+                             hkv=self._chunk << 16, d=D_l,
+                             i0=self._maxch | (w_l << 16)))
+            names.append(f"L{li}.o")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pa"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_attn")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["wgu"].data_ptr(),
+                             p3=self._gu.data_ptr(),
+                             p4=e["ln_pf"].data_ptr(), n=2 * e["I"],
+                             k=self.hidden, it=_IT_RMSNORM_G,
+                             epi=_EP_STORE, f0=self.eps))
+            names.append(f"L{li}.gateup")
+            rows.append(_row(_OP_GEMV, p1=self._gu.data_ptr(),
+                             p2=e["wdn"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=e["I"], it=_IT_GELU_GLU,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.down")
+            rows.append(_row(7, p1=tmp, p2=e["ln_ff"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps, i0=_fbits(1.0)))
+            names.append(f"L{li}.nr_ffw")
+            rows.append(_row(_OP_GEMV, p1=hid, p2=e["plig"].data_ptr(),
+                             p3=self._t256.data_ptr(),
+                             p6=self._ple.data_ptr() + li * PL * 2,
+                             n=PL, k=self.hidden, it=_IT_NONE,
+                             epi=_EP_GELU_PLE))
+            names.append(f"L{li}.plig")
+            rows.append(_row(_OP_GEMV, p1=self._t256.data_ptr(),
+                             p2=e["plpr"].data_ptr(), p3=tmp,
+                             n=self.hidden, k=PL, it=_IT_NONE,
+                             epi=_EP_STORE))
+            names.append(f"L{li}.plproj")
+            rows.append(_row(7, p1=tmp, p2=e["ln_pl"].data_ptr(), p3=hid,
+                             n=self.hidden, f0=self.eps,
+                             i0=_fbits(e["lscale"])))
+            names.append(f"L{li}.nr_ple")
+        rows.append(_row(_OP_GEMV, p1=hid, p2=embed.data_ptr(),
+                         p3=self._logits.data_ptr(),
+                         p4=fnorm.data_ptr(),
+                         p6=self._parts.data_ptr(), n=self.V,
+                         k=self.hidden, it=_IT_RMSNORM_G,
+                         epi=_EP_F32_AMAX_CAP, hkv=1, f0=self.eps,
+                         i0=_fbits(cap)))
+        names.append("lm_head")
+        rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
+                         p3=self._token.data_ptr(),
+                         p4=self._tokens_out.data_ptr()))
+        names.append("argmax.fin")
+        self._prog = torch.tensor(rows, dtype=torch.int64, device=dev)
+        self.phase_names = names
+        self._gemma = True
+        self._batch = 0
+        return self
+    # ------------------------------------------------------------------
+    def decode_step(self, token: int, pos: int, phased: bool = False):
+        """One decode step; returns the live fp32 logits buffer [V]."""
+        if pos >= self.max_seq:
+            raise ValueError("pos exceeds max_seq")
+        self._token.fill_(int(token))
+        if phased:
+            ops.mak_run_phased(self._prog, pos, self.max_gen - 1, False,
+                               self._maxk)
+        else:
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+        return self._logits[:self.V]
+    def decode_step_timed(self, token: int, pos: int):
+        """One phase-per-launch step; returns (logits, per-phase ms)."""
+        self._token.fill_(int(token))
+        ms = ops.mak_run_phased(self._prog, pos, self.max_gen - 1, True,
+                                self._maxk)
+        return self._logits[:self.V], ms
+    def prefill(self, ids: Sequence[int], phased: bool = False):
+        """Consume a prompt; returns the last position's logits. Runs in
+        chunks of up to chunk_m tokens with weight reads amortized across
+        the chunk (bitwise identical to token-by-token consumption)."""
+        P = len(ids)
+        if phased:
+            for i, t in enumerate(ids):
+                self.decode_step(int(t), i, phased=True)
+            return self._logits
+        self._prompt_in[:P].copy_(torch.tensor(ids, dtype=torch.int32))
+        ops.mak_run_seq(self._prog, self._bar, 0, P, 1 - P, P, self._maxk,
+                        self._chunk_m)
+        return self._logits[:self.V]
+    def generate(self, prompt_ids: Sequence[int], max_new: int,
+                 single_launch: bool = True) -> List[int]:
+        """Greedy generation. With single_launch (default) the entire call,
+        prompt consumption included, is one kernel launch: the in-kernel step
+        loop reads prompt tokens from a staged device buffer, then feeds each
+        step's argmax to the next step's embedding read. With
+        single_launch=False the same computation runs as one launch per
+        token; the two paths produce identical tokens.
+        """
+        P = len(prompt_ids)
+        if max_new < 1 or max_new > self.max_gen:
+            raise ValueError("max_new out of range")
+        if P < 1 or P + max_new > self.max_seq:
+            raise ValueError("sequence exceeds max_seq")
+        if single_launch:
+            self._prompt_in[:P].copy_(
+                torch.tensor(prompt_ids, dtype=torch.int32))
+            ops.mak_run_seq(self._prog, self._bar, 0, P + max_new - 1,
+                            1 - P, P, self._maxk, self._chunk_m)
+            return self._tokens_out[:max_new].cpu().tolist()
+        pos = 0
+        for t in prompt_ids[:-1]:
+            self._token.fill_(int(t))
+            ops.mak_run(self._prog, self._bar, pos, self.max_gen - 1,
+                        self._maxk)
+            pos += 1
+        self._token.fill_(int(prompt_ids[-1]))
+        ops.mak_run(self._prog, self._bar, pos, 0, self._maxk)
+        pos += 1
+        if max_new > 1:
+            ops.mak_run_steps(self._prog, self._bar, pos, max_new - 1, 1,
+                              self._maxk)
+        return self._tokens_out[:max_new].cpu().tolist()

build/torch213-cxx11-cu132-x86_64-linux/_model_as_a_kernel_cuda_e33875a.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:c0e8576073d666697bab258e0cc2ff5a784464ee621c1f5ff29403e2c6d987ca
+size 6374912

build/torch213-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

+import torch
+from . import _model_as_a_kernel_cuda_e33875a
+ops = torch.ops._model_as_a_kernel_cuda_e33875a
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_model_as_a_kernel_cuda_e33875a::{op_name}"

build/torch213-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,38 @@

+{
+  "name": "model-as-a-kernel",
+  "id": "_model_as_a_kernel_cuda_e33875a",
+  "version": 1,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "/YX5/zzAOyJacQ7H+j8zaq+xjWyy8w6r0j+ugE1EUs8=",
+      "_model_as_a_kernel_cuda_e33875a.abi3.so": "wOhXYHPWZml7qyWODML/WnhEZO5iHB9f8pQD4sbZh8o=",
+      "_ops.py": "Pqj4WKtwPyeDG9grQnzLget4g1jsGUdR7ZMa44e4rhs=",
+      "model_as_a_kernel/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "05f1ed487f599a320311aa26076ec0ad290c3f68",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "e33875ae4eba42fd136a89c388e806281cd4759c",
+      "dirty": false
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/model_as_a_kernel/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))