cuda_ops.cu

// vllm.cpp original (vt runtime, inventory deviation §9.1); no upstream mirror.
// CUDA baseline kernels for rmsnorm / silu_and_mul / embedding / rope_neox.
// Correctness-grade (M0.6): plain grid-stride / one-block-per-row kernels, f32
// accumulation, double-precision RoPE angles matching the CPU reference.
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include <math_constants.h>

#include <cub/cub.cuh>

#include <cstdint>
#include <cstdlib>
#include <stdexcept>
#include <string>
#include <type_traits>

#include "vt/cuda/rmsnorm_decode_fast.h"
#include "vt/ops.h"

namespace vt::cuda {
namespace {

constexpr int kBlock = 256;

void Check(cudaError_t err, const char* what) {
  if (err != cudaSuccess) {
    throw std::runtime_error(std::string("vt cuda: ") + what + ": " + cudaGetErrorString(err));
  }
}

cudaStream_t AsStream(const Queue& q) { return static_cast<cudaStream_t>(q.handle); }

unsigned GridFor(int64_t n) {
  const int64_t blocks = (n + kBlock - 1) / kBlock;
  return static_cast<unsigned>(blocks < 4096 ? blocks : 4096);
}

// f32 load/store overloads: bf16 converts on the way in/out, math is f32.
__device__ inline float Load(const float* p, int64_t i) { return p[i]; }
__device__ inline float Load(const __nv_bfloat16* p, int64_t i) { return __bfloat162float(p[i]); }
__device__ inline void Store(float* p, int64_t i, float v) { p[i] = v; }
__device__ inline void Store(__nv_bfloat16* p, int64_t i, float v) {
  p[i] = __float2bfloat16(v);  // round-to-nearest-even, same as host F32ToBF16
}

// Round-trip a f32 value through the residual store dtype so the variance below
// squares the SAME rounded value that gets written back to the residual stream —
// mirrors vLLM fused_add_rms_norm, whose bf16 residual add (`z += residual`) rounds
// to the model dtype before the f32 variance (layernorm_kernels.cu). Identity for a
// f32 residual, so the previous f32-residual path stays byte-for-byte unchanged.
template <typename Tres> __device__ inline float ResRound(float v);
template <> __device__ inline float ResRound<float>(float v) { return v; }
template <> __device__ inline float ResRound<__nv_bfloat16>(float v) {
  return __bfloat162float(__float2bfloat16(v));
}

// ---------------------------------------------------------------------------
// rmsnorm: one block per row, shared-memory f32 tree reduction.
// Upstream csrc counterpart: csrc/layernorm_kernels.cu (rms_norm_kernel / fused_add_rms_norm_kernel) — align signatures post-MVP.

template <typename Tin, typename Tout, typename Tres>
__global__ void RmsNormRowKernel(Tout* out, const Tin* x, const Tin* w, Tres* residual,
                                 int64_t h, float eps, bool gemma) {
  const int64_t row = blockIdx.x;
  const Tin* xrow = x + row * h;
  Tout* orow = out + row * h;
  Tres* rrow = residual == nullptr ? nullptr : residual + row * h;

  __shared__ float partial[kBlock];
  float acc = 0.0f;
  for (int64_t j = threadIdx.x; j < h; j += kBlock) {
    float v = Load(xrow, j);
    if (rrow != nullptr) {
      v = ResRound<Tres>(v + Load(rrow, j));  // new residual stream: f32 add, round to Tres
      Store(rrow, j, v);                       // updated in place (f32 or bf16)
    }
    acc += v * v;
  }
  partial[threadIdx.x] = acc;
  __syncthreads();
  for (int s = kBlock / 2; s > 0; s /= 2) {
    if (static_cast<int>(threadIdx.x) < s) partial[threadIdx.x] += partial[threadIdx.x + s];
    __syncthreads();
  }
  const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(h) + eps);
  for (int64_t j = threadIdx.x; j < h; j += kBlock) {
    const float v = rrow != nullptr ? Load(rrow, j) : Load(xrow, j);
    float wj = Load(w, j);
    if (gemma) wj += 1.0f;
    Store(orow, j, v * inv * wj);
  }
}

// ---------------------------------------------------------------------------
// Decode-fast rmsnorm variant (VT_RMSNORM_DECODE_FAST). Re-expressed 2026-07-17
// (KERNEL-EW-NORM-ACT numerics rework, CLAIM-EW-NORM-ACT-2) to be BIT-IDENTICAL
// to the shipped RmsNormRowKernel above (the through-stack 235/235 bit-reference
// that matches vLLM's production greedy stream), not merely ≤1-ulp close.
//
// WHY BIT-IDENTICAL (not ≤1-ulp): the 27B greedy token 6 is a RAZOR near-tie
// (198 "\n" prod vs 271 "\n\n" emu; the logit gap is ~zero). The shipped
// RmsNormRowKernel + GDN cubin = 235/235 (198). The prior fast kernel differed
// from shipped by ≤1 ulp (residual-add rounding + a different variance reduction
// ORDER + rsqrtf); accumulated over 64 layers alongside the GDN cubin's own
// ≤1-ulp perturbation, that tipped the tie to 271 (233/235), forcing the default
// back OFF (a875397). Making the fast output the SAME BITS as shipped removes the
// perturbation by construction: fast+cubin ≡ shipped+cubin ≡ 198 always.
//
// The three divergences vs shipped RmsNormRowKernel (cuda_ops.cu:62-93) and the
// fix that makes each bit-exact:
//   (1) residual add. shipped:62-79 does v = ResRound<bf16>(f32(x)+f32(res)) — a
//       f32 add of the two bf16 operands then a SINGLE round to bf16 (double
//       rounding through f32). The prior fast used __hadd2 (single-round bf16
//       add), which differs from the shipped double-round on rare carries. Fixed:
//       __float2bfloat16(f32(x)+f32(res)) reproduces ResRound exactly.
//   (2) variance sum ORDER. shipped:70-85 uses kBlock=256 threads, per-thread
//       partial acc = Σ_m sq[tid + 256*m] (increasing m), then the shared-memory
//       binary tree `for s=128; s>0; s>>=1`. f32 add is non-associative, so the
//       sum's bits depend on this exact 256-partial + tree structure. The prior
//       fast used cub::BlockReduce<float,1024> over 1024 threads (a different
//       thread count AND a different reduction tree) => a different f32 variance.
//       Fixed: this kernel launches with kBlock(=256) threads and reproduces
//       shipped's scalar-strided Pass 1 and shared-tree byte-for-byte.
//   (3) inv. shipped:86 uses `1.0f / sqrtf(...)` (correctly-rounded sqrt+div);
//       the prior fast used rsqrtf (a ≤2-ulp reciprocal-sqrt approximation).
//       Fixed: `1.0f / sqrtf(partial[0]/h + eps)` verbatim.
//
// The ONLY thing that legitimately differs from shipped — and the whole source of
// any speedup — is Pass 2 (normalize): out = bf16((f32(res)*inv)*(f32(w)+gemma))
// is ELEMENT-INDEPENDENT, so vectorizing it with 16-byte (4×bf162) loads/stores
// changes memory traffic, NOT the arithmetic, and stays bit-identical. Pass 1
// (the variance) is order-sensitive and is kept byte-for-byte shipped.
//
// Scope: bf16-in / bf16-out / bf16-residual, H%8==0, H>=1024 (the 129
// input/post-attn/final residual RMSNorm launches, H=5120 on the 27B / H=2048 on
// the 35B). Other dtype/residual/small-H keep RmsNormRowKernel. The q/k head norms
// take no residual => not this path. 16-byte load = 4 packed bf162.
struct alignas(16) RmsNormBf16x8 {
  __nv_bfloat162 d[4];
};

// Launch geometry. The variance REDUCTION is byte-for-byte the shipped
// RmsNormRowKernel's (kBlock=256 partials + tree), but the memory passes run with
// kFastBlock=1024 threads: at decode the RMSNorm launches only `rows` (= num
// decode tokens, ~16 at c16) blocks, so the GPU is block-starved and thread-level
// parallelism per block — not occupancy — hides the memory latency. That
// thread-count is the ORIGINAL fast kernel's win; this rework keeps it while
// making the arithmetic bit-identical to shipped. kFastMaxH bounds the f32
// square-buffer below (guard rejects larger H); 8192 covers the 27B H=5120 / 35B
// H=2048 decode norms with headroom, at 32 KB static shared (< the 48 KB no-opt-in
// cap, and free here because only ~16 blocks are resident).
constexpr int kFastBlock = 1024;
constexpr int kFastMaxH = 8192;

// One block per row. Pass 1 (vectorized, kFastBlock threads): residual =
// bf16(f32(x)+f32(res)) [== shipped ResRound], stored bf16, with each element's
// f32 square written to the shared buffer ssq. Reduction (kBlock=256 threads):
// p_i = Σ_m ssq[i + 256*m] then the shared-memory binary tree — the EXACT shipped
// summation ORDER (cuda_ops.cu:70-86), reading the same bf16-rounded squares, so
// the f32 variance is bit-identical despite the vectorized loads. inv =
// 1.0f/sqrtf (shipped:86, not rsqrtf). Pass 2 (vectorized): normalize — element-
// independent, so identical bits at higher bandwidth. Every op matches shipped
// bit-for-bit; only the memory access WIDTH and thread COUNT differ.
__global__ void RmsNormRowFastKernel(__nv_bfloat16* __restrict__ out,
                                     const __nv_bfloat16* __restrict__ x,
                                     const __nv_bfloat16* __restrict__ w,
                                     __nv_bfloat16* __restrict__ residual, int h, float eps,
                                     bool gemma) {
  const int tid = static_cast<int>(threadIdx.x);
  const int64_t base = static_cast<int64_t>(blockIdx.x) * h;
  const int vh = h / 8;
  const RmsNormBf16x8* xv = reinterpret_cast<const RmsNormBf16x8*>(x + base);
  RmsNormBf16x8* rv = reinterpret_cast<RmsNormBf16x8*>(residual + base);
  const RmsNormBf16x8* wv = reinterpret_cast<const RmsNormBf16x8*>(w);
  RmsNormBf16x8* ov = reinterpret_cast<RmsNormBf16x8*>(out + base);

  __shared__ float ssq[kFastMaxH];   // per-element v^2 (v = the bf16-rounded residual)
  __shared__ float partial[kBlock];  // 256 partials for shipped's exact tree order

  // Pass 1 — vectorized residual add + store + per-element square into ssq. The
  // residual add is bf16(f32(x)+f32(res)) (== ResRound<bf16>, shipped:75) and the
  // square is of that bf16-rounded value (shipped:78), so ssq[j] is byte-for-byte
  // shipped's v*v term for every element.
  for (int vi = tid; vi < vh; vi += kFastBlock) {
    RmsNormBf16x8 t = xv[vi];
    RmsNormBf16x8 r = rv[vi];
    RmsNormBf16x8 nr;
#pragma unroll
    for (int k = 0; k < 4; k++) {
      float2 xf = __bfloat1622float2(t.d[k]);
      float2 rf = __bfloat1622float2(r.d[k]);
      nr.d[k] = __floats2bfloat162_rn(xf.x + rf.x, xf.y + rf.y);  // per-lane == __float2bfloat16
      float2 nf = __bfloat1622float2(nr.d[k]);                    // bf16 value back to f32
      ssq[vi * 8 + 2 * k] = nf.x * nf.x;
      ssq[vi * 8 + 2 * k + 1] = nf.y * nf.y;
    }
    rv[vi] = nr;  // residual store (bf16), bit-identical to shipped
  }
  __syncthreads();  // publish ssq (and residual) before the strided reduction reads them

  // Reduction — BYTE-FOR-BYTE shipped (cuda_ops.cu:80-86). Threads >=256 idle.
  if (tid < kBlock) {
    float acc = 0.0f;
    for (int j = tid; j < h; j += kBlock) acc += ssq[j];  // p_i = Σ_m ssq[i+256m], increasing m
    partial[tid] = acc;
  }
  __syncthreads();
  for (int s = kBlock / 2; s > 0; s >>= 1) {
    if (tid < s) partial[tid] += partial[tid + s];
    __syncthreads();
  }
  const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(h) + eps);

  // Pass 2 — vectorized normalize (bit-identical to shipped:87-92; each output
  // element is (f32(res)*inv)*(f32(w)+gemma) rounded to bf16, independent of the
  // others). Each thread reloads the residual vector it wrote in Pass 1.
  for (int vi = tid; vi < vh; vi += kFastBlock) {
    RmsNormBf16x8 r = rv[vi];
    RmsNormBf16x8 wr = wv[vi];
    RmsNormBf16x8 o;
#pragma unroll
    for (int k = 0; k < 4; k++) {
      float2 rf = __bfloat1622float2(r.d[k]);
      float2 wf = __bfloat1622float2(wr.d[k]);
      float w0 = wf.x, w1 = wf.y;
      if (gemma) {
        w0 += 1.0f;
        w1 += 1.0f;
      }
      o.d[k] = __floats2bfloat162_rn(rf.x * inv * w0, rf.y * inv * w1);
    }
    ov[vi] = o;
  }
}

// f32-residual sibling of RmsNormRowFastKernel — BYTE-IDENTICAL to the shipped
// RmsNormRowKernel<float,float,float> (residual!=null, f32 in/out/residual). The Laguna
// NVFP4 decode's post-attn residual RMSNorm runs this shipped kernel as a <<<rows=1,256>>>
// single block (ncu: launch__waves_per_multiprocessor≈0.00, sm__throughput≈0.06% — one SM
// of ~100+, latency-bound). Same fix as the bf16 fast kernel: float4-vectorized memory
// passes over kFastBlock(=1024) threads hide the per-block memory latency (decode launches
// only `rows` blocks), while the variance REDUCTION reproduces the shipped kBlock(=256)
// strided partials + tree over the SAME per-element squares in the SAME Σ_m ssq[i+256m]
// order → bit-identical f32 variance. residual add is f32 x+res (ResRound<float> is
// identity, shipped:75), inv = 1.0f/sqrtf (shipped:86), normalize out=res*inv*(w[+gemma])
// element-independent — every bit matches shipped. Scope guard (h%4==0, 1024<=h<=kFastMaxH,
// 16B-aligned) in TryLaunchRmsNormDecodeFastF32.
__global__ void RmsNormRowFastF32Kernel(float* __restrict__ out, const float* __restrict__ x,
                                        const float* __restrict__ w, float* __restrict__ residual,
                                        int h, float eps, bool gemma) {
  const int tid = static_cast<int>(threadIdx.x);
  const int64_t base = static_cast<int64_t>(blockIdx.x) * h;
  const int vh = h >> 2;
  const float4* xv = reinterpret_cast<const float4*>(x + base);
  float4* rv = reinterpret_cast<float4*>(residual + base);
  const float4* wv = reinterpret_cast<const float4*>(w);
  float4* ov = reinterpret_cast<float4*>(out + base);
  __shared__ float sv[kFastMaxH];     // residual VALUE v (not v²): the reduction squares with
  __shared__ float partial[kBlock];   // shipped's `acc += v*v` expression so nvcc emits the
                                      // SAME fma (f32 v² is NOT exact — pre-squaring the value
                                      // rounds it and diverges by ≤1 ulp; the bf16 sibling can
                                      // pre-square because a bf16² is exactly representable).

  // Pass 1 — residual = x + res (f32, == shipped ResRound<float> identity), store residual + v.
  for (int vi = tid; vi < vh; vi += kFastBlock) {
    float4 t = xv[vi], r = rv[vi], nr;
    nr.x = t.x + r.x;
    nr.y = t.y + r.y;
    nr.z = t.z + r.z;
    nr.w = t.w + r.w;
    rv[vi] = nr;
    const int e = vi << 2;
    sv[e] = nr.x;
    sv[e + 1] = nr.y;
    sv[e + 2] = nr.z;
    sv[e + 3] = nr.w;
  }
  __syncthreads();
  // Reduction — BYTE-FOR-BYTE shipped (cuda_ops.cu:80-86): thread t squares+accumulates
  // v[t],v[t+256],… with the IDENTICAL `acc += v*v` expression (same nvcc fma). Threads >=256 idle.
  if (tid < kBlock) {
    float acc = 0.0f;
    for (int j = tid; j < h; j += kBlock) acc += sv[j] * sv[j];
    partial[tid] = acc;
  }
  __syncthreads();
  for (int s = kBlock / 2; s > 0; s >>= 1) {
    if (tid < s) partial[tid] += partial[tid + s];
    __syncthreads();
  }
  const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(h) + eps);
  // Pass 2 — normalize (element-independent, bit-identical to shipped:87-92).
  for (int vi = tid; vi < vh; vi += kFastBlock) {
    float4 r = rv[vi], wr = wv[vi], o;
    float w0 = wr.x, w1 = wr.y, w2 = wr.z, w3 = wr.w;
    if (gemma) {
      w0 += 1.0f;
      w1 += 1.0f;
      w2 += 1.0f;
      w3 += 1.0f;
    }
    o.x = r.x * inv * w0;
    o.y = r.y * inv * w1;
    o.z = r.z * inv * w2;
    o.w = r.w * inv * w3;
    ov[vi] = o;
  }
}

