model-as-a-kernel v1 source

Browse files

Files changed (5) hide show

mak_cuda/mak_kernels.h +3 -3
mak_cuda/megakernel.cu +322 -265
torch-ext/model_as_a_kernel/__init__.py +30 -11
torch-ext/torch_binding.cpp +16 -11
torch-ext/torch_binding.h +1 -1

mak_cuda/mak_kernels.h CHANGED Viewed

@@ -8,9 +8,9 @@
 int mak_grid_blocks_impl(int smem_bytes);
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
-                            int steps, int slot0, int prompt_len,
-                            int* bar_cnt, int* bar_sense, int grid,
-                            int smem_bytes, cudaStream_t stream);
 cudaError_t mak_launch_phase(const long long* prog, int phase, int pos,
                              int step_slot, int grid, int smem_bytes,

 int mak_grid_blocks_impl(int smem_bytes);
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
+                            int total, int slot0, int prompt_len,
+                            int chunk_m, int* bar_cnt, int* bar_sense,
+                            int grid, int smem_bytes, cudaStream_t stream);
 cudaError_t mak_launch_phase(const long long* prog, int phase, int pos,
                              int step_slot, int grid, int smem_bytes,

mak_cuda/megakernel.cu CHANGED Viewed

@@ -37,6 +37,7 @@ enum PhaseOp : int {
   OP_ATTN = 3,
   OP_ARGMAX_PART = 4,
   OP_ARGMAX_FIN = 5,
 };
 enum InputTransform : int {
@@ -130,7 +131,7 @@ __device__ unsigned long long block_max_u64(unsigned long long v,
 // and any block with nothing to do simply falls through to the barrier.
 // ---------------------------------------------------------------------------
-__device__ void ph_embed(const long long* P, int pos, int prompt_len) {
   const bf16* tab = reinterpret_cast<const bf16*>(P[2]);
   bf16* out = reinterpret_cast<bf16*>(P[3]);
   const int* prompt = reinterpret_cast<const int*>(P[5]);
@@ -138,13 +139,16 @@ __device__ void ph_embed(const long long* P, int pos, int prompt_len) {
   const long long K = P[8];
   // Positions inside the prompt read from the staged prompt buffer; beyond
   // it, from the feedback token written by the previous step's argmax.
-  const long long t = (pos < prompt_len && prompt != nullptr)
-                          ? (long long)prompt[pos]
-                          : (long long)tok[0];
-  const bf16* row = tab + t * K;
-  for (long long k = (long long)blockIdx.x * blockDim.x + threadIdx.x; k < K;
-       k += (long long)gridDim.x * blockDim.x)
-    out[k] = row[k];
 }
 __device__ __forceinline__ unsigned long long amax_key(float v, unsigned idx) {
@@ -157,10 +161,10 @@ __device__ __forceinline__ unsigned long long amax_key(float v, unsigned idx) {
 // Finalize one element of split-S attention: combine per-chunk partials
 // (max, sum, weighted V) into the bf16 attention output, matching eager's
 // bf16 rounding of the attention output before the O projection.
-__device__ __forceinline__ float attn_finalize(const float* part, int h,
                                                int d, int D, int maxch,
                                                int nch) {
-  const float* base = part + ((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]);
@@ -197,78 +201,72 @@ __device__ __forceinline__ float mk_load_x(const bf16* x, const bf16* gamma,
 // Reduced (or plain) bf16 input element k for a consumer phase: sum the
 // producer's fp32 k-slice partials in fixed order, add the residual, round
 // to bf16, optionally write the materialized vector back.
-__device__ __forceinline__ bf16 consume_x(const long long* P, long long k,
-                                          long long K, int itflags) {
-  if (itflags & ITF_REDUCE) {
-    const float* pp = reinterpret_cast<const float*>(P[1]);
-    float s = pp[k];
-#pragma unroll
-    for (int sl = 1; sl < NSLICE; ++sl) s += pp[(long long)sl * K + k];
-    const bf16* resid = reinterpret_cast<const bf16*>(P[5]);
-    if (resid != nullptr) s = bf2f(resid[k]) + s;
-    const bf16 v = f2bf(s);
-    if (itflags & ITF_WRITEBACK) reinterpret_cast<bf16*>(P[6])[k] = v;
-    return v;
-  }
-  return reinterpret_cast<const bf16*>(P[1])[k];
-}
-__device__ void ph_gemv(const long long* P, int pos) {
   const bf16* W = reinterpret_cast<const bf16*>(P[2]);
   const bf16* gamma = reinterpret_cast<const bf16*>(P[4]);
   const bf16* resid = reinterpret_cast<const bf16*>(P[5]);
   const long long N = P[7], K = P[8];
-  const int itflags = (int)P[9];
-  const int itrans = itflags & ITF_BASE_MASK;
   const int epi = (int)P[10];
   const float eps = bits2f(P[14]);
   const float* part = reinterpret_cast<const float*>(P[1]);
   const int D = (int)P[13];
   const int maxch = (int)P[15];
-  const int nch = (pos + ATTN_CHUNK) / ATTN_CHUNK;  // ceil((pos+1)/CHUNK)
-  // For reduce inputs, the vector length that carries the reduction is the
-  // producer's row count: K for most transforms, 2K for swiglu pairs.
-  const long long KR = (itrans == IT_SWIGLU) ? 2 * K : K;
-  // Stage the transformed input once per block: every transform value is a
-  // bf16-rounded quantity by semantics, so bf16 staging is exact. This turns
-  // per-row-visit transform work (swiglu exp, attention finalize) into
-  // once-per-block work and makes the dot loop pure shared-memory reads.
-  // RMSNORM stages in two passes (raw materialized value first, then the
-  // norm applied in place) so reduce+resid+writeback consumption happens
-  // exactly once per element.
   extern __shared__ bf16 xs[];
   __shared__ float s_red[32];
   __shared__ float s_scale;
-  if (itrans == IT_RMSNORM) {
-    float ss = 0.f;
-    for (long long k = threadIdx.x; k < K; k += blockDim.x) {
-      const bf16 xv = consume_x(P, k, KR, itflags);
-      xs[k] = xv;
-      const float v = bf2f(xv);
-      ss += v * v;
-    }
-    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 k = threadIdx.x; k < K; k += blockDim.x)
-      xs[k] = f2bf(bf2f(gamma[k]) * bf2f(f2bf(bf2f(xs[k]) * inv)));
-  } else {
-    for (long long k = threadIdx.x; k < K; k += blockDim.x) {
-      float v;
-      if (itrans == IT_SWIGLU) {
-        const float g = bf2f(consume_x(P, k, KR, itflags));
         const float sig = 1.f / (1.f + expf(-g));
         const float sg = bf2f(f2bf(g * sig));
-        v = bf2f(f2bf(sg * bf2f(consume_x(P, K + k, KR, itflags))));
-      } else if (itrans == IT_ATTNFIN) {
-        v = attn_finalize(part, (int)(k / D), (int)(k % D), D, maxch, nch);
-      } else {
-        v = bf2f(consume_x(P, k, KR, itflags));
       }
-      xs[k] = f2bf(v);
     }
   }
   __syncthreads();
@@ -277,57 +275,61 @@ __device__ void ph_gemv(const long long* P, int pos) {
   const long long warp =
       ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
   const long long nwarps = ((long long)gridDim.x * blockDim.x) >> 5;
-  __shared__ unsigned long long s64[32];
   unsigned long long best = amax_key(-INFINITY, 0xFFFFFFFFu);
-  // EP_PARTIAL producers split each row across NSLICE k-slices handled by
-  // independent warps: NSLICE times the parallelism and chains a quarter as
-  // long, with the fixed-order reduction done by the consumer's staging.
-  const bool partial = (epi == EP_PARTIAL);
-  const long long V = partial ? N * NSLICE : N;
-  const long long kslice = K / NSLICE;
-  for (long long v = warp; v < V; v += nwarps) {
-    const long long row = partial ? v / NSLICE : v;
-    const int slice = partial ? (int)(v % NSLICE) : 0;
-    const long long klo = partial ? (long long)slice * kslice : 0;
-    const long long khi = partial ? klo + kslice : K;
     const bf16* wr = W + row * K;
-    float acc = 0.f;
-    // 16-byte loads, next tile prefetched so the weight fetch for iteration
-    // i+1 overlaps the fma chain of iteration i.
-    long long k = klo + (long long)lane * 8;
-    if (k + 7 < khi) {
       int4 wv4 = *reinterpret_cast<const int4*>(wr + k);
-      for (; k + 7 < khi; k += 256) {
         int4 nxt;
-        const bool more = (k + 256 + 7 < khi);
         if (more) nxt = *reinterpret_cast<const int4*>(wr + k + 256);
         const __nv_bfloat162* wp =
             reinterpret_cast<const __nv_bfloat162*>(&wv4);
-        const __nv_bfloat162* xp =
-            reinterpret_cast<const __nv_bfloat162*>(xs + k);
 #pragma unroll
-        for (int j = 0; j < 4; ++j) {
-          acc = fmaf(__low2float(wp[j]), __low2float(xp[j]), acc);
-          acc = fmaf(__high2float(wp[j]), __high2float(xp[j]), acc);
         }
         if (more) wv4 = nxt;
       }
     }
-    acc = warp_sum(acc);
-    if (lane == 0) {
-      if (partial)
-        reinterpret_cast<float*>(P[3])[(long long)slice * N + row] = acc;
-      else if (epi == EP_RESID)
-        reinterpret_cast<bf16*>(P[3])[row] = f2bf(bf2f(resid[row]) + acc);
-      else if (epi == EP_F32)
-        reinterpret_cast<float*>(P[3])[row] = acc;
-      else if (epi == EP_F32_AMAX) {
-        reinterpret_cast<float*>(P[3])[row] = acc;
-        const unsigned long long kk = amax_key(bf2f(f2bf(acc)),
-                                               (unsigned)row);
-        best = (kk > best) ? kk : best;
-      } else {
-        reinterpret_cast<bf16*>(P[3])[row] = f2bf(acc);
       }
     }
   }
@@ -339,16 +341,58 @@ __device__ void ph_gemv(const long long* P, int pos) {
   }
 }
-// Attention with rope fused in and cached positions split across chunks.
-// Slots: p1 qkv, p2 kcache, p3 vcache, p4 fp32 partials
-// [Hq, maxch, D+2] (weighted V, chunk max, chunk sum), p5 q-norm w,
-// p6 k-norm w, 7 max_seq, 8 qk flag, 9 eps bits, 10 theta bits, 11 Hq,
-// 12 Hkv, 13 D, 14 scale, 15 maxch. Each (head, chunk) block ropes its own
-// copy of q (identical bits across chunks); the chunk owning `pos` also
 // norms/ropes K and appends K/V for its group (duplicate identical writes
-// across GQA members are benign). The O-projection finalizes the partials.
-__device__ void ph_attn(const long long* P, int pos) {
   const long long max_seq = P[7];
   const int qk_flag = (int)P[8];
   const float eps = bits2f(P[9]);
@@ -356,29 +400,11 @@ __device__ void ph_attn(const long long* P, int pos) {
   const int Hq = (int)P[11], Hkv = (int)P[12], D = (int)P[13];
   const float scale = bits2f(P[14]);
   const int maxch = (int)P[15];
-  const int S = pos + 1;
-  const int nch = (S + ATTN_CHUNK - 1) / ATTN_CHUNK;
-  const int own_ch = pos / ATTN_CHUNK;
   const int rep = Hq / Hkv;
   const bf16* qnw = reinterpret_cast<const bf16*>(P[5]);
   const bf16* knw = reinterpret_cast<const bf16*>(P[6]);
-  const bool red = (qk_flag & 2) != 0;
-  const long long KT = (long long)(Hq + 2 * Hkv) * D;
-  const float* qp = reinterpret_cast<const float*>(P[1]);
-  const bf16* qb = reinterpret_cast<const bf16*>(P[1]);
-  // qkv element: reduce the producer's fp32 k-slice partials (fixed order)
-  // and round to bf16, reproducing the monolithic store's value semantics.
-  auto qkv_at = [&](long long idx) -> float {
-    if (red) {
-      float s = qp[idx];
-#pragma unroll
-      for (int sl = 1; sl < NSLICE; ++sl) s += qp[(long long)sl * KT + idx];
-      return bf2f(f2bf(s));
-    }
-    return bf2f(qb[idx]);
-  };
   __shared__ float qs[MAX_HEAD_DIM];
   __shared__ float hbuf[MAX_HEAD_DIM];
@@ -386,137 +412,150 @@ __device__ void ph_attn(const long long* P, int pos) {
   __shared__ float s_red[32];
   __shared__ float s_inv, s_m;
-  const int total = Hq * nch;
-  for (int v = blockIdx.x; v < total; v += gridDim.x) {
-    const int head = v / nch;
-    const int ch = v % nch;
-    const int kvh = head / rep;
-    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;
-    float* part = reinterpret_cast<float*>(P[4]) +
-                  ((long long)head * maxch + ch) * (D + 2);
-    const int half = D >> 1;
-    // stage rope'd q (bf16-rounded values) into shared
-    {
-      const long long qbase = (long long)head * D;
-      if (qk_flag & 1) {
-        float ss = 0.f;
-        for (int d = threadIdx.x; d < D; d += blockDim.x) {
-          const float x = qkv_at(qbase + 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)
-          hbuf[d] = bf2f(f2bf(bf2f(qnw[d]) *
-                              bf2f(f2bf(qkv_at(qbase + d) * s_inv))));
-      } else {
-        for (int d = threadIdx.x; d < D; d += blockDim.x)
-          hbuf[d] = qkv_at(qbase + d);
       }
-      __syncthreads();
-      for (int d = threadIdx.x; d < half; d += blockDim.x) {
-        const float inv_freq = 1.f / powf(theta, (float)(2 * d) / (float)D);
-        const float ang = (float)pos * inv_freq;
-        const float c = bf2f(f2bf(cosf(ang)));
-        const float s = bf2f(f2bf(sinf(ang)));
-        const float x1 = hbuf[d], x2 = hbuf[d + half];
-        qs[d] = bf2f(f2bf(bf2f(f2bf(x1 * c)) + bf2f(f2bf(-x2 * s))));
-        qs[d + half] = bf2f(f2bf(bf2f(f2bf(x2 * c)) + bf2f(f2bf(x1 * s))));
       }
-      __syncthreads();
-    }
-    // the chunk owning `pos` norms/ropes K and appends K/V for its group
-    if (ch == own_ch) {
-      const long long kbase = (long long)(Hq + kvh) * D;
-      const long long vbase = (long long)(Hq + Hkv + kvh) * D;
-      if (qk_flag & 1) {
-        float ss = 0.f;
-        for (int d = threadIdx.x; d < D; d += blockDim.x) {
-          const float x = qkv_at(kbase + 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)
-          hbuf[d] = bf2f(f2bf(bf2f(knw[d]) *
-                              bf2f(f2bf(qkv_at(kbase + d) * s_inv))));
-      } else {
-        for (int d = threadIdx.x; d < D; d += blockDim.x)
-          hbuf[d] = qkv_at(kbase + d);
       }
       __syncthreads();
-      bf16* krow = kc + (long long)pos * D;
-      bf16* vrow = vc + (long long)pos * D;
-      for (int d = threadIdx.x; d < half; d += blockDim.x) {
-        const float inv_freq = 1.f / powf(theta, (float)(2 * d) / (float)D);
-        const float ang = (float)pos * inv_freq;
-        const float c = bf2f(f2bf(cosf(ang)));
-        const float s = bf2f(f2bf(sinf(ang)));
-        const float x1 = hbuf[d], x2 = hbuf[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] = f2bf(qkv_at(vbase + d));
-      __threadfence_block();
       __syncthreads();
-    }
-    // score this chunk: eager rounds QK^T to bf16, scales as a bf16 op,
-    // upcasts to fp32 for the softmax
-    const int s0 = ch * ATTN_CHUNK;
-    const int s1 = (s0 + ATTN_CHUNK < S) ? s0 + ATTN_CHUNK : S;
-    const int cn = s1 - s0;
-    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 = kc + (long long)(s0 + si) * D;
-      float acc = 0.f;
-      for (int d = lane; d < D; d += 32) acc = fmaf(qs[d], bf2f(kr[d]), acc);
-      acc = warp_sum(acc);
-      if (lane == 0) {
-        float x = bf2f(f2bf(acc));
-        sc[si] = bf2f(f2bf(x * scale));
       }
     }
-    __syncthreads();
-    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;
-    }
-    lsum = block_sum(lsum, s_red);
-    __syncthreads();
-    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(vc[(long long)(s0 + si) * D + d]), acc);
-      part[d] = acc;
-    }
-    if (threadIdx.x == 0) {
-      part[D] = s_m;
-      part[D + 1] = lsum;
-    }
-    __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]);
@@ -555,13 +594,15 @@ __device__ void ph_argmax_fin(const long long* P, int step_slot) {
 }
 __device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
-                                               int step_slot, int prompt_len) {
   switch ((int)P[0]) {
-    case OP_EMBED: ph_embed(P, pos, prompt_len); break;
-    case OP_GEMV: ph_gemv(P, pos); break;
-    case OP_ATTN: ph_attn(P, pos); break;
     case OP_ARGMAX_PART: ph_argmax_part(P); break;
     case OP_ARGMAX_FIN: ph_argmax_fin(P, step_slot); break;
     default: break;
   }
 }
@@ -592,29 +633,44 @@ __device__ void grid_barrier(int* cnt, int* sense, int* local_sense) {
   __syncthreads();
 }
-// One launch runs `steps` decode steps; the grid barrier carries across the
-// step boundary, so the argmax of step i feeds the embedding read of step
-// i+1 without leaving the kernel. steps=1 is a single decode step.
 __global__ void mega_kernel(const long long* __restrict__ prog, int n_phases,
-                            int pos0, int steps, int slot0, int prompt_len,
-                            int* bar_cnt, int* bar_sense) {
   __shared__ int s_sense;
   if (threadIdx.x == 0) s_sense = atomicAdd(bar_sense, 0);
   __syncthreads();
-  for (int st = 0; st < steps; ++st) {
-    const int pos = pos0 + st;
-    const int slot = slot0 + st;
     for (int p = 0; p < n_phases; ++p) {
-      dispatch_phase(prog + (long long)p * 16, pos, slot, prompt_len);
-      const bool last = (st + 1 == steps) && (p + 1 == n_phases);
       if (!last) grid_barrier(bar_cnt, bar_sense, &s_sense);
     }
   }
 }
 __global__ void phase_kernel(const long long* __restrict__ prog, int p,
                              int pos, int step_slot, int prompt_len) {
-  dispatch_phase(prog + (long long)p * 16, pos, step_slot, prompt_len);
 }
 }  // namespace mak
@@ -648,11 +704,12 @@ int mak_grid_blocks_impl(int smem_bytes) {
 }
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
-                            int steps, int slot0, int prompt_len,
-                            int* bar_cnt, int* bar_sense, int grid,
-                            int smem_bytes, cudaStream_t stream) {
   mak::mega_kernel<<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
-      prog, n_phases, pos0, steps, slot0, prompt_len, bar_cnt, bar_sense);
   return cudaGetLastError();
 }

   OP_ATTN = 3,
   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 {
 // and any block with nothing to do simply falls through to the barrier.
 // ---------------------------------------------------------------------------
+__device__ void ph_embed(const long long* P, int pos, int M, int prompt_len) {
   const bf16* tab = reinterpret_cast<const bf16*>(P[2]);
   bf16* out = reinterpret_cast<bf16*>(P[3]);
   const int* prompt = reinterpret_cast<const int*>(P[5]);
   const long long K = P[8];
   // Positions inside the prompt read from the staged prompt buffer; beyond
   // it, from the feedback token written by the previous step's argmax.
+  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 k = i % K;
+    const int pm = pos + m;
+    const long long t = (pm < prompt_len && prompt != nullptr)
+                            ? (long long)prompt[pm]
+                            : (long long)tok[0];
+    out[i] = tab[t * K + k];
+  }
 }
 __device__ __forceinline__ unsigned long long amax_key(float v, unsigned idx) {
 // Finalize one element of split-S attention: combine per-chunk partials
 // (max, sum, weighted V) into the bf16 attention output, matching eager's
 // 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]);
 // Reduced (or plain) bf16 input element k for a consumer phase: sum the
 // producer's fp32 k-slice partials in fixed order, add the residual, round
 // to bf16, optionally write the materialized vector back.
+// GEMV over an M-token chunk. Weights are read once per chunk: each warp's
+// 16-byte weight tile feeds M accumulators against M staged input rows, so
+// prefill weight traffic is amortized by the chunk size. mmode (slot 12)
+// restricts the phase to the chunk's last row (LM head). Buffers are laid
+// out [M][len]; the residual stream and outputs address row m.
+__device__ void ph_gemv(const long long* P, int pos, int M) {
   const bf16* W = reinterpret_cast<const bf16*>(P[2]);
   const bf16* gamma = reinterpret_cast<const bf16*>(P[4]);
   const bf16* resid = reinterpret_cast<const bf16*>(P[5]);
   const long long N = P[7], K = P[8];
+  const int itrans = ((int)P[9]) & ITF_BASE_MASK;
   const int epi = (int)P[10];
+  const int mmode = (int)P[12];
   const float eps = bits2f(P[14]);
   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
+  // Stage transformed inputs for all chunk rows: [Ma][K] in shared, bf16
+  // (every transform value is bf16-rounded by semantics, so staging is
+  // exact). RMSNORM computes one inverse rms per row.
   extern __shared__ bf16 xs[];
   __shared__ float s_red[32];
   __shared__ float s_scale;
+  __shared__ unsigned long long s64[32];
+  for (int m = 0; m < Ma; ++m) {
+    const int srcm = m0 + m;
+    if (itrans == IT_RMSNORM) {
+      const bf16* xr = xin + (long long)srcm * K;
+      float ss = 0.f;
+      for (long long k = threadIdx.x; k < K; k += blockDim.x) {
+        const float v = bf2f(xr[k]);
+        ss += v * v;
+      }
+      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 k = threadIdx.x; k < K; k += blockDim.x)
+        xs[(long long)m * K + k] =
+            f2bf(bf2f(gamma[k]) * bf2f(f2bf(bf2f(xr[k]) * inv)));
+      __syncthreads();
+    } else if (itrans == IT_SWIGLU) {
+      const bf16* xr = xin + (long long)srcm * 2 * K;
+      for (long long k = threadIdx.x; k < K; k += blockDim.x) {
+        const float g = bf2f(xr[k]);
         const float sig = 1.f / (1.f + expf(-g));
         const float sg = bf2f(f2bf(g * sig));
+        xs[(long long)m * K + k] = f2bf(bf2f(f2bf(sg * bf2f(xr[K + k]))));
       }
+    } else if (itrans == IT_ATTNFIN) {
+      const int pm = pos + srcm;
+      const int nch = (pm + ATTN_CHUNK) / ATTN_CHUNK;
+      const float* mpart =
+          part + (long long)srcm * ((long long)P[11] * maxch * (D + 2));
+      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;
+      for (long long k = threadIdx.x; k < K; k += blockDim.x)
+        xs[(long long)m * K + k] = xr[k];
     }
   }
   __syncthreads();
   const long long warp =
       ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
   const long long nwarps = ((long long)gridDim.x * blockDim.x) >> 5;
   unsigned long long best = amax_key(-INFINITY, 0xFFFFFFFFu);
+  for (long long row = warp; row < N; row += nwarps) {
     const bf16* wr = W + row * K;
+    float acc[8];
+#pragma unroll
+    for (int m = 0; m < 8; ++m) acc[m] = 0.f;
+    // 16-byte weight loads, next tile prefetched; each tile feeds all
+    // staged rows.
+    long long k = (long long)lane * 8;
+    if (k + 7 < K) {
       int4 wv4 = *reinterpret_cast<const int4*>(wr + k);
+      for (; k + 7 < K; k += 256) {
         int4 nxt;
+        const bool more = (k + 256 + 7 < K);
         if (more) nxt = *reinterpret_cast<const int4*>(wr + k + 256);
         const __nv_bfloat162* wp =
             reinterpret_cast<const __nv_bfloat162*>(&wv4);
 #pragma unroll
+        for (int m = 0; m < 8; ++m) {
+          if (m < Ma) {
+            const __nv_bfloat162* xp =
+                reinterpret_cast<const __nv_bfloat162*>(xs +
+                                                        (long long)m * K + k);
+#pragma unroll
+            for (int j = 0; j < 4; ++j) {
+              acc[m] = fmaf(__low2float(wp[j]), __low2float(xp[j]), acc[m]);
+              acc[m] = fmaf(__high2float(wp[j]), __high2float(xp[j]),
+                            acc[m]);
+            }
+          }
         }
         if (more) wv4 = nxt;
       }
     }
+#pragma unroll
+    for (int m = 0; m < 8; ++m) {
+      if (m < Ma) {
+        const float a = warp_sum(acc[m]);
+        if (lane == 0) {
+          const int srcm = m0 + m;
+          if (epi == EP_RESID)
+            reinterpret_cast<bf16*>(P[3])[(long long)srcm * N + row] =
+                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_F32_AMAX) {
+            reinterpret_cast<float*>(P[3])[row] = a;
+            const unsigned long long kk = amax_key(bf2f(f2bf(a)),
+                                                   (unsigned)row);
+            best = (kk > best) ? kk : best;
+          } else {
+            reinterpret_cast<bf16*>(P[3])[(long long)srcm * N + row] =
+                f2bf(a);
+          }
+        }
       }
     }
   }
   }
 }
+// Norm/rope K and copy V into the cache for one (kv-head, position);
+// block-cooperative, semantics identical to the eager reference.
+__device__ void kv_norm_rope_append(const bf16* qkv, bf16* kc, bf16* vc,
+                                    const bf16* knw, int qk_flag, float eps,
+                                    float theta, 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;
+  if (qk_flag & 1) {
+    float ss = 0.f;
+    for (int d = threadIdx.x; d < D; d += blockDim.x) {
+      const float x = bf2f(ksrc[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)
+      hbuf[d] = bf2f(f2bf(bf2f(knw[d]) *
+                          bf2f(f2bf(bf2f(ksrc[d]) * (*s_inv)))));
+  } else {
+    for (int d = threadIdx.x; d < D; d += blockDim.x) hbuf[d] = bf2f(ksrc[d]);
+  }
+  __syncthreads();
+  bf16* krow = kc + (long long)pm * D;
+  bf16* vrow = vc + (long long)pm * D;
+  for (int d = threadIdx.x; d < half; d += blockDim.x) {
+    const float inv_freq = 1.f / powf(theta, (float)(2 * d) / (float)D);
+    const float ang = (float)pm * inv_freq;
+    const float c = bf2f(f2bf(cosf(ang)));
+    const float s = bf2f(f2bf(sinf(ang)));
+    const float x1 = hbuf[d], x2 = hbuf[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();
+}
+// Attention with rope fused in, cached positions split across chunks, and
+// an M-token prefill loop. Slots: p1 qkv [M rows], p2 kcache, p3 vcache,
+// p4 fp32 partials [M][Hq][maxch][D+2], p5 q-norm w, p6 k-norm w,
+// 7 max_seq, 8 qk flag, 9 eps bits, 10 theta bits, 11 Hq, 12 Hkv, 13 D,
+// 14 scale, 15 maxch. In decode (M == 1) the chunk owning `pos` also
 // norms/ropes K and appends K/V for its group (duplicate identical writes
+// across GQA members are benign); in prefill chunks the OP_KV_APPEND phase
+// has already appended all M positions, so causal reads within the chunk
+// 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 int Hq = (int)P[11], Hkv = (int)P[12], D = (int)P[13];
   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;
   const bf16* qnw = reinterpret_cast<const bf16*>(P[5]);
   const bf16* knw = reinterpret_cast<const bf16*>(P[6]);
   __shared__ float qs[MAX_HEAD_DIM];
   __shared__ float hbuf[MAX_HEAD_DIM];
   __shared__ float s_red[32];
   __shared__ float s_inv, s_m;
+  // (m, head, chunk) flattened into one block-stride space so prefill
+  // chunks use M times the parallelism; per-m chunk counts vary, so
+  // out-of-range (m, ch) pairs skip uniformly.
+  const int nch_max =
+      (pos + M - 1 + 1 + ATTN_CHUNK - 1) / ATTN_CHUNK;
+  {
+    const int total = M * Hq * nch_max;
+    for (int v = blockIdx.x; v < total; v += gridDim.x) {
+      const int m = v / (Hq * nch_max);
+      const int rem = v % (Hq * nch_max);
+      const int head = rem / nch_max;
+      const int ch = rem % nch_max;
+      const int pm = pos + m;
+      const int S = pm + 1;
+      const int nch = (S + ATTN_CHUNK - 1) / ATTN_CHUNK;
+      if (ch >= nch) continue;
+      const int own_ch = pm / ATTN_CHUNK;
+      const bf16* qkv =
+          reinterpret_cast<const bf16*>(P[1]) + (long long)m * KT;
+      float* partm = reinterpret_cast<float*>(P[4]) +
+                     (long long)m * ((long long)Hq * maxch * (D + 2));
+      const int kvh = head / rep;
+      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;
+      float* part = partm + ((long long)head * maxch + ch) * (D + 2);
+      const int half = D >> 1;
+      // stage rope'd q (bf16-rounded values) into shared
+      {
+        const bf16* q = qkv + (long long)head * D;
+        if (qk_flag & 1) {
+          float ss = 0.f;
+          for (int d = threadIdx.x; d < D; d += blockDim.x) {
+            const float x = bf2f(q[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)
+            hbuf[d] =
+                bf2f(f2bf(bf2f(qnw[d]) * bf2f(f2bf(bf2f(q[d]) * s_inv))));
+        } else {
+          for (int d = threadIdx.x; d < D; d += blockDim.x)
+            hbuf[d] = bf2f(q[d]);
+        }
+        __syncthreads();
+        for (int d = threadIdx.x; d < half; d += blockDim.x) {
+          const float inv_freq = 1.f / powf(theta, (float)(2 * d) / (float)D);
+          const float ang = (float)pm * inv_freq;
+          const float c = bf2f(f2bf(cosf(ang)));
+          const float s = bf2f(f2bf(sinf(ang)));
+          const float x1 = hbuf[d], x2 = hbuf[d + half];
+          qs[d] = bf2f(f2bf(bf2f(f2bf(x1 * c)) + bf2f(f2bf(-x2 * s))));
+          qs[d + half] = bf2f(f2bf(bf2f(f2bf(x2 * c)) + bf2f(f2bf(x1 * s))));
         }
         __syncthreads();
       }
+      // decode only: the chunk owning `pos` appends K/V for its group
+      if (M == 1 && ch == own_ch) {
+        kv_norm_rope_append(qkv, kc, vc, knw, qk_flag, eps, theta, Hq, Hkv,
+                            kvh, D, pm, s_red, &s_inv, hbuf);
       }
+      // score this chunk: eager rounds QK^T to bf16, scales as a bf16 op,
+      // upcasts to fp32 for the softmax
+      const int s0 = ch * ATTN_CHUNK;
+      const int s1 = (s0 + ATTN_CHUNK < S) ? s0 + ATTN_CHUNK : S;
+      const int cn = s1 - s0;
+      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 = kc + (long long)(s0 + si) * D;
+        float acc = 0.f;
+        for (int d = lane; d < D; d += 32) acc = fmaf(qs[d], bf2f(kr[d]), acc);
+        acc = warp_sum(acc);
+        if (lane == 0) {
+          float x = bf2f(f2bf(acc));
+          sc[si] = bf2f(f2bf(x * scale));
         }
       }
       __syncthreads();
+      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;
       }
+      lsum = block_sum(lsum, s_red);
       __syncthreads();
+      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(vc[(long long)(s0 + si) * D + d]), acc);
+        part[d] = acc;
+      }
+      if (threadIdx.x == 0) {
+        part[D] = s_m;
+        part[D + 1] = lsum;
       }
+      __syncthreads();
     }
+  }
+}
+// Prefill chunks: norm/rope K and copy V into the cache for every position
+// of the chunk, one (m, kv-head) pair per block iteration. Runs before the
+// 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 theta = bits2f(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]);
+  __shared__ float s_red[32];
+  __shared__ float s_inv;
+  __shared__ float hbuf[MAX_HEAD_DIM];
+  const int total = M * Hkv;
+  for (int v = blockIdx.x; v < total; v += gridDim.x) {
+    const int m = v / Hkv;
+    const int kvh = v % Hkv;
+    const int pm = pos + m;
+    const bf16* qkv = reinterpret_cast<const bf16*>(P[1]) + (long long)m * KT;
+    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, theta, 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]);
 }
 __device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
+                                               int M, int step_slot,
+                                               int prompt_len) {
   switch ((int)P[0]) {
+    case OP_EMBED: ph_embed(P, pos, M, prompt_len); break;
+    case OP_GEMV: ph_gemv(P, pos, M); break;
+    case OP_ATTN: ph_attn(P, pos, M); break;
     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;
   }
 }
   __syncthreads();
 }
