// model-as-a-kernel: an entire llama-family decode step as one kernel launch.
//
// The kernel is a phase interpreter. A "program" (int64 [n_phases, 16] array)
// lists every operation of the forward pass: embedding read, fused
// rmsnorm+GEMV, qk-norm+rope+cache-append, attention over the KV cache,
// swiglu+GEMV, lm-head GEMV, argmax. All blocks execute each phase
// cooperatively, then cross a software grid barrier; one launch therefore runs
// the whole model. A per-phase launch mode executes the identical device code
// one phase per kernel for debugging and cudaEvent timing; it is bitwise
// identical to the fused mode by construction (same grid, same arithmetic).
//
// Arithmetic mirrors transformers eager bf16 semantics: bf16 storage with
// fp32 accumulation and a single rounding at each op boundary, fp32 softmax
// with the probability cast to bf16 before the PV product, bf16 rope
// multiplies with the cos/sin table rounded to bf16, and the scores rounded
// to bf16 before and after the 1/sqrt(D) scaling as eager attention does.
// Reductions use fixed trees, so outputs are bitwise deterministic run to run.
//
// This translation unit is torch-free; see mak_kernels.h for the interface.

#include <cuda_bf16.h>
#include <cuda_pipeline_primitives.h>
#include <math.h>

#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
// latency-bandwidth product demands many groups in flight per lane); it
// is pure scheduling and does not touch arithmetic. One stage is
// 8 warps x 32 lanes x 16B = 4KB per block.
#define GEMV_RING_STAGE_BYTES ((BLOCK_THREADS / 32) * 32 * 16)

namespace mak {

using bf16 = __nv_bfloat16;

enum PhaseOp : int {
  OP_EMBED = 0,
  OP_GEMV = 1,
  OP_QKV_POST = 2,
  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
  OP_NORMRES = 7,    // gemma post-sublayer norm + residual add
  OP_PLEMIX = 8,     // gemma per-layer-input mix (norm ctx, add tok, scale)
  // batched-decode transform phases: write the transformed [B][K] input to
  // a global scratch so the following OP_GEMV_PLAIN needs no shared staging,
  // decoupling the batch size from the projection width.
  OP_GEMV_PLAIN = 9, // GEMV over a global [B][K] input, no fused transform
  OP_NORMB = 10,     // RMSNorm (llama or gemma) [B][K] -> scratch
  OP_GLUB = 11,      // SwiGLU / gelu-glu gating [B][2K] -> scratch [B][K]
  OP_ATTNFINB = 12,  // finalize split-S attention [B] -> scratch [B][qdim]
  OP_GEMV_Q4 = 13,   // GEMV with nf4-packed weights, dequantized in-kernel
};

// NF4 codebook (bitsandbytes): the 16 quantile levels a 4-bit index maps
// to. Dequantized weight = code[nibble] * per-block absmax, rounded to
// bf16; the lookup is exact and the multiply is IEEE, so it reproduces on
// every card.
__device__ __constant__ float NF4_CODE[16] = {
    -1.0f, -0.6961928009986877f, -0.5250730514526367f,
    -0.39491748809814453f, -0.28444138169288635f, -0.18477343022823334f,
    -0.09105003625154495f, 0.0f, 0.07958029955625534f, 0.16093020141124725f,
    0.24611230194568634f, 0.33791524171829224f, 0.44070982933044434f,
    0.5626170039176941f, 0.7229568362236023f, 1.0f};

enum InputTransform : int {
  IT_NONE = 0,
  IT_RMSNORM = 1,
  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
#define ITF_REDUCE 8      // input is fp32 k-slice partials: sum 4 slices
#define ITF_WRITEBACK 16  // write the reduced bf16 vector to p6

enum Epilogue : int {
  EP_STORE = 0,
  EP_RESID = 1,
  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
#define NSLICE 4

// Phase slots (int64 each):
//  0 op | 1 in | 2 w | 3 out | 4 a | 5 b | 6 c
//  7 N/max_seq | 8 K | 9 itrans/flags | 10 epi | 11 Hq | 12 Hkv | 13 D
// 14 f0 (fp32 bits in low 32) | 15 i0 / f1 (fp32 bits)

__device__ __forceinline__ float bits2f(long long v) {
  return __int_as_float((int)v);
}
__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));
  return r;
}

__device__ __forceinline__ float warp_sum(float v) {
  for (int o = 16; o > 0; o >>= 1) v += __shfl_down_sync(0xffffffffu, v, o);
  return v;
}
__device__ __forceinline__ float warp_max(float v) {
  for (int o = 16; o > 0; o >>= 1)
    v = fmaxf(v, __shfl_down_sync(0xffffffffu, v, o));
  return v;
}
__device__ __forceinline__ unsigned long long warp_max_u64(unsigned long long v) {
  for (int o = 16; o > 0; o >>= 1) {
    unsigned long long w = __shfl_down_sync(0xffffffffu, v, o);
    v = (w > v) ? w : v;
  }
  return v;
}

// Valid on thread 0 after return; safe to call repeatedly (trailing sync).
__device__ float block_sum(float v, float* s_red) {
  const int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
  v = warp_sum(v);
  if (lane == 0) s_red[w] = v;
  __syncthreads();
  const int nw = blockDim.x >> 5;
  v = (threadIdx.x < nw) ? s_red[threadIdx.x] : 0.f;
  if (w == 0) v = warp_sum(v);
  __syncthreads();
  return v;
}
__device__ float block_max(float v, float* s_red) {
  const int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
  v = warp_max(v);
  if (lane == 0) s_red[w] = v;
  __syncthreads();
  const int nw = blockDim.x >> 5;
  v = (threadIdx.x < nw) ? s_red[threadIdx.x] : -INFINITY;
  if (w == 0) v = warp_max(v);
  __syncthreads();
  return v;
}
__device__ unsigned long long block_max_u64(unsigned long long v,
                                            unsigned long long* s_red) {
  const int lane = threadIdx.x & 31, w = threadIdx.x >> 5;
  v = warp_max_u64(v);
  if (lane == 0) s_red[w] = v;
  __syncthreads();
  const int nw = blockDim.x >> 5;
  v = (threadIdx.x < nw) ? s_red[threadIdx.x] : 0ull;
  if (w == 0) v = warp_max_u64(v);
  __syncthreads();
  return v;
}