// Runtime predicate + launch for the decode-fast path. Returns true iff it ran.
// Guard: bf16 in/out/weight/residual, 16-byte-aligned pointers (vectorized loads),
// H%8==0, 1024<=H<=kFastMaxH (H>=1024 scopes to the big residual RMSNorm launches;
// H<=kFastMaxH bounds the f32 square-buffer shared array). The launch uses
// kFastBlock threads for the memory passes, while the reduction reproduces the
// shipped RmsNormRowKernel's kBlock=256 partial + tree order, so the output is
// bit-identical. Out-of-scope shapes keep RmsNormRowKernel.
inline bool TryLaunchRmsNormDecodeFast(cudaStream_t s, Tensor& out, const Tensor& x,
                                       const Tensor& w, const RmsNormArgs& args,
                                       Tensor* residual, unsigned rows, int64_t h) {
  if (!RmsNormDecodeFastFlagIsOn(std::getenv("VT_RMSNORM_DECODE_FAST"))) return false;
  if (out.dtype != DType::kBF16 || x.dtype != DType::kBF16 || w.dtype != DType::kBF16)
    return false;
  if (residual == nullptr || residual->dtype != DType::kBF16) return false;
  if (h % 8 != 0 || h < 1024 || h > kFastMaxH) return false;
  auto aligned16 = [](const void* p) {
    return (reinterpret_cast<std::uintptr_t>(p) & 0xF) == 0;
  };
  if (!aligned16(out.data) || !aligned16(x.data) || !aligned16(w.data) ||
      !aligned16(residual->data))
    return false;
  RmsNormRowFastKernel<<<rows, kFastBlock, 0, s>>>(out.Ptr<__nv_bfloat16>(), x.Ptr<__nv_bfloat16>(),
                                             w.Ptr<__nv_bfloat16>(),
                                             residual->Ptr<__nv_bfloat16>(),
                                             static_cast<int>(h), args.eps, args.gemma);
  return true;
}

// f32-residual sibling of TryLaunchRmsNormDecodeFast. Same VT_RMSNORM_DECODE_FAST contract
// (default ON, '0' = rollback), same bit-identity guarantee (RmsNormRowFastF32Kernel). Scope:
// f32 in/out/weight AND a f32 residual (the add+RMSNorm decode launch — the Laguna NVFP4
// post-attn residual norm), h%4==0, 1024<=h<=kFastMaxH, 16-byte-aligned. Every other case
// keeps RmsNormRowKernel. float4 vectorization needs h%4==0 (vs the bf16 path's h%8==0).
inline bool TryLaunchRmsNormDecodeFastF32(cudaStream_t s, Tensor& out, const Tensor& x,
                                          const Tensor& w, const RmsNormArgs& args,
                                          Tensor* residual, unsigned rows, int64_t h) {
  if (!RmsNormDecodeFastFlagIsOn(std::getenv("VT_RMSNORM_DECODE_FAST"))) return false;
  if (out.dtype != DType::kF32 || x.dtype != DType::kF32 || w.dtype != DType::kF32) return false;
  if (residual == nullptr || residual->dtype != DType::kF32) return false;
  if (h % 4 != 0 || h < 1024 || h > kFastMaxH) return false;
  auto aligned16 = [](const void* p) {
    return (reinterpret_cast<std::uintptr_t>(p) & 0xF) == 0;
  };
  if (!aligned16(out.data) || !aligned16(x.data) || !aligned16(w.data) ||
      !aligned16(residual->data))
    return false;
  RmsNormRowFastF32Kernel<<<rows, kFastBlock, 0, s>>>(out.Ptr<float>(), x.Ptr<float>(),
                                                      w.Ptr<float>(), residual->Ptr<float>(),
                                                      static_cast<int>(h), args.eps, args.gemma);
  return true;
}

// Dispatch the residual store dtype (f32 or bf16). A bf16 residual mirrors vLLM's
// bf16 model dtype (model_config.dtype=bfloat16): the residual stream is bf16, only
// the variance/normalize accumulation below stays f32. A f32 residual (or none)
// takes the byte-identical previous path.
template <typename Tin, typename Tout>
void LaunchRmsNormRes(cudaStream_t s, Tensor& out, const Tensor& x, const Tensor& w,
                      const RmsNormArgs& args, Tensor* residual, unsigned rows, int64_t h) {
  if (residual != nullptr && residual->dtype == DType::kBF16) {
    RmsNormRowKernel<Tin, Tout, __nv_bfloat16><<<rows, kBlock, 0, s>>>(
        out.Ptr<Tout>(), x.Ptr<Tin>(), w.Ptr<Tin>(), residual->Ptr<__nv_bfloat16>(), h,
        args.eps, args.gemma);
  } else {
    float* res = residual == nullptr ? nullptr : residual->Ptr<float>();
    RmsNormRowKernel<Tin, Tout, float><<<rows, kBlock, 0, s>>>(
        out.Ptr<Tout>(), x.Ptr<Tin>(), w.Ptr<Tin>(), res, h, args.eps, args.gemma);
  }
}

template <typename Tin>
void LaunchRmsNorm(cudaStream_t s, Tensor& out, const Tensor& x, const Tensor& w,
                   const RmsNormArgs& args, Tensor* residual) {
  const int64_t t = x.shape[0], h = x.shape[1];
  if (t == 0 || h == 0) return;
  const unsigned rows = static_cast<unsigned>(t);
  // Decode-fast path (VT_RMSNORM_DECODE_FAST, default ON; '0' = rollback): only
  // engages for the bf16 add+RMSNorm decode launches; every other case keeps
  // RmsNormRowKernel. Output is bit-identical to RmsNormRowKernel by construction.
  if constexpr (std::is_same_v<Tin, __nv_bfloat16>) {
    if (TryLaunchRmsNormDecodeFast(s, out, x, w, args, residual, rows, h)) {
      Check(cudaGetLastError(), "rmsnorm fast launch");
      return;
    }
  }
  // f32-residual decode-fast path (Laguna NVFP4 post-attn residual norm). Bit-identical
  // to RmsNormRowKernel<float,float,float>; same VT_RMSNORM_DECODE_FAST contract.
  if constexpr (std::is_same_v<Tin, float>) {
    if (TryLaunchRmsNormDecodeFastF32(s, out, x, w, args, residual, rows, h)) {
      Check(cudaGetLastError(), "rmsnorm fast-f32 launch");
      return;
    }
  }
  switch (out.dtype) {
    case DType::kF32:
      LaunchRmsNormRes<Tin, float>(s, out, x, w, args, residual, rows, h);
      break;
    case DType::kBF16:
      LaunchRmsNormRes<Tin, __nv_bfloat16>(s, out, x, w, args, residual, rows, h);
      break;
    default: VT_CHECK(false, "cuda rmsnorm: unsupported out dtype");
  }
  Check(cudaGetLastError(), "rmsnorm launch");
}