+// One launch consumes `total` positions: prompt positions in chunks of up
+// to `chunk_m` (weights amortized across the chunk), generated positions
+// one at a time with the argmax of each step feeding the next embedding
+// read. The grid barrier carries across every boundary. OP_KV_APPEND
+// phases run only inside multi-token chunks; in decode both the dispatch
+// and the barrier are skipped uniformly, so they cost nothing.
 __global__ void mega_kernel(const long long* __restrict__ prog, int n_phases,
+                            int pos0, int total, int slot0, int prompt_len,
+                            int chunk_m, int* bar_cnt, int* bar_sense) {
   __shared__ int s_sense;
   if (threadIdx.x == 0) s_sense = atomicAdd(bar_sense, 0);
   __syncthreads();
+  int pos = pos0;
+  const int pos_end = pos0 + total;
+  while (pos < pos_end) {
+    int M = 1;
+    if (pos < prompt_len) {
+      M = prompt_len - pos;
+      if (M > chunk_m) M = chunk_m;
+      if (pos + M > pos_end) M = pos_end - pos;
+    }
+    const int slot = slot0 + (pos + M - 1 - pos0);
     for (int p = 0; p < n_phases; ++p) {
+      const long long* P = prog + (long long)p * 16;
+      if ((int)P[0] == OP_KV_APPEND && M == 1) continue;  // uniform skip
+      dispatch_phase(P, pos, M, slot, prompt_len);
+      const bool last = (pos + M >= pos_end) && (p + 1 == n_phases);
       if (!last) grid_barrier(bar_cnt, bar_sense, &s_sense);
     }
+    pos += M;
   }
 }
 __global__ void phase_kernel(const long long* __restrict__ prog, int p,
                              int pos, int step_slot, int prompt_len) {
+  const long long* P = prog + (long long)p * 16;
+  if ((int)P[0] == OP_KV_APPEND) return;
+  dispatch_phase(P, pos, 1, step_slot, prompt_len);
 }
 }  // namespace mak
 }
 cudaError_t mak_launch_mega(const long long* prog, int n_phases, int pos0,
+                            int total, int slot0, int prompt_len,
+                            int chunk_m, int* bar_cnt, int* bar_sense,
+                            int grid, int smem_bytes, cudaStream_t stream) {
   mak::mega_kernel<<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
+      prog, n_phases, pos0, total, slot0, prompt_len, chunk_m, bar_cnt,
+      bar_sense);
   return cudaGetLastError();
 }

