cpu_ops.cpp

// vllm.cpp original (vt runtime, inventory deviation §9.1); no upstream mirror.
// EXCEPT the parallel dispatch (QUANT-GGUF-CPU-THREADPOOL): the GEMM chunk
// policy and the row/batch chunking are ported 1:1 from llama.cpp (local fork)
// ggml/src/ggml-cpu/ggml-cpu.c:1155-1443 and ggml-cpu/ops.cpp:9070-9126 @
// 237ad9b96 — see cpu_threadpool.h and the per-kernel anchors below. Every
// kernel keeps its exact per-element math and per-output sequential reduction
// order; parallelism partitions OUTPUT elements only, so results are
// bit-identical to single-thread by construction (spec § Dispatch behavior).
#include "vt/ops.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <vector>

#include "cpu_threadpool.h"

namespace vt::cpu {
namespace {

// Row/batch-chunked dispatch through the process pool (or the test-swapped
// pool). body(r0, r1) produces output rows [r0, r1) exactly once each.
inline void ForRows(int64_t nr, const std::function<void(int64_t, int64_t)>& body) {
  ParallelForRows(CurrentThreadpool(), nr, body);
}

float LoadF32(const Tensor& t, int64_t elem_offset) {
  switch (t.dtype) {
    case DType::kF32: return t.Ptr<float>()[elem_offset];
    case DType::kF16: return F16ToF32(t.Ptr<uint16_t>()[elem_offset]);
    case DType::kBF16: return BF16ToF32(t.Ptr<uint16_t>()[elem_offset]);
    default: VT_CHECK(false, "LoadF32: unsupported dtype"); return 0.0f;
  }
}

// Mirror of LoadF32 for outputs: reduced-width formats round to storage dtype.
void StoreF32(const Tensor& t, int64_t elem_offset, float v) {
  switch (t.dtype) {
    case DType::kF32: t.Ptr<float>()[elem_offset] = v; break;
    case DType::kF16: t.Ptr<uint16_t>()[elem_offset] = F32ToF16(v); break;
    case DType::kBF16: t.Ptr<uint16_t>()[elem_offset] = F32ToBF16(v); break;
    default: VT_CHECK(false, "StoreF32: unsupported dtype");
  }
}

// GEMM chunk worker — 16x16 block tiling inside a chunk, ported from
// ggml_compute_forward_mul_mat_one_chunk (ggml-cpu.c:1155-1243; empty-chunk
// yield :1181-1184, blck_0/blck_1 = 16 :1192-1194). ggml's vec_dot per output
// element is our per-element K loop: byte-identical accumulation (sequential
// over K, f32, -ffp-contract=off pinned) to the pre-threadpool kernels.
// ir0 indexes output COLUMNS j (ggml nr0 = src0/weight rows = N), ir1 indexes
// output ROWS i (ggml nr1 = src1 rows = M). kBT selects the [N,K] row-major
// weight orientation (MatmulBT) vs [K,N] (Matmul).
template <bool kBT>
void MatmulOneChunk(Tensor& out, const Tensor& a, const Tensor& b, int64_t k, int64_t n,
                    int64_t ir0_start, int64_t ir0_end, int64_t ir1_start, int64_t ir1_end) {
  // MLA campaign W6: the activation may be ROW-STRIDED (a column slice of a
  // wider buffer — see vt::MatmulBT). For a contiguous activation `a_rs == k`,
  // so the offsets are integer-identical to the pre-W6 `i * k + p` form and
  // every existing model is bit-identical by construction.
  const int64_t a_rs = a.stride[0];
  // threads with no work simply yield
  if (ir0_start >= ir0_end || ir1_start >= ir1_end) {
    return;
  }

  // block-tiling attempt
  const int64_t blck_0 = 16;
  const int64_t blck_1 = 16;

  for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) {
    for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) {
      for (int64_t i = iir1; i < iir1 + blck_1 && i < ir1_end; ++i) {
        for (int64_t j = iir0; j < iir0 + blck_0 && j < ir0_end; ++j) {
          float acc = 0.0f;
          for (int64_t p = 0; p < k; ++p) {
            acc += LoadF32(a, i * a_rs + p) * LoadF32(b, kBT ? j * k + p : p * n + j);
          }
          StoreF32(out, i * n + j, acc);
        }
      }
    }
  }
}

// GEMM chunking policy + atomic work stealing, ported from
// ggml_compute_forward_mul_mat (ggml-cpu.c:1245-1443): thread 0 seeds the
// steal cursor at nth and a barrier publishes it (:1350-1355); chunk_size 16,
// 64 for vector shapes (:1388-1393); nchunk0 x nchunk1 grid (:1398-1399);
// re-chunk per-thread when the grid is < nth*4 or NUMA (:1404-1408, IsNuma()
// stubbed false); each thread starts at chunk ith then steals via the atomic
// cursor (:1415-1442). num_rows_per_vec_dot is 1 for our scalar dot (no mmla).
template <bool kBT>
void MatmulChunked(Tensor& out, const Tensor& a, const Tensor& b) {
  const int64_t m = a.shape[0], k = a.shape[1];
  const int64_t n = kBT ? b.shape[0] : b.shape[1];
  // ggml nr0 = ne0 (dst dim0 = weight rows) -> our N; nr1 = ne1*ne2*ne3
  // (src1 rows) -> our M.
  const int64_t nr0 = n;
  const int64_t nr1 = m;

  Threadpool& tp = CurrentThreadpool();
  tp.Run([&](int ith, int nth) {
    if (ith == 0) {
      // Every thread starts at ith, so the first unprocessed chunk is nth.
      tp.ChunkSet(nth);
    }

    tp.Barrier();

    // Now select a reasonable chunk size.
    int chunk_size = 16;

    // We need to step up the size if it's small
    if (nr0 == 1 || nr1 == 1) {
      chunk_size = 64;
    }

    // distribute the work across the inner or outer loop based on which one is larger
    int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size;
    int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size;

    // If the chunking is poor for the number of threads on this setup, scrap
    // the whole plan. Re-chunk it by thread.
    if (nchunk0 * nchunk1 < nth * 4 || IsNuma()) {
      nchunk0 = nr0 > nr1 ? nth : 1;  // parallelize by weight rows (N)
      nchunk1 = nr0 > nr1 ? 1 : nth;  // parallelize by src1 rows (M)
    }

    // The number of elements in each chunk
    const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0;
    const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1;

    // The first chunk comes from our thread_id, the rest will get auto-assigned.
    int64_t current_chunk = ith;

    while (current_chunk < nchunk0 * nchunk1) {
      const int64_t ith0 = current_chunk % nchunk0;
      const int64_t ith1 = current_chunk / nchunk0;

      const int64_t ir0_start = dr0 * ith0;
      const int64_t ir0_end = std::min(ir0_start + dr0, nr0);

      const int64_t ir1_start = dr1 * ith1;
      const int64_t ir1_end = std::min(ir1_start + dr1, nr1);

      MatmulOneChunk<kBT>(out, a, b, k, n, ir0_start, ir0_end, ir1_start, ir1_end);

      if (nth >= nchunk0 * nchunk1) {
        break;
      }

      current_chunk = tp.ChunkAdd(1);
    }
  });
}

void MatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) {
  MatmulChunked<false>(out, a, b);
}

// b is the torch Linear weight [N,K] row-major (see vt::MatmulBT); identical
// accumulation order to MatmulKernel (sequential over K), so on CPU the two
// orientations are bit-identical for the same logical weight.
void MatmulBTKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) {
  MatmulChunked<true>(out, a, b);
}

// vt::BatchedMatmul (`torch.bmm`) CPU reference — out[G,M,N] = a[G,M,K] @
// b[G,K,N]. Sequential f32 accumulation over K, exactly like MatmulOneChunk's
// per-element dot, so the reference and the cuBLASLt CUDA path share the same
// numeric contract (f32 accumulate, round on store). Every operand is addressed
// through its STRIDES: the MLA absorption call sites (mla_attention.py:789,
// :1034) pass transposed views whose batch axis is not the outermost storage
// axis. Parallelized over the flattened (batch, row) output space, which leaves
// each output element's K reduction on one thread — bit-identical to a serial
// run and run-to-run reproducible.
void BatchedMatmulKernel(Queue&, Tensor& out, const Tensor& a, const Tensor& b) {
  const int64_t g = out.shape[0], m = out.shape[1], n = out.shape[2];
  const int64_t k = a.shape[2];
  const int64_t rows = g * m;
  if (rows == 0 || n == 0) return;
  ForRows(rows, [&](int64_t r0, int64_t r1) {
    for (int64_t r = r0; r < r1; ++r) {
      const int64_t bi = r / m, i = r % m;
      const int64_t a_row = bi * a.stride[0] + i * a.stride[1];
      const int64_t b_base = bi * b.stride[0];
      const int64_t o_row = bi * out.stride[0] + i * out.stride[1];
      for (int64_t j = 0; j < n; ++j) {
        float acc = 0.0f;
        for (int64_t p = 0; p < k; ++p) {
          acc += LoadF32(a, a_row + p) * LoadF32(b, b_base + p * b.stride[1] + j);
        }
        StoreF32(out, o_row + j, acc);
      }
    }
  });
}

// vt::ConcatMlaNopeRope CPU reference — the scalar, width-generic form of
// upstream's `ConcatMLAQKernel` (csrc/libtorch_stable/concat_mla_q.cuh) and of
// the `_concat_k_nope_k_pe` slice assignments (mla_attention.py:2085-2090).
// Pure copy: no arithmetic, so it is exact for every dtype. `rope.shape[1] == 1`
// with more output heads is the broadcast form the prefill K concat needs.
void ConcatMlaNopeRopeKernel(Queue&, Tensor& out, const Tensor& nope, const Tensor& rope) {
  const int64_t tokens = out.shape[0], heads = out.shape[1];
  const int64_t dn = nope.shape[2], dr = rope.shape[2];
  const bool rope_broadcast = rope.shape[1] == 1 && heads > 1;
  ForRows(tokens, [&](int64_t t0, int64_t t1) {
    for (int64_t t = t0; t < t1; ++t) {
      for (int64_t h = 0; h < heads; ++h) {
        const int64_t o = t * out.stride[0] + h * out.stride[1];
        const int64_t n = t * nope.stride[0] + h * nope.stride[1];
        const int64_t r =
            t * rope.stride[0] + (rope_broadcast ? 0 : h) * rope.stride[1];
        for (int64_t d = 0; d < dn; ++d) StoreF32(out, o + d, LoadF32(nope, n + d));
        for (int64_t d = 0; d < dr; ++d) StoreF32(out, o + dn + d, LoadF32(rope, r + d));
      }
    }
  });
}

void RmsNormKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& w,
                   const RmsNormArgs& args, Tensor* residual) {
  const int64_t t = x.shape[0], h = x.shape[1];
  // Row-chunked over tokens (ops.cpp:9070-9126 pattern); each row's f32
  // variance reduction stays sequential on one thread — bit-identical.
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const int64_t rbase = i * h;
    float sumsq = 0.0f;
    for (int64_t j = 0; j < h; ++j) {
      float v = LoadF32(x, i * h + j);
      if (residual) {
        v += LoadF32(*residual, rbase + j);   // add in f32
        StoreF32(*residual, rbase + j, v);     // new residual stream (rounds to its dtype)
        v = LoadF32(*residual, rbase + j);     // re-read rounded value (bf16-faithful)
      }
      sumsq += v * v;
    }
    float inv = 1.0f / std::sqrt(sumsq / static_cast<float>(h) + args.eps);
    for (int64_t j = 0; j < h; ++j) {
      float v = residual ? LoadF32(*residual, rbase + j) : LoadF32(x, i * h + j);
      float wj = LoadF32(w, j);
      if (args.gemma) wj += 1.0f;
      StoreF32(out, i * h + j, v * inv * wj);
    }
  }
  });
}