// ---------------------------------------------------------------------------
// Phase bodies. Each is grid-cooperative: work is striped over all blocks,
// and any block with nothing to do simply falls through to the barrier.
// ---------------------------------------------------------------------------

__device__ unsigned long long block_max_u64(unsigned long long v,
                                            unsigned long long* s_red);
__device__ __forceinline__ unsigned long long amax_key(float v, unsigned idx);

// Positions inside the prompt read from the staged prompt buffer. Beyond it,
// with argmax partials available (closed loop), every block reduces the
// previous step's per-block partials itself, which removes the separate
// argmax-combine phase from all but the final step; the derived token is
// also recorded to the output buffer. prompt_len == -1 selects the external
// token buffer (single-step decode with a caller-supplied token).
// Row position: prompt chunks use consecutive positions of one sequence;
// batch mode (pos_b != nullptr) gives every row its own sequence position.
__device__ __forceinline__ int pm_of(int pos, int m, const int* pos_b,
                                     int step) {
  return (pos_b != nullptr) ? (pos_b[m] + step) : (pos + m);
}

__device__ void ph_embed(const long long* P, int pos, int M, int prompt_len,
                         int step_slot, const int* pos_b) {
  const bf16* tab = reinterpret_cast<const bf16*>(P[2]);
  bf16* out = reinterpret_cast<bf16*>(P[3]);
  const unsigned long long* parts =
      reinterpret_cast<const unsigned long long*>(P[4]);
  const int* prompt = reinterpret_cast<const int*>(P[5]);
  const int* tok = reinterpret_cast<const int*>(P[6]);
  int* toks_out = reinterpret_cast<int*>(P[7]);
  const long long K = P[8];

  __shared__ int s_tok;
  __shared__ int s_tokb[8];
  const bool need_feedback = (pos + M - 1 >= prompt_len);
  if (pos_b != nullptr) {
    // batch decode: every row reads its own sequence's fed-back token
    if (threadIdx.x < M) s_tokb[threadIdx.x] = tok[threadIdx.x];
    __syncthreads();
  } else if (need_feedback) {
    if (prompt_len >= 0 && parts != nullptr) {
      unsigned long long best = amax_key(-INFINITY, 0xFFFFFFFFu);
      for (int i = threadIdx.x; i < (int)gridDim.x; i += blockDim.x) {
        const unsigned long long kk = parts[i];
        best = (kk > best) ? kk : best;
      }
      __shared__ unsigned long long s64[32];
      best = block_max_u64(best, s64);
      if (threadIdx.x == 0) {
        const int idx = (int)(0xFFFFFFFFu - (unsigned)(best & 0xFFFFFFFFull));
        s_tok = idx;
        const int prev_slot = step_slot - M;
        if (toks_out != nullptr && prev_slot >= 0 && blockIdx.x == 0)
          toks_out[prev_slot] = idx;
      }
    } else {
      if (threadIdx.x == 0) s_tok = tok[0];
    }
    __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 k = i % K;
    const int pm = pos + m;
    const long long t =
        (pos_b != nullptr)
            ? (long long)s_tokb[m]
            : ((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);
  }
}

__device__ __forceinline__ unsigned long long amax_key(float v, unsigned idx) {
  unsigned u = __float_as_uint(v);
  u = (u & 0x80000000u) ? ~u : (u | 0x80000000u);
  return ((unsigned long long)u << 32) |
         (unsigned long long)(0xFFFFFFFFu - 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, 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);
    acc = fmaf(pc[d], w, acc);
  }
  return bf2f(f2bf(acc / L));
}

__device__ __forceinline__ float mk_load_x(const bf16* x, const bf16* gamma,
                                           long long k, long long K, int itrans,
                                           float inv, const float* part,
                                           int D, int maxch, int nch) {
  if (itrans == IT_RMSNORM)
    // transformers RMSNorm: (x_f32 * invrms) rounds to bf16 first, then the
    // bf16 weight multiply computes in fp32 and rounds again.
    return bf2f(f2bf(bf2f(gamma[k]) * bf2f(f2bf(bf2f(x[k]) * inv))));
  if (itrans == IT_SWIGLU) {
    // silu(gate) rounded to bf16, then bf16 multiply by up, rounded again.
    const float g = bf2f(x[k]);
    const float sig = 1.f / (1.f + expf(-g));
    const float sg = bf2f(f2bf(g * sig));
    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]);
}

// One GEMV output element for the non-argmax epilogues.
__device__ __forceinline__ void gemv_store(const long long* P, int epi,
                                           long long N, const bf16* resid,
                                           int srcm, long long row, float a) {
  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])[(long long)srcm * N + 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);
}