torch-ext/model_as_a_kernel/__init__.py CHANGED Viewed

@@ -111,14 +111,18 @@ class MegaModel:
         # Working buffers. Pointers to these are baked into the program, so
         # they (like the packed weights) live for the model's lifetime.
-        self._hidden = torch.empty(self.hidden, dtype=dt, device=dev)
-        self._qkv = torch.empty(qdim + 2 * kvdim, dtype=dt, device=dev)
-        self._gu = torch.empty(2 * self.I, dtype=dt, device=dev)
         self._logits = torch.empty(self.V, dtype=torch.float32, device=dev)
         self._chunk = 256  # ATTN_CHUNK in the kernel
         self._maxch = (self.max_seq + self._chunk - 1) // self._chunk
         self._partials = torch.empty(
-            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)
@@ -131,7 +135,6 @@ class MegaModel:
                                        device=dev)
         self._bar = torch.zeros(2, dtype=torch.int32, device=dev)
-        self._maxk = max(self.hidden, qdim, self.I)
         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)
@@ -153,6 +156,12 @@ class MegaModel:
                              n=qdim + 2 * kvdim, k=self.hidden,
                              it=_IT_RMSNORM, epi=_EP_STORE, f0=self.eps))
             names.append(f"L{li}.qkv")
             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,