void RmsNormKernelCuda(Queue& q, Tensor& out, const Tensor& x, const Tensor& w,
                       const RmsNormArgs& args, Tensor* residual) {
  VT_CHECK(w.dtype == x.dtype, "cuda rmsnorm: weight dtype must match x");
  switch (x.dtype) {
    case DType::kF32: LaunchRmsNorm<float>(AsStream(q), out, x, w, args, residual); break;
    case DType::kBF16:
      LaunchRmsNorm<__nv_bfloat16>(AsStream(q), out, x, w, args, residual);
      break;
    default: VT_CHECK(false, "cuda rmsnorm: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// rmsnorm + static fp8 quant, fused: emit the fp8 activation (and optionally the
// bf16 normed activation) directly from the RMSNorm's normalize loop, so the
// standalone QuantFp8Static pass + its bf16 round-trip disappear. Mirror of
// vLLM's Inductor fused_add_rms_norm_static_fp8_quant
// (vllm/compilation/passes/fusion/rms_quant_fusion.py:124). Same reduction as
// RmsNormRowKernel; BIT-IDENTICAL to RmsNorm(bf16)+QuantFp8Static because the
// fp8 is taken from the SAME bf16-rounded value the split path quantizes
// (F32ToFp8Dev(__bfloat162float(bf16(n)) * inv_scale)).
__device__ __forceinline__ uint8_t RmsNormF32ToFp8Dev(float f) {
  return static_cast<uint8_t>(__nv_cvt_float_to_fp8(f, __NV_SATFINITE, __NV_E4M3));
}

template <typename Tin, typename Tres>
__global__ void RmsNormQuantFp8RowKernel(uint8_t* out_fp8, __nv_bfloat16* out_bf16, const Tin* x,
                                         const Tin* w, Tres* residual, int64_t h, float eps,
                                         bool gemma, float inv_scale) {
  const int64_t row = blockIdx.x;
  const Tin* xrow = x + row * h;
  uint8_t* orow = out_fp8 + row * h;
  __nv_bfloat16* brow = out_bf16 == nullptr ? nullptr : out_bf16 + row * h;
  Tres* rrow = residual == nullptr ? nullptr : residual + row * h;

  __shared__ float partial[kBlock];
  float acc = 0.0f;
  for (int64_t j = threadIdx.x; j < h; j += kBlock) {
    float v = Load(xrow, j);
    if (rrow != nullptr) {
      v = ResRound<Tres>(v + Load(rrow, j));  // new residual stream: f32 add, round to Tres
      Store(rrow, j, v);                       // updated in place (f32 or bf16)
    }
    acc += v * v;
  }
  partial[threadIdx.x] = acc;
  __syncthreads();
  for (int s = kBlock / 2; s > 0; s /= 2) {
    if (static_cast<int>(threadIdx.x) < s) partial[threadIdx.x] += partial[threadIdx.x + s];
    __syncthreads();
  }
  const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(h) + eps);
  for (int64_t j = threadIdx.x; j < h; j += kBlock) {
    const float v = rrow != nullptr ? Load(rrow, j) : Load(xrow, j);
    float wj = Load(w, j);
    if (gemma) wj += 1.0f;
    // bf16-intermediate (matches RmsNorm's bf16 store then QuantFp8Static's bf16 load).
    const __nv_bfloat16 nb = __float2bfloat16(v * inv * wj);
    if (brow != nullptr) brow[j] = nb;
    orow[j] = RmsNormF32ToFp8Dev(__bfloat162float(nb) * inv_scale);
  }
}

template <typename Tin>
void LaunchRmsNormQuantFp8(cudaStream_t s, Tensor& out_fp8, Tensor* out_bf16, const Tensor& x,
                           const Tensor& w, const RmsNormArgs& args, Tensor* residual,
                           float input_scale) {
  const int64_t t = x.shape[0], h = x.shape[1];
  if (t == 0 || h == 0) return;
  const unsigned rows = static_cast<unsigned>(t);
  const float inv_scale = 1.0f / input_scale;
  __nv_bfloat16* bf16 = out_bf16 == nullptr ? nullptr : out_bf16->Ptr<__nv_bfloat16>();
  if (residual != nullptr && residual->dtype == DType::kBF16) {
    RmsNormQuantFp8RowKernel<Tin, __nv_bfloat16><<<rows, kBlock, 0, s>>>(
        out_fp8.Ptr<uint8_t>(), bf16, x.Ptr<Tin>(), w.Ptr<Tin>(),
        residual->Ptr<__nv_bfloat16>(), h, args.eps, args.gemma, inv_scale);
  } else {
    float* res = residual == nullptr ? nullptr : residual->Ptr<float>();
    RmsNormQuantFp8RowKernel<Tin, float><<<rows, kBlock, 0, s>>>(
        out_fp8.Ptr<uint8_t>(), bf16, x.Ptr<Tin>(), w.Ptr<Tin>(), res, h, args.eps, args.gemma,
        inv_scale);
  }
  Check(cudaGetLastError(), "rmsnorm_quant_fp8 launch");
}

void RmsNormQuantFp8KernelCuda(Queue& q, Tensor& out_fp8, Tensor* out_bf16, const Tensor& x,
                               const Tensor& w, const RmsNormArgs& args, Tensor* residual,
                               float input_scale) {
  VT_CHECK(w.dtype == x.dtype, "cuda rmsnorm_quant_fp8: weight dtype must match x");
  switch (x.dtype) {
    case DType::kF32:
      LaunchRmsNormQuantFp8<float>(AsStream(q), out_fp8, out_bf16, x, w, args, residual,
                                   input_scale);
      break;
    case DType::kBF16:
      LaunchRmsNormQuantFp8<__nv_bfloat16>(AsStream(q), out_fp8, out_bf16, x, w, args, residual,
                                           input_scale);
      break;
    default: VT_CHECK(false, "cuda rmsnorm_quant_fp8: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// silu_and_mul: grid-stride over the T*D output elements.
// Upstream csrc counterpart: csrc/activation_kernels.cu (act_and_mul_kernel<silu>) — align post-MVP.

template <typename Tin, typename Tout>
__global__ void SiluAndMulKernel(Tout* out, const Tin* x, int64_t n, int64_t d) {
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step) {
    const int64_t i = idx / d;
    const int64_t j = idx - i * d;
    const float gate = Load(x, i * 2 * d + j);
    const float up = Load(x, i * 2 * d + d + j);
    const float silu = gate / (1.0f + expf(-gate));
    Store(out, idx, silu * up);
  }
}

template <typename Tin>
void LaunchSiluAndMul(cudaStream_t s, Tensor& out, const Tensor& x) {
  const int64_t t = x.shape[0], d = x.shape[1] / 2;
  const int64_t n = t * d;
  if (n == 0) return;
  switch (out.dtype) {
    case DType::kF32:
      SiluAndMulKernel<Tin, float>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<float>(), x.Ptr<Tin>(), n, d);
      break;
    case DType::kBF16:
      SiluAndMulKernel<Tin, __nv_bfloat16>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<__nv_bfloat16>(), x.Ptr<Tin>(), n, d);
      break;
    default: VT_CHECK(false, "cuda silu_and_mul: unsupported out dtype");
  }
  Check(cudaGetLastError(), "silu_and_mul launch");
}

void SiluAndMulKernelCuda(Queue& q, Tensor& out, const Tensor& x) {
  switch (x.dtype) {
    case DType::kF32: LaunchSiluAndMul<float>(AsStream(q), out, x); break;
    case DType::kBF16: LaunchSiluAndMul<__nv_bfloat16>(AsStream(q), out, x); break;
    default: VT_CHECK(false, "cuda silu_and_mul: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// gelu_and_mul (Gemma GeGLU): out = gelu_tanh(gate) * up. gelu_tanh is the
// exact `gelu_pytorch_tanh` / F.gelu(approximate="tanh") — computed in f32 then
// stored, mirroring vLLM GeluAndMul(approximate="tanh").

template <typename Tin, typename Tout>
__global__ void GeluAndMulKernel(Tout* out, const Tin* x, int64_t n, int64_t d) {
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step) {
    const int64_t i = idx / d;
    const int64_t j = idx - i * d;
    const float g = Load(x, i * 2 * d + j);
    const float up = Load(x, i * 2 * d + d + j);
    // 0.5*g*(1 + tanh( sqrt(2/pi) * (g + 0.044715*g^3) )); sqrt(2/pi)=0.7978845608...
    const float inner = 0.7978845608028654f * (g + 0.044715f * g * g * g);
    const float gelu = 0.5f * g * (1.0f + tanhf(inner));
    Store(out, idx, gelu * up);
  }
}

template <typename Tin>
void LaunchGeluAndMul(cudaStream_t s, Tensor& out, const Tensor& x) {
  const int64_t t = x.shape[0], d = x.shape[1] / 2;
  const int64_t n = t * d;
  if (n == 0) return;
  switch (out.dtype) {
    case DType::kF32:
      GeluAndMulKernel<Tin, float>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<float>(), x.Ptr<Tin>(), n, d);
      break;
    case DType::kBF16:
      GeluAndMulKernel<Tin, __nv_bfloat16>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<__nv_bfloat16>(), x.Ptr<Tin>(), n, d);
      break;
    default: VT_CHECK(false, "cuda gelu_and_mul: unsupported out dtype");
  }
  Check(cudaGetLastError(), "gelu_and_mul launch");
}

void GeluAndMulKernelCuda(Queue& q, Tensor& out, const Tensor& x) {
  switch (x.dtype) {
    case DType::kF32: LaunchGeluAndMul<float>(AsStream(q), out, x); break;
    case DType::kBF16: LaunchGeluAndMul<__nv_bfloat16>(AsStream(q), out, x); break;
    default: VT_CHECK(false, "cuda gelu_and_mul: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// mul_scalar: out[i] = x[i] * scalar (f32 compute, out-dtype store). The Gemma
// embedding normalizer `embed * sqrt(hidden_size)`.

template <typename Tin, typename Tout>
__global__ void MulScalarKernel(Tout* out, const Tin* x, int64_t n, float scalar) {
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step)
    Store(out, idx, Load(x, idx) * scalar);
}

template <typename Tin>
void LaunchMulScalar(cudaStream_t s, Tensor& out, const Tensor& x, float scalar) {
  const int64_t n = x.Numel();
  if (n == 0) return;
  switch (out.dtype) {
    case DType::kF32:
      MulScalarKernel<Tin, float>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<float>(), x.Ptr<Tin>(), n, scalar);
      break;
    case DType::kBF16:
      MulScalarKernel<Tin, __nv_bfloat16>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<__nv_bfloat16>(), x.Ptr<Tin>(), n, scalar);
      break;
    default: VT_CHECK(false, "cuda mul_scalar: unsupported out dtype");
  }
  Check(cudaGetLastError(), "mul_scalar launch");
}

void MulScalarKernelCuda(Queue& q, Tensor& out, const Tensor& x, double scalar) {
  const float s = static_cast<float>(scalar);
  switch (x.dtype) {
    case DType::kF32: LaunchMulScalar<float>(AsStream(q), out, x, s); break;
    case DType::kBF16: LaunchMulScalar<__nv_bfloat16>(AsStream(q), out, x, s); break;
    default: VT_CHECK(false, "cuda mul_scalar: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// soft_cap: out[i] = cap * tanh(x[i] / cap) (f32 compute, out-dtype store). The
// Gemma-2 final logit soft-cap (gemma2.py:344-345). Mirrors torch
// logits.div_(cap).tanh_().mul_(cap).

template <typename Tin, typename Tout>
__global__ void SoftCapKernel(Tout* out, const Tin* x, int64_t n, float cap) {
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step)
    Store(out, idx, cap * tanhf(static_cast<float>(Load(x, idx)) / cap));
}

template <typename Tin>
void LaunchSoftCap(cudaStream_t s, Tensor& out, const Tensor& x, float cap) {
  const int64_t n = x.Numel();
  if (n == 0) return;
  switch (out.dtype) {
    case DType::kF32:
      SoftCapKernel<Tin, float>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<float>(), x.Ptr<Tin>(), n, cap);
      break;
    case DType::kBF16:
      SoftCapKernel<Tin, __nv_bfloat16>
          <<<GridFor(n), kBlock, 0, s>>>(out.Ptr<__nv_bfloat16>(), x.Ptr<Tin>(), n, cap);
      break;
    default: VT_CHECK(false, "cuda soft_cap: unsupported out dtype");
  }
  Check(cudaGetLastError(), "soft_cap launch");
}

void SoftCapKernelCuda(Queue& q, Tensor& out, const Tensor& x, double cap) {
  const float c = static_cast<float>(cap);
  switch (x.dtype) {
    case DType::kF32: LaunchSoftCap<float>(AsStream(q), out, x, c); break;
    case DType::kBF16: LaunchSoftCap<__nv_bfloat16>(AsStream(q), out, x, c); break;
    default: VT_CHECK(false, "cuda soft_cap: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// embedding: grid-stride gather. Ids live on the device, so bounds are checked
// in-kernel: bad ids are clamped for the gather (no OOB read) and the first bad
// id is recorded in a device-side flag via atomicCAS. The host wrapper
// synchronizes the stream, reads the flag back, and throws — CUDA Embedding is
// synchronizing for now (M0.6 decision, see ops.h; revisit for full async in
// M0.9/M2).
// No direct csrc counterpart (upstream uses torch embedding); keep vt-native.

struct EmbeddingErr {
  int status;    // 0 = ok, 1 = bad id recorded
  int pad;       // keep `id` naturally aligned
  long long id;  // first out-of-range id seen (valid when status != 0)
};

template <typename Tin, typename Tout, typename Tid>
__global__ void EmbeddingKernel(Tout* out, const Tin* table, const Tid* ids, int64_t n,
                                int64_t h, int64_t v, EmbeddingErr* err) {
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step) {
    const int64_t i = idx / h;
    const int64_t j = idx - i * h;
    int64_t id = static_cast<int64_t>(ids[i]);
    if (id < 0 || id >= v) {
      if (atomicCAS(&err->status, 0, 1) == 0) err->id = static_cast<long long>(id);
      id = id < 0 ? 0 : v - 1;  // clamp: keep the gather in-bounds
    }
    Store(out, idx, Load(table, id * h + j));
  }
}

template <typename Tin, typename Tout>
cudaError_t LaunchEmbedding(cudaStream_t s, Tensor& out, const Tensor& table,
                            const Tensor& ids, EmbeddingErr* err) {
  const int64_t t = ids.shape[0], h = table.shape[1], v = table.shape[0];
  const int64_t n = t * h;
  if (ids.dtype == DType::kI32) {
    EmbeddingKernel<Tin, Tout, int32_t><<<GridFor(n), kBlock, 0, s>>>(
        out.Ptr<Tout>(), table.Ptr<Tin>(), ids.Ptr<int32_t>(), n, h, v, err);
  } else {
    EmbeddingKernel<Tin, Tout, int64_t><<<GridFor(n), kBlock, 0, s>>>(
        out.Ptr<Tout>(), table.Ptr<Tin>(), ids.Ptr<int64_t>(), n, h, v, err);
  }
  return cudaGetLastError();
}

template <typename Tin>
cudaError_t LaunchEmbeddingIn(cudaStream_t s, Tensor& out, const Tensor& table,
                              const Tensor& ids, EmbeddingErr* err) {
  if (out.dtype == DType::kF32) return LaunchEmbedding<Tin, float>(s, out, table, ids, err);
  return LaunchEmbedding<Tin, __nv_bfloat16>(s, out, table, ids, err);
}

// Out-of-range reporting WITHOUT a per-call barrier.
//
// The original shape was cudaMalloc + kernel + D2H + cudaStreamSynchronize +
// cudaFree on EVERY call. All three of those driver calls SYNCHRONIZE the device,
// and embedding runs once per engine step at the very front of the forward, so
// the sequence is a hard barrier between consecutive steps. It costs almost
// nothing while the engine is serialized anyway (measured 23.7 us/call), and it
// costs the ENTIRE overlap the moment the engine stops being serialized — which
// is exactly what ENG-ASYNC-SCHED is for. See
// .agents/specs/async-discrete-device-combine.md W4e.
//
// The replacement is a small RING of persistent slots. Each slot owns a device
// flag, a PINNED host mirror and an event; a call takes the next slot, resets and
// launches into it, and records the event — no sync, no allocation. The check is
// then DEFERRED: before reusing a slot, its event is consumed. If the event has
// already completed the check is free (cudaEventQuery); only a genuinely
// in-flight slot blocks, which bounds the ring's memory without ever making the
// common path wait.
//
// Semantics change in exactly one way, deliberately: a bad id is reported up to
// kSlots calls later than it was, instead of on the offending call. It is still
// LOUD (the same exception, the same message, on the same queue) and it is still
// memory-safe on the offending call itself, because the kernel clamps the gather
// — an out-of-range id has never produced an out-of-bounds read, only a wrong
// row. The one report the ring cannot deliver is an error in the final embedding
// of a process that then exits without another embedding; that residue is
// covered by the token-exactness gates, which compare the produced ids.
struct EmbeddingErrSlot {
  EmbeddingErr* dev = nullptr;   // device-side flag the kernel atomically sets
  EmbeddingErr* host = nullptr;  // pinned mirror the async D2H lands in
  cudaEvent_t done = nullptr;    // completion of that D2H
  bool pending = false;          // event recorded and not yet consumed
  int64_t vocab = 0;             // table rows of the call that armed the slot
};

// One embedding call per engine step per queue, so a handful of slots covers any
// realistic overlap depth; the ring only has to outlive the in-flight window.
constexpr int kEmbeddingErrSlots = 4;

struct EmbeddingErrRing {
  EmbeddingErrSlot slots[kEmbeddingErrSlots];
  int next = 0;
};

// One ring per process: the forward runs on a single device and a single host
// thread (the same assumption DevicePool and the resident-weight caches make).
EmbeddingErrRing& ErrRing() {
  static EmbeddingErrRing ring;
  return ring;
}

// Consume a slot's outstanding result, blocking only if `force`. Returns the
// error to report, if any, WITHOUT throwing: the caller decides when to throw so
// a partially-armed slot is never left behind.
bool ConsumeEmbeddingErr(EmbeddingErrSlot& slot, bool force, EmbeddingErr* err_out,
                         int64_t* vocab_out) {
  if (!slot.pending) return false;
  if (!force) {
    const cudaError_t q = cudaEventQuery(slot.done);
    if (q == cudaErrorNotReady) return false;  // still in flight: check it later
    if (q != cudaSuccess) Check(q, "embedding flag event query");
  } else {
    Check(cudaEventSynchronize(slot.done), "embedding flag event sync");
  }
  slot.pending = false;
  if (slot.host->status == 0) return false;
  *err_out = *slot.host;
  *vocab_out = slot.vocab;
  return true;
}

void EmbeddingKernelCuda(Queue& q, Tensor& out, const Tensor& table, const Tensor& ids) {
  // Validate dtypes before touching the ring so a throw cannot leave a slot armed.
  VT_CHECK(table.dtype == DType::kF32 || table.dtype == DType::kBF16,
           "cuda embedding: unsupported table dtype (f32/bf16 only)");
  VT_CHECK(out.dtype == DType::kF32 || out.dtype == DType::kBF16,
           "cuda embedding: unsupported out dtype");
  const int64_t n = ids.shape[0] * table.shape[1];
  if (n == 0) return;
  // v == 0 with nonempty ids can never gather anything valid, and the in-kernel
  // clamp (v - 1) would go out of bounds — throw loudly before launching.
  VT_CHECK(table.shape[0] > 0, "cuda embedding: empty table (vocab 0) with nonempty ids");
  cudaStream_t s = AsStream(q);

  EmbeddingErrRing& ring = ErrRing();
  EmbeddingErrSlot& slot = ring.slots[ring.next];
  ring.next = (ring.next + 1) % kEmbeddingErrSlots;

  if (slot.dev == nullptr) {
    // First use of this slot: the ONLY allocation this path ever makes. Pinned
    // host memory so the D2H is a real async copy rather than a staged one.
    Check(cudaMalloc(&slot.dev, sizeof(EmbeddingErr)), "cudaMalloc embedding flag");
    Check(cudaHostAlloc(reinterpret_cast<void**>(&slot.host), sizeof(EmbeddingErr),
                        cudaHostAllocDefault),
          "cudaHostAlloc embedding flag mirror");
    Check(cudaEventCreateWithFlags(&slot.done, cudaEventDisableTiming),
          "cudaEventCreate embedding flag");
    slot.host->status = 0;
  }

  // Reusing this slot means its previous result must be consumed first. Force the
  // wait: the slot is about to be overwritten, so skipping the check here would
  // DROP a report rather than defer it. With kSlots slots in the ring this only
  // blocks when more than kSlots embeddings are genuinely in flight.
  EmbeddingErr prev{};
  int64_t prev_vocab = 0;
  const bool had_prev = ConsumeEmbeddingErr(slot, /*force=*/true, &prev, &prev_vocab);

  cudaError_t st = cudaMemsetAsync(slot.dev, 0, sizeof(EmbeddingErr), s);
  if (st == cudaSuccess) {
    st = table.dtype == DType::kF32
             ? LaunchEmbeddingIn<float>(s, out, table, ids, slot.dev)
             : LaunchEmbeddingIn<__nv_bfloat16>(s, out, table, ids, slot.dev);
  }
  if (st == cudaSuccess) {
    st = cudaMemcpyAsync(slot.host, slot.dev, sizeof(EmbeddingErr), cudaMemcpyDeviceToHost, s);
  }
  if (st == cudaSuccess) st = cudaEventRecord(slot.done, s);
  if (st == cudaSuccess) {
    slot.pending = true;
    slot.vocab = table.shape[0];
  }

  // A launch/copy failure is reported before a deferred out-of-range id: it is
  // the more immediate fault, and it may be why the older flag never arrived.
  Check(st, "embedding");
  if (had_prev) {
    throw std::runtime_error("vt cuda: embedding: id " + std::to_string(prev.id) +
                             " out of range [0, " + std::to_string(prev_vocab) + ")");
  }

  // Opportunistically drain every OTHER slot whose copy has already landed, so a
  // bad id surfaces at the next call rather than only when its slot comes round
  // again. Free: cudaEventQuery on a completed event does not block.
  for (int i = 0; i < kEmbeddingErrSlots; ++i) {
    EmbeddingErrSlot& other = ring.slots[i];
    if (&other == &slot) continue;
    EmbeddingErr err{};
    int64_t vocab = 0;
    if (ConsumeEmbeddingErr(other, /*force=*/false, &err, &vocab)) {
      throw std::runtime_error("vt cuda: embedding: id " + std::to_string(err.id) +
                               " out of range [0, " + std::to_string(vocab) + ")");
    }
  }
}

// ---------------------------------------------------------------------------
// rope_neox: grid-stride over (token, head, rotation pair) across q and k.
// Angle math in double (pow/cos/sin) to match the CPU reference numerics.
// Upstream csrc counterpart: csrc/pos_encoding_kernels.cu (rotary_embedding_kernel) — align post-MVP.

// Llama-3 rope frequency rescale (vLLM Llama3RotaryEmbedding._compute_inv_freq,
// rotary_embedding/llama3_rope.py:33-54). `freq` is the base inv_freq
// (base^(-2i/rot)); when scaling_factor <= 0 this is a no-op (plain RoPE). The
// piecewise low/high wavelength interpolation matches vLLM's torch.where ladder.
__device__ __host__ inline double Llama3ScaleFreq(double freq, double scaling_factor,
                                                  double low_ff, double high_ff,
                                                  double orig_max) {
  if (!(scaling_factor > 0.0)) return freq;
  constexpr double kTwoPi = 6.283185307179586476925286766559;
  const double low_freq_wavelen = orig_max / low_ff;
  const double high_freq_wavelen = orig_max / high_ff;
  const double wave_len = kTwoPi / freq;
  double smooth = 0.0;
  if (low_ff != high_ff)
    smooth = (orig_max / wave_len - low_ff) / (high_ff - low_ff);
  if (wave_len < high_freq_wavelen) return freq;                  // high-freq: keep
  if (wave_len > low_freq_wavelen) return freq / scaling_factor;  // low-freq: scale
  return (1.0 - smooth) * freq / scaling_factor + smooth * freq;  // mid: interpolate
}

template <typename T, typename Tid>
__global__ void RopeNeoxKernel(T* qs, T* ks, const Tid* pos, int64_t hq, int64_t hk,
                               int64_t d, int64_t half, int rot, double base,
                               double l3_sf, double l3_lo, double l3_hi, double l3_omax,
                               int64_t n) {
  const int64_t heads = hq + hk;
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step) {
    const int64_t pair = idx % half;
    const int64_t head = (idx / half) % heads;
    const int64_t tok = idx / (half * heads);
    T* ptr;
    int64_t off;
    if (head < hq) {
      ptr = qs;
      off = (tok * hq + head) * d;
    } else {
      ptr = ks;
      off = (tok * hk + (head - hq)) * d;
    }
    const int64_t p = static_cast<int64_t>(pos[tok]);
    double freq = pow(base, -2.0 * static_cast<double>(pair) / static_cast<double>(rot));
    freq = Llama3ScaleFreq(freq, l3_sf, l3_lo, l3_hi, l3_omax);
    const double angle = static_cast<double>(p) * freq;
    const float c = static_cast<float>(cos(angle));
    const float sn = static_cast<float>(sin(angle));
    const float x = Load(ptr, off + pair);
    const float y = Load(ptr, off + pair + half);
    Store(ptr, off + pair, x * c - y * sn);
    Store(ptr, off + pair + half, x * sn + y * c);
  }
}

template <typename T>
void LaunchRope(cudaStream_t s, Tensor& qs, Tensor& ks, const Tensor& pos,
                const RopeArgs& args) {
  const int64_t t = qs.shape[0], hq = qs.shape[1], hk = ks.shape[1], d = qs.shape[2];
  const int64_t half = args.rotary_dim / 2;
  const int64_t n = t * (hq + hk) * half;
  if (n == 0) return;
  const double base = static_cast<double>(args.base);
  const double l3_sf = static_cast<double>(args.llama3_scaling_factor);
  const double l3_lo = static_cast<double>(args.llama3_low_freq_factor);
  const double l3_hi = static_cast<double>(args.llama3_high_freq_factor);
  const double l3_omax = static_cast<double>(args.llama3_orig_max_position);
  if (pos.dtype == DType::kI32) {
    RopeNeoxKernel<T, int32_t><<<GridFor(n), kBlock, 0, s>>>(
        qs.Ptr<T>(), ks.Ptr<T>(), pos.Ptr<int32_t>(), hq, hk, d, half, args.rotary_dim, base,
        l3_sf, l3_lo, l3_hi, l3_omax, n);
  } else {
    RopeNeoxKernel<T, int64_t><<<GridFor(n), kBlock, 0, s>>>(
        qs.Ptr<T>(), ks.Ptr<T>(), pos.Ptr<int64_t>(), hq, hk, d, half, args.rotary_dim, base,
        l3_sf, l3_lo, l3_hi, l3_omax, n);
  }
  Check(cudaGetLastError(), "rope_neox launch");
}

void RopeNeoxKernelCuda(Queue& q, Tensor& qs, Tensor& ks, const Tensor& pos,
                        const RopeArgs& args) {
  switch (qs.dtype) {
    case DType::kF32: LaunchRope<float>(AsStream(q), qs, ks, pos, args); break;
    case DType::kBF16: LaunchRope<__nv_bfloat16>(AsStream(q), qs, ks, pos, args); break;
    default: VT_CHECK(false, "cuda rope: unsupported dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// Supplied-cache RoPE. Ported from pinned vLLM base.py:160-252,
// csrc/libtorch_stable/pos_encoding_kernels.cu:8-200, and the 3-axis selection
// in mrope.py:14-187,263-375 @ e24d1b24fe96. This hot kernel performs only
// cache lookup + rotation; YaRN formula construction happens once on the host.

__device__ inline int MropeAxisForPair(int64_t pair, int section_t,
                                       int section_h, int section_w,
                                       bool interleaved) {
  if (interleaved) {
    if (pair % 3 == 1 && pair <= 3LL * section_h) return 1;
    if (pair % 3 == 2 && pair <= 3LL * section_w) return 2;
    return 0;
  }
  if (pair < section_t) return 0;
  if (pair < static_cast<int64_t>(section_t) + section_h) return 1;
  return 2;
}

template <typename T, typename Tid>
__global__ void RopeFromCacheKernel(
    T* qs, T* ks, const Tid* positions, const T* cache, int64_t cache_rows,
    int64_t tokens, int64_t hq, int64_t hk, int64_t q_tok_stride,
    int64_t q_head_stride, int64_t k_tok_stride, int64_t k_head_stride,
    int rotary_dim, int64_t half, bool is_neox_style, bool is_mrope,
    int section_t, int section_h, int section_w, bool mrope_interleaved,
    int64_t n) {
  const int64_t heads = hq + hk;
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x +
                         threadIdx.x;
       idx < n; idx += step) {
    const int64_t pair = idx % half;
    const int64_t head = (idx / half) % heads;
    const int64_t token = idx / (half * heads);
    const int axis =
        is_mrope ? MropeAxisForPair(pair, section_t, section_h, section_w,
                                   mrope_interleaved)
                 : 0;
    const int64_t position_offset =
        is_mrope ? static_cast<int64_t>(axis) * tokens + token : token;
    const int64_t position = static_cast<int64_t>(positions[position_offset]);
    // The public contract, like upstream's custom op, requires positions to be
    // valid cache rows. Avoid an out-of-bounds read if a broken caller violates
    // it; CPU validation reports the exact error in reference tests.
    if (position < 0 || position >= cache_rows) continue;
    const int64_t cache_offset = position * rotary_dim;
    const float c = Load(cache, cache_offset + pair);
    const float sn = Load(cache, cache_offset + half + pair);

    // MLA campaign W6: STRIDE-driven addressing. DeepSeek's decoupled RoPE
    // rotates only the trailing rope slice of each query head and its k_pe is a
    // column block of the fused kv_a projection — both are strided views. For a
    // contiguous tensor these offsets are integer-identical to the previous
    // (token * heads + head) * head_dim formula, so existing callers are
    // bit-identical by construction.
    T* states = head < hq ? qs : ks;
    const int64_t local_head = head < hq ? head : head - hq;
    const int64_t row = head < hq
                            ? token * q_tok_stride + local_head * q_head_stride
                            : token * k_tok_stride + local_head * k_head_stride;
    const int64_t first = is_neox_style ? pair : pair * 2;
    const int64_t second = is_neox_style ? pair + half : pair * 2 + 1;
    const float x = Load(states, row + first);
    const float y = Load(states, row + second);
    Store(states, row + first, x * c - y * sn);
    Store(states, row + second, x * sn + y * c);
  }
}

template <typename T, typename Tid>
void LaunchRopeFromCacheTyped(cudaStream_t stream, Tensor& qs, Tensor* ks,
                              const Tensor& positions, const Tensor& cache,
                              const RopeArgs& args) {
  const int64_t tokens = qs.shape[0];
  const int64_t hq = qs.shape[1];
  const int64_t hk = ks == nullptr ? 0 : ks->shape[1];
  const int64_t half = args.rotary_dim / 2;
  const int64_t n = tokens * (hq + hk) * half;
  if (n == 0) return;
  RopeFromCacheKernel<T, Tid><<<GridFor(n), kBlock, 0, stream>>>(
      qs.Ptr<T>(), ks == nullptr ? nullptr : ks->Ptr<T>(),
      positions.Ptr<Tid>(), cache.Ptr<T>(), cache.shape[0], tokens, hq, hk,
      qs.stride[0], qs.stride[1], ks == nullptr ? 0 : ks->stride[0],
      ks == nullptr ? 0 : ks->stride[1], args.rotary_dim, half,
      args.is_neox_style, positions.rank == 2, args.mrope_section[0],
      args.mrope_section[1], args.mrope_section[2], args.mrope_interleaved, n);
}

template <typename T>
void LaunchRopeFromCache(cudaStream_t stream, Tensor& qs, Tensor* ks,
                         const Tensor& positions, const Tensor& cache,
                         const RopeArgs& args) {
  if (positions.dtype == DType::kI32) {
    LaunchRopeFromCacheTyped<T, int32_t>(stream, qs, ks, positions, cache,
                                         args);
  } else {
    LaunchRopeFromCacheTyped<T, int64_t>(stream, qs, ks, positions, cache,
                                         args);
  }
  Check(cudaGetLastError(), "rope_from_cache launch");
}

void RopeFromCacheKernelCuda(Queue& q, Tensor& qs, Tensor* ks,
                             const Tensor& positions, const Tensor& cache,
                             const RopeArgs& args) {
  switch (qs.dtype) {
    case DType::kF32:
      LaunchRopeFromCache<float>(AsStream(q), qs, ks, positions, cache, args);
      break;
    case DType::kBF16:
      LaunchRopeFromCache<__nv_bfloat16>(AsStream(q), qs, ks, positions,
                                         cache, args);
      break;
    default:
      VT_CHECK(false,
               "cuda rope_from_cache: unsupported dtype (f32/bf16 only)");
  }
}

// Fused MLA norm-rope (kFusedNormRope): one block per token. The RMS reduction
// over the latent slice [0,off) reuses RmsNormRowKernel's EXACT tree reduce +
// scale (cuda_ops.cu:62-93), and the decoupled-pe rotation over [off,off+rot)
// reuses RopeFromCacheKernel's EXACT cache read + rotation (cuda_ops.cu:896-943).
// The two halves address DISJOINT dims, so the fused output is BIT-IDENTICAL to
// {RmsNorm(x[:,:off]) ; RopeFromCache(x[:,off:])} — the CPU FusedNormRopeKernel
// runs exactly that composite as the golden.
template <typename T, typename Tid>
__global__ void FusedNormRopeKernel(T* latent_out, T* pe_out, const T* x, const T* w,
                                    const Tid* positions, const T* cache, int64_t cache_rows,
                                    int64_t off, int rot, int64_t half, int64_t x_row_stride,
                                    int64_t lat_row_stride, int64_t pe_row_stride,
                                    bool is_neox_style, bool gemma, float eps) {
  const int64_t row = blockIdx.x;
  const T* xrow = x + row * x_row_stride;
  T* lrow = latent_out + row * lat_row_stride;
  T* prow = pe_out + row * pe_row_stride;

  // --- latent RMSNorm over [0, off) — identical to RmsNormRowKernel. ----------
  __shared__ float partial[kBlock];
  float acc = 0.0f;
  for (int64_t j = threadIdx.x; j < off; j += kBlock) {
    const float v = Load(xrow, j);
    acc += v * v;
  }
  partial[threadIdx.x] = acc;
  __syncthreads();
  for (int s = kBlock / 2; s > 0; s /= 2) {
    if (static_cast<int>(threadIdx.x) < s) partial[threadIdx.x] += partial[threadIdx.x + s];
    __syncthreads();
  }
  const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(off) + eps);
  for (int64_t j = threadIdx.x; j < off; j += kBlock) {
    float wj = Load(w, j);
    if (gemma) wj += 1.0f;
    Store(lrow, j, Load(xrow, j) * inv * wj);
  }

  // --- decoupled-pe RopeFromCache over [off, off+rot) — identical to
  //     RopeFromCacheKernel (single vector, base rope, positions rank-1). ------
  const int64_t position = static_cast<int64_t>(positions[row]);
  if (position < 0 || position >= cache_rows) return;
  const int64_t cache_offset = position * rot;
  for (int64_t pair = threadIdx.x; pair < half; pair += kBlock) {
    const float c = Load(cache, cache_offset + pair);
    const float sn = Load(cache, cache_offset + half + pair);
    const int64_t first = is_neox_style ? pair : pair * 2;
    const int64_t second = is_neox_style ? pair + half : pair * 2 + 1;
    const float xr = Load(xrow, off + first);
    const float yr = Load(xrow, off + second);
    Store(prow, first, xr * c - yr * sn);
    Store(prow, second, xr * sn + yr * c);
  }
}

template <typename T>
void LaunchFusedNormRope(cudaStream_t stream, Tensor& latent_out, Tensor& pe_out,
                         const Tensor& x, const Tensor& w, const Tensor& positions,
                         const Tensor& cache, const RmsNormArgs& norm_args,
                         const RopeArgs& rope_args) {
  const int64_t t = x.shape[0];
  const int64_t off = w.shape[0];
  const int rot = rope_args.rotary_dim;
  const int64_t half = rot / 2;
  if (t == 0) return;
  const unsigned rows = static_cast<unsigned>(t);
  if (positions.dtype == DType::kI32) {
    FusedNormRopeKernel<T, int32_t><<<rows, kBlock, 0, stream>>>(
        latent_out.Ptr<T>(), pe_out.Ptr<T>(), x.Ptr<T>(), w.Ptr<T>(),
        positions.Ptr<int32_t>(), cache.Ptr<T>(), cache.shape[0], off, rot, half,
        x.stride[0], latent_out.stride[0], pe_out.stride[0], rope_args.is_neox_style,
        norm_args.gemma, norm_args.eps);
  } else {
    FusedNormRopeKernel<T, int64_t><<<rows, kBlock, 0, stream>>>(
        latent_out.Ptr<T>(), pe_out.Ptr<T>(), x.Ptr<T>(), w.Ptr<T>(),
        positions.Ptr<int64_t>(), cache.Ptr<T>(), cache.shape[0], off, rot, half,
        x.stride[0], latent_out.stride[0], pe_out.stride[0], rope_args.is_neox_style,
        norm_args.gemma, norm_args.eps);
  }
  Check(cudaGetLastError(), "fused_norm_rope launch");
}

void FusedNormRopeKernelCuda(Queue& q, Tensor& latent_out, Tensor& pe_out, const Tensor& x,
                             const Tensor& w, const Tensor& positions, const Tensor& cache,
                             const RmsNormArgs& norm_args, const RopeArgs& rope_args) {
  switch (x.dtype) {
    case DType::kF32:
      LaunchFusedNormRope<float>(AsStream(q), latent_out, pe_out, x, w, positions, cache,
                                 norm_args, rope_args);
      break;
    case DType::kBF16:
      LaunchFusedNormRope<__nv_bfloat16>(AsStream(q), latent_out, pe_out, x, w, positions, cache,
                                         norm_args, rope_args);
      break;
    default:
      VT_CHECK(false, "cuda fused_norm_rope: unsupported dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// rope_cos_sin_cache: precompute the batch's cos|sin ONCE per step (grid-stride
// over (token, pair)) so the fused preamble below does zero in-kernel
// transcendentals. Angle math in DOUBLE + f32 cast — bit-for-bit RopeNeoxKernel's
// c/sn. cos_sin[t, i]=cos, cos_sin[t, half+i]=sin. Mirrors vLLM's cos_sin_cache
// (RotaryEmbedding._compute_cos_sin_cache; read by fla fused_qk_norm_rope.py:95).

template <typename Tid>
__global__ void RopeCosSinCacheKernel(float* cos_sin, const Tid* pos, int64_t t, int rot,
                                      int64_t half, double base, double l3_sf, double l3_lo,
                                      double l3_hi, double l3_omax) {
  const int64_t n = t * half;
  const int64_t step = static_cast<int64_t>(gridDim.x) * blockDim.x;
  for (int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; idx < n;
       idx += step) {
    const int64_t pair = idx % half;
    const int64_t tok = idx / half;
    const int64_t p = static_cast<int64_t>(pos[tok]);
    double freq = pow(base, -2.0 * static_cast<double>(pair) / static_cast<double>(rot));
    freq = Llama3ScaleFreq(freq, l3_sf, l3_lo, l3_hi, l3_omax);
    const double angle = static_cast<double>(p) * freq;
    cos_sin[tok * rot + pair] = static_cast<float>(cos(angle));
    cos_sin[tok * rot + half + pair] = static_cast<float>(sin(angle));
  }
}

void RopeCosSinCacheKernelCuda(Queue& q, Tensor& cos_sin, const Tensor& pos, const RopeArgs& args) {
  const int64_t t = cos_sin.shape[0];
  const int64_t half = args.rotary_dim / 2;
  const int64_t n = t * half;
  if (n == 0) return;
  cudaStream_t s = AsStream(q);
  const double base = static_cast<double>(args.base);
  const double l3_sf = static_cast<double>(args.llama3_scaling_factor);
  const double l3_lo = static_cast<double>(args.llama3_low_freq_factor);
  const double l3_hi = static_cast<double>(args.llama3_high_freq_factor);
  const double l3_omax = static_cast<double>(args.llama3_orig_max_position);
  if (pos.dtype == DType::kI32) {
    RopeCosSinCacheKernel<int32_t><<<GridFor(n), kBlock, 0, s>>>(
        cos_sin.Ptr<float>(), pos.Ptr<int32_t>(), t, args.rotary_dim, half, base,
        l3_sf, l3_lo, l3_hi, l3_omax);
  } else {
    RopeCosSinCacheKernel<int64_t><<<GridFor(n), kBlock, 0, s>>>(
        cos_sin.Ptr<float>(), pos.Ptr<int64_t>(), t, args.rotary_dim, half, base,
        l3_sf, l3_lo, l3_hi, l3_omax);
  }
  Check(cudaGetLastError(), "rope_cos_sin_cache launch");
}

// ---------------------------------------------------------------------------
// attn_qk_norm_rope_gate: the fused full-attention preamble — one launch replacing
// AttnGateSplit + RmsNorm(q) + RmsNorm(k) + RopeNeox. Grid (T, Hq+Hkv); one block
// per (token, head) does the shared-mem gemma-RMSNorm tree reduction over Dh (reuses
// RmsNormRowKernel's math), then partial NeoX RoPE reading the precomputed cos_sin
// cache (no per-element pow/cos/sin), then passes the gate half through (q heads).
// Structure mirrors GdnPostConvKernel (cuda_gdn.cu). Bit-for-bit equal to the four
// composed ops for f32 out; templated on Tsrc (qgate/kf) and Tout (q/k/gate).
// Mirrors vLLM's fused_qk_rmsnorm_rope (fla fused_qk_norm_rope.py:95-102).

// gemma-RMSNorm one element: (v*inv)*(gemma ? w+1 : w) — matches RmsNormRowKernel.
__device__ inline float GemmaNormElem(float v, float inv, float w, bool gemma) {
  float wj = w;
  if (gemma) wj += 1.0f;
  return v * inv * wj;
}

// Tqk = q_out/k_out store dtype, Tgate = gate_out store dtype. All math stays
// f32; a bf16 Tqk store is the RN round of the exact f32 value, so (Tqk=bf16,
// Tgate=f32) is bit-identical to the f32 path followed by CastBf16 on q/k — the
// FA-2 prefill combo (bf16 q feeds FA-2, bf16 k feeds the bf16 KV-cache write,
// gate stays f32 because sigmoid(gate) must see the un-rounded f32 value).
//
// NOTE (2026-07-18, CLAIM-35B-FA2-FLIP-1): a "round normed q/k to bf16 BEFORE
// RoPE" tighten (mirror of vLLM fused_qk_norm_rope.py:67) was implemented and
// op-level VALIDATED bit-identical to the unfused bf16 path, but it flipped the
// 27B's known tok6 whitespace near-tie AWAY from the vLLM oracle (want_prod →
// want_emu) in COMBINATION with our other sub-ULP op diffs (a compensating-error
// interaction — the RMSNorm-saga lesson). The 35B passes 315/315 with OR without
// the round, so the round is NOT shipped: the untightened preamble keeps BOTH
// the 27B (235/235) and the 35B (315/315) token-exact to their graphed oracles.
template <typename Tsrc, typename Tqk, typename Tgate>
__global__ void AttnQkNormRopeGateKernel(Tqk* q_out, Tqk* k_out, Tgate* gate_out,
                                         const Tsrc* qgate, const Tsrc* kf, const float* q_norm,
                                         const float* k_norm, const float* cos_sin, int64_t hq,
                                         int64_t hkv, int64_t dh,
                                         int64_t qgate_row_stride,
                                         int64_t kf_row_stride, int rot,
                                         int64_t half, float eps, bool gemma) {
  const int64_t tok = blockIdx.x;
  const int64_t head = blockIdx.y;  // [0, hq+hkv)
  const bool is_q = head < hq;

  const Tsrc* src;
  const float* w;
  Tqk* out;
  int64_t gate_base = 0;  // only meaningful for q heads
  if (is_q) {
    const int64_t qrow = tok * qgate_row_stride + head * 2 * dh;
    src = qgate + qrow;       // q half [0,dh)
    gate_base = qrow + dh;    // gate half [dh,2dh)
    w = q_norm;
    out = q_out + (tok * hq + head) * dh;
  } else {
    const int64_t hk = head - hq;
    src = kf + tok * kf_row_stride + hk * dh;
    w = k_norm;
    out = k_out + (tok * hkv + hk) * dh;
  }

  // ---- gemma-RMSNorm reduction over Dh (f32 variance) ----
  __shared__ float partial[kBlock];
  float acc = 0.0f;
  for (int64_t j = threadIdx.x; j < dh; j += kBlock) {
    const float v = Load(src, j);
    acc += v * v;
  }
  partial[threadIdx.x] = acc;
  __syncthreads();
  for (int s = kBlock / 2; s > 0; s /= 2) {
    if (static_cast<int>(threadIdx.x) < s) partial[threadIdx.x] += partial[threadIdx.x + s];
    __syncthreads();
  }
  const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(dh) + eps);

  // ---- partial NeoX RoPE + write (recompute the paired normed elems; no shared
  // round-trip). Matches RopeNeoxKernel: out[i]=x*c - y*sn, out[i+half]=x*sn + y*c
  // with x=normed[i], y=normed[i+half]; dims [rot,Dh) normed but unrotated. ----
  const float* cs = cos_sin + tok * rot;
  for (int64_t j = threadIdx.x; j < dh; j += kBlock) {
    if (j < half) {
      const float ni = GemmaNormElem(Load(src, j), inv, w[j], gemma);
      const float nih = GemmaNormElem(Load(src, j + half), inv, w[j + half], gemma);
      Store(out, j, ni * cs[j] - nih * cs[half + j]);
    } else if (j < rot) {
      const int64_t i = j - half;
      const float ni = GemmaNormElem(Load(src, i), inv, w[i], gemma);
      const float nih = GemmaNormElem(Load(src, i + half), inv, w[i + half], gemma);
      Store(out, j, ni * cs[half + i] + nih * cs[i]);
    } else {
      Store(out, j, GemmaNormElem(Load(src, j), inv, w[j], gemma));
    }
  }

  // ---- gate passthrough (q heads only): the raw gate half, no norm/rope ----
  if (is_q) {
    Tgate* go = gate_out + (tok * hq + head) * dh;
    for (int64_t j = threadIdx.x; j < dh; j += kBlock) Store(go, j, Load(qgate, gate_base + j));
  }
}

template <typename Tsrc, typename Tqk, typename Tgate>
void LaunchAttnPreamble(cudaStream_t s, Tensor& q_out, Tensor& k_out, Tensor& gate_out,
                        const Tensor& qgate, const Tensor& kf, const Tensor& q_norm,
                        const Tensor& k_norm, const Tensor& cos_sin, const RmsNormArgs& na,
                        const RopeArgs& ra) {
  const int64_t t = q_out.shape[0], hq = q_out.shape[1], dh = q_out.shape[2];
  const int64_t hkv = k_out.shape[1];
  const int64_t half = ra.rotary_dim / 2;
  if (t == 0) return;
  dim3 grid(static_cast<unsigned>(t), static_cast<unsigned>(hq + hkv));
  AttnQkNormRopeGateKernel<Tsrc, Tqk, Tgate><<<grid, kBlock, 0, s>>>(
      q_out.Ptr<Tqk>(), k_out.Ptr<Tqk>(), gate_out.Ptr<Tgate>(), qgate.Ptr<Tsrc>(),
      kf.Ptr<Tsrc>(), q_norm.Ptr<float>(), k_norm.Ptr<float>(), cos_sin.Ptr<float>(), hq, hkv, dh,
      qgate.stride[0], kf.stride[0], ra.rotary_dim, half, na.eps, na.gemma);
  Check(cudaGetLastError(), "attn_qk_norm_rope_gate launch");
}

template <typename Tsrc>
void LaunchAttnPreambleOut(cudaStream_t s, Tensor& q_out, Tensor& k_out, Tensor& gate_out,
                           const Tensor& qgate, const Tensor& kf, const Tensor& q_norm,
                           const Tensor& k_norm, const Tensor& cos_sin, const RmsNormArgs& na,
                           const RopeArgs& ra) {
  // (q/k out, gate out) combos: (f32,f32) — the default token-exact path;
  // (bf16,bf16) — all-bf16; (bf16,f32) — the FA-2 prefill combo (bf16 q/k for
  // FA-2 + the bf16 KV-cache write, f32 gate for the sigmoid). Validation in
  // ops.cpp admits exactly these.
  switch (q_out.dtype) {
    case DType::kF32:
      LaunchAttnPreamble<Tsrc, float, float>(s, q_out, k_out, gate_out, qgate, kf, q_norm, k_norm,
                                             cos_sin, na, ra);
      break;
    case DType::kBF16:
      if (gate_out.dtype == DType::kF32) {
        LaunchAttnPreamble<Tsrc, __nv_bfloat16, float>(s, q_out, k_out, gate_out, qgate, kf,
                                                       q_norm, k_norm, cos_sin, na, ra);
      } else {
        LaunchAttnPreamble<Tsrc, __nv_bfloat16, __nv_bfloat16>(s, q_out, k_out, gate_out, qgate,
                                                               kf, q_norm, k_norm, cos_sin, na, ra);
      }
      break;
    default: VT_CHECK(false, "cuda attn_qk_norm_rope_gate: unsupported out dtype");
  }
}

void AttnQkNormRopeGateKernelCuda(Queue& q, Tensor& q_out, Tensor& k_out, Tensor& gate_out,
                                  const Tensor& qgate, const Tensor& kf, const Tensor& q_norm,
                                  const Tensor& k_norm, const Tensor& cos_sin,
                                  const RmsNormArgs& na, const RopeArgs& ra) {
  cudaStream_t s = AsStream(q);
  switch (qgate.dtype) {
    case DType::kF32:
      LaunchAttnPreambleOut<float>(s, q_out, k_out, gate_out, qgate, kf, q_norm, k_norm, cos_sin,
                                   na, ra);
      break;
    case DType::kBF16:
      LaunchAttnPreambleOut<__nv_bfloat16>(s, q_out, k_out, gate_out, qgate, kf, q_norm, k_norm,
                                           cos_sin, na, ra);
      break;
    default: VT_CHECK(false, "cuda attn_qk_norm_rope_gate: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// attention: one block per (query i, q-head h); block threads cooperate over
// the head_dim and stream the keys with an online (flash-style) softmax. The
// online update is algebraically identical to the CPU two-pass reference
// (qwen36-forward-notes.md §5); f32 accumulation. Correctness-grade (M0.9).

template <typename Tin, typename Tout>
__global__ void AttentionKernel(Tout* out, const Tin* query, const Tin* key, const Tin* value,
                                int64_t hq, int64_t hk, int64_t d, int64_t t, float scale,
                                bool causal) {
  const int64_t i = blockIdx.x;  // query position
  const int64_t h = blockIdx.y;  // q-head
  const int64_t g = h / (hq / hk);
  const int64_t jmax = causal ? i : t - 1;
  const int64_t qoff = (i * hq + h) * d;

  extern __shared__ float smem[];
  float* acc = smem;                    // [d] running output accumulator
  float* red = smem + d;                // [blockDim.x] reduction scratch
  __shared__ float s_score, s_m, s_l;   // block-wide score / running max / denom
  for (int64_t e = threadIdx.x; e < d; e += blockDim.x) acc[e] = 0.0f;
  if (threadIdx.x == 0) {
    s_m = -CUDART_INF_F;
    s_l = 0.0f;
  }
  __syncthreads();

  for (int64_t j = 0; j <= jmax; ++j) {
    const int64_t koff = (j * hk + g) * d;
    float part = 0.0f;
    for (int64_t e = threadIdx.x; e < d; e += blockDim.x)
      part += Load(query, qoff + e) * Load(key, koff + e);
    red[threadIdx.x] = part;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
      if (threadIdx.x < stride) red[threadIdx.x] += red[threadIdx.x + stride];
      __syncthreads();
    }
    if (threadIdx.x == 0) s_score = red[0] * scale;
    __syncthreads();

    const float s = s_score;
    const float m_new = fmaxf(s_m, s);
    const float corr = expf(s_m - m_new);  // 0 on the first key (s_m == -inf)
    const float p = expf(s - m_new);
    const int64_t voff = (j * hk + g) * d;
    for (int64_t e = threadIdx.x; e < d; e += blockDim.x)
      acc[e] = acc[e] * corr + p * Load(value, voff + e);
    __syncthreads();
    if (threadIdx.x == 0) {
      s_l = s_l * corr + p;
      s_m = m_new;
    }
    __syncthreads();
  }

  const float inv = 1.0f / s_l;
  for (int64_t e = threadIdx.x; e < d; e += blockDim.x) Store(out, qoff + e, acc[e] * inv);
}

template <typename Tin>
void LaunchAttention(cudaStream_t s, Tensor& out, const Tensor& query, const Tensor& key,
                     const Tensor& value, const AttentionArgs& args) {
  const int64_t t = query.shape[0], hq = query.shape[1], d = query.shape[2];
  const int64_t hk = key.shape[1];
  if (t == 0 || hq == 0 || d == 0) return;
  const dim3 grid(static_cast<unsigned>(t), static_cast<unsigned>(hq));
  const size_t shmem = (static_cast<size_t>(d) + kBlock) * sizeof(float);
  switch (out.dtype) {
    case DType::kF32:
      AttentionKernel<Tin, float><<<grid, kBlock, shmem, s>>>(
          out.Ptr<float>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), hq, hk, d, t,
          args.scale, args.causal);
      break;
    case DType::kBF16:
      AttentionKernel<Tin, __nv_bfloat16><<<grid, kBlock, shmem, s>>>(
          out.Ptr<__nv_bfloat16>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), hq, hk,
          d, t, args.scale, args.causal);
      break;
    default: VT_CHECK(false, "cuda attention: unsupported out dtype");
  }
  Check(cudaGetLastError(), "attention launch");
}

void AttentionKernelCuda(Queue& q, Tensor& out, const Tensor& query, const Tensor& key,
                         const Tensor& value, const AttentionArgs& args) {
  switch (query.dtype) {
    case DType::kF32: LaunchAttention<float>(AsStream(q), out, query, key, value, args); break;
    case DType::kBF16:
      LaunchAttention<__nv_bfloat16>(AsStream(q), out, query, key, value, args);
      break;
    default: VT_CHECK(false, "cuda attention: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// DFlash in-block attention (SPEC-DFLASH D2, DF-DRAFT-MODEL) — the project's FIRST
// non-causal / bidirectional attention CUDA kernel. Mirrors AttentionKernel's
// one-block-per-(query,head) block-reduction online-softmax recurrence EXACTLY
// (same f32 accumulation), generalized to (a) per-request BLOCK bounds so a query
// attends only within its own (1+k) block and (b) a BIDIRECTIONAL (non-causal) or
// causal-within-window mask per DFlashQwen3Attention (_resolve_layer_attention,
// qwen3_dflash.py:86-146). Kept a SEPARATE op so the causal kAttention used by the
// text/audio paths is byte-identical. `cu` is a device copy of cu_seqlens
// (num_reqs+1); each block linearly finds its request (num_reqs is small).
template <typename Tin, typename Tout>
__global__ void DFlashBlockAttentionKernel(Tout* out, const Tin* query, const Tin* key,
                                           const Tin* value, const int32_t* cu, int num_reqs,
                                           int64_t hq, int64_t hk, int64_t d, float scale,
                                           bool causal, int64_t window) {
  const int64_t i = blockIdx.x;  // GLOBAL query row
  const int64_t h = blockIdx.y;  // q-head
  const int64_t g = h / (hq / hk);
  // Find the request block owning row i (small num_reqs; the uniform DFlash case
  // is num_reqs blocks of 1+k). qs/qe are GLOBAL row bounds of the block.
  int64_t qs = 0, qe = 0;
  for (int r = 0; r < num_reqs; ++r) {
    if (i >= cu[r] && i < cu[r + 1]) {
      qs = cu[r];
      qe = cu[r + 1];
      break;
    }
  }
  const int64_t ii = i - qs;                       // intra-block query offset
  const int64_t jhi = causal ? ii : (qe - qs - 1);  // last visible intra-block key
  int64_t jlo = 0;
  if (causal && window > 0) jlo = ii - (window - 1) > 0 ? ii - (window - 1) : 0;
  const int64_t qoff = (i * hq + h) * d;

  extern __shared__ float smem[];
  float* acc = smem;                   // [d] running output accumulator
  float* red = smem + d;               // [blockDim.x] reduction scratch
  __shared__ float s_score, s_m, s_l;  // block-wide score / running max / denom
  for (int64_t e = threadIdx.x; e < d; e += blockDim.x) acc[e] = 0.0f;
  if (threadIdx.x == 0) {
    s_m = -CUDART_INF_F;
    s_l = 0.0f;
  }
  __syncthreads();

  for (int64_t jj = jlo; jj <= jhi; ++jj) {
    const int64_t koff = ((qs + jj) * hk + g) * d;
    float part = 0.0f;
    for (int64_t e = threadIdx.x; e < d; e += blockDim.x)
      part += Load(query, qoff + e) * Load(key, koff + e);
    red[threadIdx.x] = part;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
      if (threadIdx.x < stride) red[threadIdx.x] += red[threadIdx.x + stride];
      __syncthreads();
    }
    if (threadIdx.x == 0) s_score = red[0] * scale;
    __syncthreads();

    const float s = s_score;
    const float m_new = fmaxf(s_m, s);
    const float corr = expf(s_m - m_new);  // 0 on the first key (s_m == -inf)
    const float p = expf(s - m_new);
    const int64_t voff = ((qs + jj) * hk + g) * d;
    for (int64_t e = threadIdx.x; e < d; e += blockDim.x)
      acc[e] = acc[e] * corr + p * Load(value, voff + e);
    __syncthreads();
    if (threadIdx.x == 0) {
      s_l = s_l * corr + p;
      s_m = m_new;
    }
    __syncthreads();
  }

  const float inv = 1.0f / s_l;
  for (int64_t e = threadIdx.x; e < d; e += blockDim.x) Store(out, qoff + e, acc[e] * inv);
}

template <typename Tin>
void LaunchDFlashBlockAttention(cudaStream_t s, Tensor& out, const Tensor& query,
                                const Tensor& key, const Tensor& value,
                                const DFlashBlockAttentionArgs& args) {
  const int64_t t = query.shape[0], hq = query.shape[1], d = query.shape[2];
  const int64_t hk = key.shape[1];
  if (t == 0 || hq == 0 || d == 0) return;
  // Upload cu_seqlens (host, num_reqs+1) to a stream-ordered device scratch.
  const size_t cub = static_cast<size_t>(args.num_reqs + 1) * sizeof(int32_t);
  int32_t* d_cu = nullptr;
  Check(cudaMallocAsync(&d_cu, cub, s), "dflash-block-attn cu malloc");
  Check(cudaMemcpyAsync(d_cu, args.cu_seqlens, cub, cudaMemcpyHostToDevice, s),
        "dflash-block-attn cu upload");
  const dim3 grid(static_cast<unsigned>(t), static_cast<unsigned>(hq));
  const size_t shmem = (static_cast<size_t>(d) + kBlock) * sizeof(float);
  switch (out.dtype) {
    case DType::kF32:
      DFlashBlockAttentionKernel<Tin, float><<<grid, kBlock, shmem, s>>>(
          out.Ptr<float>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), d_cu,
          args.num_reqs, hq, hk, d, args.scale, args.causal, args.sliding_window);
      break;
    case DType::kBF16:
      DFlashBlockAttentionKernel<Tin, __nv_bfloat16><<<grid, kBlock, shmem, s>>>(
          out.Ptr<__nv_bfloat16>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), d_cu,
          args.num_reqs, hq, hk, d, args.scale, args.causal, args.sliding_window);
      break;
    default: VT_CHECK(false, "cuda dflash-block-attn: unsupported out dtype");
  }
  Check(cudaGetLastError(), "dflash-block-attn launch");
  Check(cudaFreeAsync(d_cu, s), "dflash-block-attn cu free");
}

void DFlashBlockAttentionKernelCuda(Queue& q, Tensor& out, const Tensor& query, const Tensor& key,
                                    const Tensor& value, const DFlashBlockAttentionArgs& args) {
  switch (query.dtype) {
    case DType::kF32:
      LaunchDFlashBlockAttention<float>(AsStream(q), out, query, key, value, args);
      break;
    case DType::kBF16:
      LaunchDFlashBlockAttention<__nv_bfloat16>(AsStream(q), out, query, key, value, args);
      break;
    default: VT_CHECK(false, "cuda dflash-block-attn: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// DFlash PAGED in-block attention (SPEC-DFLASH D12 Part B) — the CAPTURE-SAFE form
// of DFlashBlockAttentionKernel. Same block-cooperative f32 online-softmax
// recurrence, but each (1+k) block query attends over [PAGED context ; its own
// (1+k) block] with the context supplied as DATA: a paged K/V cache
// ([pages, block_size, Hkv, D]) + per-request seq_lens (context length C_r) +
// block_table (logical page -> physical page), mirroring PagedAttentionKernel's
// paged read. Grid is STATIC over the fixed Nq=(1+k)*num_reqs query rows; EVERY
// metadata input is a persistent device pointer read in place — NO cudaMallocAsync
// / cudaMemcpyAsync of a function-local host buffer (the eager DFlash launcher's
// cu_seqlens upload was the cudagraph-capture-bakes-stack-addresses UAF class).
// The D2 mask is applied over the COMBINED index: context rows occupy combined
// positions [0,C_r) (position-ordered) and block rows [C_r, C_r+blen_r), so a
// query's combined offset is C_r+ii and the mask/order match the materialized
// combined buffer bit-for-bit.
template <typename Tin, typename Tout>
__global__ void DFlashPagedBlockAttentionKernel(
    Tout* out, const Tin* query, const Tin* block_key, const Tin* block_value,
    const Tin* ctx_key, const Tin* ctx_value, const int32_t* cu, const int32_t* slen,
    const int32_t* btbl, int num_reqs, int64_t hq, int64_t hk, int64_t d, int64_t block_size,
    int64_t max_pages, int64_t ck_blk, int64_t ck_pg, int64_t ck_hd, float scale, bool causal,
    int64_t window) {
  const int64_t i = blockIdx.x;  // GLOBAL block-query row
  const int64_t h = blockIdx.y;  // q-head
  const int64_t g = h / (hq / hk);
  int req = -1;
  int64_t qs = 0, qe = 0;
  for (int r = 0; r < num_reqs; ++r) {
    if (i >= cu[r] && i < cu[r + 1]) {
      qs = cu[r];
      qe = cu[r + 1];
      req = r;
      break;
    }
  }
  if (req < 0) return;  // padding row beyond the last request
  const int64_t blen = qe - qs;
  const int64_t C = slen[req];
  const int64_t N = C + blen;  // combined key length
  const int64_t ii_comb = C + (i - qs);
  const int64_t jhi = causal ? ii_comb : (N - 1);
  int64_t jlo = 0;
  if (causal && window > 0) jlo = ii_comb - (window - 1) > 0 ? ii_comb - (window - 1) : 0;
  const int64_t qoff = (i * hq + h) * d;

  extern __shared__ float smem[];
  float* acc = smem;                   // [d] running output accumulator
  float* red = smem + d;               // [blockDim.x] reduction scratch
  __shared__ float s_score, s_m, s_l;  // block-wide score / running max / denom
  for (int64_t e = threadIdx.x; e < d; e += blockDim.x) acc[e] = 0.0f;
  if (threadIdx.x == 0) {
    s_m = -CUDART_INF_F;
    s_l = 0.0f;
  }
  __syncthreads();

  for (int64_t cj = jlo; cj <= jhi; ++cj) {
    int64_t koff;
    const Tin* ksrc;
    const Tin* vsrc;
    if (cj < C) {  // paged context key/value (k and v share the [pages,bs,Hkv,D] layout)
      const int64_t page = btbl[req * max_pages + cj / block_size];
      const int64_t off = cj % block_size;
      koff = page * ck_blk + off * ck_pg + g * ck_hd;
      ksrc = ctx_key;
      vsrc = ctx_value;
    } else {  // contiguous block key/value
      const int64_t brow = qs + (cj - C);
      koff = (brow * hk + g) * d;
      ksrc = block_key;
      vsrc = block_value;
    }
    float part = 0.0f;
    for (int64_t e = threadIdx.x; e < d; e += blockDim.x)
      part += Load(query, qoff + e) * Load(ksrc, koff + e);
    red[threadIdx.x] = part;
    __syncthreads();
    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
      if (threadIdx.x < stride) red[threadIdx.x] += red[threadIdx.x + stride];
      __syncthreads();
    }
    if (threadIdx.x == 0) s_score = red[0] * scale;
    __syncthreads();

    const float s = s_score;
    const float m_new = fmaxf(s_m, s);
    const float corr = expf(s_m - m_new);
    const float p = expf(s - m_new);
    for (int64_t e = threadIdx.x; e < d; e += blockDim.x)
      acc[e] = acc[e] * corr + p * Load(vsrc, koff + e);
    __syncthreads();
    if (threadIdx.x == 0) {
      s_l = s_l * corr + p;
      s_m = m_new;
    }
    __syncthreads();
  }

  const float inv = 1.0f / s_l;
  for (int64_t e = threadIdx.x; e < d; e += blockDim.x) Store(out, qoff + e, acc[e] * inv);
}

// D14 WARP-scoped variant of DFlashPagedBlockAttentionKernel — the SAME paged/
// block combined-index read, GQA broadcast, and causal/SWA mask, but ONE WARP
// (not a whole kBlock=256 block) owns each (block-query row, q-head): the
// head_dim dot product is a butterfly `__shfl_xor` reduction (no __syncthreads),
// the output accumulator lives in registers, and the online-softmax stats are
// per-lane. This removes the block kernel's per-key __syncthreads storm (2 barriers
// + a 256-wide shared-mem tree reduction PER context key, over C~500-640 keys) that
// nsys attributed ~1.8% of the graphed DFlash step to (median ~460 us/call vs
// vLLM's fused flash draft-attn ~0.15%) — exactly the AttentionDenseFast fix applied
// to the ViT tower. It is NOT bit-identical to the block kernel (the head_dim
// partial-sum grouping over 32 lanes differs), but is the SAME f32-online-softmax
// math within the bf16 envelope (CUDA==CPU gates f32 eps 1e-4 / bf16 eps 3e-2), and
// spec-decode output is exact by construction (the target verify is unchanged; only
// which draft proposals are accepted can shift, within the ratified ±4 gate). Reads
// only persistent device pointers (capture-safe, no shared mem, no host uploads).
// Grounded in the shipped AttentionWarpKernel + the online-softmax recurrence
// (src/vt/cuda/flash_attn/).
template <typename Tin, typename Tout>
__global__ void DFlashPagedBlockAttentionWarpKernel(
    Tout* out, const Tin* query, const Tin* block_key, const Tin* block_value,
    const Tin* ctx_key, const Tin* ctx_value, const int32_t* cu, const int32_t* slen,
    const int32_t* btbl, int num_reqs, int64_t hq, int64_t hk, int64_t d, int64_t block_size,
    int64_t max_pages, int64_t ck_blk, int64_t ck_pg, int64_t ck_hd, float scale, bool causal,
    int64_t window) {
  constexpr int kMaxPerLane = 8;  // head_dim up to 256
  const int64_t i = blockIdx.x;   // GLOBAL block-query row
  const int64_t h = blockIdx.y;   // q-head
  const int64_t g = h / (hq / hk);
  const int lane = static_cast<int>(threadIdx.x);  // 0..31
  int req = -1;
  int64_t qs = 0, qe = 0;
  for (int r = 0; r < num_reqs; ++r) {
    if (i >= cu[r] && i < cu[r + 1]) {
      qs = cu[r];
      qe = cu[r + 1];
      req = r;
      break;
    }
  }
  if (req < 0) return;  // padding row beyond the last request
  const int64_t blen = qe - qs;
  const int64_t C = slen[req];
  const int64_t N = C + blen;  // combined key length
  const int64_t ii_comb = C + (i - qs);
  const int64_t jhi = causal ? ii_comb : (N - 1);
  int64_t jlo = 0;
  if (causal && window > 0) jlo = ii_comb - (window - 1) > 0 ? ii_comb - (window - 1) : 0;
  const int64_t qoff = (i * hq + h) * d;
  const int npl = static_cast<int>((d + 31) / 32);
  float qreg[kMaxPerLane];
  float acc[kMaxPerLane];
#pragma unroll
  for (int k = 0; k < kMaxPerLane; ++k) {
    qreg[k] = 0.0f;
    acc[k] = 0.0f;
  }
  for (int k = 0; k < npl; ++k) {
    const int e = lane + 32 * k;
    if (e < d) qreg[k] = Load(query, qoff + e);
  }
  float m = -CUDART_INF_F, l = 0.0f;
  for (int64_t cj = jlo; cj <= jhi; ++cj) {
    int64_t koff;
    const Tin* ksrc;
    const Tin* vsrc;
    if (cj < C) {  // paged context key/value (k and v share the [pages,bs,Hkv,D] layout)
      const int64_t page = btbl[req * max_pages + cj / block_size];
      const int64_t off = cj % block_size;
      koff = page * ck_blk + off * ck_pg + g * ck_hd;
      ksrc = ctx_key;
      vsrc = ctx_value;
    } else {  // contiguous block key/value
      const int64_t brow = qs + (cj - C);
      koff = (brow * hk + g) * d;
      ksrc = block_key;
      vsrc = block_value;
    }
    float part = 0.0f;
#pragma unroll
    for (int k = 0; k < kMaxPerLane; ++k) {
      const int e = lane + 32 * k;
      if (k < npl && e < d) part += qreg[k] * Load(ksrc, koff + e);
    }
#pragma unroll
    for (int off = 16; off > 0; off >>= 1) part += __shfl_xor_sync(0xffffffffu, part, off);
    const float s = part * scale;
    const float m_new = fmaxf(m, s);
    const float corr = expf(m - m_new);  // 0 on the first key (m == -inf)
    const float p = expf(s - m_new);
#pragma unroll
    for (int k = 0; k < kMaxPerLane; ++k) {
      const int e = lane + 32 * k;
      if (k < npl && e < d) acc[k] = acc[k] * corr + p * Load(vsrc, koff + e);
    }
    l = l * corr + p;
    m = m_new;
  }
  const float inv = 1.0f / l;
  for (int k = 0; k < npl; ++k) {
    const int e = lane + 32 * k;
    if (e < d) Store(out, qoff + e, acc[k] * inv);
  }
}

// Select the block kernel (bit-identical, the D12/D13 reference) only when
// VT_DFLASH_ATTN_BLOCK=1; the WARP kernel is the D14 default (same math, bf16
// envelope, ~no __syncthreads storm). Cached once (getenv is not hot-path-safe).
inline bool UseDflashAttnBlockKernel() {
  static const bool on = [] {
    const char* e = std::getenv("VT_DFLASH_ATTN_BLOCK");
    return e != nullptr && e[0] == '1';
  }();
  return on;
}

template <typename Tin>
void LaunchDFlashPagedBlockAttention(cudaStream_t s, Tensor& out, const Tensor& query,
                                     const Tensor& block_key, const Tensor& block_value,
                                     const Tensor& ctx_key, const Tensor& ctx_value,
                                     const Tensor& cu_seqlens, const Tensor& seq_lens,
                                     const Tensor& block_table,
                                     const DFlashPagedBlockAttentionArgs& args) {
  const int64_t nq = query.shape[0], hq = query.shape[1], d = query.shape[2];
  const int64_t hk = block_key.shape[1];
  if (nq == 0 || hq == 0 || d == 0) return;
  const int64_t block_size = args.block_size;
  const int64_t max_pages = block_table.shape[1];
  // Contiguous paged cache strides: [num_pages, block_size, Hkv, D].
  const int64_t ck_blk = block_size * hk * d, ck_pg = hk * d, ck_hd = d;
  const dim3 grid(static_cast<unsigned>(nq), static_cast<unsigned>(hq));
  if (!UseDflashAttnBlockKernel()) {
    VT_CHECK(d <= 256, "cuda dflash-paged-block-attn(warp): head_dim <= 256 only");
    switch (out.dtype) {
      case DType::kF32:
        DFlashPagedBlockAttentionWarpKernel<Tin, float><<<grid, 32, 0, s>>>(
            out.Ptr<float>(), query.Ptr<Tin>(), block_key.Ptr<Tin>(), block_value.Ptr<Tin>(),
            ctx_key.Ptr<Tin>(), ctx_value.Ptr<Tin>(), cu_seqlens.Ptr<int32_t>(),
            seq_lens.Ptr<int32_t>(), block_table.Ptr<int32_t>(), args.num_reqs, hq, hk, d,
            block_size, max_pages, ck_blk, ck_pg, ck_hd, args.scale, args.causal,
            args.sliding_window);
        break;
      case DType::kBF16:
        DFlashPagedBlockAttentionWarpKernel<Tin, __nv_bfloat16><<<grid, 32, 0, s>>>(
            out.Ptr<__nv_bfloat16>(), query.Ptr<Tin>(), block_key.Ptr<Tin>(),
            block_value.Ptr<Tin>(), ctx_key.Ptr<Tin>(), ctx_value.Ptr<Tin>(),
            cu_seqlens.Ptr<int32_t>(), seq_lens.Ptr<int32_t>(), block_table.Ptr<int32_t>(),
            args.num_reqs, hq, hk, d, block_size, max_pages, ck_blk, ck_pg, ck_hd, args.scale,
            args.causal, args.sliding_window);
        break;
      default: VT_CHECK(false, "cuda dflash-paged-block-attn: unsupported out dtype");
    }
    Check(cudaGetLastError(), "dflash-paged-block-attn(warp) launch");
    return;
  }
  const size_t shmem = (static_cast<size_t>(d) + kBlock) * sizeof(float);
  switch (out.dtype) {
    case DType::kF32:
      DFlashPagedBlockAttentionKernel<Tin, float><<<grid, kBlock, shmem, s>>>(
          out.Ptr<float>(), query.Ptr<Tin>(), block_key.Ptr<Tin>(), block_value.Ptr<Tin>(),
          ctx_key.Ptr<Tin>(), ctx_value.Ptr<Tin>(), cu_seqlens.Ptr<int32_t>(),
          seq_lens.Ptr<int32_t>(), block_table.Ptr<int32_t>(), args.num_reqs, hq, hk, d, block_size,
          max_pages, ck_blk, ck_pg, ck_hd, args.scale, args.causal, args.sliding_window);
      break;
    case DType::kBF16:
      DFlashPagedBlockAttentionKernel<Tin, __nv_bfloat16><<<grid, kBlock, shmem, s>>>(
          out.Ptr<__nv_bfloat16>(), query.Ptr<Tin>(), block_key.Ptr<Tin>(),
          block_value.Ptr<Tin>(), ctx_key.Ptr<Tin>(), ctx_value.Ptr<Tin>(),
          cu_seqlens.Ptr<int32_t>(), seq_lens.Ptr<int32_t>(), block_table.Ptr<int32_t>(),
          args.num_reqs, hq, hk, d, block_size, max_pages, ck_blk, ck_pg, ck_hd, args.scale,
          args.causal, args.sliding_window);
      break;
    default: VT_CHECK(false, "cuda dflash-paged-block-attn: unsupported out dtype");
  }
  Check(cudaGetLastError(), "dflash-paged-block-attn launch");
}

void DFlashPagedBlockAttentionKernelCuda(Queue& q, Tensor& out, const Tensor& query,
                                         const Tensor& block_key, const Tensor& block_value,
                                         const Tensor& ctx_key, const Tensor& ctx_value,
                                         const Tensor& cu_seqlens, const Tensor& seq_lens,
                                         const Tensor& block_table,
                                         const DFlashPagedBlockAttentionArgs& args) {
  switch (query.dtype) {
    case DType::kF32:
      LaunchDFlashPagedBlockAttention<float>(AsStream(q), out, query, block_key, block_value,
                                             ctx_key, ctx_value, cu_seqlens, seq_lens, block_table,
                                             args);
      break;
    case DType::kBF16:
      LaunchDFlashPagedBlockAttention<__nv_bfloat16>(AsStream(q), out, query, block_key,
                                                     block_value, ctx_key, ctx_value, cu_seqlens,
                                                     seq_lens, block_table, args);
      break;
    default:
      VT_CHECK(false, "cuda dflash-paged-block-attn: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// AttentionDenseFast — a WARP-scoped variant of the AttentionKernel above with
// the IDENTICAL online-softmax recurrence (flash-style, f32 accumulation), for
// dense full/causal attention over small head_dim (the Qwen3-VL vision tower:
// head_dim 72, non-causal). One WARP (not a whole block) owns each (query,head):
// the head_dim reduction is a butterfly `__shfl_xor` (NO __syncthreads), the
// output accumulator lives in registers, and the softmax stats are tracked per
// lane. This removes the naive kernel's per-key block-reduction __syncthreads
// storm — nsys attributed ~99% of the ViT tower forward to that kernel (56 ms
// per block at 784 tokens). It is NOT bit-identical to AttentionKernel (the
// head_dim partial-sum grouping over 32 lanes differs from the block version),
// but is the same f32-online-softmax math within the tower's bf16 envelope.
// Registered as a SEPARATE op (kAttentionDenseFast) so kAttention — used by the
// text/audio paths — is byte-identical. Grounded in the online-softmax /
// FlashAttention recurrence we already ship (src/vt/cuda/flash_attn/).
template <typename Tin, typename Tout>
__global__ void AttentionWarpKernel(Tout* out, const Tin* query, const Tin* key, const Tin* value,
                                    int64_t hq, int64_t hk, int64_t d, int64_t t, float scale,
                                    bool causal) {
  constexpr int kMaxPerLane = 8;  // head_dim up to 256
  const int64_t i = blockIdx.x;   // query position
  const int64_t h = blockIdx.y;   // q-head
  const int64_t g = h / (hq / hk);
  const int lane = static_cast<int>(threadIdx.x);  // 0..31
  const int64_t jmax = causal ? i : t - 1;
  const int64_t qoff = (i * hq + h) * d;
  const int npl = static_cast<int>((d + 31) / 32);  // elements this lane owns
  float qreg[kMaxPerLane];
  float acc[kMaxPerLane];
#pragma unroll
  for (int k = 0; k < kMaxPerLane; ++k) {
    qreg[k] = 0.0f;
    acc[k] = 0.0f;
  }
  for (int k = 0; k < npl; ++k) {
    const int e = lane + 32 * k;
    if (e < d) qreg[k] = Load(query, qoff + e);
  }
  float m = -CUDART_INF_F, l = 0.0f;
  for (int64_t j = 0; j <= jmax; ++j) {
    const int64_t koff = (j * hk + g) * d;
    float part = 0.0f;
#pragma unroll
    for (int k = 0; k < kMaxPerLane; ++k) {
      const int e = lane + 32 * k;
      if (k < npl && e < d) part += qreg[k] * Load(key, koff + e);
    }
#pragma unroll
    for (int off = 16; off > 0; off >>= 1) part += __shfl_xor_sync(0xffffffffu, part, off);
    const float s = part * scale;
    const float m_new = fmaxf(m, s);
    const float corr = expf(m - m_new);  // 0 on the first key (m == -inf)
    const float p = expf(s - m_new);
    const int64_t voff = (j * hk + g) * d;
#pragma unroll
    for (int k = 0; k < kMaxPerLane; ++k) {
      const int e = lane + 32 * k;
      if (k < npl && e < d) acc[k] = acc[k] * corr + p * Load(value, voff + e);
    }
    l = l * corr + p;
    m = m_new;
  }
  const float inv = 1.0f / l;
  for (int k = 0; k < npl; ++k) {
    const int e = lane + 32 * k;
    if (e < d) Store(out, qoff + e, acc[k] * inv);
  }
}

template <typename Tin>
void LaunchAttentionWarp(cudaStream_t s, Tensor& out, const Tensor& query, const Tensor& key,
                         const Tensor& value, const AttentionArgs& args) {
  const int64_t t = query.shape[0], hq = query.shape[1], d = query.shape[2];
  const int64_t hk = key.shape[1];
  if (t == 0 || hq == 0 || d == 0) return;
  VT_CHECK(d <= 256, "cuda attention-dense-fast: head_dim <= 256 only");
  const dim3 grid(static_cast<unsigned>(t), static_cast<unsigned>(hq));
  switch (out.dtype) {
    case DType::kF32:
      AttentionWarpKernel<Tin, float><<<grid, 32, 0, s>>>(
          out.Ptr<float>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), hq, hk, d, t,
          args.scale, args.causal);
      break;
    case DType::kBF16:
      AttentionWarpKernel<Tin, __nv_bfloat16><<<grid, 32, 0, s>>>(
          out.Ptr<__nv_bfloat16>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), hq, hk,
          d, t, args.scale, args.causal);
      break;
    default: VT_CHECK(false, "cuda attention-dense-fast: unsupported out dtype");
  }
  Check(cudaGetLastError(), "attention-dense-fast launch");
}

void AttentionDenseFastKernelCuda(Queue& q, Tensor& out, const Tensor& query, const Tensor& key,
                                  const Tensor& value, const AttentionArgs& args) {
  switch (query.dtype) {
    case DType::kF32:
      LaunchAttentionWarp<float>(AsStream(q), out, query, key, value, args);
      break;
    case DType::kBF16:
      LaunchAttentionWarp<__nv_bfloat16>(AsStream(q), out, query, key, value, args);
      break;
    default:
      VT_CHECK(false, "cuda attention-dense-fast: unsupported input dtype (f32/bf16 only)");
  }
}

// AttentionDenseFlash — a SHARED-MEMORY-TILED flash variant of AttentionWarpKernel
// with the BIT-IDENTICAL per-warp online-softmax recurrence, structured so that a
// block of `kFlashBr` query-warps SHARES each streamed K/V tile out of shared
// memory instead of re-reading all K/V from global once PER (query,head). This is
// the classic FlashAttention K/V-tiling: K and V are streamed into shared memory in
// tiles of `kFlashBc` columns and reused across the whole query block, killing the
// O(t^2) redundant global K/V reads that make AttentionWarpKernel memory-bound
// (§13.5: ~21 ms/layer of K reads alone at 1500 frames, no cross-query tile reuse).
// Ported 1:1 in STRUCTURE from the vendored FlashAttention-2 forward kernel
// (src/vt/cuda/flash_attn/src/flash_fwd_kernel.h compute_attn_1rowblock :52 — sK/sV
// shared tiles :163-165, the `for (int n_block ...)` K/V-tile stream + online
// rescale), and cross-checked to vLLM's non-causal encoder attention dispatch
// (vllm/model_executor/models/whisper.py WhisperEncoderAttention:255). Unlike the
// cute/MMA FA2 (head_dim {128,192,256}, paged-KV), this is a scalar warp-per-query
// kernel for the DENSE single-request head_dim-64 encoder layout — so the ARITHMETIC
// is byte-for-byte the AttentionWarpKernel recurrence (same per-lane head_dim
// grouping lane+32k, same butterfly __shfl_xor, same sequential j-order 0..t-1, same
// f32 accumulation) with K/V bytes sourced from shared memory rather than global.
// Because every float op and its order are unchanged, the output is BIT-IDENTICAL to
// AttentionDenseFast ⇒ token-identical by construction; kAttention (text decode) is
// untouched. One q-head per CTA (all warps share the same GQA kv-head g).
constexpr int kFlashBr = 16;  // query-warps per CTA (= K/V global-read reuse factor)
constexpr int kFlashBc = 64;  // key/value columns streamed per shared-memory tile

template <typename Tin, typename Tout>
__global__ void AttentionDenseFlashKernel(Tout* out, const Tin* query, const Tin* key,
                                          const Tin* value, int64_t hq, int64_t hk, int64_t d,
                                          int64_t t, float scale, bool causal) {
  constexpr int kMaxPerLane = 8;  // head_dim up to 256
  extern __shared__ __align__(16) char flash_smem[];
  Tin* sK = reinterpret_cast<Tin*>(flash_smem);
  Tin* sV = sK + static_cast<int64_t>(kFlashBc) * d;

  const int warp = static_cast<int>(threadIdx.x >> 5);
  const int lane = static_cast<int>(threadIdx.x & 31);
  const int nthreads = kFlashBr * 32;
  const int64_t h = blockIdx.y;             // q-head (one per CTA)
  const int64_t g = h / (hq / hk);          // shared kv-head for every warp in CTA
  const int64_t qi = static_cast<int64_t>(blockIdx.x) * kFlashBr + warp;  // this warp's query
  const bool active = qi < t;
  const int npl = static_cast<int>((d + 31) / 32);  // head_dim elements this lane owns

  // This warp's query row, in registers (identical layout to AttentionWarpKernel).
  float qreg[kMaxPerLane];
  float acc[kMaxPerLane];
#pragma unroll
  for (int k = 0; k < kMaxPerLane; ++k) {
    qreg[k] = 0.0f;
    acc[k] = 0.0f;
  }
  if (active) {
    const int64_t qoff = (qi * hq + h) * d;
    for (int k = 0; k < npl; ++k) {
      const int e = lane + 32 * k;
      if (e < d) qreg[k] = Load(query, qoff + e);
    }
  }
  float m = -CUDART_INF_F, l = 0.0f;

  // Non-causal: every warp scans keys [0, t). Causal: up to the max query in this
  // block; each warp then stops at its own qi (same 0..jmax order as the warp kernel).
  const int64_t block_qmax =
      min(static_cast<int64_t>(blockIdx.x) * kFlashBr + (kFlashBr - 1), t - 1);
  const int64_t key_end = causal ? (block_qmax + 1) : t;

  for (int64_t c0 = 0; c0 < key_end; c0 += kFlashBc) {
    const int tile = static_cast<int>(min(static_cast<int64_t>(kFlashBc), key_end - c0));
    // Cooperative load of this K/V tile into shared memory (all warps participate).
    __syncthreads();
    for (int idx = static_cast<int>(threadIdx.x); idx < tile * static_cast<int>(d);
         idx += nthreads) {
      const int jj = idx / static_cast<int>(d);
      const int e = idx % static_cast<int>(d);
      const int64_t off = ((c0 + jj) * hk + g) * d + e;
      sK[idx] = key[off];
      sV[idx] = value[off];
    }
    __syncthreads();
    if (!active) continue;
    // This warp's online-softmax update over the tile — bit-identical math to
    // AttentionWarpKernel, K/V now read from shared memory.
    const int64_t jstop = causal ? min(static_cast<int64_t>(tile), qi - c0 + 1)
                                 : static_cast<int64_t>(tile);
    for (int64_t j = 0; j < jstop; ++j) {
      const int64_t base = j * d;
      float part = 0.0f;
#pragma unroll
      for (int k = 0; k < kMaxPerLane; ++k) {
        const int e = lane + 32 * k;
        if (k < npl && e < d) part += qreg[k] * Load(sK, base + e);
      }
#pragma unroll
      for (int off = 16; off > 0; off >>= 1) part += __shfl_xor_sync(0xffffffffu, part, off);
      const float s = part * scale;
      const float m_new = fmaxf(m, s);
      const float corr = expf(m - m_new);
      const float p = expf(s - m_new);
#pragma unroll
      for (int k = 0; k < kMaxPerLane; ++k) {
        const int e = lane + 32 * k;
        if (k < npl && e < d) acc[k] = acc[k] * corr + p * Load(sV, base + e);
      }
      l = l * corr + p;
      m = m_new;
    }
  }
  if (!active) return;
  const int64_t qoff = (qi * hq + h) * d;
  const float inv = 1.0f / l;
  for (int k = 0; k < npl; ++k) {
    const int e = lane + 32 * k;
    if (e < d) Store(out, qoff + e, acc[k] * inv);
  }
}

template <typename Tin>
void LaunchAttentionDenseFlash(cudaStream_t s, Tensor& out, const Tensor& query, const Tensor& key,
                               const Tensor& value, const AttentionArgs& args) {
  const int64_t t = query.shape[0], hq = query.shape[1], d = query.shape[2];
  const int64_t hk = key.shape[1];
  if (t == 0 || hq == 0 || d == 0) return;
  VT_CHECK(d <= 256, "cuda attention-dense-flash: head_dim <= 256 only");
  const unsigned nblk = static_cast<unsigned>((t + kFlashBr - 1) / kFlashBr);
  const dim3 grid(nblk, static_cast<unsigned>(hq));
  const size_t shmem = static_cast<size_t>(2) * kFlashBc * d * sizeof(Tin);  // sK + sV
  switch (out.dtype) {
    case DType::kF32:
      AttentionDenseFlashKernel<Tin, float><<<grid, kFlashBr * 32, shmem, s>>>(
          out.Ptr<float>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), hq, hk, d, t,
          args.scale, args.causal);
      break;
    case DType::kBF16:
      AttentionDenseFlashKernel<Tin, __nv_bfloat16><<<grid, kFlashBr * 32, shmem, s>>>(
          out.Ptr<__nv_bfloat16>(), query.Ptr<Tin>(), key.Ptr<Tin>(), value.Ptr<Tin>(), hq, hk, d,
          t, args.scale, args.causal);
      break;
    default: VT_CHECK(false, "cuda attention-dense-flash: unsupported out dtype");
  }
  Check(cudaGetLastError(), "attention-dense-flash launch");
}

void AttentionDenseFlashKernelCuda(Queue& q, Tensor& out, const Tensor& query, const Tensor& key,
                                   const Tensor& value, const AttentionArgs& args) {
  switch (query.dtype) {
    case DType::kF32:
      LaunchAttentionDenseFlash<float>(AsStream(q), out, query, key, value, args);
      break;
    case DType::kBF16:
      LaunchAttentionDenseFlash<__nv_bfloat16>(AsStream(q), out, query, key, value, args);
      break;
    default:
      VT_CHECK(false, "cuda attention-dense-flash: unsupported input dtype (f32/bf16 only)");
  }
}

// ---------------------------------------------------------------------------
// fused_chain (TDR): the Tier-1 single-pass INTERPRETER over the canonical
// (out, x, weight, residual) 4-operand shape. The Tier-0 composite is device-
// agnostic (ops.cpp), dispatching each opcode to the standalone vt:: op. This
// kernel is reached ONLY for Tier-1-able recipes (steps in {kAdd,kMul,kSilu,
// kSigmoid,kRmsNorm}) when VT_FUSED_TIER=1; the general FusedChain wrapper
// resolves the canonical operand order [0=x,1=weight,2=residual,3=out] and
// forwards here. One block per row, shared-mem tree reduction (the
// RmsNormRowKernel skeleton). Bit-for-bit equal to the CUDA RmsNorm(residual)
// golden: same f32 tree reduction, gemma (1+w), Tres-rounded residual add.
// Operand indices resolve to typed pointers: 0->Tin x, 1->Tin w, 2->Tres res,
// 3->Tout out. The recipe POD is passed BY VALUE into the kernel (small, no heap).

template <typename Tin, typename Tout, typename Tres>
struct FusedCtx {
  const Tin* x;
  const Tin* w;
  Tres* res;
  Tout* out;
  int64_t h;
};

template <typename Tin, typename Tout, typename Tres>
__device__ inline float FusedLoadDev(uint8_t idx, const FusedCtx<Tin, Tout, Tres>& c, int64_t row,
                                     int64_t j) {
  switch (idx) {
    case 0: return Load(c.x, row * c.h + j);
    case 1: return Load(c.w, j);
    case 2: return Load(c.res, row * c.h + j);
    case 3: return Load(c.out, row * c.h + j);
    default: return 0.0f;  // read-only / out-of-range (validated host-side)
  }
}

template <typename Tin, typename Tout, typename Tres>
__device__ inline void FusedStoreDev(uint8_t idx, const FusedCtx<Tin, Tout, Tres>& c, int64_t row,
                                     int64_t j, float v) {
  switch (idx) {
    case 2: Store(c.res, row * c.h + j, v); break;
    case 3: Store(c.out, row * c.h + j, v); break;
    default: break;  // read-only operands (validated host-side)
  }
}

__device__ inline float FSigmoidDev(float x) { return 1.0f / (1.0f + expf(-x)); }

// Tier 1 — interpreter: one block per row walks the whole recipe.
template <typename Tin, typename Tout, typename Tres>
__global__ void FusedChainInterpKernel(FusedCtx<Tin, Tout, Tres> c, float eps, FusedRecipe r) {
  const int64_t row = blockIdx.x;
  const int64_t h = c.h;
  __shared__ float partial[kBlock];
  for (int s = 0; s < r.n; ++s) {
    const FStep st = r.steps[s];
    if (st.op == FOp::kRmsNorm) {
      float acc = 0.0f;
      for (int64_t j = threadIdx.x; j < h; j += kBlock) {
        const float v = FusedLoadDev(st.in[0], c, row, j);
        acc += v * v;  // kMeanSquare, f32
      }
      partial[threadIdx.x] = acc;
      __syncthreads();
      for (int stride = kBlock / 2; stride > 0; stride /= 2) {
        if (static_cast<int>(threadIdx.x) < stride)
          partial[threadIdx.x] += partial[threadIdx.x + stride];
        __syncthreads();
      }
      const float inv = 1.0f / sqrtf(partial[0] / static_cast<float>(h) + eps);
      for (int64_t j = threadIdx.x; j < h; j += kBlock) {
        const float v = FusedLoadDev(st.in[0], c, row, j);
        float wj = FusedLoadDev(st.in[1], c, row, j);
        if (st.gemma) wj += 1.0f;
        FusedStoreDev(st.out, c, row, j, v * inv * wj);
      }
      __syncthreads();
    } else if (st.op == FOp::kSilu || st.op == FOp::kSigmoid) {  // unary elementwise
      for (int64_t j = threadIdx.x; j < h; j += kBlock) {
        const float a = FusedLoadDev(st.in[0], c, row, j);
        FusedStoreDev(st.out, c, row, j, st.op == FOp::kSilu ? a * FSigmoidDev(a) : FSigmoidDev(a));
      }
      __syncthreads();
    } else {  // kAdd / kMul — binary elementwise over the row
      for (int64_t j = threadIdx.x; j < h; j += kBlock) {
        const float a = FusedLoadDev(st.in[0], c, row, j);
        const float b = FusedLoadDev(st.in[1], c, row, j);
        FusedStoreDev(st.out, c, row, j, st.op == FOp::kAdd ? a + b : a * b);
      }
      __syncthreads();  // writes (e.g. residual) visible before the next step reads
    }
  }
}

template <typename Tin, typename Tout, typename Tres>
void LaunchFusedInterp(cudaStream_t s, Tensor& out, const Tensor& x, const Tensor& weight,
                       Tensor* residual, const FusedRecipe& r, float eps) {
  const int64_t t = x.shape[0], h = x.shape[1];
  if (t == 0 || h == 0) return;
  FusedCtx<Tin, Tout, Tres> c{x.Ptr<Tin>(), weight.Ptr<Tin>(),
                              residual == nullptr ? nullptr : residual->Ptr<Tres>(),
                              out.Ptr<Tout>(), h};
  FusedChainInterpKernel<Tin, Tout, Tres><<<static_cast<unsigned>(t), kBlock, 0, s>>>(c, eps, r);
  Check(cudaGetLastError(), "fused_chain interp launch");
}

// Resolve the (Tin, Tout, Tres) triple and launch the interpreter. Tres follows
// the residual dtype (f32 or bf16); with no residual it is unused (pass float).
template <typename Tin, typename Tout>
void LaunchFusedInterpRes(cudaStream_t s, Tensor& out, const Tensor& x, const Tensor& weight,
                          Tensor* residual, const FusedRecipe& r, float eps) {
  if (residual != nullptr && residual->dtype == DType::kBF16) {
    LaunchFusedInterp<Tin, Tout, __nv_bfloat16>(s, out, x, weight, residual, r, eps);
  } else {
    LaunchFusedInterp<Tin, Tout, float>(s, out, x, weight, residual, r, eps);
  }
}

// Registered kernel: the Tier-1 interpreter (the general wrapper only dispatches
// here for Tier-1-able recipes; Tier-0 composite is device-agnostic in ops.cpp).
void FusedChainKernelCuda(Queue& q, Tensor& out, const Tensor& x, const Tensor& weight,
                          Tensor* residual, const FusedRecipe& r, float eps) {
  VT_CHECK(weight.dtype == x.dtype, "cuda fused_chain: weight dtype must match x");
  cudaStream_t s = AsStream(q);
  switch (x.dtype) {
    case DType::kF32:
      switch (out.dtype) {
        case DType::kF32:
          LaunchFusedInterpRes<float, float>(s, out, x, weight, residual, r, eps);
          break;
        case DType::kBF16:
          LaunchFusedInterpRes<float, __nv_bfloat16>(s, out, x, weight, residual, r, eps);
          break;
        default: VT_CHECK(false, "cuda fused_chain: unsupported out dtype");
      }
      break;
    case DType::kBF16:
      switch (out.dtype) {
        case DType::kF32:
          LaunchFusedInterpRes<__nv_bfloat16, float>(s, out, x, weight, residual, r, eps);
          break;
        case DType::kBF16:
          LaunchFusedInterpRes<__nv_bfloat16, __nv_bfloat16>(s, out, x, weight, residual, r, eps);
          break;
        default: VT_CHECK(false, "cuda fused_chain: unsupported out dtype");
      }
      break;
    default: VT_CHECK(false, "cuda fused_chain: unsupported input dtype (f32/bf16 only)");
  }
}

// Registers the CUDA kernels during static init (pre-main, like the CPU ops).
// Filling the op table is harmless on machines without a GPU: the kCUDA
// backend never registers there, so no CUDA queue can exist to dispatch with.
struct Registrar {
  Registrar() {
    RegisterOp(OpId::kRmsNorm, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<RmsNormFn>(&RmsNormKernelCuda)));
    RegisterOp(OpId::kRmsNormQuantFp8, DeviceType::kCUDA,
               reinterpret_cast<void*>(
                   static_cast<RmsNormQuantFp8Fn>(&RmsNormQuantFp8KernelCuda)));
    RegisterOp(OpId::kSiluAndMul, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<SiluAndMulFn>(&SiluAndMulKernelCuda)));
    RegisterOp(OpId::kGeluAndMul, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<GeluAndMulFn>(&GeluAndMulKernelCuda)));
    RegisterOp(OpId::kMulScalar, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<MulScalarFn>(&MulScalarKernelCuda)));
    RegisterOp(OpId::kSoftCap, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<SoftCapFn>(&SoftCapKernelCuda)));
    RegisterOp(OpId::kEmbedding, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<EmbeddingFn>(&EmbeddingKernelCuda)));
    RegisterOp(OpId::kRopeNeox, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<RopeFn>(&RopeNeoxKernelCuda)));
    RegisterOp(
        OpId::kRopeFromCache, DeviceType::kCUDA,
        reinterpret_cast<void*>(
            static_cast<RopeFromCacheFn>(&RopeFromCacheKernelCuda)));
    RegisterOp(
        OpId::kFusedNormRope, DeviceType::kCUDA,
        reinterpret_cast<void*>(static_cast<FusedNormRopeFn>(&FusedNormRopeKernelCuda)));
    RegisterOp(
        OpId::kRopeCosSinCache, DeviceType::kCUDA,
        reinterpret_cast<void*>(static_cast<RopeCosSinCacheFn>(&RopeCosSinCacheKernelCuda)));
    RegisterOp(OpId::kAttnQkNormRopeGate, DeviceType::kCUDA,
               reinterpret_cast<void*>(
                   static_cast<AttnQkNormRopeGateFn>(&AttnQkNormRopeGateKernelCuda)));
    RegisterOp(OpId::kAttention, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<AttentionFn>(&AttentionKernelCuda)));
    RegisterOp(OpId::kAttentionDenseFast, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<AttentionFn>(&AttentionDenseFastKernelCuda)));
    RegisterOp(OpId::kAttentionDenseFlash, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<AttentionFn>(&AttentionDenseFlashKernelCuda)));
    RegisterOp(OpId::kDFlashBlockAttention, DeviceType::kCUDA,
               reinterpret_cast<void*>(
                   static_cast<DFlashBlockAttentionFn>(&DFlashBlockAttentionKernelCuda)));
    RegisterOp(OpId::kDFlashPagedBlockAttention, DeviceType::kCUDA,
               reinterpret_cast<void*>(
                   static_cast<DFlashPagedBlockAttentionFn>(&DFlashPagedBlockAttentionKernelCuda)));
    RegisterOp(OpId::kFusedChain, DeviceType::kCUDA,
               reinterpret_cast<void*>(static_cast<FusedChainFn>(&FusedChainKernelCuda)));
  }
} registrar;

}  // namespace
}  // namespace vt::cuda