// 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.
template <int NB>
__device__ void ph_gemv(const long long* P, int pos, int M, int rs,
                        const int* pos_b, int step) {
  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 bool tiled = ((P[9] >> 20) & 1) != 0;  // row-tile-of-8 weights
  const int epi = (int)P[10];
  const int mmode = ((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 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] & 0xFFFF);
  const int fin_win = (int)((P[15] >> 16) & 0xFFFF);  // 0 = unwindowed

  // mmode restricts to the chunk's last row (prompt LM head); in batch
  // mode every row is a live sequence needing logits.
  const int m0 = (mmode && pos_b == nullptr) ? (M - 1) : 0;
  const int Ma = M - m0;  // staged rows

  // dynamic shared: [cp.async weight ring (rs stages) | staged inputs xs]
  extern __shared__ unsigned char smem_raw[];
  int4* rng = reinterpret_cast<int4*>(smem_raw);
  bf16* xs = reinterpret_cast<bf16*>(
      smem_raw + (long long)rs * GEMV_RING_STAGE_BYTES);
  __shared__ float s_red[32];
  __shared__ float s_scale;
  __shared__ unsigned long long s64[32];
  __shared__ float s_fine[1024 + 128];  // attn-finalize combine weights

  const int lane = threadIdx.x & 31;
  const int wib = threadIdx.x >> 5;

  // Weight streaming starts before input staging: the ring touches only
  // global weight rows, so the first GEMV_NSTAGE tile groups fly while
  // the staging transform computes.
  const long long gwarp =
      ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  const long long nwarps = ((long long)gridDim.x * blockDim.x) >> 5;
  const int rmask = rs - 1;  // rs is a power of two
  int4* ringw = rng + (wib * rs) * 32 + lane;
  const long long kbase = (long long)lane * 8;
  const long long ntiles = (kbase < K) ? (K - kbase + 255) / 256 : 0;
  long long irow = gwarp, itile = 0;  // next (row, tile) to issue
  long long ig = 0, cg = 0;           // groups issued / consumed
  // With tiled weights ([N/8][K/8][8][8]) a block's eight warps hold
  // eight consecutive rows, so their simultaneous 16-byte loads form one
  // contiguous 4KB stream per block (HBM row-buffer locality); the plain
  // layout keeps per-warp sequential streams, which GDDR parts prefer.
  // Same bytes per lane in the same order either way; addresses only.
  const long long K8 = K >> 3;
  auto issue_one = [&]() {
    if (irow < N && ntiles > 0) {
      const bf16* src =
          tiled ? W + (((irow >> 3) * K8 + (kbase >> 3) + itile * 32) << 6) +
                      ((irow & 7) << 3)
                : W + irow * K + kbase + itile * 256;
      __pipeline_memcpy_async(ringw + ((int)ig & rmask) * 32, src, 16);
      __pipeline_commit();
      ++ig;
      if (++itile == ntiles) {
        itile = 0;
        irow += nwarps;
      }
    }
  };
  for (int i = 0; i < rs; ++i) issue_one();

  // Stage transformed inputs for all chunk rows: [Ma][K] in shared, bf16
  // (every transform value is bf16-rounded by semantics, so staging is
  // exact). Staging reads are 16-byte vectorized: scalar bf16 loads make
  // the pre-GEMV serialization a chain of tens of dependent global reads
  // per thread, which costs more than the GEMV itself on short phases.
  // RMSNORM lands the raw row in shared on the first pass and transforms
  // it in place, so global x is read once instead of twice.
  {
    for (int m = 0; m < Ma; ++m) {
      const int srcm = m0 + m;
      const long long G = K / 8;  // 16-byte groups (K % 8 == 0 by contract)
      if (itrans == IT_RMSNORM) {
        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(gp[j]) *
                           bf2f(f2bf(__low2float(xp[j]) * inv)));
            xp[j].y = f2bf(__high2float(gp[j]) *
                           bf2f(f2bf(__high2float(xp[j]) * inv)));
          }
          *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) {
          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 sig = 1.f / (1.f + expf(-gv));
              const float sg = bf2f(f2bf(gv * sig));
              op[j].x = f2bf(sg * __low2float(up[j]));
            }
            {
              const float gv = __high2float(gp[j]);
              const float sig = 1.f / (1.f + expf(-gv));
              const float sg = bf2f(f2bf(gv * sig));
              op[j].y = f2bf(sg * __high2float(up[j]));
            }
          }
          *reinterpret_cast<int4*>(xs + (long long)m * K + g * 8) = o4;
        }
      } else if (itrans == IT_ATTNFIN) {
        const int pm = pm_of(pos, srcm, pos_b, step);
        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];
        if (Hq_ * nch <= 1024 && Hq_ <= 128) {
          // Per-(head, chunk) softmax-combine weights computed once per
          // block: identical expf inputs and identical per-element fma
          // chains as attn_finalize, so the staged values are
          // bit-identical while the exp count drops from K*nch to Hq*nch.
          float* ew = s_fine;
          float* lw = s_fine + 1024;
          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;
              L = fmaf(pc[D + 1], w, L);
            }
            lw[h] = L;
          }
          __syncthreads();
          if (D % 4 == 0) {
            // four elements per thread with 8-byte partial loads: the
            // scalar walk costs a 32-byte sector per 4-byte read; the
            // per-element combine chains are unchanged.
            for (long long g = threadIdx.x; g < K / 4; g += blockDim.x) {
              const long long k = g * 4;
              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 float w = ew[h * nch + c];
                a0 = fmaf(pa.x, w, a0);
                a1 = fmaf(pa.y, w, a1);
                a2 = fmaf(pb.x, w, a2);
                a3 = fmaf(pb.y, w, a3);
              }
              const float L = lw[h];
              bf16* xp = xs + (long long)m * K + k;
              xp[0] = f2bf(bf2f(f2bf(a0 / L)));
              xp[1] = f2bf(bf2f(f2bf(a1 / L)));
              xp[2] = f2bf(bf2f(f2bf(a2 / L)));
              xp[3] = f2bf(bf2f(f2bf(a3 / L)));
            }
          } else {
            for (long long k = threadIdx.x; k < K; k += blockDim.x) {
              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])));
            }
          }
          __syncthreads();
        } 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;
        for (long long g = threadIdx.x; g < G; g += blockDim.x)
          *reinterpret_cast<int4*>(xs + (long long)m * K + g * 8) =
              *reinterpret_cast<const int4*>(xr + g * 8);
      }
    }
    __syncthreads();
  }

  // Warp-per-row GEMV with weight tiles streamed through the per-warp
  // cp.async ring primed above. A register prefetch pipeline serializes
  // on write-after-read hazards (a buffer reload cannot issue until the
  // previous tile's fma chain drains), which caps effective memory depth
  // near one on long-latency parts; the ring keeps rs tile groups in
  // flight per lane with no register dependence, and the (row, tile)
  // stream is continuous across rows so the pipeline never drains inside
  // a phase. Each lane consumes its own tiles in order, so the fma chain
  // and reduction tree are exactly those of the register schedule:
  // results are bit-identical.
  unsigned long long best = amax_key(-INFINITY, 0xFFFFFFFFu);
  {
    for (long long row = gwarp; row < N; row += nwarps) {
      float acc[NB];
#pragma unroll
      for (int m = 0; m < NB; ++m) acc[m] = 0.f;
      auto consume = [&](const int4& w4, long long t) {
        const long long k = kbase + t * 256;
        const __nv_bfloat162* wp =
            reinterpret_cast<const __nv_bfloat162*>(&w4);
#pragma unroll
        for (int m = 0; m < NB; ++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]);
            }
          }
        }
      };
      for (long long t = 0; t < ntiles; ++t) {
        __pipeline_wait_prior((int)(ig - cg - 1));
        const int4 w4 = ringw[((int)cg & rmask) * 32];
        ++cg;
        issue_one();
        consume(w4, t);
      }