@@ -165,7 +174,8 @@ class MegaModel:
             rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
                              p2=lw["wo"].data_ptr(), p3=hid, p5=hid,
                              n=self.hidden, k=qdim, it=_IT_ATTNFIN,
-                             epi=_EP_RESID, d=self.D, i0=self._maxch))
             names.append(f"L{li}.o")
             rows.append(_row(_OP_GEMV, p1=hid, p2=lw["wgu"].data_ptr(),
                              p3=self._gu.data_ptr(),
@@ -182,7 +192,8 @@ class MegaModel:
                          p3=self._logits.data_ptr(),
                          p4=final_norm.data_ptr(),
                          p6=self._parts.data_ptr(), n=self.V, k=self.hidden,
-                         it=_IT_RMSNORM, epi=_EP_F32_AMAX, f0=self.eps))
         names.append("lm_head")
         rows.append(_row(_OP_ARGMAX_FIN, p1=self._parts.data_ptr(),
                          p3=self._token.data_ptr(),
@@ -275,9 +286,17 @@ class MegaModel:
         return self._logits, ms
     def prefill(self, ids: Sequence[int], phased: bool = False):
-        """Feed a prompt token by token; returns the last position's logits."""
-        for i, t in enumerate(ids):
-            self.decode_step(int(t), i, phased=phased)
         return self._logits
     def generate(self, prompt_ids: Sequence[int], max_new: int,
@@ -298,7 +317,7 @@ class MegaModel:
             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)
             return self._tokens_out[:max_new].cpu().tolist()
         pos = 0
         for t in prompt_ids[:-1]:

         # 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
+        self._chunk_m = max(1, min(8, 40960 // (2 * 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)
         self._chunk = 256  # ATTN_CHUNK in the kernel
         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)
                                        device=dev)
         self._bar = torch.zeros(2, 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)
                              n=qdim + 2 * kvdim, k=self.hidden,
                              it=_IT_RMSNORM, epi=_EP_STORE, f0=self.eps))
             names.append(f"L{li}.qkv")
