batched decode; batch-invariance gates

Browse files

Files changed (4) hide show

README.md +5 -3
mak_cuda/megakernel.cu +73 -0
tests/test_mak.py +35 -0
torch-ext/model_as_a_kernel/__init__.py +141 -43

README.md CHANGED Viewed

@@ -179,6 +179,8 @@ to 1673 at batch 16 (6.5x aggregate).
 One CUDA device, greedy sampling in-kernel (route the fp32 logits to an
 external sampler for anything else). Llama-family architectures without
-rope attention scaling or attention biases. Prompts prefill in chunks of
-at most eight tokens. The program bakes device pointers, so the
-`MegaModel` owns its weights and buffers for its lifetime.

 One CUDA device, greedy sampling in-kernel (route the fp32 logits to an
 external sampler for anything else). Llama-family architectures without
+rope attention scaling or attention biases. Weights are dense bf16 or
+bitsandbytes nf4 (4-bit checkpoints load with their weights kept packed
+and dequantized in the kernel). Prompts prefill in chunks of at most eight
+tokens. The program bakes device pointers, so the `MegaModel` owns its
+weights and buffers for its lifetime.

mak_cuda/megakernel.cu CHANGED Viewed

@@ -55,8 +55,20 @@ enum PhaseOp : int {
   OP_NORMB = 10,     // RMSNorm (llama or gemma) [B][K] -> scratch
   OP_GLUB = 11,      // SwiGLU / gelu-glu gating [B][2K] -> scratch [B][K]
   OP_ATTNFINB = 12,  // finalize split-S attention [B] -> scratch [B][qdim]
 };
 enum InputTransform : int {
   IT_NONE = 0,
   IT_RMSNORM = 1,
@@ -1326,6 +1338,66 @@ __device__ void ph_gemv_plain(const long long* P, int M, int rs) {
   }
 }
 template <int NB>
 __device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
                                                int M, int step_slot,
@@ -1343,6 +1415,7 @@ __device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
     case OP_NORMRES: ph_normres(P); break;
     case OP_PLEMIX: ph_plemix(P); break;
     case OP_GEMV_PLAIN: ph_gemv_plain<NB>(P, M, ring_stages); break;
     case OP_NORMB: ph_normb(P, M); break;
     case OP_GLUB: ph_glub(P, M); break;
     case OP_ATTNFINB: ph_attnfinb(P, M, pos_b, step); break;

   OP_NORMB = 10,     // RMSNorm (llama or gemma) [B][K] -> scratch
   OP_GLUB = 11,      // SwiGLU / gelu-glu gating [B][2K] -> scratch [B][K]
   OP_ATTNFINB = 12,  // finalize split-S attention [B] -> scratch [B][qdim]
+  OP_GEMV_Q4 = 13,   // GEMV with nf4-packed weights, dequantized in-kernel
 };
+// NF4 codebook (bitsandbytes): the 16 quantile levels a 4-bit index maps
+// to. Dequantized weight = code[nibble] * per-block absmax, rounded to
+// bf16; the lookup is exact and the multiply is IEEE, so it reproduces on
+// every card.
+__device__ __constant__ float NF4_CODE[16] = {
+    -1.0f, -0.6961928009986877f, -0.5250730514526367f,
+    -0.39491748809814453f, -0.28444138169288635f, -0.18477343022823334f,
+    -0.09105003625154495f, 0.0f, 0.07958029955625534f, 0.16093020141124725f,
+    0.24611230194568634f, 0.33791524171829224f, 0.44070982933044434f,
+    0.5626170039176941f, 0.7229568362236023f, 1.0f};
 enum InputTransform : int {
   IT_NONE = 0,
   IT_RMSNORM = 1,
   }
 }