void SiluAndMulKernel(Queue&, Tensor& out, const Tensor& x) {
  const int64_t t = x.shape[0], d = x.shape[1] / 2;
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    for (int64_t j = 0; j < d; ++j) {
      float gate = LoadF32(x, i * 2 * d + j);
      float up = LoadF32(x, i * 2 * d + d + j);
      float silu = gate / (1.0f + std::exp(-gate));
      StoreF32(out, i * d + j, silu * up);
    }
  }
  });
}

void MoeSiluMulKernel(Queue&, Tensor& out, const Tensor& gate, const Tensor& up) {
  const int64_t n = out.Numel();
  // Elementwise: partition the flat output range.
  ForRows(n, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const float g = LoadF32(gate, i);
    const float silu = g / (1.0f + std::exp(-g));
    StoreF32(out, i, silu * LoadF32(up, i));
  }
  });
}

// --- TRUE W4A4 (fp4xfp4) helpers + kernels (notes §7). Self-contained fp8/fp4
// codec (vt does not depend on vllm), bit-matching vllm::F8E4M3ToF32 /
// F32ToF8E4M3 / CastToFp4 / kE2M1Lut so the op equals vllm::RunNvfp4Emulation.
constexpr float kFp4Max = 6.0F;    // E2M1 max magnitude
constexpr float kFp8Max = 448.0F;  // fp8-e4m3fn max finite
constexpr float kE2M1[8] = {0.0F, 0.5F, 1.0F, 1.5F, 2.0F, 3.0F, 4.0F, 6.0F};

inline float ClampF(float x, float lo, float hi) { return x < lo ? lo : (x > hi ? hi : x); }
inline float RecipF(float x) { return x == 0.0F ? 0.0F : 1.0F / x; }

// IEEE fp8-e4m3fn byte -> f32 (bit-matches vllm::F8E4M3ToF32).
float Fp8ToF32(uint8_t byte) {
  const uint32_t sign = static_cast<uint32_t>(byte >> 7) & 0x1U;
  const uint32_t exp = static_cast<uint32_t>(byte >> 3) & 0xFU;
  const uint32_t mant = static_cast<uint32_t>(byte) & 0x7U;
  const float sm = sign ? -1.0F : 1.0F;
  if (exp == 0xFU && mant == 0x7U) return std::numeric_limits<float>::quiet_NaN();
  if (exp == 0U) return sm * (static_cast<float>(mant) * (1.0F / 512.0F));
  const float mantissa = 1.0F + static_cast<float>(mant) * (1.0F / 8.0F);
  return sm * std::ldexp(mantissa, static_cast<int>(exp) - 7);
}

// f32 -> fp8-e4m3fn byte, round-to-nearest-even saturating (bit-matches
// vllm::F32ToF8E4M3).
uint8_t F32ToFp8(float f) {
  if (std::isnan(f)) return 0x7FU;
  const uint8_t sign = std::signbit(f) ? 0x80U : 0x00U;
  const float a = std::fabs(f);
  if (!std::isfinite(a) || a >= kFp8Max) return static_cast<uint8_t>(sign | 0x7EU);
  if (a == 0.0F) return sign;
  int e2 = 0;
  const float frac = std::frexp(a, &e2);
  int exp_field = (e2 - 1) + 7;
  if (exp_field <= 0) {
    const double qd = static_cast<double>(a) * 512.0;
    const int qi = static_cast<int>(std::nearbyint(qd));
    if (qi <= 0) return sign;
    if (qi < 8) return static_cast<uint8_t>(sign | static_cast<uint8_t>(qi));
    return static_cast<uint8_t>(sign | (1U << 3));
  }
  const double sig = static_cast<double>(frac) * 2.0;
  int mi = static_cast<int>(std::nearbyint(sig * 8.0));
  if (mi == 16) {
    mi = 8;
    exp_field += 1;
  }
  const int mant = mi - 8;
  if (exp_field > 15 || (exp_field == 15 && mant >= 7)) {
    return static_cast<uint8_t>(sign | 0x7EU);
  }
  return static_cast<uint8_t>(sign | (static_cast<uint8_t>(exp_field) << 3) |
                              static_cast<uint8_t>(mant));
}

// Fused fp8 RMSNorm -> static per-tensor quant (mirror vLLM Inductor
// fused_add_rms_norm_static_fp8_quant, rms_quant_fusion.py:124). Same reduction
// order as RmsNormKernel; the fp8 is taken from the SAME bf16-rounded normed value
// the split RmsNorm(bf16)+QuantFp8Static path quantizes (bf16-intermediate form),
// so the two are bit-identical. out_bf16 (optional) is the normed activation in
// bf16 (for a coexisting bf16 consumer, e.g. GDN in_proj_a/b). CUDA has the hot
// path; this CPU kernel keeps the op available on the host backend.
void RmsNormQuantFp8Kernel(Queue&, 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];
  const float inv_scale = 1.0F / input_scale;
  uint8_t* op = out_fp8.Ptr<uint8_t>();
  uint16_t* bp = out_bf16 == nullptr ? nullptr : out_bf16->Ptr<uint16_t>();
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const int64_t rbase = i * h;
    float sumsq = 0.0F;
    for (int64_t j = 0; j < h; ++j) {
      float v = LoadF32(x, rbase + j);
      if (residual) {
        v += LoadF32(*residual, rbase + j);   // add in f32
        StoreF32(*residual, rbase + j, v);     // new residual stream (rounds to its dtype)
        v = LoadF32(*residual, rbase + j);     // re-read rounded value (bf16-faithful)
      }
      sumsq += v * v;
    }
    const float inv = 1.0F / std::sqrt(sumsq / static_cast<float>(h) + args.eps);
    for (int64_t j = 0; j < h; ++j) {
      float v = residual ? LoadF32(*residual, rbase + j) : LoadF32(x, rbase + j);
      float wj = LoadF32(w, j);
      if (args.gemma) wj += 1.0F;
      // bf16-intermediate: round the normed value to bf16 (as RmsNorm's bf16 store),
      // then quant from that bf16 (as QuantFp8Static's bf16 load).
      const uint16_t nb = F32ToBF16(v * inv * wj);
      if (bp) bp[rbase + j] = nb;
      op[rbase + j] = F32ToFp8(BF16ToF32(nb) * inv_scale);
    }
  }
  });
}

// f32 -> E2M1 nibble (bit-matches vllm::CastToFp4 + Fp4ToNibble). Input pre-scaled.
uint8_t F32ToFp4Nibble(float x) {
  const float a = std::fabs(x);
  uint8_t idx = 7;  // 6.0
  if (a <= 0.25F) idx = 0;
  else if (a < 0.75F) idx = 1;
  else if (a <= 1.25F) idx = 2;
  else if (a < 1.75F) idx = 3;
  else if (a <= 2.5F) idx = 4;
  else if (a < 3.5F) idx = 5;
  else if (a <= 5.0F) idx = 6;
  if (idx == 0) return 0;
  return static_cast<uint8_t>((x < 0.0F ? 0x8U : 0x0U) | idx);
}

inline float Nibble(uint8_t nib) {
  return kE2M1[nib & 0x7U] * ((nib & 0x8U) ? -1.0F : 1.0F);
}

// ScaledFp4Quant CPU kernel: x [M,K] float -> out_packed [M,K/2] i8 + out_scale
// [M,K/16] i8. Per-token, per-16-group; equals vllm::RefScaledFp4Quant.
void ScaledFp4QuantKernel(Queue&, Tensor& out_packed, Tensor& out_scale, const Tensor& x,
                          float input_global_scale_inv,
                          Fp4ScaleLayout scale_layout) {
  const int64_t m = x.shape[0], k = x.shape[1];
  constexpr int kBS = 16;
  const int64_t groups = k / kBS;
  const float gs_recip = RecipF(input_global_scale_inv);
  auto* packed = out_packed.Ptr<uint8_t>();
  auto* scale = out_scale.Ptr<uint8_t>();
  const int64_t scale_cols = out_scale.shape[1];
  if (scale_layout == Fp4ScaleLayout::kCutlassSwizzled) {
    std::fill_n(scale, out_scale.Numel(), uint8_t{0});
  }
  ForRows(m, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    for (int64_t g = 0; g < groups; ++g) {
      const int64_t base = g * kBS;
      float vec_max = 0.0F;
      for (int j = 0; j < kBS; ++j) vec_max = std::fmax(vec_max, std::fabs(LoadF32(x, i * k + base + j)));
      float sc = ClampF(input_global_scale_inv * (vec_max * (1.0F / kFp4Max)), -kFp8Max, kFp8Max);
      const uint8_t sc_f8 = F32ToFp8(sc);
      if (scale_layout == Fp4ScaleLayout::kLinear) {
        scale[i * groups + g] = sc_f8;
      } else {
        const int64_t m_tile = i / 128;
        const int64_t outer_m = i % 32;
        const int64_t inner_m = (i % 128) / 32;
        const int64_t k_tile = g / 4;
        const int64_t inner_k = g % 4;
        const int64_t scale_offset =
            ((((m_tile * (scale_cols / 4) + k_tile) * 32 + outer_m) * 4 +
               inner_m) *
                  4 +
              inner_k);
        scale[scale_offset] = sc_f8;
      }
      const float block_scale = Fp8ToF32(sc_f8) * gs_recip;
      const float out_scale_v = RecipF(block_scale);
      for (int j = 0; j < kBS; j += 2) {
        const float lo = ClampF(LoadF32(x, i * k + base + j) * out_scale_v, -kFp4Max, kFp4Max);
        const float hi = ClampF(LoadF32(x, i * k + base + j + 1) * out_scale_v, -kFp4Max, kFp4Max);
        packed[(i * k + base + j) / 2] =
            static_cast<uint8_t>(F32ToFp4Nibble(lo) | (F32ToFp4Nibble(hi) << 4));
      }
    }
  }
  });
}

// SiluMulFp4Quant CPU fallback = the exact composite (bf16 intermediate then
// quant) — which IS the definition of correctness for the CUDA fused kernel. The
// bf16 scratch reproduces the round-through-bf16 the CUDA kernel folds in.
void SiluMulFp4QuantKernel(Queue& q, Tensor& out_packed, Tensor& out_scale, const Tensor& gate,
                           const Tensor& up, float input_global_scale_inv,
                           Fp4ScaleLayout scale_layout) {
  const int64_t m = gate.shape[0], i = gate.shape[1];
  std::vector<uint16_t> tmp(static_cast<size_t>(m) * static_cast<size_t>(i));
  Tensor act = Tensor::Contiguous(tmp.data(), DType::kBF16, gate.device, {m, i});
  MoeSiluMulKernel(q, act, gate, up);
  ScaledFp4QuantKernel(q, out_packed, out_scale, act, input_global_scale_inv,
                       scale_layout);
}

// CPU definition of vLLM's one-input silu_and_mul_nvfp4_quant custom op. Keep
// this visibly composite: it is the correctness oracle for the CUDA single-pass
// producer and preserves the BF16 store/load boundary exactly.
void SiluAndMulFp4QuantKernel(Queue& q, Tensor& out_packed, Tensor& out_scale,
                              const Tensor& gate_up,
                              float input_global_scale_inv,
                              Fp4ScaleLayout scale_layout) {
  const int64_t m = gate_up.shape[0], i = gate_up.shape[1] / 2;
  std::vector<uint16_t> tmp(static_cast<size_t>(m) * static_cast<size_t>(i));
  Tensor act = Tensor::Contiguous(tmp.data(), DType::kBF16, gate_up.device, {m, i});
  SiluAndMulKernel(q, act, gate_up);
  ScaledFp4QuantKernel(q, out_packed, out_scale, act,
                       input_global_scale_inv, scale_layout);
}