+            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=_fbits(self.theta),
+                             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,
             rows.append(_row(_OP_GEMV, p1=self._partials.data_ptr(),
                              p2=lw["wo"].data_ptr(), p3=hid, p5=hid,
                              n=self.hidden, k=qdim, it=_IT_ATTNFIN,
+                             epi=_EP_RESID, hq=self.Hq, d=self.D,
+                             i0=self._maxch))
             names.append(f"L{li}.o")
             rows.append(_row(_OP_GEMV, p1=hid, p2=lw["wgu"].data_ptr(),
                              p3=self._gu.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, 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(),
         return self._logits, 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
     def generate(self, prompt_ids: Sequence[int], max_new: int,
             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]:

torch-ext/torch_binding.cpp CHANGED Viewed

@@ -1,4 +1,4 @@
-#include <torch/library.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -24,9 +24,11 @@ void check_barrier(const at::Tensor& barrier) {
               "barrier must be a CUDA int32 tensor with >= 2 elements");
 }
-int smem_for(int64_t max_k) {
-  TORCH_CHECK(max_k > 0 && max_k <= 16384, "max_k out of range");
-  return (int)(max_k * 2);
 }
 }  // namespace
@@ -46,7 +48,7 @@ void mak_run(at::Tensor prog, at::Tensor barrier, int64_t pos,
   const int smem = smem_for(max_k);
   const cudaError_t err = mak_launch_mega(
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
-      (int)prog.size(0), (int)pos, 1, (int)step_slot, 0, bar, bar + 1,
       mak_grid_blocks_impl(smem), smem, stream.stream());
   TORCH_CHECK(err == cudaSuccess, "mak_run launch failed: ",
               cudaGetErrorString(err));
@@ -66,8 +68,8 @@ void mak_run_steps(at::Tensor prog, at::Tensor barrier, int64_t pos0,
   const int grid = mak_grid_blocks_impl(smem);
   for (int64_t i = 0; i < steps; ++i) {
     const cudaError_t err =
-        mak_launch_mega(pp, n, (int)(pos0 + i), 1, (int)(slot0 + i), 0, bar,
-                        bar + 1, grid, smem, stream.stream());
     TORCH_CHECK(err == cudaSuccess, "mak_run_steps launch failed: ",
                 cudaGetErrorString(err));
   }
@@ -75,18 +77,20 @@ void mak_run_steps(at::Tensor prog, at::Tensor barrier, int64_t pos0,
 void mak_run_seq(at::Tensor prog, at::Tensor barrier, int64_t pos0,
                  int64_t steps, int64_t slot0, int64_t prompt_len,
-                 int64_t max_k) {
   check_prog(prog);
   check_barrier(barrier);
   TORCH_CHECK(steps >= 1, "steps must be >= 1");
   const c10::cuda::CUDAGuard guard(prog.device());
   auto stream = at::cuda::getCurrentCUDAStream();
   int* bar = barrier.data_ptr<int>();
-  const int smem = smem_for(max_k);
   const cudaError_t err = mak_launch_mega(
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
       (int)prog.size(0), (int)pos0, (int)steps, (int)slot0, (int)prompt_len,
-      bar, bar + 1, mak_grid_blocks_impl(smem), smem, stream.stream());
   TORCH_CHECK(err == cudaSuccess, "mak_run_seq launch failed: ",
               cudaGetErrorString(err));
 }
@@ -136,7 +140,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
   ops.def("mak_run_steps(Tensor prog, Tensor barrier, int pos0, int steps, "
           "int slot0, int max_k) -> ()");
   ops.def("mak_run_seq(Tensor prog, Tensor barrier, int pos0, int steps, "
-          "int slot0, int prompt_len, int max_k) -> ()");
   ops.def("mak_run_phased(Tensor prog, int pos, int step_slot, bool timed, "
           "int max_k) -> Tensor");
   ops.def("mak_num_blocks(Tensor prog, int max_k) -> int");
@@ -151,3 +155,4 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
 }
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

+#include <torch/library.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
               "barrier must be a CUDA int32 tensor with >= 2 elements");
 }