#pragma unroll
      for (int m = 0; m < NB; ++m) {
        if (m < Ma) {
          const float a = warp_sum(acc[m]);
          if (lane == 0) {
            const int srcm = m0 + m;
            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 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)
      reinterpret_cast<unsigned long long*>(P[6])[blockIdx.x] = best;
  }
}

// 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,
                                    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 (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] = (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
  if (qk_flag & 1) __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 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();
}

// 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 inv_freq table ptr, 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 int* pos_b, int step, long long kv_bstride) {
  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 D = (int)P[13];
  // chunk length rides in slot 12 bits 16+ (0 = compile-time default);
  // shared staging is sized for the maximum, so CH <= ATTN_CHUNK.
  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 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 sc[ATTN_CHUNK];
  __shared__ float pv[BLOCK_THREADS];
  __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.
  int pm_max = pos + M - 1;
  if (pos_b != nullptr) {
    int mx = pos_b[0];
    for (int m = 1; m < M; ++m) mx = (pos_b[m] > mx) ? pos_b[m] : mx;
    pm_max = mx + step;
  }
  const int nch_max = (pm_max + 1 + CH - 1) / CH;
  {
    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 = pm_of(pos, m, pos_b, step);
      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;
      float* partm = reinterpret_cast<float*>(P[4]) +
                     (long long)m * ((long long)Hq * maxch * (D + 2));
      const int kvh = head / rep;
      const long long kvb = (pos_b != nullptr) ? (long long)m * kv_bstride : 0;
      bf16* kc = reinterpret_cast<bf16*>(P[2]) + kvb +
                 (long long)kvh * max_seq * D;
      bf16* vc = reinterpret_cast<bf16*>(P[3]) + kvb +
                 (long long)kvh * max_seq * D;
      float* part = partm + ((long long)head * maxch + ch) * (D + 2);
      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 || pos_b != nullptr) && 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
      // as a chain of serial global K-row reads per warp. cp.async copies
      // overlap the q staging below; values are unchanged, so both paths
      // compute identical bits.
      extern __shared__ unsigned char smem_raw[];
      bf16* ks_sh = reinterpret_cast<bf16*>(smem_raw);
      bf16* vs_sh = ks_sh + (long long)CH * D;
      const bool kv_staged =
          (D % 8 == 0) &&
          (2u * (unsigned)CH * (unsigned)D * 2u <= dyn_smem_size());
      if (kv_staged) {
        const int ncpy = cn * D / 8;
        for (int i = threadIdx.x; i < ncpy; i += blockDim.x)
          __pipeline_memcpy_async(
              ks_sh + (long long)i * 8,
              kc + (long long)s0 * D + (long long)i * 8, 16);
        for (int i = threadIdx.x; i < ncpy; i += blockDim.x)
          __pipeline_memcpy_async(
              vs_sh + (long long)i * 8,
              vc + (long long)s0 * D + (long long)i * 8, 16);
        __pipeline_commit();
      }

      // 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] = (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]);
          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))));
        }
        if (kv_staged) __pipeline_wait_prior(0);
        __syncthreads();
      }

      // score this chunk: eager rounds QK^T to bf16, scales as a bf16 op,
      // upcasts to fp32 for the softmax
      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;
        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)
        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;
      }
      lsum = block_sum(lsum, s_red);
      __syncthreads();

      // PV: long chunks split positions across NG thread groups whose fp32
      // partials combine in group order — the tree depends only on cn and
      // D, so it is identical on every device and in every launch mode.
      const int NG = (BLOCK_THREADS / D < 8) ? BLOCK_THREADS / D : 8;
      if (NG > 1 && cn >= 64) {
        const int g = threadIdx.x / D;
        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],
                       bf2f(kv_staged ? vs_sh[(long long)si * D + d]
                                      : vc[(long long)(s0 + si) * D + d]),
                       acc);
          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;
        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]) & 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]);

  __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, 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, int M) {
  const float* lg = reinterpret_cast<const float*>(P[1]);
  unsigned long long* parts = reinterpret_cast<unsigned long long*>(P[3]);
  const long long V = P[7];
  __shared__ unsigned long long s64[32];
  for (int m = 0; m < M; ++m) {
    unsigned long long best = amax_key(-INFINITY, 0xFFFFFFFFu);
    for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
         i < V; i += (long long)gridDim.x * blockDim.x) {
      // Greedy parity with transformers: argmax over the bf16-rounded
      // logits, lowest index on ties (torch.argmax convention).
      const unsigned long long k =
          amax_key(bf2f(f2bf(lg[(long long)m * V + i])), (unsigned)i);
      best = (k > best) ? k : best;
    }
    best = block_max_u64(best, s64);
    if (threadIdx.x == 0)
      parts[(long long)m * gridDim.x + blockIdx.x] = best;
    if (m + 1 < M) __syncthreads();
  }
}

// Row m finalizes on block m; slot 8 carries the per-row output stride
// (0 for the single-sequence programs, whose only row is m = 0).
__device__ void ph_argmax_fin(const long long* P, int step_slot, int M) {
  if ((int)blockIdx.x >= M) return;
  const int m = (int)blockIdx.x;
  const unsigned long long* parts =
      reinterpret_cast<const unsigned long long*>(P[1]) +
      (long long)m * gridDim.x;
  int* tok = reinterpret_cast<int*>(P[3]);
  int* toks_out = reinterpret_cast<int*>(P[4]);
  const long long tstride = P[8];
  unsigned long long best = amax_key(-INFINITY, 0xFFFFFFFFu);
  for (int i = threadIdx.x; i < gridDim.x; i += blockDim.x) {
    const unsigned long long k = parts[i];
    best = (k > best) ? k : best;
  }
  __shared__ unsigned long long s64[32];
  best = block_max_u64(best, s64);
  if (threadIdx.x == 0) {
    const int idx = (int)(0xFFFFFFFFu - (unsigned)(best & 0xFFFFFFFFull));
    tok[m] = idx;
    if (toks_out != nullptr && step_slot >= 0)
      toks_out[(long long)m * tstride + step_slot] = idx;
  }
}