+// GEMV with nf4-packed weights [N][K/2] (two 4-bit indices per byte, row
+// major) and per-64-element fp32 absmax. Each lane dequantizes its own
+// weight tile (code[nibble] * absmax, rounded to bf16) and accumulates in
+// the exact order of the dense path, so the result is bit-identical to a
+// dense GEMV over the same dequantized weights and reproduces on every
+// card. Input is dense bf16 in a global buffer (a transform phase writes
+// it). Weights stay packed, so the quantized memory footprint is kept.
+template <int NB>
+__device__ void ph_gemv_q4(const long long* P, int M) {
+  const unsigned char* W = reinterpret_cast<const unsigned char*>(P[2]);
+  const float* absmax = reinterpret_cast<const float*>(P[6]);
+  const bf16* resid = reinterpret_cast<const bf16*>(P[5]);
+  const bf16* xin = reinterpret_cast<const bf16*>(P[1]);
+  const long long N = P[7], K = P[8];
+  const int epi = (int)P[10];
+  const int lane = threadIdx.x & 31;
+  const long long gwarp =
+      ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
+  const long long nwarps = ((long long)gridDim.x * blockDim.x) >> 5;
+  const long long kbase = (long long)lane * 8;
+  const long long ntiles = (kbase < K) ? (K - kbase + 255) / 256 : 0;
+  const long long Kh = K >> 1;  // packed bytes per row
+  for (long long row = gwarp; row < N; row += nwarps) {
+    float acc[NB];
+#pragma unroll
+    for (int m = 0; m < NB; ++m) acc[m] = 0.f;
+    for (long long t = 0; t < ntiles; ++t) {
+      const long long k = kbase + t * 256;
+      const unsigned int pw =
+          *reinterpret_cast<const unsigned int*>(W + row * Kh + (k >> 1));
+      const float am = absmax[(row * K + k) >> 6];
+      float w8[8];
+#pragma unroll
+      for (int b = 0; b < 4; ++b) {
+        const unsigned int by = (pw >> (b * 8)) & 0xFFu;
+        // bitsandbytes packs the earlier element in the high nibble
+        w8[2 * b] = bf2f(f2bf(NF4_CODE[by >> 4] * am));
+        w8[2 * b + 1] = bf2f(f2bf(NF4_CODE[by & 0xF] * am));
+      }
+#pragma unroll
+      for (int m = 0; m < NB; ++m) {
+        if (m < M) {
+          const bf16* xp = xin + (long long)m * K + k;
+#pragma unroll
+          for (int j = 0; j < 8; ++j)
+            acc[m] = fmaf(w8[j], bf2f(xp[j]), acc[m]);
+        }
+      }
+    }
+#pragma unroll
+    for (int m = 0; m < NB; ++m) {
+      if (m < M) {
+        const float a = warp_sum(acc[m]);
+        if (lane == 0) gemv_store(P, epi, N, resid, m, row, a);
+      }
+    }
+  }
+}
 template <int NB>
 __device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
                                                int M, int step_slot,
     case OP_NORMRES: ph_normres(P); break;
     case OP_PLEMIX: ph_plemix(P); break;
     case OP_GEMV_PLAIN: ph_gemv_plain<NB>(P, M, ring_stages); break;
+    case OP_GEMV_Q4: ph_gemv_q4<NB>(P, M); break;
     case OP_NORMB: ph_normb(P, M); break;
     case OP_GLUB: ph_glub(P, M); break;
     case OP_ATTNFINB: ph_attnfinb(P, M, pos_b, step); break;

tests/test_mak.py CHANGED Viewed

@@ -212,6 +212,41 @@ def test_batch_invariance(models):
     assert all(perm[j] == solo[order[j]] for j in range(B))
 @requires_cuda
 @pytest.mark.parametrize("model_id", ["Qwen/Qwen3-0.6B",
                                       "TinyLlama/TinyLlama-1.1B-Chat-v1.0"])

     assert all(perm[j] == solo[order[j]] for j in range(B))
+@requires_cuda
+def test_nf4_packed_matches_dequant():
+    """A bitsandbytes nf4 checkpoint loads with weights kept packed and
+    dequantized in the kernel, bitwise identical to the same weights
+    dequantized to bf16, and deterministic."""
+    bnb = pytest.importorskip("bitsandbytes")
+    transformers = pytest.importorskip("transformers")
+    from transformers import BitsAndBytesConfig, AutoModelForCausalLM
+    mid = "HuggingFaceTB/SmolLM2-135M"
+    try:
+        q = AutoModelForCausalLM.from_pretrained(
+            mid, quantization_config=BitsAndBytesConfig(
+                load_in_4bit=True, bnb_4bit_quant_type="nf4",
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_use_double_quant=False))
+    except Exception as e:
+        pytest.skip(f"bitsandbytes 4-bit unavailable here: {e}")
+    mm = mak.MegaModel.from_pretrained(q, max_seq=512, max_gen=64)
+    assert mm._has_q4 and mm._layers[0]["wdown"]["packed"].dtype == torch.uint8
+    ref = AutoModelForCausalLM.from_pretrained(mid, dtype=torch.bfloat16)
+    for name, mod in q.named_modules():
+        if isinstance(mod, bnb.nn.Linear4bit):
+            w = bnb.functional.dequantize_4bit(mod.weight.data,
+                                               mod.weight.quant_state)
+            ref.get_submodule(name).weight.data = \
+                w.to(torch.bfloat16).cpu().clone()
+    mr = mak.MegaModel.from_pretrained(ref.cuda().eval(), max_seq=512,
+                                       max_gen=64)
+    prompt = [5, 9, 13, 21, 33, 41, 40, 32]
+    assert torch.equal(mm.decode_step(prompt[0], 0),
+                       mr.decode_step(prompt[0], 0))
+    g = mm.generate(prompt, 32)
+    assert g == mr.generate(prompt, 32) == mm.generate(prompt, 32)
 @requires_cuda
 @pytest.mark.parametrize("model_id", ["Qwen/Qwen3-0.6B",
                                       "TinyLlama/TinyLlama-1.1B-Chat-v1.0"])

torch-ext/model_as_a_kernel/__init__.py CHANGED Viewed

@@ -28,6 +28,7 @@ _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
 _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
@@ -124,7 +125,7 @@ class MegaModel:
         self._keep: List[torch.Tensor] = []
         embed = pack(weights["embed"])
         final_norm = pack(weights["norm"])