void SigmoidGateBf16Kernel(Queue&, Tensor& out, const Tensor& attn,
                           const Tensor& gate);  // defined below

// SigmoidGateFp4Quant CPU fallback = the exact composite (bf16 intermediate then
// quant) — the definition of correctness for the CUDA fused kernel. The bf16
// scratch reproduces the round-through-bf16 the CUDA kernel folds in.
void SigmoidGateFp4QuantKernel(Queue& q, Tensor& out_packed, Tensor& out_scale,
                               const Tensor& attn, const Tensor& gate,
                               float input_global_scale_inv, Fp4ScaleLayout scale_layout) {
  const int64_t m = attn.shape[0], i = attn.shape[1];
  std::vector<uint16_t> tmp(static_cast<size_t>(m) * static_cast<size_t>(i));
  Tensor act = Tensor::Contiguous(tmp.data(), DType::kBF16, attn.device, {m, i});
  SigmoidGateBf16Kernel(q, act, attn, gate);
  ScaledFp4QuantKernel(q, out_packed, out_scale, act, input_global_scale_inv,
                       scale_layout);
}

// MatmulNvfp4Fp4 CPU kernel: out[m,n] = alpha * Σ_k (a_fp4·f8(a_scale))·(b_fp4·
// f8(b_scale)). Equals vllm::RunNvfp4Emulation up to K-reduction order.
void MatmulNvfp4Fp4Kernel(Queue&, Tensor& out, const Tensor& a_packed, const Tensor& a_scale,
                          const Tensor& b_packed, const Tensor& b_scale, float alpha) {
  const int64_t m = a_packed.shape[0], k = a_packed.shape[1] * 2, n = b_packed.shape[0];
  constexpr int kBS = 16;
  const int64_t groups = k / kBS;
  const auto* ap = a_packed.Ptr<uint8_t>();
  const auto* as = a_scale.Ptr<uint8_t>();
  const auto* bp = b_packed.Ptr<uint8_t>();
  const auto* bs = b_scale.Ptr<uint8_t>();
  // Row-chunked over M (each output row + its arow decode owned by one
  // thread); per-column K-group reduction order unchanged.
  ForRows(m, [&](int64_t r0, int64_t r1) {
  std::vector<float> arow(static_cast<size_t>(k));
  for (int64_t i = r0; i < r1; ++i) {
    // Decode a_fp4·a_scale_fp8 for this row once (reused across N columns).
    for (int64_t g = 0; g < groups; ++g) {
      const float asf = Fp8ToF32(as[i * groups + g]);
      for (int j = 0; j < kBS / 2; ++j) {
        const uint8_t byte = ap[(i * k + g * kBS) / 2 + j];
        arow[static_cast<size_t>(g * kBS + 2 * j)] = Nibble(byte & 0x0FU) * asf;
        arow[static_cast<size_t>(g * kBS + 2 * j + 1)] = Nibble(byte >> 4) * asf;
      }
    }
    for (int64_t col = 0; col < n; ++col) {
      float acc = 0.0F;
      for (int64_t g = 0; g < groups; ++g) {
        const float bsf = Fp8ToF32(bs[col * groups + g]);
        for (int j = 0; j < kBS / 2; ++j) {
          const uint8_t byte = bp[(col * k + g * kBS) / 2 + j];
          acc += arow[static_cast<size_t>(g * kBS + 2 * j)] * (Nibble(byte & 0x0FU) * bsf);
          acc += arow[static_cast<size_t>(g * kBS + 2 * j + 1)] * (Nibble(byte >> 4) * bsf);
        }
      }
      StoreF32(out, i * n + col, alpha * acc);
    }
  }
  });
}

void EmbeddingKernel(Queue&, Tensor& out, const Tensor& table, const Tensor& ids) {
  const int64_t t = ids.shape[0], h = table.shape[1], v = table.shape[0];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    int64_t id = ids.dtype == DType::kI32 ? ids.Ptr<int32_t>()[i] : ids.Ptr<int64_t>()[i];
    VT_CHECK(id >= 0 && id < v, "embedding: id out of range");
    for (int64_t j = 0; j < h; ++j) {
      StoreF32(out, i * h + j, LoadF32(table, id * h + j));
    }
  }
  });
}

// In-place rotation of one head starting at element head_off; f32 math,
// stores round back to the tensor's dtype (f32 or bf16).
void RopeRotateHead(const Tensor& t, int64_t head_off, int rot, double base, int64_t pos) {
  const int half = rot / 2;
  for (int i = 0; i < half; ++i) {
    double freq = std::pow(base, -2.0 * i / rot);
    double angle = static_cast<double>(pos) * freq;
    float c = static_cast<float>(std::cos(angle));
    float s = static_cast<float>(std::sin(angle));
    float x = LoadF32(t, head_off + i);
    float y = LoadF32(t, head_off + i + half);
    StoreF32(t, head_off + i, x * c - y * s);
    StoreF32(t, head_off + i + half, x * s + y * c);
  }
}

void RopeNeoxKernel(Queue&, 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];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    int64_t p = pos.dtype == DType::kI32 ? pos.Ptr<int32_t>()[i] : pos.Ptr<int64_t>()[i];
    for (int64_t hh = 0; hh < hq; ++hh) {
      RopeRotateHead(qs, (i * hq + hh) * d, args.rotary_dim, static_cast<double>(args.base), p);
    }
    for (int64_t hh = 0; hh < hk; ++hh) {
      RopeRotateHead(ks, (i * hk + hh) * d, args.rotary_dim, static_cast<double>(args.base), p);
    }
  }
  });
}

// Ported from vLLM's supplied-cache rotary path:
//   base.py:160-252; common.py:145-185; mrope.py:14-187,263-375
// @ e24d1b24fe96. Formula construction stays outside this hot apply loop.
int MropeAxisForPair(int64_t pair, const RopeArgs& args) {
  if (args.mrope_interleaved) {
    if (pair % 3 == 1 &&
        pair <= 3LL * static_cast<int64_t>(args.mrope_section[1])) {
      return 1;
    }
    if (pair % 3 == 2 &&
        pair <= 3LL * static_cast<int64_t>(args.mrope_section[2])) {
      return 2;
    }
    return 0;
  }
  if (pair < args.mrope_section[0]) return 0;
  if (pair < static_cast<int64_t>(args.mrope_section[0]) +
                 args.mrope_section[1]) {
    return 1;
  }
  return 2;
}

void RopeFromCacheKernel(Queue&, 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 bool is_mrope = positions.rank == 2;
  // MLA campaign W6: q/k are addressed through their STRIDES, not a contiguous
  // (token * heads + head) * head_dim formula. DeepSeek's DECOUPLED RoPE rotates
  // only the trailing qk_rope_head_dim slice of the query head
  // (deepseek_v2.py:580-595 / mla.py:160-167 pass `q[..., qk_nope_head_dim:]`),
  // and its k_pe is the trailing column block of the single fused
  // kv_a_proj_with_mqa output — both are STRIDED VIEWS. For a contiguous tensor
  // the strided offsets are integer-identical to the old formula, so every
  // existing caller is bit-identical by construction.
  ForRows(tokens, [&](int64_t row_start, int64_t row_end) {
    for (int64_t token = row_start; token < row_end; ++token) {
      for (int64_t pair = 0; pair < half; ++pair) {
        const int axis = is_mrope ? MropeAxisForPair(pair, args) : 0;
        const int64_t pos_offset =
            is_mrope ? static_cast<int64_t>(axis) * tokens + token : token;
        const int64_t position =
            positions.dtype == DType::kI32
                ? static_cast<int64_t>(positions.Ptr<int32_t>()[pos_offset])
                : positions.Ptr<int64_t>()[pos_offset];
        VT_CHECK(position >= 0 && position < cache.shape[0],
                 "rope_from_cache: position outside cache");
        const int64_t cache_off = position * args.rotary_dim;
        const float c = LoadF32(cache, cache_off + pair);
        const float s = LoadF32(cache, cache_off + half + pair);
        const int64_t first = args.is_neox_style ? pair : pair * 2;
        const int64_t second =
            args.is_neox_style ? pair + half : pair * 2 + 1;
        for (int64_t head = 0; head < hq; ++head) {
          const int64_t off = token * qs.stride[0] + head * qs.stride[1];
          const float x = LoadF32(qs, off + first);
          const float y = LoadF32(qs, off + second);
          StoreF32(qs, off + first, x * c - y * s);
          StoreF32(qs, off + second, x * s + y * c);
        }
        if (ks != nullptr) {
          for (int64_t head = 0; head < hk; ++head) {
            const int64_t off = token * ks->stride[0] + head * ks->stride[1];
            const float x = LoadF32(*ks, off + first);
            const float y = LoadF32(*ks, off + second);
            StoreF32(*ks, off + first, x * c - y * s);
            StoreF32(*ks, off + second, x * s + y * c);
          }
        }
      }
    }
  });
}

float Silu(float x) { return x / (1.0f + std::exp(-x)); }

// Per-step RoPE cos|sin cache fill (fused-attn-preamble prep). cos_sin[T,rot] f32:
// cols [0,half)=cos, [half,rot)=sin. Angle math in DOUBLE + f32 cast, matching
// RopeRotateHead/RopeNeoxKernel element-for-element so the cache reproduces the
// inline rotation bit-for-bit.
void RopeCosSinCacheKernel(Queue&, Tensor& cos_sin, const Tensor& positions, const RopeArgs& args) {
  const int64_t t = cos_sin.shape[0];
  const int rot = args.rotary_dim;
  const int64_t half = rot / 2;
  const double base = static_cast<double>(args.base);
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const int64_t p =
        positions.dtype == DType::kI32 ? positions.Ptr<int32_t>()[i] : positions.Ptr<int64_t>()[i];
    for (int64_t pair = 0; pair < half; ++pair) {
      const double freq = std::pow(base, -2.0 * static_cast<double>(pair) / static_cast<double>(rot));
      const double angle = static_cast<double>(p) * freq;
      StoreF32(cos_sin, i * rot + pair, static_cast<float>(std::cos(angle)));
      StoreF32(cos_sin, i * rot + half + pair, static_cast<float>(std::sin(angle)));
    }
  }
  });
}

// gemma-RMSNorm one element: (v*inv)*(gemma ? w+1 : w) — matches RmsNormKernel's
// `v * inv * wj` (wj = w [+1 if gemma]) grouping and order exactly.
float GemmaNormElem(float v, float inv, float w, bool gemma) {
  float wj = w;
  if (gemma) wj += 1.0f;
  return v * inv * wj;
}