// ---- batched-decode transform phases: transformed [B][K] input -> scratch
// The reductions and per-element op chains reproduce ph_gemv's fused
// staging exactly, so the scratch bytes equal the shared-staged bytes and
// the following OP_GEMV_PLAIN is bit-identical to a fused OP_GEMV. This is
// what makes a batched sequence's logits identical to its batch-1 run.

// RMSNorm one row per block (llama IT_RMSNORM or gemma IT_RMSNORM_G).
__device__ void ph_normb(const long long* P, int M) {
  const int m = blockIdx.x;
  if (m >= M) return;
  const long long K = P[8];
  const bf16* xin = reinterpret_cast<const bf16*>(P[1]) + (long long)m * K;
  const bf16* gamma = reinterpret_cast<const bf16*>(P[2]);
  bf16* xg = reinterpret_cast<bf16*>(P[3]) + (long long)m * K;
  const int variant = (int)P[9] & ITF_BASE_MASK;
  const float eps = bits2f(P[14]);
  __shared__ float s_red[32];
  __shared__ float s_scale;
  const long long G = K / 8;  // K % 8 == 0 by contract
  float ss = 0.f;
  for (long long g = threadIdx.x; g < G; g += blockDim.x) {
    const int4 x4 = *reinterpret_cast<const int4*>(xin + g * 8);
    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) {
    const int4 x4 = *reinterpret_cast<const int4*>(xin + g * 8);
    const int4 g4 = *reinterpret_cast<const int4*>(gamma + g * 8);
    const __nv_bfloat162* xp = reinterpret_cast<const __nv_bfloat162*>(&x4);
    const __nv_bfloat162* gp = reinterpret_cast<const __nv_bfloat162*>(&g4);
    int4 o4;
    __nv_bfloat162* op = reinterpret_cast<__nv_bfloat162*>(&o4);
#pragma unroll
    for (int j = 0; j < 4; ++j) {
      if (variant == IT_RMSNORM_G) {
        op[j].x = f2bf((__low2float(xp[j]) * inv) * __low2float(gp[j]));
        op[j].y = f2bf((__high2float(xp[j]) * inv) * __high2float(gp[j]));
      } else {
        op[j].x = f2bf(__low2float(gp[j]) *
                       bf2f(f2bf(__low2float(xp[j]) * inv)));
        op[j].y = f2bf(__high2float(gp[j]) *
                       bf2f(f2bf(__high2float(xp[j]) * inv)));
      }
    }
    *reinterpret_cast<int4*>(xg + g * 8) = o4;
  }
}

// SwiGLU / gelu-glu gating: xg[m][k] = act(gate[m][k]) * up[m][k].
__device__ void ph_glub(const long long* P, int M) {
  const long long K = P[8];
  const bf16* xin = reinterpret_cast<const bf16*>(P[1]);  // [M][2K]
  bf16* xg = reinterpret_cast<bf16*>(P[3]);               // [M][K]
  const int variant = (int)P[9] & ITF_BASE_MASK;
  const long long tot = (long long)M * K;
  for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
       i < tot; i += (long long)gridDim.x * blockDim.x) {
    const int m = (int)(i / K);
    const long long k = i % K;
    const bf16* gate = xin + (long long)m * 2 * K;
    const float g = bf2f(gate[k]);
    const float up = bf2f(gate[K + k]);
    if (variant == IT_GELU_GLU) {
      const float act = bf2f(f2bf(gelu_tanh_f(g)));
      xg[i] = f2bf(act * up);
    } else {
      const float sig = 1.f / (1.f + expf(-g));
      const float sg = bf2f(f2bf(g * sig));
      xg[i] = f2bf(sg * up);
    }
  }
}

// Finalize split-S attention for the O-projection input.
__device__ void ph_attnfinb(const long long* P, int M, const int* pos_b,
                            int step) {
  const long long K = P[8];  // qdim = Hq * D
  const float* part = reinterpret_cast<const float*>(P[1]);
  bf16* xg = reinterpret_cast<bf16*>(P[3]);
  const int Hq = (int)P[11];
  const int D = (int)P[13];
  const int chbits = (int)(P[12] >> 16) & 0xFFFF;
  const int CH = (chbits > 0 && chbits <= ATTN_CHUNK) ? chbits : ATTN_CHUNK;
  const int maxch = (int)(P[15] & 0xFFFF);
  const int fin_win = (int)((P[15] >> 16) & 0xFFFF);
  const long long tot = (long long)M * K;
  for (long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
       i < tot; i += (long long)gridDim.x * blockDim.x) {
    const int m = (int)(i / K);
    const long long k = i % K;
    const int pm = pm_of(0, m, pos_b, step);
    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)m * ((long long)Hq * maxch * (D + 2));
    xg[i] = f2bf(attn_finalize(mpart, (int)(k / D), (int)(k % D), D, maxch,
                               nch, nch_lo));
  }
}

// GEMV over a global [M][K] input already transformed into scratch: no
// shared staging, so the batch width M is bounded only by the NB register
// accumulators, not by shared memory. The weight ring, per-lane tile order,
// and warp reduction are exactly ph_gemv's, so results are bit-identical.
template <int NB>
__device__ void ph_gemv_plain(const long long* P, int M, int rs) {
  const bf16* W = reinterpret_cast<const bf16*>(P[2]);
  const bf16* resid = reinterpret_cast<const bf16*>(P[5]);
  const bf16* xin = reinterpret_cast<const bf16*>(P[1]);
  const long long N = P[7], K = P[8];
  const bool tiled = ((P[9] >> 20) & 1) != 0;
  const int epi = (int)P[10];

  extern __shared__ unsigned char smem_raw[];
  int4* rng = reinterpret_cast<int4*>(smem_raw);
  const int lane = threadIdx.x & 31, wib = threadIdx.x >> 5;
  const long long gwarp =
      ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  const long long nwarps = ((long long)gridDim.x * blockDim.x) >> 5;
  const int rmask = rs - 1;
  int4* ringw = rng + (wib * rs) * 32 + lane;
  const long long kbase = (long long)lane * 8;
  const long long ntiles = (kbase < K) ? (K - kbase + 255) / 256 : 0;
  const long long K8 = K >> 3;
  long long irow = gwarp, itile = 0, ig = 0, cg = 0;
  auto issue_one = [&]() {
    if (irow < N && ntiles > 0) {
      const bf16* src =
          tiled ? W + (((irow >> 3) * K8 + (kbase >> 3) + itile * 32) << 6) +
                      ((irow & 7) << 3)
                : W + irow * K + kbase + itile * 256;
      __pipeline_memcpy_async(ringw + ((int)ig & rmask) * 32, src, 16);
      __pipeline_commit();
      ++ig;
      if (++itile == ntiles) {
        itile = 0;
        irow += nwarps;
      }
    }
  };
  for (int i = 0; i < rs; ++i) issue_one();

  for (long long row = gwarp; row < N; row += nwarps) {
    float acc[NB];
#pragma unroll
    for (int m = 0; m < NB; ++m) acc[m] = 0.f;
    for (long long t = 0; t < ntiles; ++t) {
      __pipeline_wait_prior((int)(ig - cg - 1));
      const int4 w4 = ringw[((int)cg & rmask) * 32];
      ++cg;
      issue_one();
      const long long k = kbase + t * 256;
      const __nv_bfloat162* wp = reinterpret_cast<const __nv_bfloat162*>(&w4);
#pragma unroll
      for (int m = 0; m < NB; ++m) {
        if (m < M) {
          const __nv_bfloat162* xp =
              reinterpret_cast<const __nv_bfloat162*>(xin +
                                                      (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]);
          }
        }
      }
    }
#pragma unroll
    for (int m = 0; m < NB; ++m) {
      if (m < M) {
        const float a = warp_sum(acc[m]);
        if (lane == 0) gemv_store(P, epi, N, resid, m, row, a);
      }
    }
  }
}

