model-as-a-kernel v1 source

Browse files

Files changed (3) hide show

mak_cuda/megakernel.cu +265 -48
tests/test_mak.py +42 -0
torch-ext/model_as_a_kernel/__init__.py +250 -0

mak_cuda/megakernel.cu CHANGED Viewed

@@ -25,7 +25,7 @@
 #include "mak_kernels.h"
 #define BLOCK_THREADS 256
-#define MAX_HEAD_DIM 256
 // GEMV weight tiles stream through a per-warp cp.async shared ring. The
 // stage count is a power of two chosen per launch from the device's
 // shared-memory budget (deep on high-bandwidth parts, where the
@@ -46,13 +46,17 @@ enum PhaseOp : int {
   OP_ARGMAX_PART = 4,
   OP_ARGMAX_FIN = 5,
   OP_KV_APPEND = 6,  // prefill chunks only: rope+append K/V for all m
 };
 enum InputTransform : int {
   IT_NONE = 0,
   IT_RMSNORM = 1,
   IT_SWIGLU = 2,
-  IT_ATTNFIN = 3,  // finalize split-S attention partials on the fly
 };
 // slot-9 flag bits above the base transform
 #define ITF_BASE_MASK 7
@@ -65,6 +69,8 @@ enum Epilogue : int {
   EP_F32 = 2,
   EP_F32_AMAX = 3,   // fp32 store + per-block argmax partial (lm head)
   EP_PARTIAL = 4,    // fp32 k-slice partials [NSLICE][N], no rounding
 };
 #define ATTN_CHUNK 128