// Fused full-attention preamble: split q|gate + gemma qk-RMSNorm(Dh) + partial
// NeoX RoPE (from the cos_sin cache) + gate passthrough, in one pass. Bit-for-bit
// equal (f32 out) to AttnGateSplit + RmsNorm(q) + RmsNorm(k) + RopeNeox composed:
// the variance is f32, the weight is applied as (1+w), and the rotation reuses the
// same f32 c/sn the cache holds (x*c - y*sn / x*sn + y*c). Tail dims [rot,Dh) are
// normed but unrotated.
void AttnQkNormRopeGateKernel(Queue&, 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 int rot = ra.rotary_dim;
  const int64_t half = rot / 2;
  const bool gemma = na.gemma;
  // Normalize one head row (src..src+Dh) into out.., applying partial NeoX RoPE
  // from cs[0..rot). Recomputes normed[i]/normed[i+half] where paired.
  auto do_head = [&](const Tensor& src, int64_t src_off, const Tensor& w, const Tensor& out,
                     int64_t out_off, const float* cs) {
    float ss = 0.0f;
    for (int64_t j = 0; j < dh; ++j) {
      const float v = LoadF32(src, src_off + j);
      ss += v * v;
    }
    const float inv = 1.0f / std::sqrt(ss / static_cast<float>(dh) + na.eps);
    for (int64_t j = 0; j < dh; ++j) {
      if (j < half) {
        const float ni = GemmaNormElem(LoadF32(src, src_off + j), inv, LoadF32(w, j), gemma);
        const float nih =
            GemmaNormElem(LoadF32(src, src_off + j + half), inv, LoadF32(w, j + half), gemma);
        StoreF32(out, out_off + j, ni * cs[j] - nih * cs[half + j]);
      } else if (j < rot) {
        const int64_t i = j - half;
        const float ni = GemmaNormElem(LoadF32(src, src_off + i), inv, LoadF32(w, i), gemma);
        const float nih =
            GemmaNormElem(LoadF32(src, src_off + i + half), inv, LoadF32(w, i + half), gemma);
        StoreF32(out, out_off + j, ni * cs[half + i] + nih * cs[i]);
      } else {
        StoreF32(out, out_off + j,
                 GemmaNormElem(LoadF32(src, src_off + j), inv, LoadF32(w, j), gemma));
      }
    }
  };
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t tok = r0; tok < r1; ++tok) {
    const float* cs = cos_sin.Ptr<float>() + tok * rot;
    for (int64_t h = 0; h < hq; ++h) {
      const int64_t qrow = tok * qgate.stride[0] + h * 2 * dh;
      do_head(qgate, qrow, q_norm, q_out, (tok * hq + h) * dh, cs);
      // gate passthrough: the second Dh of each (t,h) q|gate pair (no norm/rope).
      const int64_t gbase = qrow + dh;
      const int64_t gout = (tok * hq + h) * dh;
      for (int64_t j = 0; j < dh; ++j) StoreF32(gate_out, gout + j, LoadF32(qgate, gbase + j));
    }
    for (int64_t h = 0; h < hkv; ++h) {
      do_head(kf, tok * kf.stride[0] + h * dh, k_norm, k_out,
              (tok * hkv + h) * dh, cs);
    }
  }
  });
}

// GDN CPU reference kernels. Formulas: .agents/specs/gdn-semantics.md (§ cited per
// kernel); scalar f32 math throughout, states f32 in place.

// §2 causal_conv1d_fn. Per sequence s (tokens [qsl[s], qsl[s+1])), channel c,
// token t: window[j] = x token t-(K-1-j), falling back to
// conv_state[c, (K-1)+(t-i)] (init state) or 0 before the sequence start.
// Write-back: last K-1 RAW x tokens, left-padded with zeros / shifted old
// state when T < K-1. Outputs read the OLD state, so the row is buffered
// before overwrite.
void CausalConv1dFwdKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& w,
                           const Tensor* bias, Tensor& conv_state, const Tensor& qsl,
                           const Tensor& his, const CausalConv1dArgs& args) {
  const int64_t total = x.shape[0], c_dim = x.shape[1], k = w.shape[1], width = k - 1;
  const int64_t n = conv_state.shape[0];
  // x may be a padded-row (inner-contiguous) view of the merged qkvz output;
  // out/conv_state stay contiguous so they keep the c_dim row stride.
  const int64_t x_rs = x.stride[0];
  const int32_t* qslp = qsl.Ptr<int32_t>();
  VT_CHECK(qslp[0] == 0 && qslp[n] == total, "causal_conv1d_fwd: bad query_start_loc bounds");
  for (int64_t s = 0; s < n; ++s) {
    VT_CHECK(qslp[s + 1] >= qslp[s] && qslp[s] >= 0,
             "causal_conv1d_fwd: query_start_loc not monotonic");
  }
  // Row-chunked over (sequence, channel) pairs: each pair owns its out column
  // slice and its conv_state row — independent outputs (spec W3 "conv1d").
  ForRows(n * c_dim, [&](int64_t r0, int64_t r1) {
  std::vector<float> old_row(static_cast<size_t>(width));
  for (int64_t r = r0; r < r1; ++r) {
    const int64_t s = r / c_dim, c = r % c_dim;
    const int64_t begin = qslp[s], end = qslp[s + 1], t_len = end - begin;
    const bool init = his.dtype == DType::kI8 ? his.Ptr<int8_t>()[s] != 0
                                               : his.Ptr<int32_t>()[s] != 0;
    float* srow_base = conv_state.Ptr<float>() + s * c_dim * width;
    {
      float* srow = srow_base + c * width;
      for (int64_t j = 0; j < width; ++j) old_row[static_cast<size_t>(j)] = srow[j];
      const float b = bias != nullptr ? LoadF32(*bias, c) : 0.0f;
      for (int64_t t = 0; t < t_len; ++t) {
        float acc = b;
        for (int64_t j = 0; j < k; ++j) {
          const int64_t ti = t - (k - 1 - j);  // token index of window[j]
          float v = 0.0f;
          if (ti >= 0) {
            v = LoadF32(x, (begin + ti) * x_rs + c);
          } else if (init) {
            v = old_row[static_cast<size_t>(width + ti)];  // state col (K-1)+(t-i)
          }
          acc += LoadF32(w, c * k + j) * v;
        }
        StoreF32(out, (begin + t) * c_dim + c, args.silu_activation ? Silu(acc) : acc);
      }
      for (int64_t j = 0; j < width; ++j) {
        const int64_t tj = t_len - width + j;  // new state col j holds token tj
        float v = 0.0f;
        if (tj >= 0) {
          v = LoadF32(x, (begin + tj) * x_rs + c);
        } else if (init) {
          v = old_row[static_cast<size_t>(width + tj)];  // shifted old state
        }
        srow[j] = v;
      }
    }
  }
  });
}

// §3 causal_conv1d_update (seqlen==1): read-old-then-roll. conv_state_indices
// (optional; mirrors mamba conv_state_indices): token bt's row is cache slot
// idx[bt] (idx<0 == NULL block → skip); null => compact row == bt.
void CausalConv1dUpdateKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& w,
                              const Tensor* bias, Tensor& conv_state,
                              const Tensor* conv_state_indices,
                              const CausalConv1dArgs& args) {
  const int64_t batch = x.shape[0], c_dim = x.shape[1], k = w.shape[1], width = k - 1;
  // x may be a padded-row view of the merged qkvz output; out is contiguous.
  const int64_t x_rs = x.stride[0];
  const int32_t* cache_idx =
      conv_state_indices != nullptr ? conv_state_indices->Ptr<int32_t>() : nullptr;
  // Row-chunked over (batch, channel) pairs: each pair owns its out element and
  // its conv_state row slice (batch rows map to distinct cache slots).
  ForRows(batch * c_dim, [&](int64_t r0, int64_t r1) {
  for (int64_t r = r0; r < r1; ++r) {
    const int64_t bt = r / c_dim, c = r % c_dim;
    int64_t srow_row = bt;
    if (cache_idx != nullptr) {
      if (cache_idx[bt] < 0) continue;  // NULL block
      srow_row = cache_idx[bt];
    }
    float* srow_base = conv_state.Ptr<float>() + srow_row * c_dim * width;
    {
      float* srow = srow_base + c * width;
      const float xt = LoadF32(x, bt * x_rs + c);
      float acc = bias != nullptr ? LoadF32(*bias, c) : 0.0f;
      for (int64_t j = 0; j < width; ++j) acc += LoadF32(w, c * k + j) * srow[j];
      acc += LoadF32(w, c * k + width) * xt;
      StoreF32(out, bt * c_dim + c, args.silu_activation ? Silu(acc) : acc);
      for (int64_t j = 0; j + 1 < width; ++j) srow[j] = srow[j + 1];  // roll left
      if (width > 0) srow[width - 1] = xt;                            // raw x
    }
  }
  });
}

// §4 l2norm_fwd: y = x * rsqrt(sum(x^2) + eps) over the last dim (plain SUM).
void L2NormKernel(Queue&, Tensor& out, const Tensor& x, const L2NormArgs& args) {
  const int64_t d = x.shape[x.rank - 1];
  const int64_t rows = x.Numel() / d;
  ForRows(rows, [&](int64_t r0, int64_t r1) {
  for (int64_t r = r0; r < r1; ++r) {
    float sumsq = 0.0f;
    for (int64_t j = 0; j < d; ++j) {
      const float v = LoadF32(x, r * d + j);
      sumsq += v * v;
    }
    const float inv = 1.0f / std::sqrt(sumsq + args.eps);
    for (int64_t j = 0; j < d; ++j) StoreF32(out, r * d + j, LoadF32(x, r * d + j) * inv);
  }
  });
}

// §5 RMSNormGated (norm_before_gate=True, group_size=None):
// out = x * rsqrt(mean(x^2) + eps) * w * act(gate); act = silu or sigmoid.
void RmsNormGatedKernel(Queue&, Tensor& out, const Tensor& x, const Tensor& gate,
                        const Tensor& w, const RmsNormGatedArgs& args) {
  const int64_t d = x.shape[x.rank - 1];
  const int64_t t = x.Numel() / d;
  // x/out are contiguous (flat row i at i*d). The gate may be a padded-row
  // rank-3 [T,Hv,D] view of the merged qkvz z slice: map super-row i to
  // (token=i/Hv, head=i%Hv) and honor the token stride. Contiguous rank-2 gate
  // degenerates to i*gate.stride[0] == i*d (group == 1), byte-identical.
  const int64_t gate_group = gate.rank == 3 ? gate.shape[1] : 1;
  const int64_t gate_outer = gate.stride[0];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    float sumsq = 0.0f;
    for (int64_t j = 0; j < d; ++j) {
      const float v = LoadF32(x, i * d + j);
      sumsq += v * v;
    }
    const float inv = 1.0f / std::sqrt(sumsq / static_cast<float>(d) + args.eps);
    const int64_t gbase = (i / gate_group) * gate_outer + (i % gate_group) * d;
    for (int64_t j = 0; j < d; ++j) {
      const float z = LoadF32(gate, gbase + j);
      const float act = args.sigmoid_gate ? 1.0f / (1.0f + std::exp(-z)) : Silu(z);
      StoreF32(out, i * d + j, LoadF32(x, i * d + j) * inv * LoadF32(w, j) * act);
    }
  }
  });
}

// RmsNormGated + static fp8 quant, fused (RmsNormGatedQuantFp8). Composite of
// RmsNormGatedKernel (bf16-rounded output) + QuantFp8Static: bit-identical to the
// split path because the fp8 is taken from the SAME bf16-rounded gated-norm value.
// CUDA has the hot path; this keeps the op available on the host backend for tests.
void RmsNormGatedQuantFp8Kernel(Queue&, Tensor& out_fp8, const Tensor& x, const Tensor& gate,
                                const Tensor& w, const RmsNormGatedArgs& args, float input_scale) {
  const int64_t d = x.shape[x.rank - 1];
  const int64_t t = x.Numel() / d;
  const float inv_scale = 1.0F / input_scale;
  uint8_t* op = out_fp8.Ptr<uint8_t>();
  const int64_t gate_group = gate.rank == 3 ? gate.shape[1] : 1;
  const int64_t gate_outer = gate.stride[0];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    float sumsq = 0.0F;
    for (int64_t j = 0; j < d; ++j) {
      const float v = LoadF32(x, i * d + j);
      sumsq += v * v;
    }
    const float inv = 1.0F / std::sqrt(sumsq / static_cast<float>(d) + args.eps);
    const int64_t gbase = (i / gate_group) * gate_outer + (i % gate_group) * d;
    for (int64_t j = 0; j < d; ++j) {
      const float z = LoadF32(gate, gbase + j);
      const float act = args.sigmoid_gate ? 1.0f / (1.0f + std::exp(-z)) : Silu(z);
      // bf16-intermediate (matches RmsNormGated's bf16 store then QuantFp8Static's load).
      const uint16_t nb = F32ToBF16(LoadF32(x, i * d + j) * inv * LoadF32(w, j) * act);
      op[i * d + j] = F32ToFp8(BF16ToF32(nb) * inv_scale);
    }
  }
  });
}