// GEMV with nf4-packed weights [N][K/2] (two 4-bit indices per byte, row
// major) and per-64-element fp32 absmax. Each lane dequantizes its own
// weight tile (code[nibble] * absmax, rounded to bf16) and accumulates in
// the exact order of the dense path, so the result is bit-identical to a
// dense GEMV over the same dequantized weights and reproduces on every
// card. Input is dense bf16 in a global buffer (a transform phase writes
// it). Weights stay packed, so the quantized memory footprint is kept.
template <int NB>
__device__ void ph_gemv_q4(const long long* P, int M) {
  const unsigned char* W = reinterpret_cast<const unsigned char*>(P[2]);
  const float* absmax = reinterpret_cast<const float*>(P[6]);
  const bf16* resid = reinterpret_cast<const bf16*>(P[5]);
  const bf16* xin = reinterpret_cast<const bf16*>(P[1]);
  const long long N = P[7], K = P[8];
  const int epi = (int)P[10];
  const int lane = threadIdx.x & 31;
  const long long gwarp =
      ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
  const long long nwarps = ((long long)gridDim.x * blockDim.x) >> 5;
  const long long kbase = (long long)lane * 8;
  const long long ntiles = (kbase < K) ? (K - kbase + 255) / 256 : 0;
  const long long Kh = K >> 1;  // packed bytes per row

  for (long long row = gwarp; row < N; row += nwarps) {
    float acc[NB];
#pragma unroll
    for (int m = 0; m < NB; ++m) acc[m] = 0.f;
    for (long long t = 0; t < ntiles; ++t) {
      const long long k = kbase + t * 256;
      const unsigned int pw =
          *reinterpret_cast<const unsigned int*>(W + row * Kh + (k >> 1));
      const float am = absmax[(row * K + k) >> 6];
      float w8[8];
#pragma unroll
      for (int b = 0; b < 4; ++b) {
        const unsigned int by = (pw >> (b * 8)) & 0xFFu;
        // bitsandbytes packs the earlier element in the high nibble
        w8[2 * b] = bf2f(f2bf(NF4_CODE[by >> 4] * am));
        w8[2 * b + 1] = bf2f(f2bf(NF4_CODE[by & 0xF] * am));
      }
#pragma unroll
      for (int m = 0; m < NB; ++m) {
        if (m < M) {
          const bf16* xp = xin + (long long)m * K + k;
#pragma unroll
          for (int j = 0; j < 8; ++j)
            acc[m] = fmaf(w8[j], bf2f(xp[j]), acc[m]);
        }
      }
    }
#pragma unroll
    for (int m = 0; m < NB; ++m) {
      if (m < M) {
        const float a = warp_sum(acc[m]);
        if (lane == 0) gemv_store(P, epi, N, resid, m, row, a);
      }
    }
  }
}

template <int NB>
__device__ __forceinline__ void dispatch_phase(const long long* P, int pos,
                                               int M, int step_slot,
                                               int prompt_len,
                                               int ring_stages,
                                               const int* pos_b, int step,
                                               long long kv_bstride) {
  switch ((int)P[0]) {
    case OP_EMBED: ph_embed(P, pos, M, prompt_len, step_slot, pos_b); break;
    case OP_GEMV: ph_gemv<NB>(P, pos, M, ring_stages, pos_b, step); break;
    case OP_ATTN: ph_attn(P, pos, M, pos_b, step, kv_bstride); break;
    case OP_ARGMAX_PART: ph_argmax_part(P, M); break;
    case OP_ARGMAX_FIN: ph_argmax_fin(P, step_slot, M); 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;
    case OP_GEMV_PLAIN: ph_gemv_plain<NB>(P, M, ring_stages); break;
    case OP_GEMV_Q4: ph_gemv_q4<NB>(P, M); break;
    case OP_NORMB: ph_normb(P, M); break;
    case OP_GLUB: ph_glub(P, M); break;
    case OP_ATTNFINB: ph_attnfinb(P, M, pos_b, step); break;
    default: break;
  }
}

// Sense-reversing grid barrier over co-resident blocks. The launcher sizes
// the grid from the occupancy API, so every block is resident and the spin
// cannot deadlock. The count returns to zero at each crossing and the sense
// word persists across launches; blocks read it at kernel entry.
// Barrier state layout (int32): [0..31] striped arrival counters (one
// cacheline apart via 32-int stride would be ideal; adjacent ints already
// split traffic 32 ways), [32] level-2 counter, [33] sense word.
#define BAR_STRIPES 32