+int smem_for(int64_t stage_elems) {
+  // bf16 staging must fit the default 48KB dynamic shared memory window
+  TORCH_CHECK(stage_elems > 0 && stage_elems <= 24576,
+              "staging size out of range");
+  return (int)(stage_elems * 2);
 }
 }  // namespace
   const int smem = smem_for(max_k);
   const cudaError_t err = mak_launch_mega(
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
+      (int)prog.size(0), (int)pos, 1, (int)step_slot, 0, 1, bar, bar + 1,
       mak_grid_blocks_impl(smem), smem, stream.stream());
   TORCH_CHECK(err == cudaSuccess, "mak_run launch failed: ",
               cudaGetErrorString(err));
   const int grid = mak_grid_blocks_impl(smem);
   for (int64_t i = 0; i < steps; ++i) {
     const cudaError_t err =
+        mak_launch_mega(pp, n, (int)(pos0 + i), 1, (int)(slot0 + i), 0, 1,
+                        bar, bar + 1, grid, smem, stream.stream());
     TORCH_CHECK(err == cudaSuccess, "mak_run_steps launch failed: ",
                 cudaGetErrorString(err));
   }
 void mak_run_seq(at::Tensor prog, at::Tensor barrier, int64_t pos0,
                  int64_t steps, int64_t slot0, int64_t prompt_len,
+                 int64_t max_k, int64_t chunk_m) {
   check_prog(prog);
   check_barrier(barrier);
   TORCH_CHECK(steps >= 1, "steps must be >= 1");
+  TORCH_CHECK(chunk_m >= 1 && chunk_m <= 8, "chunk_m must be in [1, 8]");
   const c10::cuda::CUDAGuard guard(prog.device());
   auto stream = at::cuda::getCurrentCUDAStream();
   int* bar = barrier.data_ptr<int>();
+  const int smem = smem_for(max_k * chunk_m);
   const cudaError_t err = mak_launch_mega(
       reinterpret_cast<const long long*>(prog.data_ptr<int64_t>()),
       (int)prog.size(0), (int)pos0, (int)steps, (int)slot0, (int)prompt_len,
+      (int)chunk_m, bar, bar + 1, mak_grid_blocks_impl(smem), smem,
+      stream.stream());
   TORCH_CHECK(err == cudaSuccess, "mak_run_seq launch failed: ",
               cudaGetErrorString(err));
 }
   ops.def("mak_run_steps(Tensor prog, Tensor barrier, int pos0, int steps, "
           "int slot0, int max_k) -> ()");
   ops.def("mak_run_seq(Tensor prog, Tensor barrier, int pos0, int steps, "
+          "int slot0, int prompt_len, int max_k, int chunk_m) -> ()");
   ops.def("mak_run_phased(Tensor prog, int pos, int step_slot, bool timed, "
           "int max_k) -> Tensor");
   ops.def("mak_num_blocks(Tensor prog, int max_k) -> int");
 }
 REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h CHANGED Viewed

@@ -10,7 +10,7 @@ void mak_run_steps(at::Tensor prog, at::Tensor barrier, int64_t pos0,
 void mak_run_seq(at::Tensor prog, at::Tensor barrier, int64_t pos0,
                  int64_t steps, int64_t slot0, int64_t prompt_len,
-                 int64_t max_k);
 at::Tensor mak_run_phased(at::Tensor prog, int64_t pos, int64_t step_slot,
                           bool timed, int64_t max_k);

 void mak_run_seq(at::Tensor prog, at::Tensor barrier, int64_t pos0,
                  int64_t steps, int64_t slot0, int64_t prompt_len,
+                 int64_t max_k, int64_t chunk_m);
 at::Tensor mak_run_phased(at::Tensor prog, int64_t pos, int64_t step_slot,
                           bool timed, int64_t max_k);