// §7 gated-delta-rule token step, shared by prefill and decode. state points
// at this sequence's [Hv,Dv,Dk] f32 block; tok indexes the packed q/k/v/g/beta
// rows. GQA broadcast: v-head hv reads q/k head hv / (Hv/Hk).
//   q' = q * scale;  S *= exp(g[hv]);  v' = (v - S @ k) * beta[hv];
//   S += outer(v', k);  out = S @ q'      (k is NOT scaled)
void GdnTokenStep(Tensor& out, const Tensor& q_in, const Tensor& k_in, const Tensor& v_in,
                  const Tensor& g, const Tensor& beta, float* state, int64_t tok, float scale,
                  std::vector<float>& qbuf, std::vector<float>& kbuf,
                  std::vector<float>& vbuf) {
  const int64_t hk_n = q_in.shape[1], dk = q_in.shape[2];
  const int64_t hv_n = v_in.shape[1], dv = v_in.shape[2];
  const int64_t ratio = hv_n / hk_n;
  for (int64_t hv = 0; hv < hv_n; ++hv) {
    const int64_t hk = hv / ratio;
    float* s_head = state + hv * dv * dk;  // [Dv, Dk]
    const float g_t = g.Ptr<float>()[tok * hv_n + hv];
    const float beta_t = beta.Ptr<float>()[tok * hv_n + hv];
    const float decay = std::exp(g_t);
    for (int64_t i = 0; i < dk; ++i) {
      qbuf[static_cast<size_t>(i)] = LoadF32(q_in, (tok * hk_n + hk) * dk + i) * scale;
      kbuf[static_cast<size_t>(i)] = LoadF32(k_in, (tok * hk_n + hk) * dk + i);
    }
    for (int64_t vi = 0; vi < dv; ++vi) {
      float* s_row = s_head + vi * dk;
      float dot = 0.0f;  // (S * exp(g)) @ k, fused with the decay pass
      for (int64_t ki = 0; ki < dk; ++ki) {
        s_row[ki] *= decay;
        dot += s_row[ki] * kbuf[static_cast<size_t>(ki)];
      }
      vbuf[static_cast<size_t>(vi)] =
          (LoadF32(v_in, (tok * hv_n + hv) * dv + vi) - dot) * beta_t;
    }
    for (int64_t vi = 0; vi < dv; ++vi) {
      float* s_row = s_head + vi * dk;
      float o = 0.0f;  // (S + outer(v',k)) @ q', fused with the rank-1 update
      for (int64_t ki = 0; ki < dk; ++ki) {
        s_row[ki] += vbuf[static_cast<size_t>(vi)] * kbuf[static_cast<size_t>(ki)];
        o += s_row[ki] * qbuf[static_cast<size_t>(ki)];
      }
      StoreF32(out, (tok * hv_n + hv) * dv + vi, o);
    }
  }
}

void GdnPrefillKernel(Queue&, Tensor& out, const Tensor& q_in, const Tensor& k, const Tensor& v,
                      const Tensor& g, const Tensor& beta, Tensor& state, const Tensor& qsl,
                      const GdnArgs& args) {
  const int64_t n = state.shape[0], hv_n = state.shape[1], dv = state.shape[2],
                dk = state.shape[3];
  const int32_t* qslp = qsl.Ptr<int32_t>();
  VT_CHECK(qslp[0] == 0 && qslp[n] == q_in.shape[0],
           "gdn_prefill: bad query_start_loc bounds");
  for (int64_t s = 0; s < n; ++s) {
    VT_CHECK(qslp[s + 1] >= qslp[s], "gdn_prefill: query_start_loc not monotonic");
  }
  // Row-chunked over SEQUENCES (spec W3): each sequence owns its state block
  // and its token range's outputs; the in-sequence recurrence stays sequential.
  ForRows(n, [&](int64_t r0, int64_t r1) {
  std::vector<float> qbuf(static_cast<size_t>(dk)), kbuf(static_cast<size_t>(dk)),
      vbuf(static_cast<size_t>(dv));
  for (int64_t s = r0; s < r1; ++s) {
    float* s_state = state.Ptr<float>() + s * hv_n * dv * dk;
    for (int64_t t = qslp[s]; t < qslp[s + 1]; ++t)
      GdnTokenStep(out, q_in, k, v, g, beta, s_state, t, args.scale, qbuf, kbuf, vbuf);
  }
  });
}

void GdnDecodeKernel(Queue&, Tensor& out, const Tensor& q_in, const Tensor& k, const Tensor& v,
                     const Tensor& g, const Tensor& beta, Tensor& state,
                     const Tensor* state_idx, const GdnArgs& args) {
  const int64_t batch = q_in.shape[0], hv_n = state.shape[1], dv = state.shape[2],
                dk = state.shape[3];
  const int32_t* sidx = state_idx != nullptr ? state_idx->Ptr<int32_t>() : nullptr;
  // Row-chunked over the BATCH (spec W3): each token owns its state slot
  // (distinct cache rows) and its output row.
  ForRows(batch, [&](int64_t r0, int64_t r1) {
  std::vector<float> qbuf(static_cast<size_t>(dk)), kbuf(static_cast<size_t>(dk)),
      vbuf(static_cast<size_t>(dv));
  for (int64_t bt = r0; bt < r1; ++bt) {
    // state_idx != null => in-place on the FULL cache at slot sidx[bt] (fla
    // ssm_state_indices); sidx[bt]<0 == NULL block → zero out; null => row == bt.
    int64_t srow = bt;
    if (sidx != nullptr) {
      if (sidx[bt] < 0) {
        for (int64_t hv = 0; hv < hv_n; ++hv)
          for (int64_t d = 0; d < dv; ++d)
            StoreF32(out, (bt * hv_n + hv) * dv + d, 0.0f);
        continue;
      }
      srow = sidx[bt];
    }
    float* s_state = state.Ptr<float>() + srow * hv_n * dv * dk;
    GdnTokenStep(out, q_in, k, v, g, beta, s_state, bt, args.scale, qbuf, kbuf, vbuf);
  }
  });
}

// Pure packed decode ported from vLLM v0.25.0
// vllm/model_executor/layers/fla/ops/fused_recurrent.py:255-336 @ 702f4814.
// Unlike GdnDecodeKernel, raw q/k are normalized here in f32 and never
// materialized through the activation dtype. sigmoid(b) is rounded through
// b.dtype before recurrence, exactly at the upstream store/load boundary.
void GdnPackedDecodeKernel(Queue&, Tensor& out, const Tensor& mixed_qkv,
                           const Tensor& a, const Tensor& b,
                           const Tensor& a_log, const Tensor& dt_bias,
                           Tensor& state, const Tensor& state_idx,
                           const GdnArgs& args) {
  const int64_t batch = mixed_qkv.shape[0];
  const int64_t hv_n = state.shape[1];
  const int64_t dv = state.shape[2];
  const int64_t dk = state.shape[3];
  const int64_t qk_dim = mixed_qkv.shape[1] - hv_n * dv;
  const int64_t hk_n = qk_dim / (2 * dk);
  const int64_t ratio = hv_n / hk_n;
  const int32_t* indices = state_idx.Ptr<int32_t>();

  ForRows(batch, [&](int64_t r0, int64_t r1) {
    std::vector<float> qv(static_cast<size_t>(dk));
    std::vector<float> kv(static_cast<size_t>(dk));
    std::vector<float> state_row(static_cast<size_t>(dk));
    for (int64_t bt = r0; bt < r1; ++bt) {
      const int32_t slot = indices[bt];
      if (slot < 0) {
        for (int64_t i = 0; i < hv_n * dv; ++i)
          StoreF32(out, bt * hv_n * dv + i, 0.0f);
        continue;
      }
      for (int64_t hv = 0; hv < hv_n; ++hv) {
        const int64_t hk = hv / ratio;
        float q_sumsq = 0.0f;
        float k_sumsq = 0.0f;
        for (int64_t ki = 0; ki < dk; ++ki) {
          qv[static_cast<size_t>(ki)] =
              LoadF32(mixed_qkv, bt * mixed_qkv.stride[0] + hk * dk + ki);
          kv[static_cast<size_t>(ki)] =
              LoadF32(mixed_qkv, bt * mixed_qkv.stride[0] + hk_n * dk + hk * dk + ki);
          q_sumsq += qv[static_cast<size_t>(ki)] * qv[static_cast<size_t>(ki)];
          k_sumsq += kv[static_cast<size_t>(ki)] * kv[static_cast<size_t>(ki)];
        }
        const float q_inv = 1.0f / std::sqrt(q_sumsq + 1e-6f);
        const float k_inv = 1.0f / std::sqrt(k_sumsq + 1e-6f);
        for (int64_t ki = 0; ki < dk; ++ki) {
          qv[static_cast<size_t>(ki)] *= q_inv * args.scale;
          kv[static_cast<size_t>(ki)] *= k_inv;
        }

        const float av = LoadF32(a, bt * a.stride[0] + hv);
        const float bv = LoadF32(b, bt * b.stride[0] + hv);
        const float x = av + LoadF32(dt_bias, hv);
        const float softplus =
            x <= 20.0f ? std::log1p(std::exp(std::min(x, 20.0f))) : x;
        const float g = -std::exp(LoadF32(a_log, hv)) * softplus;
        float beta = 1.0f / (1.0f + std::exp(-bv));
        if (b.dtype == DType::kF16)
          beta = F16ToF32(F32ToF16(beta));
        else if (b.dtype == DType::kBF16)
          beta = BF16ToF32(F32ToBF16(beta));
        const float decay = std::exp(g);

        for (int64_t vi = 0; vi < dv; ++vi) {
          const int64_t state_base =
              ((static_cast<int64_t>(slot) * hv_n + hv) * dv + vi) * dk;
          float dot = 0.0f;
          for (int64_t ki = 0; ki < dk; ++ki) {
            state_row[static_cast<size_t>(ki)] =
                LoadF32(state, state_base + ki) * decay;
            dot += state_row[static_cast<size_t>(ki)] * kv[static_cast<size_t>(ki)];
          }
          const int64_t v_offset = bt * mixed_qkv.stride[0] +
                                   2 * hk_n * dk + hv * dv + vi;
          const float vp = (LoadF32(mixed_qkv, v_offset) - dot) * beta;
          float output = 0.0f;
          for (int64_t ki = 0; ki < dk; ++ki) {
            const float updated = state_row[static_cast<size_t>(ki)] +
                                  vp * kv[static_cast<size_t>(ki)];
            StoreF32(state, state_base + ki, updated);
            output += updated * qv[static_cast<size_t>(ki)];
          }
          StoreF32(out, (bt * hv_n + hv) * dv + vi, output);
        }
      }
    }
  });
}

// Indexed cache boundary for mixed GDN prefill. This is the CPU executable
// reference for the fused CUDA gather/scatter kernels: cache rows may be f32 or
// bf16, while the compact recurrence/conv working state is always f32.
void GdnStateGatherKernel(Queue&, Tensor& working, const Tensor& cache,
                          const Tensor& state_idx,
                          const Tensor* has_initial_state) {
  const int64_t rows = state_idx.shape[0];
  if (rows == 0) return;
  const int64_t row_elems = working.Numel() / rows;
  const int32_t* idx = state_idx.Ptr<int32_t>();
  for (int64_t r = 0; r < rows; ++r) {
    VT_CHECK(idx[r] >= 0 && idx[r] < cache.shape[0],
             "gdn_state_gather: state_idx out of range");
  }
  ForRows(rows, [&](int64_t r0, int64_t r1) {
    for (int64_t r = r0; r < r1; ++r) {
      bool keep = true;
      if (has_initial_state != nullptr) {
        keep = has_initial_state->dtype == DType::kI8
                   ? has_initial_state->Ptr<int8_t>()[r] != 0
                   : has_initial_state->Ptr<int32_t>()[r] != 0;
      }
      const int64_t src = static_cast<int64_t>(idx[r]) * row_elems;
      const int64_t dst = r * row_elems;
      for (int64_t e = 0; e < row_elems; ++e) {
        StoreF32(working, dst + e, keep ? LoadF32(cache, src + e) : 0.0f);
      }
    }
  });
}