__device__ void grid_barrier(int* bar, int* local_sense) {
  __syncthreads();
  if (threadIdx.x == 0) {
    const int ls = *local_sense ^ 1;
    *local_sense = ls;
    __threadfence();
    // Two-level arrival: hundreds of serialized RMWs on one line cost ~2us;
    // striping arrivals across 32 counters with a 32-wide second level cuts
    // the serialization to tens of nanoseconds per line.
    const int nb = (int)gridDim.x;
    const int stripe = (int)blockIdx.x & (BAR_STRIPES - 1);
    const int in_stripe =
        nb / BAR_STRIPES + ((stripe < (nb % BAR_STRIPES)) ? 1 : 0);
    int* sense = bar + BAR_STRIPES + 1;
    if (atomicAdd(bar + stripe, 1) == in_stripe - 1) {
      atomicExch(bar + stripe, 0);
      __threadfence();
      const int nstripes = (nb < BAR_STRIPES) ? nb : BAR_STRIPES;
      if (atomicAdd(bar + BAR_STRIPES, 1) == nstripes - 1) {
        atomicExch(bar + BAR_STRIPES, 0);
        __threadfence();
        atomicExch(sense, ls);
      } else {
        const volatile int* vs = sense;
        while (*vs != ls) {
        }
      }
    } else {
      const volatile int* vs = sense;
        while (*vs != ls) {
        }
    }
    __threadfence();
  }
  __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.
// Two blocks per SM everywhere: the phase interpreter's GEMV needs every
// co-resident warp on high-bandwidth parts (H200 at one block per SM
// measures 1.5x slower), so the register budget is pinned to 128. The
// cp.async ring keeps the hot loop's register frame small enough that the
// bound does not force spills.
template <int NB>
__global__ void __launch_bounds__(BLOCK_THREADS, 2)
    mega_kernel(const long long* __restrict__ prog, int n_phases,
                int pos0, int total, int slot0, int prompt_len,
                int chunk_m, int ring_stages, int* bar,
                long long* __restrict__ ts,
                const int* __restrict__ pos_b, int B,
                long long kv_bstride) {
  __shared__ int s_sense;
  if (threadIdx.x == 0) s_sense = atomicAdd(bar + BAR_STRIPES + 1, 0);
  __syncthreads();
  // diagnostic: block 0 stamps its SM clock at entry and after each
  // barrier crossing of the first step (pure observation, no data effect)
  if (ts != nullptr && blockIdx.x == 0 && threadIdx.x == 0)
    ts[0] = (long long)clock64();
  const bool batch = (pos_b != nullptr);
  int pos = pos0;
  const int pos_end = pos0 + total;
  while (pos < pos_end) {
    int M = 1;
    if (batch) {
      M = B;  // one decode step for every sequence; positions per row
    } else 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 step = pos - pos0;
    const int slot = batch ? (slot0 + step) : (slot0 + (pos + M - 1 - pos0));
    const bool final_step = batch ? (pos + 1 >= pos_end)
                                  : (pos + M >= pos_end);
    for (int p = 0; p < n_phases; ++p) {
      const long long* P = prog + (long long)p * 16;
      const int op = (int)P[0];
      if (op == OP_KV_APPEND && (M == 1 || batch)) continue;  // uniform skip
      if (op == OP_ARGMAX_FIN && !final_step && !batch)
        continue;  // embed derives it; batch feeds tokens through it
      dispatch_phase<NB>(P, pos, M, slot, batch ? -1 : prompt_len,
                         ring_stages, pos_b, step, kv_bstride);
      if (!(final_step && p + 1 == n_phases)) {
        // Warm the next GEMV's first weight tiles into L2 while the grid
        // barrier settles. Weights are constant, so the prefetch is safe
        // regardless of phase state, and it has no memory semantics.
        for (int q = p + 1; q < n_phases; ++q) {
          const long long* Pn = prog + (long long)q * 16;
          const int opn = (int)Pn[0];
          if (opn == OP_KV_APPEND && (M == 1 || batch)) continue;
          if (opn == OP_GEMV || opn == OP_GEMV_PLAIN) {
            const long long w =
                ((long long)blockIdx.x * blockDim.x + threadIdx.x) >> 5;
            const long long kb = (long long)(threadIdx.x & 31) * 8;
            if (w < Pn[7] && kb < Pn[8]) {
              const long long off =
                  ((Pn[9] >> 20) & 1)
                      ? (((w >> 3) * (Pn[8] >> 3) + (kb >> 3)) << 6) +
                            ((w & 7) << 3)
                      : w * Pn[8] + kb;
              const char* wr =
                  reinterpret_cast<const char*>(Pn[2]) + off * 2;
              asm volatile("prefetch.global.L2 [%0];" ::"l"(wr));
              asm volatile("prefetch.global.L2 [%0];" ::"l"(wr + 512));
            }
          }
          break;
        }
        grid_barrier(bar, &s_sense);
      }
      if (ts != nullptr && pos == pos0 && blockIdx.x == 0 &&
          threadIdx.x == 0)
        ts[p + 1] = (long long)clock64();
    }
    pos += batch ? 1 : M;
  }
}

__global__ void __launch_bounds__(BLOCK_THREADS, 2)
    phase_kernel(const long long* __restrict__ prog, int p,
                 int pos, int step_slot, int prompt_len,
                 int ring_stages) {
  const long long* P = prog + (long long)p * 16;
  if ((int)P[0] == OP_KV_APPEND) return;
  dispatch_phase<8>(P, pos, 1, step_slot, -1, ring_stages, nullptr, 0, 0);
}

// Batch width of the global-scratch decode kernel: acc[MAK_MAXB_BIG] per
// lane. 16 keeps the register frame under the launch_bounds(256, 2) budget.
#define MAK_MAXB_BIG 16

}  // namespace mak

#include <stdlib.h>