@@ -81,6 +87,13 @@ __device__ __forceinline__ float bits2f(long long v) {
 __device__ __forceinline__ float bf2f(bf16 v) { return __bfloat162float(v); }
 __device__ __forceinline__ bf16 f2bf(float v) { return __float2bfloat16(v); }
 __device__ __forceinline__ unsigned dyn_smem_size() {
   unsigned r;
   asm("mov.u32 %0, %%dynamic_smem_size;" : "=r"(r));
@@ -190,6 +203,7 @@ __device__ void ph_embed(const long long* P, int pos, int M, int prompt_len,
     __syncthreads();
   }
   for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
        i < (long long)M * K; i += (long long)gridDim.x * blockDim.x) {
     const int m = (int)(i / K);
@@ -198,7 +212,23 @@ __device__ void ph_embed(const long long* P, int pos, int M, int prompt_len,
     const long long t = (pm < prompt_len && prompt != nullptr)
                             ? (long long)prompt[pm]
                             : (long long)s_tok;
-    out[i] = tab[t * K + k];
   }
 }
@@ -214,13 +244,13 @@ __device__ __forceinline__ unsigned long long amax_key(float v, unsigned idx) {
 // bf16 rounding of the attention output before the O projection.
 __device__ __forceinline__ float attn_finalize(const float* mpart, int h,
                                                int d, int D, int maxch,
-                                               int nch) {
   const float* base = mpart + ((long long)h * maxch) * (D + 2);
   float M = -INFINITY;
-  for (int c = 0; c < nch; ++c)
     M = fmaxf(M, base[(long long)c * (D + 2) + D]);
   float L = 0.f, acc = 0.f;
-  for (int c = 0; c < nch; ++c) {
     const float* pc = base + (long long)c * (D + 2);
     const float w = expf(pc[D] - M);
     L = fmaf(pc[D + 1], w, L);
@@ -245,7 +275,8 @@ __device__ __forceinline__ float mk_load_x(const bf16* x, const bf16* gamma,
     return bf2f(f2bf(sg * bf2f(x[K + k])));
   }
   if (itrans == IT_ATTNFIN)
-    return attn_finalize(part, (int)(k / D), (int)(k % D), D, maxch, nch);
   return bf2f(x[k]);
 }
@@ -258,7 +289,15 @@ __device__ __forceinline__ void gemv_store(const long long* P, int epi,
         f2bf(bf2f(resid[(long long)srcm * N + row]) + a);
   else if (epi == EP_F32)
     reinterpret_cast<float*>(P[3])[row] = a;
-  else
     reinterpret_cast<bf16*>(P[3])[(long long)srcm * N + row] = f2bf(a);
 }
@@ -282,7 +321,8 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
   const float* part = reinterpret_cast<const float*>(P[1]);
   const bf16* xin = reinterpret_cast<const bf16*>(P[1]);
   const int D = (int)P[13];
-  const int maxch = (int)P[15];
   const int m0 = mmode ? (M - 1) : 0;
   const int Ma = M - m0;  // staged rows
@@ -383,6 +423,70 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
           *reinterpret_cast<int4*>(xsp) = x4;
         }
         __syncthreads();
       } else if (itrans == IT_SWIGLU) {
         const bf16* xr = xin + (long long)srcm * 2 * K;
         for (long long g = threadIdx.x; g < G; g += blockDim.x) {
@@ -414,6 +518,9 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
       } else if (itrans == IT_ATTNFIN) {
         const int pm = pos + srcm;
         const int nch = (pm + CH) / CH;
         const float* mpart =
             part + (long long)srcm * ((long long)P[11] * maxch * (D + 2));
         const int Hq_ = (int)P[11];
@@ -427,10 +534,10 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
           for (int h = threadIdx.x; h < Hq_; h += blockDim.x) {
             const float* base = mpart + ((long long)h * maxch) * (D + 2);
             float Mx = -INFINITY;
-            for (int c = 0; c < nch; ++c)
               Mx = fmaxf(Mx, base[(long long)c * (D + 2) + D]);
             float L = 0.f;
-            for (int c = 0; c < nch; ++c) {
               const float* pc = base + (long long)c * (D + 2);
               const float w = expf(pc[D] - Mx);
               ew[h * nch + c] = w;
@@ -448,7 +555,7 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
               const int h = (int)(k / D), d0 = (int)(k % D);
               const float* base = mpart + ((long long)h * maxch) * (D + 2);
               float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
-              for (int c = 0; c < nch; ++c) {
                 const float* pc = base + (long long)c * (D + 2) + d0;
                 const float2 pa = *reinterpret_cast<const float2*>(pc);
                 const float2 pb = *reinterpret_cast<const float2*>(pc + 2);
@@ -470,7 +577,7 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
               const int h = (int)(k / D), d = (int)(k % D);
               const float* base = mpart + ((long long)h * maxch) * (D + 2);
               float acc = 0.f;
-              for (int c = 0; c < nch; ++c)
                 acc = fmaf(base[(long long)c * (D + 2) + d],
                            ew[h * nch + c], acc);
               xs[(long long)m * K + k] = f2bf(bf2f(f2bf(acc / lw[h])));
@@ -480,7 +587,7 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
         } else {
           for (long long k = threadIdx.x; k < K; k += blockDim.x)
             xs[(long long)m * K + k] = f2bf(attn_finalize(
-                mpart, (int)(k / D), (int)(k % D), D, maxch, nch));
         }
       } else {
         const bf16* xr = xin + (long long)srcm * K;
@@ -545,6 +652,16 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
               const unsigned long long kk = amax_key(bf2f(f2bf(a)),
                                                      (unsigned)row);
               best = (kk > best) ? kk : best;
             } else {
               gemv_store(P, epi, N, resid, srcm, row, a);
             }
@@ -553,7 +670,7 @@ __device__ void ph_gemv(const long long* P, int pos, int M, int rs) {
       }
     }
   }
-  if (epi == EP_F32_AMAX) {
     __syncthreads();
     best = block_max_u64(best, s64);
     if (threadIdx.x == 0)
@@ -567,7 +684,7 @@ __device__ void kv_norm_rope_append(const bf16* qkv, bf16* kc, bf16* vc,
                                     const bf16* knw, int qk_flag, float eps,
                                     const float* invf, int Hq, int Hkv, int kvh,
                                     int D, int pm, float* s_red,
-                                    float* s_inv, float* hbuf) {
   const int half = D >> 1;
   const bf16* ksrc = qkv + (long long)(Hq + kvh) * D;
   const bf16* vsrc = qkv + (long long)(Hq + Hkv + kvh) * D;
@@ -581,8 +698,10 @@ __device__ void kv_norm_rope_append(const bf16* qkv, bf16* kc, bf16* vc,
     if (threadIdx.x == 0) *s_inv = rsqrtf(ss / (float)D + eps);
     __syncthreads();
     for (int d = threadIdx.x; d < D; d += blockDim.x)
-      hbuf[d] = bf2f(f2bf(bf2f(knw[d]) *
-                          bf2f(f2bf(bf2f(ksrc[d]) * (*s_inv)))));
   }
   // without k-norm the rope reads the source row directly (identical
   // values), skipping the staging pass and its barrier
@@ -591,14 +710,28 @@ __device__ void kv_norm_rope_append(const bf16* qkv, bf16* kc, bf16* vc,
   bf16* vrow = vc + (long long)pm * D;
   for (int d = threadIdx.x; d < half; d += blockDim.x) {
     const float ang = (float)pm * invf[d];
-    const float c = bf2f(f2bf(cosf(ang)));
-    const float s = bf2f(f2bf(sinf(ang)));
     const float x1 = (qk_flag & 1) ? hbuf[d] : bf2f(ksrc[d]);
     const float x2 = (qk_flag & 1) ? hbuf[d + half] : bf2f(ksrc[d + half]);
     krow[d] = f2bf(bf2f(f2bf(x1 * c)) + bf2f(f2bf(-x2 * s)));
     krow[d + half] = f2bf(bf2f(f2bf(x2 * c)) + bf2f(f2bf(x1 * s)));
   }
-  for (int d = threadIdx.x; d < D; d += blockDim.x) vrow[d] = vsrc[d];
   __threadfence_block();
   __syncthreads();
 }
@@ -615,7 +748,10 @@ __device__ void kv_norm_rope_append(const bf16* qkv, bf16* kc, bf16* vc,
 // are safe. The O-projection finalizes the partials.
 __device__ void ph_attn(const long long* P, int pos, int M) {
   const long long max_seq = P[7];
-  const int qk_flag = (int)P[8];
   const float eps = bits2f(P[9]);
   const float* invf = reinterpret_cast<const float*>(P[10]);
   const int Hq = (int)P[11], Hkv = ((int)P[12]) & 0xFFFF;
@@ -625,7 +761,10 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
   const int chbits = (int)(P[12] >> 16) & 0xFFFF;
   const int CH = (chbits > 0 && chbits <= ATTN_CHUNK) ? chbits : ATTN_CHUNK;
   const float scale = bits2f(P[14]);
-  const int maxch = (int)P[15];
   const int rep = Hq / Hkv;
   const long long KT = (long long)(Hq + 2 * Hkv) * D;
@@ -654,6 +793,8 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
       const int S = pm + 1;
       const int nch = (S + CH - 1) / CH;
       if (ch >= nch) continue;
       const int own_ch = pm / CH;
       const bf16* qkv =
           reinterpret_cast<const bf16*>(P[1]) + (long long)m * KT;
@@ -666,15 +807,18 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
       const int half = D >> 1;
       // decode only: the chunk owning `pos` appends K/V for its group,
-      // before the chunk copy below so the copy sees the appended row
-      if (M == 1 && ch == own_ch) {
         kv_norm_rope_append(qkv, kc, vc, knw, qk_flag, eps, invf, Hq, Hkv,
-                            kvh, D, pm, s_red, &s_inv, hbuf);
       }
       const int s0 = ch * CH;
       const int s1 = (s0 + CH < S) ? s0 + CH : S;
       const int cn = s1 - s0;
       // Copy this chunk's K/V into the dynamic-shared region (idle during
       // attention) when it fits: the score loop otherwise walks the chunk
@@ -713,16 +857,18 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
           if (threadIdx.x == 0) s_inv = rsqrtf(ss / (float)D + eps);
           __syncthreads();
           for (int d = threadIdx.x; d < D; d += blockDim.x)
-            hbuf[d] =
-                bf2f(f2bf(bf2f(qnw[d]) * bf2f(f2bf(bf2f(q[d]) * s_inv))));
         }
         // without q-norm the rope reads q directly (identical values),
         // skipping the staging pass and its barrier
         if (qk_flag & 1) __syncthreads();
         for (int d = threadIdx.x; d < half; d += blockDim.x) {
           const float ang = (float)pm * invf[d];
-          const float c = bf2f(f2bf(cosf(ang)));
-          const float s = bf2f(f2bf(sinf(ang)));
           const float x1 = (qk_flag & 1) ? hbuf[d] : bf2f(q[d]);
           const float x2 =
               (qk_flag & 1) ? hbuf[d + half] : bf2f(q[d + half]);
@@ -738,6 +884,7 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
       const int wid = threadIdx.x >> 5, lane = threadIdx.x & 31;
       const int nw = blockDim.x >> 5;
       for (int si = wid; si < cn; si += nw) {
         const bf16* kr = kv_staged ? ks_sh + (long long)si * D
                                    : kc + (long long)(s0 + si) * D;
         float acc = 0.f;
@@ -752,13 +899,14 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
       float mx = -INFINITY;
       for (int si = threadIdx.x; si < cn; si += blockDim.x)
-        mx = fmaxf(mx, sc[si]);
       mx = block_max(mx, s_red);
       if (threadIdx.x == 0) s_m = mx;
       __syncthreads();
       float lsum = 0.f;
       for (int si = threadIdx.x; si < cn; si += blockDim.x) {
         const float p = expf(sc[si] - s_m);
         sc[si] = p;
         lsum += p;
@@ -775,8 +923,9 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
         const int d = threadIdx.x - g * D;
         if (g < NG) {
           const int csz = (cn + NG - 1) / NG;
-          const int a0 = g * csz;
-          const int a1 = (a0 + csz < cn) ? a0 + csz : cn;
           float acc = 0.f;
           for (int si = a0; si < a1; ++si)
             acc = fmaf(sc[si],
@@ -786,20 +935,21 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
           pv[g * D + d] = acc;
         }
         __syncthreads();
-        if (threadIdx.x < D) {
           float acc = 0.f;
-          for (int g2 = 0; g2 < NG; ++g2) acc += pv[g2 * D + threadIdx.x];
-          part[threadIdx.x] = acc;
         }
-      } else if (threadIdx.x < D) {
-        const int d = threadIdx.x;
-        float acc = 0.f;
-        for (int si = 0; si < cn; ++si)
-          acc = fmaf(sc[si],
-                     bf2f(kv_staged ? vs_sh[(long long)si * D + d]
-                                    : vc[(long long)(s0 + si) * D + d]),
-                     acc);
-        part[d] = acc;
       }
       if (threadIdx.x == 0) {
         part[D] = s_m;
@@ -815,10 +965,13 @@ __device__ void ph_attn(const long long* P, int pos, int M) {
 // attention phase so causal reads inside the chunk see complete K/V.
 __device__ void ph_kv_append(const long long* P, int pos, int M) {
   const long long max_seq = P[7];
-  const int qk_flag = (int)P[8];
   const float eps = bits2f(P[9]);
   const float* invf = reinterpret_cast<const float*>(P[10]);
-  const int Hq = (int)P[11], Hkv = (int)P[12], D = (int)P[13];
   const long long KT = (long long)(Hq + 2 * Hkv) * D;
   const bf16* knw = reinterpret_cast<const bf16*>(P[6]);
@@ -835,9 +988,71 @@ __device__ void ph_kv_append(const long long* P, int pos, int M) {
     bf16* kc = reinterpret_cast<bf16*>(P[2]) + (long long)kvh * max_seq * D;
     bf16* vc = reinterpret_cast<bf16*>(P[3]) + (long long)kvh * max_seq * D;
     kv_norm_rope_append(qkv, kc, vc, knw, qk_flag, eps, invf, Hq, Hkv, kvh,
-                        D, pm, s_red, &s_inv, hbuf);
   }
 }
 __device__ void ph_argmax_part(const long long* P) {
   const float* lg = reinterpret_cast<const float*>(P[1]);
   unsigned long long* parts = reinterpret_cast<unsigned long long*>(P[3]);
@@ -886,6 +1101,8 @@ __device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
     case OP_ARGMAX_PART: ph_argmax_part(P); break;
     case OP_ARGMAX_FIN: ph_argmax_fin(P, step_slot); break;
     case OP_KV_APPEND: ph_kv_append(P, pos, M); break;
     default: break;
   }
 }

 #include "mak_kernels.h"
 #define BLOCK_THREADS 256
+#define MAX_HEAD_DIM 512
 // GEMV weight tiles stream through a per-warp cp.async shared ring. The
 // stage count is a power of two chosen per launch from the device's
 // shared-memory budget (deep on high-bandwidth parts, where the
   OP_ARGMAX_PART = 4,
   OP_ARGMAX_FIN = 5,
   OP_KV_APPEND = 6,  // prefill chunks only: rope+append K/V for all m
+  OP_NORMRES = 7,    // gemma post-sublayer norm + residual add
+  OP_PLEMIX = 8,     // gemma per-layer-input mix (norm ctx, add tok, scale)
 };
 enum InputTransform : int {
   IT_NONE = 0,
   IT_RMSNORM = 1,
   IT_SWIGLU = 2,
+  IT_ATTNFIN = 3,   // finalize split-S attention partials on the fly
+  IT_RMSNORM_G = 4, // gemma norm: fp32 throughout, raw weight, one rounding
+  IT_GELU_GLU = 5,  // gelu_tanh(gate) * up
 };
 // slot-9 flag bits above the base transform
 #define ITF_BASE_MASK 7
   EP_F32 = 2,
   EP_F32_AMAX = 3,   // fp32 store + per-block argmax partial (lm head)
   EP_PARTIAL = 4,    // fp32 k-slice partials [NSLICE][N], no rounding
+  EP_GELU_PLE = 5,   // gemma: gelu_tanh(out) * per-layer-input element
+  EP_F32_AMAX_CAP = 6,  // lm head with final logit softcapping
 };
 #define ATTN_CHUNK 128
 __device__ __forceinline__ float bf2f(bf16 v) { return __bfloat162float(v); }
 __device__ __forceinline__ bf16 f2bf(float v) { return __float2bfloat16(v); }
+// torch gelu(approximate="tanh") in fp32
+__device__ __forceinline__ float gelu_tanh_f(float x) {
+  const float c = 0.7978845608028654f;  // sqrt(2/pi)
+  const float t = tanhf(c * (x + 0.044715f * x * x * x));
+  return 0.5f * x * (1.f + t);
+}
 __device__ __forceinline__ unsigned dyn_smem_size() {
   unsigned r;
   asm("mov.u32 %0, %%dynamic_smem_size;" : "=r"(r));
     __syncthreads();
   }
+  const float hsc = bits2f(P[14]);  // gemma embed scale (0 bits = off)
   for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
        i < (long long)M * K; i += (long long)gridDim.x * blockDim.x) {
     const int m = (int)(i / K);
     const long long t = (pm < prompt_len && prompt != nullptr)
                             ? (long long)prompt[pm]
                             : (long long)s_tok;
+    out[i] = (hsc != 0.f) ? f2bf(bf2f(tab[t * K + k]) * hsc)
+                          : tab[t * K + k];
+  }
+  // gemma: gather this token's per-layer-input row, scaled (decode only;
+  // the gemma builder pins chunk_m to 1)
+  const bf16* ptab = reinterpret_cast<const bf16*>(P[1]);
+  bf16* pout = reinterpret_cast<bf16*>(P[9]);
+  if (ptab != nullptr && pout != nullptr && M == 1) {
+    const long long PL = P[12];
+    const float esc = bits2f(P[15]);
+    const long long t = (pos < prompt_len && prompt != nullptr)
+                            ? (long long)prompt[pos]
+                            : (long long)s_tok;
+    for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+         i < PL; i += (long long)gridDim.x * blockDim.x)
+      pout[i] = f2bf(bf2f(ptab[t * PL + i]) * esc);
   }
 }
 // bf16 rounding of the attention output before the O projection.
 __device__ __forceinline__ float attn_finalize(const float* mpart, int h,
                                                int d, int D, int maxch,
+                                               int nch, int nch_lo) {
   const float* base = mpart + ((long long)h * maxch) * (D + 2);
   float M = -INFINITY;
+  for (int c = nch_lo; c < nch; ++c)
     M = fmaxf(M, base[(long long)c * (D + 2) + D]);
   float L = 0.f, acc = 0.f;
+  for (int c = nch_lo; c < nch; ++c) {
     const float* pc = base + (long long)c * (D + 2);
     const float w = expf(pc[D] - M);
     L = fmaf(pc[D + 1], w, L);
     return bf2f(f2bf(sg * bf2f(x[K + k])));
   }
   if (itrans == IT_ATTNFIN)
+    return attn_finalize(part, (int)(k / D), (int)(k % D), D, maxch, nch,
+                         0);
   return bf2f(x[k]);
 }
         f2bf(bf2f(resid[(long long)srcm * N + row]) + a);
   else if (epi == EP_F32)
     reinterpret_cast<float*>(P[3])[row] = a;
+  else if (epi == EP_GELU_PLE) {
+    // gemma per-layer-input gate: bf16 linear output, gelu in fp32
+    // rounded back, multiplied by the layer's per-layer-input element
+    const bf16* ple = reinterpret_cast<const bf16*>(P[6]);
+    const float g = bf2f(f2bf(a));
+    const float act = bf2f(f2bf(gelu_tanh_f(g)));
+    reinterpret_cast<bf16*>(P[3])[(long long)srcm * N + row] =
+        f2bf(act * bf2f(ple[row]));
+  } else
     reinterpret_cast<bf16*>(P[3])[(long long)srcm * N + row] = f2bf(a);
 }
   const float* part = reinterpret_cast<const float*>(P[1]);
   const bf16* xin = reinterpret_cast<const bf16*>(P[1]);
   const int D = (int)P[13];
+  const int maxch = (int)(P[15] & 0xFFFF);
+  const int fin_win = (int)((P[15] >> 16) & 0xFFFF);  // 0 = unwindowed
   const int m0 = mmode ? (M - 1) : 0;
   const int Ma = M - m0;  // staged rows
           *reinterpret_cast<int4*>(xsp) = x4;
         }
         __syncthreads();
+      } else if (itrans == IT_RMSNORM_G) {
+        // gemma norm: fp32 throughout, raw weight, one rounding at the end
+        const bf16* xr = xin + (long long)srcm * K;
+        float ss = 0.f;
+        for (long long g = threadIdx.x; g < G; g += blockDim.x) {
+          const int4 x4 = *reinterpret_cast<const int4*>(xr + g * 8);
+          *reinterpret_cast<int4*>(xs + (long long)m * K + g * 8) = x4;
+          const __nv_bfloat162* xp =
+              reinterpret_cast<const __nv_bfloat162*>(&x4);
+#pragma unroll
+          for (int j = 0; j < 4; ++j) {
+            const float lo = __low2float(xp[j]);
+            const float hi = __high2float(xp[j]);
+            ss += lo * lo;
+            ss += hi * hi;
+          }
+        }
+        ss = block_sum(ss, s_red);
+        if (threadIdx.x == 0) s_scale = rsqrtf(ss / (float)K + eps);
+        __syncthreads();
+        const float inv = s_scale;
+        for (long long g = threadIdx.x; g < G; g += blockDim.x) {
+          bf16* xsp = xs + (long long)m * K + g * 8;
+          int4 x4 = *reinterpret_cast<const int4*>(xsp);
+          const int4 g4 = *reinterpret_cast<const int4*>(gamma + g * 8);
+          __nv_bfloat162* xp = reinterpret_cast<__nv_bfloat162*>(&x4);
+          const __nv_bfloat162* gp =
+              reinterpret_cast<const __nv_bfloat162*>(&g4);
+#pragma unroll
+          for (int j = 0; j < 4; ++j) {
+            xp[j].x = f2bf((__low2float(xp[j]) * inv) * __low2float(gp[j]));
+            xp[j].y =
+                f2bf((__high2float(xp[j]) * inv) * __high2float(gp[j]));
+          }
+          *reinterpret_cast<int4*>(xsp) = x4;
+        }
+        __syncthreads();
+      } else if (itrans == IT_GELU_GLU) {
+        // gemma MLP gating: gelu_tanh(gate) rounded to bf16, times up
+        const bf16* xr = xin + (long long)srcm * 2 * K;
+        for (long long g = threadIdx.x; g < G; g += blockDim.x) {
+          const int4 g4 = *reinterpret_cast<const int4*>(xr + g * 8);
+          const int4 u4 = *reinterpret_cast<const int4*>(xr + K + g * 8);
+          int4 o4;
+          const __nv_bfloat162* gp =
+              reinterpret_cast<const __nv_bfloat162*>(&g4);
+          const __nv_bfloat162* up =
+              reinterpret_cast<const __nv_bfloat162*>(&u4);
+          __nv_bfloat162* op = reinterpret_cast<__nv_bfloat162*>(&o4);
+#pragma unroll
+          for (int j = 0; j < 4; ++j) {
+            {
+              const float gv = __low2float(gp[j]);
+              const float act = bf2f(f2bf(gelu_tanh_f(gv)));
+              op[j].x = f2bf(act * __low2float(up[j]));
+            }
+            {
+              const float gv = __high2float(gp[j]);
+              const float act = bf2f(f2bf(gelu_tanh_f(gv)));
+              op[j].y = f2bf(act * __high2float(up[j]));
+            }
+          }
+          *reinterpret_cast<int4*>(xs + (long long)m * K + g * 8) = o4;
+        }
       } else if (itrans == IT_SWIGLU) {
         const bf16* xr = xin + (long long)srcm * 2 * K;
         for (long long g = threadIdx.x; g < G; g += blockDim.x) {
       } else if (itrans == IT_ATTNFIN) {
         const int pm = pos + srcm;
         const int nch = (pm + CH) / CH;
+        const int s_lo = (fin_win > 0 && pm + 1 > fin_win)
+                             ? pm + 1 - fin_win : 0;
+        const int nch_lo = s_lo / CH;
         const float* mpart =
             part + (long long)srcm * ((long long)P[11] * maxch * (D + 2));
         const int Hq_ = (int)P[11];
           for (int h = threadIdx.x; h < Hq_; h += blockDim.x) {
             const float* base = mpart + ((long long)h * maxch) * (D + 2);
             float Mx = -INFINITY;
+            for (int c = nch_lo; c < nch; ++c)
               Mx = fmaxf(Mx, base[(long long)c * (D + 2) + D]);
             float L = 0.f;
+            for (int c = nch_lo; c < nch; ++c) {
               const float* pc = base + (long long)c * (D + 2);
               const float w = expf(pc[D] - Mx);
               ew[h * nch + c] = w;
               const int h = (int)(k / D), d0 = (int)(k % D);
               const float* base = mpart + ((long long)h * maxch) * (D + 2);
               float a0 = 0.f, a1 = 0.f, a2 = 0.f, a3 = 0.f;
+              for (int c = nch_lo; c < nch; ++c) {
                 const float* pc = base + (long long)c * (D + 2) + d0;
                 const float2 pa = *reinterpret_cast<const float2*>(pc);
                 const float2 pb = *reinterpret_cast<const float2*>(pc + 2);
               const int h = (int)(k / D), d = (int)(k % D);
               const float* base = mpart + ((long long)h * maxch) * (D + 2);
               float acc = 0.f;
+              for (int c = nch_lo; c < nch; ++c)
                 acc = fmaf(base[(long long)c * (D + 2) + d],
                            ew[h * nch + c], acc);
               xs[(long long)m * K + k] = f2bf(bf2f(f2bf(acc / lw[h])));
         } else {
           for (long long k = threadIdx.x; k < K; k += blockDim.x)
             xs[(long long)m * K + k] = f2bf(attn_finalize(
+                mpart, (int)(k / D), (int)(k % D), D, maxch, nch, nch_lo));
         }
       } else {
         const bf16* xr = xin + (long long)srcm * K;
               const unsigned long long kk = amax_key(bf2f(f2bf(a)),
                                                      (unsigned)row);
               best = (kk > best) ? kk : best;
+            } else if (epi == EP_F32_AMAX_CAP) {
+              // final logit softcapping, mirroring the eager bf16 op chain
+              const float cap = bits2f(P[15]);
+              const bf16 t1 = f2bf(a);
+              const bf16 t2 = f2bf(bf2f(t1) / cap);
+              const bf16 t3 = f2bf(tanhf(bf2f(t2)));
+              const float v = bf2f(f2bf(bf2f(t3) * cap));
+              reinterpret_cast<float*>(P[3])[row] = v;
+              const unsigned long long kk = amax_key(v, (unsigned)row);
+              best = (kk > best) ? kk : best;
             } else {
               gemv_store(P, epi, N, resid, srcm, row, a);
             }
       }
     }
   }
+  if (epi == EP_F32_AMAX || epi == EP_F32_AMAX_CAP) {
     __syncthreads();
     best = block_max_u64(best, s64);
     if (threadIdx.x == 0)
                                     const bf16* knw, int qk_flag, float eps,
                                     const float* invf, int Hq, int Hkv, int kvh,
                                     int D, int pm, float* s_red,
+                                    float* s_inv, float* hbuf, float ascale) {
   const int half = D >> 1;
   const bf16* ksrc = qkv + (long long)(Hq + kvh) * D;
   const bf16* vsrc = qkv + (long long)(Hq + Hkv + kvh) * D;
     if (threadIdx.x == 0) *s_inv = rsqrtf(ss / (float)D + eps);
     __syncthreads();
     for (int d = threadIdx.x; d < D; d += blockDim.x)
+      hbuf[d] = (qk_flag & 2)
+                    ? bf2f(f2bf((bf2f(ksrc[d]) * (*s_inv)) * bf2f(knw[d])))
+                    : bf2f(f2bf(bf2f(knw[d]) *
+                                bf2f(f2bf(bf2f(ksrc[d]) * (*s_inv)))));
   }
   // without k-norm the rope reads the source row directly (identical
   // values), skipping the staging pass and its barrier
   bf16* vrow = vc + (long long)pm * D;
   for (int d = threadIdx.x; d < half; d += blockDim.x) {
     const float ang = (float)pm * invf[d];
+    const float c = bf2f(f2bf(cosf(ang) * ascale));
+    const float s = bf2f(f2bf(sinf(ang) * ascale));
     const float x1 = (qk_flag & 1) ? hbuf[d] : bf2f(ksrc[d]);
     const float x2 = (qk_flag & 1) ? hbuf[d + half] : bf2f(ksrc[d + half]);
     krow[d] = f2bf(bf2f(f2bf(x1 * c)) + bf2f(f2bf(-x2 * s)));
     krow[d + half] = f2bf(bf2f(f2bf(x2 * c)) + bf2f(f2bf(x1 * s)));
   }
+  if (qk_flag & 4) {
+    // gemma v-norm: weightless RMS over v, fp32, one rounding
+    float ss = 0.f;
+    for (int d = threadIdx.x; d < D; d += blockDim.x) {
+      const float x = bf2f(vsrc[d]);
+      ss += x * x;
+    }
+    ss = block_sum(ss, s_red);
+    if (threadIdx.x == 0) *s_inv = rsqrtf(ss / (float)D + eps);
+    __syncthreads();
+    for (int d = threadIdx.x; d < D; d += blockDim.x)
+      vrow[d] = f2bf(bf2f(vsrc[d]) * (*s_inv));
+  } else {
+    for (int d = threadIdx.x; d < D; d += blockDim.x) vrow[d] = vsrc[d];
+  }
   __threadfence_block();
   __syncthreads();
 }
 // are safe. The O-projection finalizes the partials.
 __device__ void ph_attn(const long long* P, int pos, int M) {
   const long long max_seq = P[7];
+  // slot 8: bit0 q/k norm, bit1 gemma norm semantics, bit2 v-norm;
+  // bits 16+ sliding window length (0 = unwindowed)
+  const int qk_flag = ((int)P[8]) & 0xFFFF;
+  const int win = (int)((P[8] >> 16) & 0xFFFFFFFFll);
   const float eps = bits2f(P[9]);
   const float* invf = reinterpret_cast<const float*>(P[10]);
   const int Hq = (int)P[11], Hkv = ((int)P[12]) & 0xFFFF;
   const int chbits = (int)(P[12] >> 16) & 0xFFFF;
   const int CH = (chbits > 0 && chbits <= ATTN_CHUNK) ? chbits : ATTN_CHUNK;
   const float scale = bits2f(P[14]);
+  const int maxch = (int)(P[15] & 0xFFFF);
+  // rope attention scaling in slot 15 bits 32+ (0 bits = 1.0)
+  const int asbits = (int)(P[15] >> 32);
+  const float ascale = (asbits != 0) ? __int_as_float(asbits) : 1.0f;
   const int rep = Hq / Hkv;
   const long long KT = (long long)(Hq + 2 * Hkv) * D;
       const int S = pm + 1;
       const int nch = (S + CH - 1) / CH;
       if (ch >= nch) continue;
+      const int s_lo = (win > 0 && S > win) ? S - win : 0;
+      if (ch < s_lo / CH) continue;  // fully outside the sliding window
       const int own_ch = pm / CH;
       const bf16* qkv =
           reinterpret_cast<const bf16*>(P[1]) + (long long)m * KT;
       const int half = D >> 1;
       // decode only: the chunk owning `pos` appends K/V for its group,
+      // before the chunk copy below so the copy sees the appended row.
+      // Shared-KV layers (flag bit 3) reuse an earlier layer's cache and
+      // append nothing.
+      if (M == 1 && ch == own_ch && !(qk_flag & 8)) {
         kv_norm_rope_append(qkv, kc, vc, knw, qk_flag, eps, invf, Hq, Hkv,
+                            kvh, D, pm, s_red, &s_inv, hbuf, ascale);
       }
       const int s0 = ch * CH;
       const int s1 = (s0 + CH < S) ? s0 + CH : S;
       const int cn = s1 - s0;
+      const int si0 = (s0 < s_lo) ? (s_lo - s0) : 0;  // window mask
       // Copy this chunk's K/V into the dynamic-shared region (idle during
       // attention) when it fits: the score loop otherwise walks the chunk
           if (threadIdx.x == 0) s_inv = rsqrtf(ss / (float)D + eps);
           __syncthreads();
           for (int d = threadIdx.x; d < D; d += blockDim.x)
+            hbuf[d] = (qk_flag & 2)
+                          ? bf2f(f2bf((bf2f(q[d]) * s_inv) * bf2f(qnw[d])))
+                          : bf2f(f2bf(bf2f(qnw[d]) *
+                                      bf2f(f2bf(bf2f(q[d]) * s_inv))));
         }
         // without q-norm the rope reads q directly (identical values),
         // skipping the staging pass and its barrier
         if (qk_flag & 1) __syncthreads();
         for (int d = threadIdx.x; d < half; d += blockDim.x) {
           const float ang = (float)pm * invf[d];
+          const float c = bf2f(f2bf(cosf(ang) * ascale));
+          const float s = bf2f(f2bf(sinf(ang) * ascale));
           const float x1 = (qk_flag & 1) ? hbuf[d] : bf2f(q[d]);
           const float x2 =
               (qk_flag & 1) ? hbuf[d + half] : bf2f(q[d + half]);
       const int wid = threadIdx.x >> 5, lane = threadIdx.x & 31;
       const int nw = blockDim.x >> 5;
       for (int si = wid; si < cn; si += nw) {
+        if (si < si0) continue;
         const bf16* kr = kv_staged ? ks_sh + (long long)si * D
                                    : kc + (long long)(s0 + si) * D;
         float acc = 0.f;
       float mx = -INFINITY;
       for (int si = threadIdx.x; si < cn; si += blockDim.x)
+        if (si >= si0) mx = fmaxf(mx, sc[si]);
       mx = block_max(mx, s_red);
       if (threadIdx.x == 0) s_m = mx;
       __syncthreads();
       float lsum = 0.f;
       for (int si = threadIdx.x; si < cn; si += blockDim.x) {
+        if (si < si0) continue;
         const float p = expf(sc[si] - s_m);
         sc[si] = p;
         lsum += p;
         const int d = threadIdx.x - g * D;
         if (g < NG) {
           const int csz = (cn + NG - 1) / NG;
+          int a0 = g * csz;
+          if (a0 < si0) a0 = si0;
+          const int a1 = ((g * csz) + csz < cn) ? (g * csz) + csz : cn;
           float acc = 0.f;
           for (int si = a0; si < a1; ++si)
             acc = fmaf(sc[si],
           pv[g * D + d] = acc;
         }
         __syncthreads();
+        for (int d = threadIdx.x; d < D; d += blockDim.x) {
           float acc = 0.f;
+          for (int g2 = 0; g2 < NG; ++g2) acc += pv[g2 * D + d];
+          part[d] = acc;
+        }
+      } else {
+        for (int d = threadIdx.x; d < D; d += blockDim.x) {
+          float acc = 0.f;
+          for (int si = si0; si < cn; ++si)
+            acc = fmaf(sc[si],
+                       bf2f(kv_staged ? vs_sh[(long long)si * D + d]
+                                      : vc[(long long)(s0 + si) * D + d]),
+                       acc);
+          part[d] = acc;
         }
       }
       if (threadIdx.x == 0) {
         part[D] = s_m;
 // attention phase so causal reads inside the chunk see complete K/V.
 __device__ void ph_kv_append(const long long* P, int pos, int M) {
   const long long max_seq = P[7];
+  const int qk_flag = ((int)P[8]) & 0xFFFF;
   const float eps = bits2f(P[9]);
   const float* invf = reinterpret_cast<const float*>(P[10]);
+  const int Hq = (int)P[11], Hkv = ((int)P[12]) & 0xFFFF;
+  const int D = (int)P[13];
+  const int asbits = (int)(P[15] >> 32);
+  const float ascale = (asbits != 0) ? __int_as_float(asbits) : 1.0f;
   const long long KT = (long long)(Hq + 2 * Hkv) * D;
   const bf16* knw = reinterpret_cast<const bf16*>(P[6]);
     bf16* kc = reinterpret_cast<bf16*>(P[2]) + (long long)kvh * max_seq * D;
     bf16* vc = reinterpret_cast<bf16*>(P[3]) + (long long)kvh * max_seq * D;
     kv_norm_rope_append(qkv, kc, vc, knw, qk_flag, eps, invf, Hq, Hkv, kvh,
+                        D, pm, s_red, &s_inv, hbuf, ascale);
   }
 }
+// gemma post-sublayer norm + residual: hid = hid + norm(t) (+ optional
+// layer scale). One vector; block 0 computes while the rest fall through.
+__device__ void ph_normres(const long long* P) {
+  if (blockIdx.x != 0) return;
+  const bf16* t = reinterpret_cast<const bf16*>(P[1]);
+  const bf16* w = reinterpret_cast<const bf16*>(P[2]);
+  bf16* hid = reinterpret_cast<bf16*>(P[3]);
+  const long long n = P[7];
+  const float eps = bits2f(P[14]);
+  const float lscale = bits2f(P[15]);
+  __shared__ float s_red[32];
+  __shared__ float s_inv;
+  float ss = 0.f;
+  for (long long k = threadIdx.x; k < n; k += blockDim.x) {
+    const float v = bf2f(t[k]);
+    ss += v * v;
+  }
+  ss = block_sum(ss, s_red);
+  if (threadIdx.x == 0) s_inv = rsqrtf(ss / (float)n + eps);
+  __syncthreads();
+  for (long long k = threadIdx.x; k < n; k += blockDim.x) {
+    const float nv = bf2f(f2bf((bf2f(t[k]) * s_inv) * bf2f(w[k])));
+    float h = bf2f(f2bf(bf2f(hid[k]) + nv));
+    if (lscale != 1.f) h = bf2f(f2bf(h * lscale));
+    hid[k] = f2bf(h);
+  }
+}
+// gemma per-layer-input mix: for each layer row, scale the context
+// projection (bf16 round), gemma-norm it over the 256-wide block, add the
+// token-identity row, and scale by 1/sqrt(2), mirroring the eager op
+// order exactly.
+__device__ void ph_plemix(const long long* P) {
+  const bf16* ctx = reinterpret_cast<const bf16*>(P[1]);
+  const bf16* w = reinterpret_cast<const bf16*>(P[2]);
+  bf16* out = reinterpret_cast<bf16*>(P[3]);
+  const bf16* tok = reinterpret_cast<const bf16*>(P[4]);
+  const long long L = P[7], C = P[8];
+  const float eps = bits2f(P[14]);
+  const float pscale = bits2f(P[15]);
+  __shared__ float s_red[32];
+  __shared__ float s_inv;
+  __shared__ float xbuf[512];
+  for (long long l = blockIdx.x; l < L; l += gridDim.x) {
+    float ss = 0.f;
+    for (int j = threadIdx.x; j < (int)C; j += blockDim.x) {
+      const float v = bf2f(f2bf(bf2f(ctx[l * C + j]) * pscale));
+      xbuf[j] = v;
+      ss += v * v;
+    }
+    ss = block_sum(ss, s_red);
+    if (threadIdx.x == 0) s_inv = rsqrtf(ss / (float)C + eps);
+    __syncthreads();
+    for (int j = threadIdx.x; j < (int)C; j += blockDim.x) {
+      const float nv = bf2f(f2bf((xbuf[j] * s_inv) * bf2f(w[j])));
+      const float mixed = bf2f(f2bf(nv + bf2f(tok[l * C + j])));
+      out[l * C + j] = f2bf(mixed * 0.70710678118654752440f);
+    }
+    __syncthreads();
+  }
+}
 __device__ void ph_argmax_part(const long long* P) {
   const float* lg = reinterpret_cast<const float*>(P[1]);
   unsigned long long* parts = reinterpret_cast<unsigned long long*>(P[3]);
     case OP_ARGMAX_PART: ph_argmax_part(P); break;
     case OP_ARGMAX_FIN: ph_argmax_fin(P, step_slot); break;
     case OP_KV_APPEND: ph_kv_append(P, pos, M); break;
+    case OP_NORMRES: ph_normres(P); break;
+    case OP_PLEMIX: ph_plemix(P); break;
     default: break;
   }
 }

tests/test_mak.py CHANGED Viewed

@@ -276,3 +276,45 @@ def test_generate_closed_loop_and_speed(models):
     speedup = dt_hf / dt_ours
     print(f"speedup vs transformers eager generate: {speedup:.2f}x")
     assert speedup > 2.0

     speedup = dt_hf / dt_ours
     print(f"speedup vs transformers eager generate: {speedup:.2f}x")
     assert speedup > 2.0
+@requires_cuda
+def test_gemma4_e2b_coverage():
+    """gemma-4-E2B-it: gemma norms, PLE pathway, shared KV, sliding
+    windows, dual head dims, softcapped head. Needs both model copies
+    resident, so it runs where memory allows."""
+    if torch.cuda.get_device_properties(0).total_memory < 30e9:
+        pytest.skip("needs ~24GB free for eager + packed copies")
+    transformers = pytest.importorskip("transformers")
+    from transformers import DynamicCache
+    ref = transformers.AutoModelForCausalLM.from_pretrained(
+        "google/gemma-4-E2B-it", dtype=torch.bfloat16,
+        attn_implementation="eager").cuda().eval()
+    mm = mak.MegaModel.from_pretrained(ref, max_seq=512, max_gen=512)
+    torch.manual_seed(5)
+    seq = torch.randint(0, mm.V, (32,)).tolist()
+    cap = ref.config.text_config.final_logit_softcapping
+    cache = DynamicCache(config=ref.config.text_config)
+    out = torch.empty(len(seq), mm.V, device="cuda")
+    with torch.no_grad():
+        for i, t_ in enumerate(seq):
+            r = ref.model.language_model(
+                input_ids=torch.tensor([[t_]], device="cuda"),
+                past_key_values=cache, use_cache=True)
+            cache = r.past_key_values
+            lg = ref.lm_head(r.last_hidden_state)[0, -1].float()
+            lg = torch.tanh(lg.bfloat16() / cap).bfloat16() * cap
+            out[i] = lg.float()
+    mine = torch.empty(len(seq), mm.V, device="cuda")
+    for i, t_ in enumerate(seq):
+        mine[i] = mm.decode_step(t_, i).clone()
+    agree = (mine.argmax(-1) == out.argmax(-1)).sum().item()
+    d = (mine.bfloat16().float() - out.bfloat16().float()).abs()
+    print(f"gemma4-E2B band mean={d.mean().item():.4f} "
+          f"agree={agree}/{len(seq)}")
+    assert agree >= len(seq) * 0.85
+    a = mm.decode_step(seq[0], 0).clone()
+    b = mm.decode_step(seq[0], 0, phased=True).clone()
+    assert torch.equal(a, b)
+    del ref, mm
+    torch.cuda.empty_cache()

torch-ext/model_as_a_kernel/__init__.py CHANGED Viewed

@@ -27,6 +27,8 @@ __all__ = ["MegaModel"]
 _OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
 _OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
 _IT_NONE, _IT_RMSNORM, _IT_SWIGLU, _IT_ATTNFIN = 0, 1, 2, 3
 _ITF_REDUCE, _ITF_WRITEBACK = 8, 16
 _EP_STORE, _EP_RESID, _EP_F32, _EP_F32_AMAX, _EP_PARTIAL = 0, 1, 2, 3, 4
 _NSLICE = 4
@@ -265,6 +267,8 @@ class MegaModel:
             from transformers import AutoModelForCausalLM
             model = AutoModelForCausalLM.from_pretrained(
                 model, torch_dtype=torch.bfloat16)
         hf = model.config
         dec = model.model
         attn0 = dec.layers[0].self_attn
@@ -324,6 +328,252 @@ class MegaModel:
         return cls(weights, config, device=device, max_seq=max_seq,
                    max_gen=max_gen)
     # ------------------------------------------------------------------
     def decode_step(self, token: int, pos: int, phased: bool = False):
         """One decode step; returns the live fp32 logits buffer [V]."""

 _OP_EMBED, _OP_GEMV, _OP_QKV_POST, _OP_ATTN = 0, 1, 2, 3
 _OP_ARGMAX_PART, _OP_ARGMAX_FIN = 4, 5
 _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
             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
         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
+        return self
     # ------------------------------------------------------------------
     def decode_step(self, token: int, pos: int, phased: bool = False):
         """One decode step; returns the live fp32 logits buffer [V]."""