void GdnStateScatterKernel(Queue&, Tensor& cache, const Tensor& working,
                           const Tensor& state_idx) {
  const int64_t rows = state_idx.shape[0];
  if (rows == 0) return;
  const int64_t row_elems = working.Numel() / rows;
  const int32_t* idx = state_idx.Ptr<int32_t>();
  for (int64_t r = 0; r < rows; ++r) {
    VT_CHECK(idx[r] >= 0 && idx[r] < cache.shape[0],
             "gdn_state_scatter: state_idx out of range");
  }
  ForRows(rows, [&](int64_t r0, int64_t r1) {
    for (int64_t r = r0; r < r1; ++r) {
      const int64_t src = r * row_elems;
      const int64_t dst = static_cast<int64_t>(idx[r]) * row_elems;
      for (int64_t e = 0; e < row_elems; ++e) {
        StoreF32(cache, dst + e, LoadF32(working, src + e));
      }
    }
  });
}

// Grouped-topk (`noaux_tc`) router — the CPU REFERENCE and the executable spec
// for the DeepSeek router. 1:1 port of
// vllm/model_executor/layers/fused_moe/router/grouped_topk_router.py:106-161
// (`grouped_topk`, the `forward_native` path) @ pin e24d1b24. Step numbering in
// the comments matches the ops.h contract. Reached ONLY when
// args.num_expert_group > 0; the ungrouped kernel below is untouched.
void MoeRouterGroupedTopKKernel(Tensor& weights, Tensor& indices, const Tensor& logits,
                                const MoeRouterTopKArgs& args, const Tensor* bias_t) {
  const int64_t t = logits.shape[0], e = logits.shape[1];
  const int k = args.top_k;
  const int64_t n_group = args.num_expert_group;
  const int64_t group_size = e / n_group;  // wrapper checked divisibility
  const float* bias = bias_t != nullptr ? bias_t->Ptr<float>() : nullptr;

  ForRows(t, [&](int64_t r0, int64_t r1) {
    std::vector<float> sel(static_cast<size_t>(e));   // SELECTION score (biased)
    std::vector<float> orig(static_cast<size_t>(e));  // WEIGHT score (unbiased)
    std::vector<float> gscore(static_cast<size_t>(n_group));
    std::vector<char> gkeep(static_cast<size_t>(n_group));
    for (int64_t row = r0; row < r1; ++row) {
      // (1) scores = softmax(logits, -1) | sigmoid(logits)   (:110-117)
      if (args.scoring_func == MoeScoringFunc::kSigmoid) {
        // ELEMENTWISE — not normalized across experts. This is the V3/R1 path.
        for (int64_t j = 0; j < e; ++j) {
          orig[static_cast<size_t>(j)] = 1.0f / (1.0f + std::exp(-LoadF32(logits, row * e + j)));
        }
      } else {
        float mx = -INFINITY;
        for (int64_t j = 0; j < e; ++j) mx = std::max(mx, LoadF32(logits, row * e + j));
        float sum = 0.0f;
        for (int64_t j = 0; j < e; ++j) {
          const float ex = std::exp(LoadF32(logits, row * e + j) - mx);
          orig[static_cast<size_t>(j)] = ex;
          sum += ex;
        }
        for (int64_t j = 0; j < e; ++j) {
          float& pj = orig[static_cast<size_t>(j)];
          pj = sum > 0.0f ? pj / sum : 0.0f;
          if (!std::isfinite(pj)) pj = 0.0f;
        }
      }
      // (2) The BIAS asymmetry (:120-124): the bias shifts the score used for
      // SELECTION only; the routing WEIGHT is always read from the unbiased
      // score. Getting this backwards is a silent accuracy bug.
      for (int64_t j = 0; j < e; ++j) {
        sel[static_cast<size_t>(j)] =
            orig[static_cast<size_t>(j)] + (bias != nullptr ? bias[j] : 0.0f);
      }
      // Group score: SUM of the top-2 in the group when a bias is present
      // (:124-126), else the group MAX (:128-131).
      for (int64_t g = 0; g < n_group; ++g) {
        const int64_t base = g * group_size;
        if (bias != nullptr) {
          float b0 = -INFINITY, b1 = -INFINITY;
          for (int64_t j = 0; j < group_size; ++j) {
            const float v = sel[static_cast<size_t>(base + j)];
            if (v > b0) {
              b1 = b0;
              b0 = v;
            } else if (v > b1) {
              b1 = v;
            }
          }
          // group_size >= 2 whenever a bias is used in practice (V3: 32); if it
          // were 1, torch's topk(2) would fail upstream too, so b1 stays -inf
          // and the sum is -inf — an honest propagation, not a silent 0.
          gscore[static_cast<size_t>(g)] = b0 + b1;
        } else {
          float m = -INFINITY;
          for (int64_t j = 0; j < group_size; ++j) {
            m = std::max(m, sel[static_cast<size_t>(base + j)]);
          }
          gscore[static_cast<size_t>(g)] = m;
        }
      }
      // (3) keep the top `topk_group` groups, mask the rest to -inf (:133-145).
      // Greedy argmax with a strict `>` ascending scan — lowest group index wins
      // an exact tie (our determinism convention; see the ops.h deviation note).
      for (int64_t g = 0; g < n_group; ++g) gkeep[static_cast<size_t>(g)] = 0;
      for (int gi = 0; gi < args.topk_group; ++gi) {
        int64_t best = -1;
        float best_v = -INFINITY;
        for (int64_t g = 0; g < n_group; ++g) {
          if (gkeep[static_cast<size_t>(g)]) continue;
          // The FIRST unkept index seeds the scan, so an all-`-inf` group row
          // still selects (the lowest unkept index) instead of falling through.
          if (best < 0 || gscore[static_cast<size_t>(g)] > best_v) {
            best_v = gscore[static_cast<size_t>(g)];
            best = g;
          }
        }
        if (best < 0) break;  // fewer groups than topk_group (wrapper forbids it)
        gkeep[static_cast<size_t>(best)] = 1;
      }
      for (int64_t g = 0; g < n_group; ++g) {
        if (gkeep[static_cast<size_t>(g)]) continue;
        for (int64_t j = 0; j < group_size; ++j) {
          sel[static_cast<size_t>(g * group_size + j)] = -INFINITY;
        }
      }
      // (4) top-k over the masked SELECTION scores; the weight comes from the
      // UNBIASED score at the selected id (:147-150).
      float denom = 0.0f;
      for (int j = 0; j < k; ++j) {
        int64_t best = -1;
        float best_v = -INFINITY;
        for (int64_t idx = 0; idx < e; ++idx) {
          if (sel[static_cast<size_t>(idx)] > best_v) {
            best_v = sel[static_cast<size_t>(idx)];
            best = idx;
          }
        }
        if (best < 0) best = 0;
        sel[static_cast<size_t>(best)] = -INFINITY;  // exclude from later rounds
        const float w = orig[static_cast<size_t>(best)];
        weights.Ptr<float>()[row * k + j] = w;
        indices.Ptr<int32_t>()[row * k + j] = static_cast<int32_t>(best);
        denom += w;
      }
      // (5) renormalize (:156-157) then routed_scaling_factor (:159-160). Order
      // matters: upstream scales AFTER the renormalize divide.
      if (args.renormalize) {
        if (!(denom > 0.0f)) denom = 1.0f;
        for (int j = 0; j < k; ++j) weights.Ptr<float>()[row * k + j] /= denom;
      }
      if (args.routed_scaling_factor != 1.0f) {
        for (int j = 0; j < k; ++j) {
          weights.Ptr<float>()[row * k + j] *= args.routed_scaling_factor;
        }
      }
    }
  });
}

// §3 router: softmax (f32, over all E) -> greedy top-k (lowest-index tie-break)
// -> optional renormalize. weights [T,K] f32, indices [T,K] i32.
void MoeRouterTopKKernel(Queue&, Tensor& weights, Tensor& indices, const Tensor& logits,
                         const MoeRouterTopKArgs& args, const Tensor* bias) {
  if (args.num_expert_group > 0) {  // W3 grouped-topk (`noaux_tc`) path
    MoeRouterGroupedTopKKernel(weights, indices, logits, args, bias);
    return;
  }
  const int64_t t = logits.shape[0], e = logits.shape[1];
  const int k = args.top_k;
  ForRows(t, [&](int64_t r0, int64_t r1) {
  std::vector<float> p(static_cast<size_t>(e));
  std::vector<char> chosen(static_cast<size_t>(e));
  for (int64_t row = r0; row < r1; ++row) {
    // softmax(logits.float()) with max-subtraction (topk_softmax_kernels.cu).
    float mx = -INFINITY;
    for (int64_t j = 0; j < e; ++j) mx = std::max(mx, LoadF32(logits, row * e + j));
    float sum = 0.0f;
    for (int64_t j = 0; j < e; ++j) {
      const float ex = std::exp(LoadF32(logits, row * e + j) - mx);
      p[static_cast<size_t>(j)] = ex;
      sum += ex;
    }
    for (int64_t j = 0; j < e; ++j) {
      float& pj = p[static_cast<size_t>(j)];
      pj = sum > 0.0f ? pj / sum : 0.0f;
      if (!std::isfinite(pj)) pj = 0.0f;  // NaN/Inf clamp (.cu:136)
      chosen[static_cast<size_t>(j)] = 0;
    }
    // Greedy argmax, k rounds; strict `>` over ascending j -> lowest index wins.
    float denom = 0.0f;
    for (int j = 0; j < k; ++j) {
      int64_t best = -1;
      float best_v = -INFINITY;
      for (int64_t idx = 0; idx < e; ++idx) {
        if (chosen[static_cast<size_t>(idx)]) continue;
        if (p[static_cast<size_t>(idx)] > best_v) {
          best_v = p[static_cast<size_t>(idx)];
          best = idx;
        }
      }
      chosen[static_cast<size_t>(best)] = 1;
      weights.Ptr<float>()[row * k + j] = best_v;
      indices.Ptr<int32_t>()[row * k + j] = static_cast<int32_t>(best);
      denom += best_v;
    }
    if (args.renormalize) {
      if (!(denom > 0.0f)) denom = 1.0f;  // (.cu:245-253) denom<=0 -> 1 guard
      for (int j = 0; j < k; ++j) weights.Ptr<float>()[row * k + j] /= denom;
    }
  }
  });
}

// §4/§6 weighted scatter-combine: out[t,:] = sum_j w[t,j]*expert_out[t,j,:]
// (f32 accumulation) + shared[t,:] (optional). Stored at out's dtype.
void MoeCombineKernel(Queue&, Tensor& out, const Tensor& expert_out, const Tensor& weights,
                      const Tensor* shared) {
  const int64_t t = out.shape[0], h = out.shape[1], k = weights.shape[1];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t row = r0; row < r1; ++row) {
    for (int64_t col = 0; col < h; ++col) {
      float acc = 0.0f;
      for (int64_t j = 0; j < k; ++j)
        acc += weights.Ptr<float>()[row * k + j] *
               LoadF32(expert_out, (row * k + j) * h + col);
      if (shared != nullptr) acc += LoadF32(*shared, row * h + col);
      StoreF32(out, row * h + col, acc);
    }
  }
  });
}