// cp.async ring depth for a launch: as many power-of-two stages as fit the
// device's shared-memory budget alongside the staging buffer while keeping
// two blocks per SM (high-bandwidth parts take the maximum of 16; a stage
// is 4KB). Pure scheduling: any value computes identical bits.
int mak_ring_stages_impl(int stage_bytes) {
  static int cached_dev = -1;
  static int cached_stage = -1;
  static int cached_rs = -1;
  int dev = 0;
  cudaGetDevice(&dev);
  if (dev != cached_dev || stage_bytes != cached_stage || cached_rs < 0) {
    int shm_sm = 0, optin = 0;
    cudaDeviceGetAttribute(&shm_sm, cudaDevAttrMaxSharedMemoryPerMultiprocessor,
                           dev);
    cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin,
                           dev);
    cudaFuncAttributes fa{};
    cudaFuncGetAttributes(&fa, mak::mega_kernel<8>);
    const int stat = (int)fa.sharedSizeBytes;
    long long by_occ = ((long long)shm_sm / 2 - stat - stage_bytes) /
                       GEMV_RING_STAGE_BYTES;
    long long by_win = ((long long)optin - stat - stage_bytes) /
                       GEMV_RING_STAGE_BYTES;
    if (by_occ < 4) by_occ = 4;  // never starve the ring for occupancy
    long long rs = (by_occ < by_win) ? by_occ : by_win;
    if (rs > 4) rs = 4;  // deeper rings measure flat-to-worse on every card
    if (rs < 1) rs = 1;
    while (rs & (rs - 1)) rs &= rs - 1;  // power-of-two floor
    const char* env = getenv("MAK_RING");
    if (env != nullptr) {
      long long r = atoll(env);
      if (r >= 1 && r <= 16 && !(r & (r - 1)) &&
          (long long)stat + stage_bytes + r * GEMV_RING_STAGE_BYTES <= optin)
        rs = r;
    }
    cached_rs = (int)rs;
    cached_dev = dev;
    cached_stage = stage_bytes;
  }
  return cached_rs;
}

int mak_grid_blocks_impl(int smem_bytes) {
  static int cached_dev = -1;
  static int cached_smem = -1;
  static int cached_grid = -1;
  int dev = 0;
  cudaGetDevice(&dev);
  if (dev != cached_dev || smem_bytes != cached_smem || cached_grid < 0) {
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, dev);
    // staging + weight ring can exceed the default 48KB dynamic window;
    // the settable maximum is the opt-in limit minus the kernel's static
    // shared allocation
    int optin = 0;
    cudaDeviceGetAttribute(&optin, cudaDevAttrMaxSharedMemoryPerBlockOptin,
                           dev);
    cudaFuncAttributes fa{};
    cudaFuncGetAttributes(&fa, mak::mega_kernel<8>);
    cudaFuncSetAttribute(mak::mega_kernel<8>,
                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                         optin - (int)fa.sharedSizeBytes);
    cudaFuncGetAttributes(&fa, mak::mega_kernel<MAK_MAXB_BIG>);
    cudaFuncSetAttribute(mak::mega_kernel<MAK_MAXB_BIG>,
                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                         optin - (int)fa.sharedSizeBytes);
    cudaFuncGetAttributes(&fa, mak::phase_kernel);
    cudaFuncSetAttribute(mak::phase_kernel,
                         cudaFuncAttributeMaxDynamicSharedMemorySize,
                         optin - (int)fa.sharedSizeBytes);
    cudaGetLastError();  // never leave a failed attribute set sticky
    int nb = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&nb, mak::mega_kernel<8>,
                                                  BLOCK_THREADS, smem_bytes);
    if (nb < 1) nb = 1;
    int grid = nb * prop.multiProcessorCount;
    const char* env = getenv("MAK_GRID");
    if (env != nullptr) {
      const int g = atoi(env);
      if (g > 0 && g <= grid) grid = g;
    }
    cached_grid = grid;
    cached_dev = dev;
    cached_smem = smem_bytes;
  }
  return cached_grid;
}

// Grid for the global-scratch batch kernel, sized from its own occupancy
// (acc[MAK_MAXB_BIG] raises the register frame, which can lower blocks per
// SM; the grid must match so every block stays co-resident for the barrier).
int mak_grid_blocks_big_impl(int smem_bytes) {
  static int cached_dev = -1;
  static int cached_smem = -1;
  static int cached_grid = -1;
  int dev = 0;
  cudaGetDevice(&dev);
  if (dev != cached_dev || smem_bytes != cached_smem || cached_grid < 0) {
    cudaDeviceProp prop;
    cudaGetDeviceProperties(&prop, dev);
    int nb = 0;
    cudaOccupancyMaxActiveBlocksPerMultiprocessor(
        &nb, mak::mega_kernel<MAK_MAXB_BIG>, BLOCK_THREADS, smem_bytes);
    if (nb < 1) nb = 1;
    int grid = nb * prop.multiProcessorCount;
    const char* env = getenv("MAK_GRID");
    if (env != nullptr) {
      const int g = atoi(env);
      if (g > 0 && g <= grid) grid = g;
    }
    cached_grid = grid;
    cached_dev = dev;
    cached_smem = smem_bytes;
  }
  return cached_grid;
}

int mak_batch_maxb_impl() { return MAK_MAXB_BIG; }

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 ring_stages, int* bar, int grid,
                            int smem_bytes, cudaStream_t stream,
                            long long* ts) {
  mak::mega_kernel<8><<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
      prog, n_phases, pos0, total, slot0, prompt_len, chunk_m, ring_stages,
      bar, ts, nullptr, 1, 0);
  return cudaGetLastError();
}

cudaError_t mak_launch_mega_batch(const long long* prog, int n_phases,
                                  int steps, int slot0, int ring_stages,
                                  int* bar, int grid, int smem_bytes,
                                  const int* pos_b, int B,
                                  long long kv_bstride,
                                  cudaStream_t stream) {
  mak::mega_kernel<8><<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
      prog, n_phases, 0, steps, slot0, -1, 1, ring_stages, bar, nullptr,
      pos_b, B, kv_bstride);
  return cudaGetLastError();
}

cudaError_t mak_launch_mega_batch_big(const long long* prog, int n_phases,
                                      int steps, int slot0, int ring_stages,
                                      int* bar, int grid, int smem_bytes,
                                      const int* pos_b, int B,
                                      long long kv_bstride,
                                      cudaStream_t stream) {
  mak::mega_kernel<MAK_MAXB_BIG><<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
      prog, n_phases, 0, steps, slot0, -1, 1, ring_stages, bar, nullptr,
      pos_b, B, kv_bstride);
  return cudaGetLastError();
}

cudaError_t mak_launch_phase(const long long* prog, int phase, int pos,
                             int step_slot, int ring_stages, int grid,
                             int smem_bytes, cudaStream_t stream) {
  mak::phase_kernel<<<grid, BLOCK_THREADS, smem_bytes, stream>>>(
      prog, phase, pos, step_slot, 0, ring_stages);
  return cudaGetLastError();
}