-        lm_head = pack_tiled(weights["lm_head"])
         inv_freq = weights.get("inv_freq")
         if inv_freq is None:
             d_idx = torch.arange(0, self.D, 2, dtype=torch.float32)
@@ -137,24 +138,49 @@ class MegaModel:
         self._invf = inv_freq
         qdim, kvdim = self.Hq * self.D, self.Hkv * self.D
         self._layers = []
         for lw in weights["layers"]:
-            assert lw["wqkv"].shape == (qdim + 2 * kvdim, self.hidden)
-            assert lw["wo"].shape == (self.hidden, qdim)
-            assert lw["wgu"].shape == (2 * self.I, self.hidden)
-            assert lw["wdown"].shape == (self.hidden, self.I)
             layer = {
-                "wqkv": pack_tiled(lw["wqkv"]), "wo": pack_tiled(lw["wo"]),
-                "wgu": pack_tiled(lw["wgu"]),
-                "wdown": pack_tiled(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._layers.append(layer)
         assert len(self._layers) == self.L
         assert embed.shape == (self.V, self.hidden)
-        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
@@ -202,6 +228,10 @@ class MegaModel:
         self._pos_b = None
         self._kv_bstride = 0
         self._gemma = False
         self._build_programs()
     # ------------------------------------------------------------------
@@ -222,7 +252,7 @@ class MegaModel:
         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() if big else 0
         B = self._batch if batch else 1
         rows, names = [], []
         staged_elems = [0]  # widest [B][K] panel any fused projection stages
@@ -233,52 +263,63 @@ class MegaModel:
         def note(K):
             staged_elems[0] = max(staged_elems[0], B * K)
         def norm_gemv(name, gamma, w, out, N, K, epi, inp, resid=0,
                       variant=_IT_RMSNORM):
-            if 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")
-                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w, p3=out, p5=resid,
-                                 n=N, k=K, it=_tbit, epi=epi))
                 names.append(name)
             else:
                 note(K)
-                rows.append(_row(_OP_GEMV, p1=inp, p2=w, 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 big and not fits(self.I):
                 rows.append(_row(_OP_GLUB, p1=gu, p3=xg, k=self.I, it=variant))
                 names.append(name + ".glu")
-                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w, p3=out, p5=resid,
-                                 n=self.hidden, k=self.I, it=_tbit,
-                                 epi=_EP_RESID))
                 names.append(name)
             else:
                 note(self.I)
-                rows.append(_row(_OP_GEMV, p1=gu, p2=w, 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 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")
-                rows.append(_row(_OP_GEMV_PLAIN, p1=xg, p2=w, p3=out, p5=resid,
-                                 n=self.hidden, k=qdim, it=_tbit,
-                                 epi=_EP_RESID))
                 names.append(name)
             else:
                 note(qdim)
                 rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
-                                 p2=w, 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))
@@ -297,7 +338,7 @@ class MegaModel:
             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"].data_ptr(), 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,
@@ -317,14 +358,13 @@ class MegaModel:
                              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"].data_ptr(), hid, hid)
             norm_gemv(f"L{li}.gateup", lw["ln2"].data_ptr(),
-                      lw["wgu"].data_ptr(), self._gu.data_ptr(), 2 * self.I,
                       self.hidden, _EP_STORE, hid)
-            glu_gemv(f"L{li}.down", self._gu.data_ptr(),
-                     lw["wdown"].data_ptr(), hid, hid)
-        if batch:
-            norm_gemv("lm_head", final_norm.data_ptr(), lm_head.data_ptr(),
                       self._logits.data_ptr(), self.V, self.hidden, _EP_F32,
                       hid)
             rows.append(_row(_OP_ARGMAX_PART, p1=self._logits.data_ptr(),
@@ -402,8 +442,9 @@ class MegaModel:
         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 for the global-scratch path (B > 8)
-        self._xg = torch.empty(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)
@@ -476,10 +517,67 @@ class MegaModel:
         return [[firsts[b]] + rest[b] for b in range(B)]
     # ------------------------------------------------------------------
     @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)."""
         if isinstance(model, str):
             from transformers import AutoModelForCausalLM
             model = AutoModelForCausalLM.from_pretrained(
@@ -513,22 +611,22 @@ class MegaModel:
             raise ValueError("rope attention_scaling != 1 is not supported")
         inv_freq = rot.inv_freq.detach().float()
         layers = []
         for lyr in dec.layers:
             a, m = lyr.self_attn, lyr.mlp
             layers.append({
-                "wqkv": torch.cat([a.q_proj.weight, a.k_proj.weight,
-                                   a.v_proj.weight], 0),
-                "wo": a.o_proj.weight,
-                "wgu": torch.cat([m.gate_proj.weight, m.up_proj.weight], 0),
-                "wdown": m.down_proj.weight,
                 "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": model.lm_head.weight, "layers": layers,
                    "inv_freq": inv_freq}
         config = {
             "num_hidden_layers": hf.num_hidden_layers,

 _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
         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)
         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
         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()
     # ------------------------------------------------------------------
         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 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))
             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,
                              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(),
         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)
         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(
             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,