// Dense causal attention (qwen36-forward-notes.md §5). Causal scaled-dot-product
// with GQA broadcast over a single packed sequence. query [T,Hq,D],
// key/value [T,Hk,D], out [T,Hq,D]. Per q-head h (kv-head g = h/(Hq/Hk)) and
// query i: softmax over keys j<=i of scale*(q·k), then weighted sum of v. f32
// softmax with online max-subtraction for numerical stability.
void AttentionKernel(Queue&, 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];
  const int64_t qpk = hq / hk;  // q-heads per kv-head (GQA ratio)
  const float scale = args.scale;
  // Row-chunked over (head, query) pairs (spec W3) — the flash-attn q-row
  // split (ops.cpp:9072-9073 nr = neq1*neq2*neq3). Same h-outer/i-inner walk.
  ForRows(hq * t, [&](int64_t r0, int64_t r1) {
  std::vector<float> probs(static_cast<size_t>(t));
  std::vector<float> acc(static_cast<size_t>(d));
  for (int64_t r = r0; r < r1; ++r) {
    const int64_t h = r / t;
    const int64_t g = h / qpk;
    {
      const int64_t i = r % t;
      const int64_t jmax = args.causal ? i : t - 1;  // causal: keys 0..i
      const int64_t qoff = (i * hq + h) * d;
      // Pass 1: scores + running max.
      float m = -std::numeric_limits<float>::infinity();
      for (int64_t j = 0; j <= jmax; ++j) {
        const int64_t koff = (j * hk + g) * d;
        float dot = 0.0f;
        for (int64_t e = 0; e < d; ++e) dot += LoadF32(query, qoff + e) * LoadF32(key, koff + e);
        dot *= scale;
        probs[static_cast<size_t>(j)] = dot;
        if (dot > m) m = dot;
      }
      // Pass 2: exp + normalization denominator.
      float denom = 0.0f;
      for (int64_t j = 0; j <= jmax; ++j) {
        float e = std::exp(probs[static_cast<size_t>(j)] - m);
        probs[static_cast<size_t>(j)] = e;
        denom += e;
      }
      const float inv = 1.0f / denom;  // denom >= 1 (j==i term is exp(0)=1)
      // Pass 3: weighted sum of v (f32 accumulation), stored at out's dtype.
      for (int64_t e = 0; e < d; ++e) acc[static_cast<size_t>(e)] = 0.0f;
      for (int64_t j = 0; j <= jmax; ++j) {
        const float p = probs[static_cast<size_t>(j)] * inv;
        const int64_t voff = (j * hk + g) * d;
        for (int64_t e = 0; e < d; ++e) acc[static_cast<size_t>(e)] += p * LoadF32(value, voff + e);
      }
      for (int64_t e = 0; e < d; ++e) StoreF32(out, qoff + e, acc[static_cast<size_t>(e)]);
    }
  }
  });
}

// --- Qwen3.6 elementwise "glue" ops (M0.9 forward). Elementwise fusions of the
// small host-side loops between the big decode ops; all math f32, dims inferred
// from the tensor shapes.

float Sigmoid(float x) { return 1.0f / (1.0f + std::exp(-x)); }

// out[i] = F32ToBF16(in[i]); out bf16, in f32, same element count.
void CastBf16Kernel(Queue&, Tensor& out, const Tensor& in) {
  const int64_t n = out.Numel();
  ForRows(n, [&](int64_t r0, int64_t r1) {
    for (int64_t i = r0; i < r1; ++i) StoreF32(out, i, LoadF32(in, i));
  });
}

// out[i] = (float)in[i]; out f32, in bf16, same element count. CPU sibling of
// the CUDA CastF32 kernel (bf16 -> f32 widen); LoadF32 reads the bf16 source as
// f32 and StoreF32 writes it to the f32 destination (widening is exact).
void CastF32Kernel(Queue&, Tensor& out, const Tensor& in) {
  const int64_t n = out.Numel();
  ForRows(n, [&](int64_t r0, int64_t r1) {
    for (int64_t i = r0; i < r1; ++i) StoreF32(out, i, LoadF32(in, i));
  });
}

// x[m,n] *= col[n]; x f32 [M,N] (inner-contiguous rows, row stride x.stride[0]),
// col f32 [N]. CPU sibling of the CUDA MulColVecF32 kernel.
void MulColVecF32Kernel(Queue&, Tensor& x, const Tensor& col) {
  const int64_t m = x.shape[0], n = x.shape[1], rs = x.stride[0];
  const float* c = col.Ptr<float>();
  ForRows(m, [&](int64_t r0, int64_t r1) {
    for (int64_t i = r0; i < r1; ++i) {
      float* row = x.Ptr<float>() + i * rs;
      for (int64_t j = 0; j < n; ++j) row[j] *= c[j];
    }
  });
}

// Split fused [T, Hq*2*Dh] q/gate projection into q_out/gate_out [T,Hq,Dh].
void AttnGateSplitKernel(Queue&, Tensor& q_out, Tensor& gate_out, const Tensor& qgate) {
  const int64_t t = q_out.shape[0], hq = q_out.shape[1], dh = q_out.shape[2];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    for (int64_t h = 0; h < hq; ++h) {
      const int64_t base = i * (hq * 2 * dh) + h * 2 * dh;  // start of (i,h) pair
      const int64_t out_off = (i * hq + h) * dh;
      for (int64_t d = 0; d < dh; ++d) {
        StoreF32(q_out, out_off + d, LoadF32(qgate, base + d));
        StoreF32(gate_out, out_off + d, LoadF32(qgate, base + dh + d));
      }
    }
  }
  });
}

// out[i] = F32ToBF16(attn[i] * sigmoid(gate[i])); out bf16, attn/gate f32.
void SigmoidGateBf16Kernel(Queue&, Tensor& out, const Tensor& attn, const Tensor& gate) {
  const int64_t n = out.Numel();
  ForRows(n, [&](int64_t r0, int64_t r1) {
    for (int64_t i = r0; i < r1; ++i)
      StoreF32(out, i, LoadF32(attn, i) * Sigmoid(LoadF32(gate, i)));
  });
}

// GDN g/beta from raw projections (gdn-semantics.md §6). softplus threshold 20.
void GdnGBetaKernel(Queue&, Tensor& g_out, Tensor& beta_out, const Tensor& araw,
                    const Tensor& braw, const Tensor& a_log, const Tensor& dt_bias) {
  const int64_t t = g_out.shape[0], hv = g_out.shape[1];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    for (int64_t h = 0; h < hv; ++h) {
      const int64_t idx = i * hv + h;
      const int64_t aidx = i * araw.stride[0] + h;
      const int64_t bidx = i * braw.stride[0] + h;
      const float x = LoadF32(araw, aidx) + LoadF32(dt_bias, h);
      const float sp = x > 20.0f ? x : std::log1p(std::exp(x));  // softplus
      StoreF32(g_out, idx, -std::exp(LoadF32(a_log, h)) * sp);
      StoreF32(beta_out, idx, Sigmoid(LoadF32(braw, bidx)));
    }
  }
  });
}

// Split GDN conv [T, 2*key_dim+value_dim] into q/k [T,key_dim] and v [T,value_dim].
void GdnConvSplitKernel(Queue&, Tensor& q_out, Tensor& k_out, Tensor& v_out, const Tensor& conv) {
  const int64_t t = conv.shape[0], key_dim = q_out.Numel() / t, value_dim = v_out.Numel() / t;
  const int64_t conv_dim = 2 * key_dim + value_dim;
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const int64_t row = i * conv_dim;
    for (int64_t j = 0; j < key_dim; ++j) {
      StoreF32(q_out, i * key_dim + j, LoadF32(conv, row + j));
      StoreF32(k_out, i * key_dim + j, LoadF32(conv, row + key_dim + j));
    }
    for (int64_t j = 0; j < value_dim; ++j)
      StoreF32(v_out, i * value_dim + j, LoadF32(conv, row + 2 * key_dim + j));
  }
  });
}

// QkvSplit: GQA merged-qkv split with INDEPENDENT q/k/v dims (q_dim != k_dim).
void QkvSplitKernel(Queue&, Tensor& q_out, Tensor& k_out, Tensor& v_out, const Tensor& qkv) {
  const int64_t t = qkv.shape[0], q_dim = q_out.Numel() / t, k_dim = k_out.Numel() / t,
                v_dim = v_out.Numel() / t;
  const int64_t total = q_dim + k_dim + v_dim;
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const int64_t row = i * total;
    for (int64_t j = 0; j < q_dim; ++j) StoreF32(q_out, i * q_dim + j, LoadF32(qkv, row + j));
    for (int64_t j = 0; j < k_dim; ++j)
      StoreF32(k_out, i * k_dim + j, LoadF32(qkv, row + q_dim + j));
    for (int64_t j = 0; j < v_dim; ++j)
      StoreF32(v_out, i * v_dim + j, LoadF32(qkv, row + q_dim + k_dim + j));
  }
  });
}

// Fused GDN post-conv prep: GdnConvSplit + L2Norm(q) + L2Norm(k) + GdnGBeta in
// one pass (mirror of fla fused_gdn_prefill_post_conv). Bit-for-bit equal to
// composing those four ops (same f32 math, same softplus threshold 20).
void GdnPostConvKernel(Queue&, Tensor& q_out, Tensor& k_out, Tensor& v_out, Tensor& g_out,
                       Tensor& beta_out, const Tensor& conv, const Tensor& araw,
                       const Tensor& braw, const Tensor& a_log, const Tensor& dt_bias,
                       const L2NormArgs& args) {
  const int64_t t = conv.shape[0];
  const int64_t hk = q_out.shape[1], dk = q_out.shape[2];
  const int64_t hv = v_out.shape[1], dv = v_out.shape[2];
  const int64_t key_dim = hk * dk, value_dim = hv * dv;
  const int64_t conv_dim = 2 * key_dim + value_dim;
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const int64_t row = i * conv_dim;
    // q/k: split then l2norm over Dk, per head (plain SUM of squares, §4).
    for (int64_t h = 0; h < hk; ++h) {
      float qss = 0.0f, kss = 0.0f;
      for (int64_t j = 0; j < dk; ++j) {
        const float qv = LoadF32(conv, row + h * dk + j);
        const float kv = LoadF32(conv, row + key_dim + h * dk + j);
        qss += qv * qv;
        kss += kv * kv;
      }
      const float qinv = 1.0f / std::sqrt(qss + args.eps);
      const float kinv = 1.0f / std::sqrt(kss + args.eps);
      for (int64_t j = 0; j < dk; ++j) {
        const int64_t o = (i * hk + h) * dk + j;
        StoreF32(q_out, o, LoadF32(conv, row + h * dk + j) * qinv);
        StoreF32(k_out, o, LoadF32(conv, row + key_dim + h * dk + j) * kinv);
      }
    }
    // v: plain copy.
    for (int64_t j = 0; j < value_dim; ++j)
      StoreF32(v_out, i * value_dim + j, LoadF32(conv, row + 2 * key_dim + j));
    // g/beta from a/b + A_log/dt_bias (§6).
    for (int64_t h = 0; h < hv; ++h) {
      const int64_t idx = i * hv + h;
      const int64_t aidx = i * araw.stride[0] + h;
      const int64_t bidx = i * braw.stride[0] + h;
      const float x = LoadF32(araw, aidx) + LoadF32(dt_bias, h);
      const float sp = x > 20.0f ? x : std::log1p(std::exp(x));  // softplus
      StoreF32(g_out, idx, -std::exp(LoadF32(a_log, h)) * sp);
      StoreF32(beta_out, idx, Sigmoid(LoadF32(braw, bidx)));
    }
  }
  });
}

// out[t,c] = F32ToBF16(sigmoid(gl[t]) * sd[t*H+c]); shared-expert sigmoid gate.
void SharedExpertGateKernel(Queue&, Tensor& out, const Tensor& sd, const Tensor& gl) {
  const int64_t t = out.shape[0], h = out.shape[1];
  ForRows(t, [&](int64_t r0, int64_t r1) {
  for (int64_t i = r0; i < r1; ++i) {
    const float g = Sigmoid(LoadF32(gl, i));
    for (int64_t c = 0; c < h; ++c) StoreF32(out, i * h + c, g * LoadF32(sd, i * h + c));
  }
  });
}

// --- Fused declarative recipe (TDR): the Tier-1 single-pass INTERPRETER over the
// canonical (out, x, weight, residual) 4-operand shape. The Tier-0 composite is
// device-agnostic and lives in ops.cpp (it dispatches each opcode to the
// standalone vt:: op). This kernel is reached ONLY for Tier-1-able recipes (all
// 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. Bit-identical to vt::RmsNorm(residual) for
// kFusedAddRmsNorm — same f32 variance accumulation, gemma (1+w), residual-add
// rounding.
//
// Bit-identity across SEPARATELY COMPILED loops (here vs RmsNormKernel) holds
// because the build pins -ffp-contract=off (top-level CMakeLists): under the
// GCC/Clang default (fast), `sumsq += v*v` may be fma-contracted in one loop
// and not the other (observed: RmsNormKernel's residual path contracted, this
// interpreter not, on FMA-capable ISAs), skewing inv by 1 ulp on ~half the rows.

// Read one operand element by its CANONICAL index. Index 1 (weight) is per-column
// ([H]); the [T,H] operands index row*h + j.
float FusedLoad(uint8_t idx, int64_t row, int64_t j, int64_t h, const Tensor& x,
                const Tensor& weight, const Tensor* residual, const Tensor& out) {
  switch (idx) {
    case 0: return LoadF32(x, row * h + j);
    case 1: return LoadF32(weight, j);
    case 2: return LoadF32(*residual, row * h + j);
    case 3: return LoadF32(out, row * h + j);
    default: VT_CHECK(false, "fused_chain interp: operand index out of canonical range"); return 0.0f;
  }
}

// Store one element into a WRITABLE operand (residual=2 / out=3 only), rounding to
// that operand's dtype (bf16 residual mirrors vLLM's model-dtype residual).
void FusedStore(uint8_t idx, int64_t row, int64_t j, int64_t h, float v, Tensor* residual,
                Tensor& out) {
  switch (idx) {
    case 2: StoreF32(*residual, row * h + j, v); break;
    case 3: StoreF32(out, row * h + j, v); break;
    default: VT_CHECK(false, "fused_chain interp: step writes a read-only operand");
  }
}

// silu(x) = x * sigmoid(x); sigmoid in f32 (matches the standalone kernels' math).
inline float FSigmoid(float x) { return 1.0f / (1.0f + std::exp(-x)); }

// Tier 1 — scalar interpreter: walk the recipe over EACH ROW in one pass.
void FusedChainInterpKernel(Queue&, 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];
  ForRows(t, [&](int64_t rr0, int64_t rr1) {
  for (int64_t row = rr0; row < rr1; ++row) {
    for (int s = 0; s < r.n; ++s) {
      const FStep& st = r.steps[s];
      switch (st.op) {
        case FOp::kAdd:
        case FOp::kMul:
          for (int64_t j = 0; j < h; ++j) {
            const float a = FusedLoad(st.in[0], row, j, h, x, weight, residual, out);
            const float b = FusedLoad(st.in[1], row, j, h, x, weight, residual, out);
            FusedStore(st.out, row, j, h, st.op == FOp::kAdd ? a + b : a * b, residual, out);
          }
          break;
        case FOp::kSilu:
        case FOp::kSigmoid:
          for (int64_t j = 0; j < h; ++j) {
            const float a = FusedLoad(st.in[0], row, j, h, x, weight, residual, out);
            FusedStore(st.out, row, j, h, st.op == FOp::kSilu ? a * FSigmoid(a) : FSigmoid(a),
                       residual, out);
          }
          break;
        case FOp::kRmsNorm: {
          VT_CHECK(st.reduce == FReduce::kMeanSquare, "fused_chain: rmsnorm needs kMeanSquare");
          float sumsq = 0.0f;
          for (int64_t j = 0; j < h; ++j) {
            const float v = FusedLoad(st.in[0], row, j, h, x, weight, residual, out);
            sumsq += v * v;  // f32 variance accumulation
          }
          const float inv = 1.0f / std::sqrt(sumsq / static_cast<float>(h) + eps);
          for (int64_t j = 0; j < h; ++j) {
            const float v = FusedLoad(st.in[0], row, j, h, x, weight, residual, out);
            float wj = FusedLoad(st.in[1], row, j, h, x, weight, residual, out);
            if (st.gemma) wj += 1.0f;
            FusedStore(st.out, row, j, h, v * inv * wj, residual, out);
          }
          break;
        }
        default:
          VT_CHECK(false, "fused_chain interp: non-Tier-1 opcode reached the interpreter");
      }
    }
  }
  });
}

// 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 FusedChainKernel(Queue& q, Tensor& out, const Tensor& x, const Tensor& weight,
                      Tensor* residual, const FusedRecipe& r, float eps) {
  FusedChainInterpKernel(q, out, x, weight, residual, r, eps);
}

struct Registrar {
  Registrar() {
    // static_cast against the ops.h aliases ties kernel signatures to the
    // registration contract at compile time.
    RegisterOp(OpId::kMatmul, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MatmulFn>(&MatmulKernel)));
    RegisterOp(OpId::kMatmulBT, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MatmulFn>(&MatmulBTKernel)));
    RegisterOp(OpId::kBatchedMatmul, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<BatchedMatmulFn>(&BatchedMatmulKernel)));
    RegisterOp(
        OpId::kConcatMlaNopeRope, DeviceType::kCPU,
        reinterpret_cast<void*>(static_cast<ConcatMlaNopeRopeFn>(&ConcatMlaNopeRopeKernel)));
    RegisterOp(OpId::kRmsNorm, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<RmsNormFn>(&RmsNormKernel)));
    RegisterOp(OpId::kRmsNormQuantFp8, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<RmsNormQuantFp8Fn>(&RmsNormQuantFp8Kernel)));
    RegisterOp(OpId::kSiluAndMul, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<SiluAndMulFn>(&SiluAndMulKernel)));
    RegisterOp(OpId::kMoeSiluMul, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MoeSiluMulFn>(&MoeSiluMulKernel)));
    RegisterOp(OpId::kScaledFp4Quant, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<ScaledFp4QuantFn>(&ScaledFp4QuantKernel)));
    RegisterOp(OpId::kSiluMulFp4Quant, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<SiluMulFp4QuantFn>(&SiluMulFp4QuantKernel)));
    RegisterOp(OpId::kSigmoidGateFp4Quant, DeviceType::kCPU,
               reinterpret_cast<void*>(
                   static_cast<SigmoidGateFp4QuantFn>(&SigmoidGateFp4QuantKernel)));
    RegisterOp(
        OpId::kSiluAndMulFp4Quant, DeviceType::kCPU,
        reinterpret_cast<void*>(static_cast<SiluAndMulFp4QuantFn>(
            &SiluAndMulFp4QuantKernel)));
    RegisterOp(OpId::kMatmulNvfp4Fp4, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MatmulNvfp4Fp4Fn>(&MatmulNvfp4Fp4Kernel)));
    RegisterOp(OpId::kEmbedding, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<EmbeddingFn>(&EmbeddingKernel)));
    RegisterOp(OpId::kRopeNeox, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<RopeFn>(&RopeNeoxKernel)));
    RegisterOp(
        OpId::kRopeFromCache, DeviceType::kCPU,
        reinterpret_cast<void*>(
            static_cast<RopeFromCacheFn>(&RopeFromCacheKernel)));
    RegisterOp(OpId::kCausalConv1dFwd, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<CausalConv1dFwdFn>(&CausalConv1dFwdKernel)));
    RegisterOp(
        OpId::kCausalConv1dUpdate, DeviceType::kCPU,
        reinterpret_cast<void*>(static_cast<CausalConv1dUpdateFn>(&CausalConv1dUpdateKernel)));
    RegisterOp(OpId::kL2Norm, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<L2NormFn>(&L2NormKernel)));
    RegisterOp(OpId::kRmsNormGated, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<RmsNormGatedFn>(&RmsNormGatedKernel)));
    RegisterOp(
        OpId::kRmsNormGatedQuantFp8, DeviceType::kCPU,
        reinterpret_cast<void*>(
            static_cast<RmsNormGatedQuantFp8Fn>(&RmsNormGatedQuantFp8Kernel)));
    RegisterOp(OpId::kGdnPrefill, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<GdnPrefillFn>(&GdnPrefillKernel)));
    RegisterOp(OpId::kGdnDecode, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<GdnDecodeFn>(&GdnDecodeKernel)));
    RegisterOp(
        OpId::kGdnPackedDecode, DeviceType::kCPU,
        reinterpret_cast<void*>(static_cast<GdnPackedDecodeFn>(
            &GdnPackedDecodeKernel)));
    RegisterOp(
        OpId::kGdnStateGather, DeviceType::kCPU,
        reinterpret_cast<void*>(
            static_cast<GdnStateGatherFn>(&GdnStateGatherKernel)));
    RegisterOp(
        OpId::kGdnStateScatter, DeviceType::kCPU,
        reinterpret_cast<void*>(
            static_cast<GdnStateScatterFn>(&GdnStateScatterKernel)));
    RegisterOp(OpId::kMoeRouterTopK, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MoeRouterTopKFn>(&MoeRouterTopKKernel)));
    RegisterOp(OpId::kMoeCombine, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MoeCombineFn>(&MoeCombineKernel)));
    RegisterOp(OpId::kAttention, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<AttentionFn>(&AttentionKernel)));
    RegisterOp(OpId::kCastBf16, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<CastBf16Fn>(&CastBf16Kernel)));
    RegisterOp(OpId::kCastF32, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<CastF32Fn>(&CastF32Kernel)));
    RegisterOp(OpId::kMulColVecF32, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<MulColVecF32Fn>(&MulColVecF32Kernel)));
    RegisterOp(OpId::kAttnGateSplit, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<AttnGateSplitFn>(&AttnGateSplitKernel)));
    RegisterOp(OpId::kSigmoidGateBf16, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<SigmoidGateBf16Fn>(&SigmoidGateBf16Kernel)));
    RegisterOp(OpId::kGdnGBeta, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<GdnGBetaFn>(&GdnGBetaKernel)));
    RegisterOp(OpId::kGdnConvSplit, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<GdnConvSplitFn>(&GdnConvSplitKernel)));
    RegisterOp(OpId::kQkvSplit, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<QkvSplitFn>(&QkvSplitKernel)));
    RegisterOp(OpId::kGdnPostConv, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<GdnPostConvFn>(&GdnPostConvKernel)));
    RegisterOp(OpId::kRopeCosSinCache, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<RopeCosSinCacheFn>(&RopeCosSinCacheKernel)));
    RegisterOp(
        OpId::kAttnQkNormRopeGate, DeviceType::kCPU,
        reinterpret_cast<void*>(static_cast<AttnQkNormRopeGateFn>(&AttnQkNormRopeGateKernel)));
    RegisterOp(OpId::kSharedExpertGate, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<SharedExpertGateFn>(&SharedExpertGateKernel)));
    RegisterOp(OpId::kFusedChain, DeviceType::kCPU,
               reinterpret_cast<void*>(static_cast<FusedChainFn>(&FusedChainKernel)));
  }
} registrar;

}  // namespace
}  // namespace vt::cpu