chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/minimax/fused_gemma_qknorm_rope.cuh

@@ -0,0 +1,200 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+// Up to this many norm "groups" are fused into one launch. A group is a
+// contiguous run of heads (within the per-token row) that share one norm
+// weight and all receive RoPE. Heads not covered by any group (e.g. the V /
+// index-V regions) are simply never assigned a job and left untouched.
+//
+// For MiniMax-M3 the groups are, in order: Q (main), K (main), index-Q,
+// index-K -- which is why 4 slots are enough. This mirrors the multi-branch
+// single-launch design of fused_store_kv_index.cuh (one kernel handles main
+// K/V + index K/V), here applied to GemmaRMSNorm + partial RoPE.
+constexpr int kMaxGroups = 4;
+
+struct FusedGemmaQKNormParams {
+  bf16_t* __restrict__ qkv;
+  const bf16_t* __restrict__ weight[kMaxGroups];  // per-group norm weight [head_dim]
+  uint32_t group_offset[kMaxGroups];              // head offset of the group in the row
+  uint32_t group_count[kMaxGroups];               // number of heads in the group
+  uint32_t num_groups;
+  uint32_t total_heads;  // sum of group_count[0..num_groups)
+  const float* __restrict__ cos_sin_cache;
+  const void* __restrict__ positions;  // dtype depends on PosT template
+  uint32_t num_tokens;
+  int64_t token_stride;
+  float eps;
+};
+
+template <typename PosT, int64_t kHeadDim, int64_t kRopeDim, bool kUsePDL>
+struct FusedTrait {
+  static_assert(kHeadDim == 128 && kRopeDim == 64, "kernel specialized for HEAD_DIM=128, ROTARY_DIM=64");
+  static constexpr uint32_t kWorkerSize = device::kWarpThreads;  // full warp per head
+  SGL_DEVICE static void forward(const FusedGemmaQKNormParams& params) {
+    using namespace device;
+    const auto tx = threadIdx.x;
+    const auto bx = blockIdx.x;
+    const auto lane_id = tx % kWorkerSize;
+    const auto work_id = (bx * blockDim.x + tx) / kWorkerSize;
+    const auto total_heads = params.total_heads;
+    if (work_id >= params.num_tokens * total_heads) return;
+    const auto token_id = work_id / total_heads;
+    const auto grouped_head = work_id % total_heads;
+
+    // Resolve which group this job belongs to (num_groups <= kMaxGroups, so a
+    // short scan is cheaper than any precomputed table and warp-uniform).
+    uint32_t g = 0, base = 0;
+    for (uint32_t i = 0; i < params.num_groups; ++i) {
+      const auto cnt = params.group_count[i];
+      if (grouped_head < base + cnt) {
+        g = i;
+        break;
+      }
+      base += cnt;
+    }
+    const auto local_head = grouped_head - base;
+    const auto head_id = params.group_offset[g] + local_head;
+    const auto weight = params.weight[g];
+    const auto input = params.qkv + token_id * params.token_stride + head_id * kHeadDim;
+
+    // prefetch weight and rope index
+    const auto idx_0 = lane_id + 0;   // rope first half  [0,32)
+    const auto idx_1 = lane_id + 32;  // rope second half [32, 64)
+    const auto idx_2 = lane_id + 64;  // pass
+    const auto idx_3 = lane_id + 96;  // pass
+    const auto w0_bf16 = weight[idx_0];
+    const auto w1_bf16 = weight[idx_1];
+    const auto w2_bf16 = weight[idx_2];
+    const auto w3_bf16 = weight[idx_3];
+    const auto rope_idx = static_cast<const PosT*>(params.positions)[token_id];
+    PDLWaitPrimary<kUsePDL>();
+
+    // load input and compute RMS, fp32 accumulation for stability
+    const auto i0_bf16 = input[idx_0];
+    const auto i1_bf16 = input[idx_1];
+    const auto i2_bf16 = input[idx_2];
+    const auto i3_bf16 = input[idx_3];
+    const auto [i0, i1] = cast<fp32x2_t>(bf16x2_t{i0_bf16, i1_bf16});
+    const auto [i2, i3] = cast<fp32x2_t>(bf16x2_t{i2_bf16, i3_bf16});
+    const auto ss = warp::reduce_sum(i0 * i0 + i1 * i1 + i2 * i2 + i3 * i3);
+    const auto inv_rms = rsqrtf(ss / static_cast<float>(kHeadDim) + params.eps);
+
+    // apply norm
+    const auto [w0, w1] = cast<fp32x2_t>(bf16x2_t{w0_bf16, w1_bf16});
+    const auto [w2, w3] = cast<fp32x2_t>(bf16x2_t{w2_bf16, w3_bf16});
+    const auto n0 = i0 * inv_rms * (1.0f + w0);
+    const auto n1 = i1 * inv_rms * (1.0f + w1);
+    const auto n2 = i2 * inv_rms * (1.0f + w2);
+    const auto n3 = i3 * inv_rms * (1.0f + w3);
+
+    const auto cs = params.cos_sin_cache + rope_idx * kRopeDim;
+    const auto cos = cs[lane_id];
+    const auto sin = cs[lane_id + kRopeDim / 2];
+
+    // apply rope to the first kRopeDim dims, and write back
+    device::PDLTriggerSecondary<kUsePDL>();
+    const auto o0 = n0 * cos - n1 * sin;
+    const auto o1 = n1 * cos + n0 * sin;
+    const auto [o0_bf16, o1_bf16] = cast<bf16x2_t>(fp32x2_t{o0, o1});
+    const auto [o2_bf16, o3_bf16] = cast<bf16x2_t>(fp32x2_t{n2, n3});
+    input[idx_0] = o0_bf16;
+    input[idx_1] = o1_bf16;
+    input[idx_2] = o2_bf16;
+    input[idx_3] = o3_bf16;
+  }
+};
+
+template <typename Trait>
+__global__ void fused_gemma_qknorm_rope_kernel(const __grid_constant__ FusedGemmaQKNormParams params) {
+  return Trait::forward(params);
+}
+
+// Multi-group fused GemmaRMSNorm + partial NeoX RoPE, in place over `qkv`.
+//
+// Up to kMaxGroups norm groups are passed as (weight, head offset, head count)
+// triples. `w0..w3` are the per-group norm weights ([head_dim] bf16 each); the
+// `offN` / `cntN` scalars give each group's head offset (within the per-token
+// row) and head count. These offsets/counts are host-known constants (passed as
+// scalars, never device tensors) so the launch stays CUDA-graph capturable.
+// Weight slots beyond `num_groups` may be dummies (e.g. == w0); the kernel
+// never reads them because `num_groups` bounds the group scan. The main Q/K and
+// index-Q/index-K heads are all normed and rotated in one launch; the V /
+// index-V heads, lying outside every group, are left untouched.
+template <typename PosT, int64_t HEAD_DIM, int64_t ROTARY_DIM, bool kUsePDL>
+void fused_gemma_qknorm_rope(
+    tvm::ffi::TensorView qkv,
+    tvm::ffi::TensorView w0,
+    tvm::ffi::TensorView w1,
+    tvm::ffi::TensorView w2,
+    tvm::ffi::TensorView w3,
+    tvm::ffi::TensorView cos_sin_cache,
+    tvm::ffi::TensorView positions,
+    int64_t off0,
+    int64_t cnt0,
+    int64_t off1,
+    int64_t cnt1,
+    int64_t off2,
+    int64_t cnt2,
+    int64_t off3,
+    int64_t cnt3,
+    int64_t num_groups,
+    double eps) {
+  using namespace host;
+  auto N = SymbolicSize{"num_tokens"};
+  auto device = SymbolicDevice{};
+  constexpr auto D = HEAD_DIM;
+  constexpr auto R = ROTARY_DIM;
+  device.set_options<kDLCUDA>();
+  TensorMatcher({N, -1}).with_dtype<bf16_t>().with_device(device).verify(qkv);
+  TensorMatcher({D}).with_dtype<bf16_t>().with_device(device).verify(w0).verify(w1).verify(w2).verify(w3);
+  TensorMatcher({-1, R}).with_dtype<fp32_t>().with_device(device).verify(cos_sin_cache);
+  TensorMatcher({N}).with_dtype<PosT>().with_device(device).verify(positions);
+
+  RuntimeCheck(num_groups >= 1 && num_groups <= kMaxGroups);
+
+  const tvm::ffi::TensorView weights[kMaxGroups] = {w0, w1, w2, w3};
+  const int64_t offsets[kMaxGroups] = {off0, off1, off2, off3};
+  const int64_t counts[kMaxGroups] = {cnt0, cnt1, cnt2, cnt3};
+
+  auto params = FusedGemmaQKNormParams{};
+  params.qkv = static_cast<bf16_t*>(qkv.data_ptr());
+  params.cos_sin_cache = static_cast<const float*>(cos_sin_cache.data_ptr());
+  params.positions = positions.data_ptr();
+  params.num_tokens = static_cast<uint32_t>(N.unwrap());
+  params.num_groups = static_cast<uint32_t>(num_groups);
+  params.token_stride = static_cast<int64_t>(qkv.stride(0));
+  params.eps = static_cast<float>(eps);
+
+  uint32_t total_heads = 0;
+  for (int64_t i = 0; i < kMaxGroups; ++i) {
+    const auto cnt = (i < num_groups) ? counts[i] : 0;
+    params.weight[i] = static_cast<const bf16_t*>(weights[i].data_ptr());
+    params.group_offset[i] = static_cast<uint32_t>(i < num_groups ? offsets[i] : 0);
+    params.group_count[i] = static_cast<uint32_t>(cnt);
+    total_heads += static_cast<uint32_t>(cnt);
+  }
+  params.total_heads = total_heads;
+
+  using Trait = FusedTrait<PosT, HEAD_DIM, ROTARY_DIM, kUsePDL>;
+  const auto needed_threads = static_cast<int64_t>(params.num_tokens) * total_heads * device::kWarpThreads;
+  RuntimeCheck(needed_threads < std::numeric_limits<uint32_t>::max());
+  if (needed_threads == 0) return;
+  const uint32_t block_size = 256u;
+  const uint32_t num_blocks = div_ceil(static_cast<uint32_t>(needed_threads), block_size);
+  LaunchKernel(num_blocks, block_size, device.unwrap())  //
+      .enable_pdl(kUsePDL)(fused_gemma_qknorm_rope_kernel<Trait>, params);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/minimax/fused_store_kv_index.cuh

@@ -0,0 +1,171 @@
+#include <sgl_kernel/tensor.h>  // For TensorMatcher, SymbolicSize/DType/Device
+#include <sgl_kernel/utils.h>   // For RuntimeCheck, div_ceil, pointer::offset
+
+#include <sgl_kernel/utils.cuh>  // For LaunchKernel, SGL_DEVICE, PDL helpers
+#include <sgl_kernel/vec.cuh>    // For AlignedVector
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+struct StoreKVIndexParams {
+  const void* __restrict__ input_ptrs[4];
+  void* __restrict__ cache_ptrs[4];
+  const void* __restrict__ indices;
+  int64_t input_stride_bytes[4];  // k stride, v stride, idx_k stride, idx_v stride
+  int64_t cache_stride_bytes[4];  // k stride, v stride, idx_k stride, idx_v stride
+  uint32_t num_k_heads;
+  uint32_t num_total_heads;  // 2 * num_k_heads + 1 + has_index_v
+  uint32_t total_jobs;       // batch_size * heads_per_token
+};
+
+template <int64_t kHeadBytes, bool kUsePDL>
+struct StoreTrait {
+  static_assert(kHeadBytes % 16 == 0, "head bytes must be a multiple of 16 (128-bit vector)");
+  using vec_t = device::AlignedVector<uint32_t, 4>;         // 16 bytes / thread
+  static constexpr uint32_t kWorkerSize = kHeadBytes / 16;  // threads per head
+
+  template <typename T>
+  SGL_DEVICE static void forward(const StoreKVIndexParams& params) {
+    using namespace device;
+    const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+    const uint32_t work_id = tid / kWorkerSize;  // which (token, head) job
+    const uint32_t lane = tid % kWorkerSize;     // which 16-byte chunk of the head
+    if (work_id >= params.total_jobs) return;
+
+    const uint32_t num_total_heads = params.num_total_heads;
+    const uint32_t token_id = work_id / num_total_heads;
+    const uint32_t slot_id = work_id % num_total_heads;
+    const uint32_t num_k_heads = params.num_k_heads;
+
+    const auto loc = static_cast<const T*>(params.indices)[token_id];
+    uint32_t head_id, which;
+    if (slot_id < num_k_heads) {
+      head_id = slot_id;
+      which = 0;  // K
+    } else if (slot_id < 2 * num_k_heads) {
+      head_id = slot_id - num_k_heads;
+      which = 1;  // V
+    } else if (slot_id == 2 * num_k_heads) {
+      head_id = 0;
+      which = 2;  // idx K
+    } else {
+      head_id = 0;
+      which = 3;  // idx V
+    }
+    const auto cache_ptr = static_cast<char*>(params.cache_ptrs[which]);
+    const auto input_ptr = static_cast<const char*>(params.input_ptrs[which]);
+    const auto cache_stride = params.cache_stride_bytes[which];
+    const auto input_stride = params.input_stride_bytes[which];
+    const auto src = pointer::offset(input_ptr, token_id * input_stride, head_id * kHeadBytes);
+    const auto dst = pointer::offset(cache_ptr, loc * cache_stride, head_id * kHeadBytes);
+    PDLWaitPrimary<kUsePDL>();
+    vec_t chunk;
+    chunk.load(src, lane);
+    chunk.store(dst, lane);
+    PDLTriggerSecondary<kUsePDL>();
+  }
+};
+
+template <typename Trait, typename T>
+__global__ void store_kv_index_kernel(const __grid_constant__ StoreKVIndexParams params) {
+  Trait::template forward<T>(params);
+}
+
+// idx_v / idx_v_cache may be dummies (== idx_k / idx_k_cache) when the layer
+// has no index value; the caller signals "no V" via heads_per_token == 2*hkv+1
+// so the index-V branch is never taken.
+template <int64_t kHeadBytes, bool kUsePDL>
+void store_kv_index(
+    tvm::ffi::TensorView k,
+    tvm::ffi::TensorView v,
+    tvm::ffi::TensorView k_cache,
+    tvm::ffi::TensorView v_cache,
+    tvm::ffi::TensorView idx_k,
+    tvm::ffi::TensorView idx_k_cache,
+    tvm::ffi::TensorView idx_v,
+    tvm::ffi::TensorView idx_v_cache,
+    tvm::ffi::TensorView indices,
+    int64_t num_kv_heads,
+    int64_t heads_per_token) {
+  using namespace host;
+  auto B = SymbolicSize{"batch"};
+  auto Mrow = SymbolicSize{"main_row"};  // num_kv_heads * head_dim
+  auto Drow = SymbolicSize{"idx_row"};   // head_dim
+  auto dtype = SymbolicDType{};
+  auto indice_dtype = SymbolicDType{};
+  auto device = SymbolicDevice{};
+  device.set_options<kDLCUDA>();
+
+  // All cache/source tensors share the same store dtype (fast-path precondition).
+  TensorMatcher({B, Mrow})  //
+      .with_strides({-1, 1})
+      .with_dtype(dtype)
+      .with_device(device)
+      .verify(k)
+      .verify(v);
+  TensorMatcher({-1, Mrow})  //
+      .with_strides({-1, 1})
+      .with_dtype(dtype)
+      .with_device(device)
+      .verify(k_cache)
+      .verify(v_cache);
+  TensorMatcher({B, Drow})  //
+      .with_strides({-1, 1})
+      .with_dtype(dtype)
+      .with_device(device)
+      .verify(idx_k)
+      .verify(idx_v);
+  TensorMatcher({-1, Drow})  //
+      .with_strides({-1, 1})
+      .with_dtype(dtype)
+      .with_device(device)
+      .verify(idx_k_cache)
+      .verify(idx_v_cache);
+  TensorMatcher({B})  //
+      .with_dtype<int32_t, int64_t>(indice_dtype)
+      .with_device(device)
+      .verify(indices);
+
+  const int64_t dsize = dtype_bytes(dtype.unwrap());
+  RuntimeCheck(kHeadBytes == Drow.unwrap() * dsize);
+  RuntimeCheck(Mrow.unwrap() == num_kv_heads * Drow.unwrap());
+
+  const auto params = StoreKVIndexParams{
+      .input_ptrs = {k.data_ptr(), v.data_ptr(), idx_k.data_ptr(), idx_v.data_ptr()},
+      .cache_ptrs = {k_cache.data_ptr(), v_cache.data_ptr(), idx_k_cache.data_ptr(), idx_v_cache.data_ptr()},
+      .indices = indices.data_ptr(),
+      .input_stride_bytes =
+          {
+              k.stride(0) * dsize,
+              v.stride(0) * dsize,
+              idx_k.stride(0) * dsize,
+              idx_v.stride(0) * dsize,
+          },
+      .cache_stride_bytes =
+          {
+              k_cache.stride(0) * dsize,
+              v_cache.stride(0) * dsize,
+              idx_k_cache.stride(0) * dsize,
+              idx_v_cache.stride(0) * dsize,
+          },
+      .num_k_heads = static_cast<uint32_t>(num_kv_heads),
+      .num_total_heads = static_cast<uint32_t>(heads_per_token),
+      .total_jobs = static_cast<uint32_t>(B.unwrap() * heads_per_token),
+  };
+  if (params.total_jobs == 0) return;
+
+  using Trait = StoreTrait<kHeadBytes, kUsePDL>;
+  constexpr uint32_t kBlockSize = 256u;
+  const uint32_t num_blocks = div_ceil(params.total_jobs * Trait::kWorkerSize, kBlockSize);
+  const auto kernel = indice_dtype.is_type<int32_t>()  //
+                          ? store_kv_index_kernel<Trait, int32_t>
+                          : store_kv_index_kernel<Trait, int64_t>;
+  LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+      .enable_pdl(kUsePDL)(kernel, params);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/minimax/minimax_decode_topk.cuh

@@ -0,0 +1,535 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cfloat>
+#include <cstdint>
+#if defined(__HIP_PLATFORM_AMD__)
+static constexpr unsigned long long kWarpSyncMask = 0xFFFFFFFFFFFFFFFFull;
+#else
+#include <math_constants.h>
+static constexpr unsigned int kWarpSyncMask = 0xFFFFFFFFu;
+#endif
+
+namespace {
+
+// Block top-k selection over a per-(head, batch) row of block scores, run by one
+// CTA of TopKTrait::kCTASize threads. Picks the `topk` highest-scoring block ids
+// (k_eff = min(topk, num_blocks)). Three size regimes, chosen by num_blocks:
+//   * <= kSmallThreshold : O(n^2) rank-by-compare (no radix).
+//   * <= kCTASize        : 4-pass 8-bit radix, one element per thread in a reg.
+//   * <= kMaxNumBlocks   : 4-pass 8-bit radix, kIters elements per thread cached
+//                          in registers (row read from global exactly once);
+//                          liveness is a uint32_t bitmask, selection is an
+//                          in-loop scatter -- nothing is cached in shared memory.
+// The trivial case num_blocks <= topk (every block selected) is handled by the
+// kernels below, outside the Trait.
+struct TopKTrait {
+  static constexpr uint32_t kMaxTopK = 32;
+  static constexpr uint32_t kCTASize = 512;
+  static constexpr uint32_t kNumWarps = kCTASize / device::kWarpThreads;
+  static constexpr uint32_t kMaxNumBlocks = 4096;  // block topk
+  static constexpr uint32_t kSmallThreshold = 8 * kNumWarps;
+  static constexpr uint32_t kRadixBits = 8;
+  static constexpr uint32_t kRadixSize = 1 << kRadixBits;
+  static constexpr float kNegInf = -std::numeric_limits<float>::infinity();
+
+  struct Smem {
+    uint32_t warp_sum[kNumWarps];
+    alignas(128) uint32_t counter;
+    alignas(128) uint32_t counter_final;
+    alignas(128) uint32_t threshold_bin;
+    uint32_t equal_count;
+    uint32_t above_count;
+    uint32_t histogram[2][kRadixSize];    // 8 bit radix
+    float small_scores[kSmallThreshold];  // small (O(n^2)) path only
+  };
+
+  SGL_DEVICE static void forward(
+      const float* __restrict__ scores,
+      const uint32_t num_blocks,
+      int32_t* __restrict__ topk_out,
+      const uint32_t topk,
+      Smem* smem) {
+    using namespace device;
+    const auto tx = threadIdx.x;
+    __builtin_assume(tx < kCTASize);
+    const auto warp_id = tx / kWarpThreads;
+    const auto lane_id = tx % kWarpThreads;
+
+    constexpr auto is_greater = [](float x, float y, int32_t delta) {
+      return (x > y) || ((x == y) && delta < 0);  // lower block id wins
+    };
+    constexpr auto warp_inclusive_sum = [](uint32_t lane_id, uint32_t val) {
+#pragma unroll
+      for (uint32_t offset = 1; offset < 32; offset *= 2) {
+        uint32_t n = __shfl_up_sync(kWarpSyncMask, val, offset, 32);
+        if (lane_id >= offset) val += n;
+      }
+      return val;
+    };
+    constexpr auto clip_nan = [](float x) { return x != x ? kNegInf : x; };
+    constexpr auto score_to_key = [](float x) {
+      uint32_t b = __float_as_uint(x);
+      return (b & 0x80000000u) ? ~b : (b | 0x80000000u);
+    };
+    // Find the radix bin holding the topk_remain-th largest of `total_active`
+    // elements currently counted in `histogram`. Writes threshold_bin (the bin),
+    // above_count (elements strictly above it), equal_count (elements in it).
+    const auto find_threshold = [&](uint32_t* histogram, uint32_t total_active, uint32_t topk_remain) {
+      using namespace device;
+      uint32_t hist_val = 0;
+      uint32_t warp_inc = 0;
+      if (tx < kRadixSize) {
+        hist_val = histogram[tx];
+        warp_inc = warp_inclusive_sum(lane_id, hist_val);
+        if (lane_id == kWarpThreads - 1) smem->warp_sum[warp_id] = warp_inc;
+      }
+      __syncthreads();
+      if (tx < kRadixSize) {
+        const auto inter = warp::reduce_sum(lane_id < warp_id ? smem->warp_sum[lane_id] : 0);
+        const auto prefix = inter + warp_inc;      // count in bins [0, tx]
+        const auto above = total_active - prefix;  // count in bins ABOVE tx
+        if (above < topk_remain && above + hist_val >= topk_remain) {
+          smem->threshold_bin = tx;
+          smem->above_count = above;
+          smem->equal_count = hist_val;
+        }
+      }
+      __syncthreads();
+    };
+
+    if (num_blocks <= kSmallThreshold) {
+      // O(n^2) compare: each block's rank = #blocks that outrank it; the ones
+      // with rank < topk are selected (rank is its position in topk_out).
+      static_assert(kSmallThreshold <= kCTASize);
+      if (tx < num_blocks) smem->small_scores[tx] = clip_nan(scores[tx]);
+      __syncthreads();
+      constexpr uint32_t kNumCandidates = kSmallThreshold / kNumWarps;
+      constexpr uint32_t kNumTargets = kSmallThreshold / kWarpThreads;
+      float candidates[kNumCandidates];
+      float target[kNumTargets];
+#pragma unroll
+      for (uint32_t i = 0; i < kNumTargets; ++i) {
+        const auto idx = lane_id + i * kWarpThreads;
+        target[i] = (idx < num_blocks) ? smem->small_scores[idx] : kNegInf;
+      }
+#pragma unroll
+      for (uint32_t i = 0; i < kNumCandidates; ++i) {
+        const auto idx = warp_id + i * kNumWarps;
+        candidates[i] = (idx < num_blocks) ? smem->small_scores[idx] : kNegInf;
+      }
+
+#pragma unroll
+      for (uint32_t i = 0; i < kNumCandidates; ++i) {
+        const int32_t idx = warp_id + i * kNumWarps;
+        if (idx >= static_cast<int32_t>(num_blocks)) break;
+        uint32_t rank = 0;
+#pragma unroll
+        for (uint32_t j = 0; j < kNumTargets; ++j) {
+          const int32_t delta = lane_id + j * kWarpThreads - idx;
+          // partial rank = how many of this lane's targets outrank the candidate
+          rank += is_greater(target[j], candidates[i], delta);
+        }
+        // full rank = sum of the per-lane partial ranks across the warp
+        rank = warp::reduce_sum(rank);
+        if (rank < topk) topk_out[rank] = idx;
+      }
+    } else if (num_blocks <= kCTASize) {
+      // 4-pass 8-bit radix select, one element per thread held in a register.
+      bool active = tx < num_blocks;
+      const auto value = active ? clip_nan(scores[tx]) : kNegInf;
+      const auto key = score_to_key(value);
+      uint32_t topk_remain = topk;
+      uint32_t write_pos = topk;  // sentinel: not selected
+      if (tx < kRadixSize) smem->histogram[0][tx] = 0;
+      if (tx == kRadixSize) smem->counter = smem->counter_final = 0;
+      __syncthreads();
+      uint32_t total_active = num_blocks;
+
+#pragma unroll
+      for (int round = 0; round < 4; round++) {
+        const uint32_t shift = 24 - round * 8;
+        const uint32_t bin = (key >> shift) & 0xFFu;
+        const auto hist_idx = round % 2;
+        const auto histogram = smem->histogram[hist_idx];
+
+        if (active) atomicAdd(&histogram[bin], 1);
+        if (round < 3 && tx < kRadixSize) smem->histogram[hist_idx ^ 1][tx] = 0;
+        __syncthreads();
+
+        find_threshold(histogram, total_active, topk_remain);
+
+        const auto threshold_bin = smem->threshold_bin;
+        const auto above_count = smem->above_count;
+        const auto equal_count = smem->equal_count;
+
+        if (round < 3) total_active = equal_count;
+        topk_remain -= above_count;
+
+        // scatter: above -> selected now; equal at the last pass -> keep the rest
+        if (active) {
+          if (bin > threshold_bin) {
+            write_pos = atomicAdd(&smem->counter, 1);
+            active = false;
+          } else if (bin < threshold_bin) {
+            active = false;
+          } else if (round == 3) {
+            write_pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1);
+          }
+          // bin == threshold && round < 3: stay active for the next pass
+        }
+
+        if (round == 3 || topk_remain == 0) break;
+      }
+
+      if (write_pos < topk) topk_out[write_pos] = tx;
+    } else {
+      // num_blocks in (kCTASize, kMaxNumBlocks]: each thread caches its (up to
+      // kIters) slice of the row in registers -- read from global exactly ONCE --
+      // then runs the same 4-pass radix select as the single-element path looped
+      // over those slots. Liveness is a uint32_t bitmask (bit i = slot i still in
+      // the running set), so there is no per-element flag array; selection is an
+      // in-loop scatter, so there is no per-element position array. Nothing is
+      // cached in shared memory beyond the histogram.
+      constexpr uint32_t kIters = kMaxNumBlocks / kCTASize;
+      static_assert(kIters <= 32, "active liveness is packed into a uint32_t");
+      uint32_t key[kIters];
+      uint32_t active = 0;
+#pragma unroll
+      for (uint32_t i = 0; i < kIters; ++i) {
+        const uint32_t idx = i * kCTASize + tx;
+        if (idx < num_blocks) {
+          key[i] = score_to_key(clip_nan(scores[idx]));
+          active |= 1u << i;
+        }
+      }
+      if (tx < kRadixSize) smem->histogram[0][tx] = 0;
+      if (tx == kRadixSize) smem->counter = smem->counter_final = 0;
+      __syncthreads();
+
+      uint32_t topk_remain = topk;
+      uint32_t total_active = num_blocks;
+
+#pragma unroll
+      for (int round = 0; round < 4; ++round) {
+        const uint32_t shift = 24 - round * 8;
+        const auto hb = round & 1;
+
+#pragma unroll
+        for (uint32_t i = 0; i < kIters; ++i)
+          if (active & (1u << i)) atomicAdd(&smem->histogram[hb][(key[i] >> shift) & 0xFFu], 1);
+        if (round < 3 && tx < kRadixSize) smem->histogram[hb ^ 1][tx] = 0;
+        __syncthreads();
+
+        find_threshold(smem->histogram[hb], total_active, topk_remain);
+        const auto threshold_bin = smem->threshold_bin;
+        const auto above_count = smem->above_count;
+        const auto equal_count = smem->equal_count;
+
+        if (round < 3) total_active = equal_count;
+        topk_remain -= above_count;
+
+#pragma unroll
+        for (uint32_t i = 0; i < kIters; ++i) {
+          if (active & (1u << i)) {
+            const uint32_t bin = (key[i] >> shift) & 0xFFu;
+            if (bin > threshold_bin) {
+              topk_out[atomicAdd(&smem->counter, 1)] = i * kCTASize + tx;
+              active &= ~(1u << i);
+            } else if (bin < threshold_bin) {
+              active &= ~(1u << i);
+            } else if (round == 3) {
+              const auto pos = topk - topk_remain + atomicAdd(&smem->counter_final, 1);
+              if (pos < topk) topk_out[pos] = i * kCTASize + tx;
+            }
+            // bin == threshold && round < 3: slot stays live for the next pass
+          }
+        }
+
+        if (round == 3 || topk_remain == 0) break;
+      }
+    }
+  }
+};
+
+// -------------------------------------------------------------------------
+// Kernels: one CTA (kCTASize threads) per (head, batch) row. The trivial case
+// num_blocks <= topk (every block selected) is special-judged here, outside the
+// Trait; otherwise the Trait selects the top-k block ids.
+// -------------------------------------------------------------------------
+
+// Block-id output: topk_idx[h, b, 0:k_eff) = selected block ids (front-packed,
+// unordered), [k_eff:topk) = -1.
+template <typename SeqLenT, bool kUsePDL>
+__global__ void minimax_decode_topk_block_kernel(
+    const float* __restrict__ score,
+    const SeqLenT* __restrict__ seq_lens,
+    int32_t* __restrict__ topk_idx,
+    int batch,
+    int num_heads,
+    int max_seqblock,
+    int block_size,
+    int topk) {
+  const int b = blockIdx.x;  // grid.x = batch
+  const int h = blockIdx.y;  // grid.y = num_heads
+  const int tx = threadIdx.x;
+
+  // seq_lens is from an earlier kernel; prefetch it (and the cheap setup) before
+  // waiting on the score producer so the prologue overlaps its tail (PDL).
+  const int64_t seq_len = static_cast<int64_t>(seq_lens[b]);
+  const int num_blocks_raw = static_cast<int>((seq_len + block_size - 1) / block_size);
+  // Never scan past the materialized score columns.
+  const int num_blocks = num_blocks_raw < max_seqblock ? num_blocks_raw : max_seqblock;
+  int32_t* __restrict__ out = topk_idx + (static_cast<int64_t>(h) * batch + b) * topk;
+  device::PDLWaitPrimary<kUsePDL>();
+
+  if (num_blocks <= topk) {  // trivial: identity, -1 padded
+    for (int i = tx; i < topk; i += TopKTrait::kCTASize)
+      out[i] = (i < num_blocks) ? i : -1;
+    return;
+  }
+
+  const float* __restrict__ row = score + (static_cast<int64_t>(h) * batch + b) * max_seqblock;
+  __shared__ TopKTrait::Smem smem;
+  TopKTrait::forward(row, static_cast<uint32_t>(num_blocks), out, static_cast<uint32_t>(topk), &smem);
+}
+
+// Page-table output: for each (batch b, kv-head h) pseudo-request emit the
+// trtllm/fa3 page table -- selected blocks sorted ascending (so the final partial
+// block's pages land last), each expanded to its ppb = block_size/page_size pages
+// via req_to_token -- plus the effective KV length seq_lens_out.
+//
+// DP attention (num_kv_heads > 1): each kv head selects its OWN blocks, so the
+// per-request page table can't be shared across heads. We flatten (b, h) into
+// num_heads*batch pseudo-requests laid out batch-major (row = b*num_heads + h,
+// matching q.view(bs, nkv, gqa, d).reshape(bs*nkv, gqa, d)). seq_lens / slot_ids /
+// req_to_token are per-batch (head-independent: a token's cache slot is the same
+// for every head). The page index is head-encoded (head-minor) as
+// base_page*num_heads + h, which is exactly the page index into an HND cache
+// [num_pages, nkv, page_size, D] reshaped to [num_pages*nkv, 1, page_size, D] (a
+// free view when the cache is contiguous HND). num_heads == 1 (h == 0) reproduces
+// the single-kv-head TP>=4 behavior (page index == base_page).
+template <typename SeqLenT, bool kUsePDL>
+__global__ void minimax_decode_topk_page_table_kernel(
+    const float* __restrict__ score,
+    const SeqLenT* __restrict__ seq_lens,
+    const int32_t* __restrict__ req_to_token,
+    const int64_t* __restrict__ slot_ids,
+    int32_t* __restrict__ page_table,
+    int32_t* __restrict__ seq_lens_out,
+    int batch,
+    int num_heads,
+    int max_seqblock,
+    int block_size,
+    int topk,
+    int page_size,
+    int r2t_stride,
+    int max_kv_len,
+    int max_sparse_pages) {
+  const int b = blockIdx.x;  // grid.x = batch
+  const int h = blockIdx.y;  // grid.y = num_heads (kv head)
+  const int tx = threadIdx.x;
+
+  // Prefetch seq_lens / slot_ids (from earlier kernels) and the cheap setup
+  // before waiting on the score producer, so the prologue overlaps its tail (PDL).
+  const int64_t seq_len = static_cast<int64_t>(seq_lens[b]);
+  const int num_blocks_raw = static_cast<int>((seq_len + block_size - 1) / block_size);
+  const int num_blocks = num_blocks_raw < max_seqblock ? num_blocks_raw : max_seqblock;
+  const int ppb = block_size / page_size;
+  const int64_t out_row = static_cast<int64_t>(b) * num_heads + h;  // flattened pseudo-request
+  int32_t* __restrict__ pt_row = page_table + out_row * max_sparse_pages;
+  const int64_t r2t_base = static_cast<int64_t>(slot_ids[b]) * r2t_stride;
+  device::PDLWaitPrimary<kUsePDL>();
+
+  if (num_blocks <= topk) {  // trivial: every block selected, all tokens valid
+    if (tx == 0) seq_lens_out[out_row] = static_cast<int>(seq_len);
+    // block id == ascending slot, so the partial final block's pages land last
+    const int total = num_blocks * ppb;
+    for (int e = tx; e < total; e += TopKTrait::kCTASize) {
+      const int slot = e / ppb;
+      const int pp = e % ppb;
+      int tok = slot * block_size + pp * page_size;
+      if (tok >= max_kv_len) tok = max_kv_len - 1;
+      pt_row[e] = req_to_token[r2t_base + tok] / page_size * num_heads + h;
+    }
+    return;
+  }
+
+  const int k_eff = topk;                                                                        // num_blocks > topk
+  const float* __restrict__ row = score + (static_cast<int64_t>(h) * batch + b) * max_seqblock;  // head-major score
+  __shared__ TopKTrait::Smem smem;
+  __shared__ int32_t s_topk[TopKTrait::kMaxTopK];
+  TopKTrait::forward(row, static_cast<uint32_t>(num_blocks), s_topk, static_cast<uint32_t>(topk), &smem);
+  __syncthreads();  // s_topk fully written before the transform reads it
+
+  // Sort the selected block ids ascending (k_eff <= kMaxTopK is tiny) so the
+  // partial final block lands last, accumulating the effective KV length in the
+  // same pass: each selected block contributes min(block_size, seq_len - c*block)
+  // valid tokens (only the final block can be partial).
+  __shared__ int32_t s_sorted[TopKTrait::kMaxTopK];
+  __shared__ int s_eff_kv;
+  if (tx == 0) s_eff_kv = 0;
+  __syncthreads();
+  for (int slot = tx; slot < k_eff; slot += TopKTrait::kCTASize) {
+    const int32_t v = s_topk[slot];
+    int rank = 0;
+    for (int j = 0; j < k_eff; ++j)
+      rank += (s_topk[j] < v);
+    s_sorted[rank] = v;
+    const int rem = static_cast<int>(seq_len - static_cast<int64_t>(v) * block_size);
+    atomicAdd(&s_eff_kv, rem < block_size ? rem : block_size);
+  }
+  __syncthreads();
+  if (tx == 0) seq_lens_out[out_row] = s_eff_kv;
+
+  // Parallel page emit: one thread per output page.
+  const int total = k_eff * ppb;
+  for (int e = tx; e < total; e += TopKTrait::kCTASize) {
+    const int slot = e / ppb;
+    const int pp = e % ppb;
+    int tok = s_sorted[slot] * block_size + pp * page_size;
+    if (tok >= max_kv_len) tok = max_kv_len - 1;
+    pt_row[e] = req_to_token[r2t_base + tok] / page_size * num_heads + h;
+  }
+}
+
+// -------------------------------------------------------------------------
+// Launchers
+// -------------------------------------------------------------------------
+template <typename SeqLenT, bool kUsePDL>
+void minimax_decode_topk(
+    tvm::ffi::TensorView score,     // [H, B, S] fp32
+    tvm::ffi::TensorView seq_lens,  // [B] int32/int64
+    tvm::ffi::TensorView topk_idx,  // [H, B, T] int32
+    int64_t block_size,
+    int64_t topk) {
+  using namespace host;
+
+  SymbolicSize H = {"num_heads"};
+  SymbolicSize B = {"batch"};
+  SymbolicSize S = {"max_seqblock"};
+  SymbolicSize T = {"topk"};
+  SymbolicDevice device_;
+  device_.set_options<kDLCUDA>();
+
+  TensorMatcher({H, B, S}).with_dtype<fp32_t>().with_device(device_).verify(score);
+  TensorMatcher({B}).with_dtype<SeqLenT>().with_device(device_).verify(seq_lens);
+  TensorMatcher({H, B, T}).with_dtype<int32_t>().with_device(device_).verify(topk_idx);
+
+  const int num_heads = static_cast<int>(H.unwrap());
+  const int batch = static_cast<int>(B.unwrap());
+  const int max_seqblock = static_cast<int>(S.unwrap());
+  const int topk_i = static_cast<int>(T.unwrap());
+  const DLDevice device = device_.unwrap();
+
+  RuntimeCheck(
+      static_cast<int64_t>(topk_i) == topk,
+      "minimax_decode_topk: topk arg (",
+      topk,
+      ") must match topk_idx last dim (",
+      topk_i,
+      ")");
+  RuntimeCheck(block_size > 0, "block_size must be > 0, got ", block_size);
+  if (batch == 0 || num_heads == 0) return;
+
+  const dim3 grid(static_cast<unsigned>(batch), static_cast<unsigned>(num_heads));
+  LaunchKernel(grid, TopKTrait::kCTASize, device, 0)
+      .enable_pdl(kUsePDL)(
+          minimax_decode_topk_block_kernel<SeqLenT, kUsePDL>,
+          static_cast<const float*>(score.data_ptr()),
+          static_cast<const SeqLenT*>(seq_lens.data_ptr()),
+          static_cast<int32_t*>(topk_idx.data_ptr()),
+          batch,
+          num_heads,
+          max_seqblock,
+          static_cast<int>(block_size),
+          topk_i);
+}
+
+// Page-table variant: emit the per-(batch, kv-head) paged page table consumed by
+// the dense backend (trtllm_mha / fa3) plus the effective KV length, instead of
+// block ids. For DP attention (num_kv_heads > 1) each kv head selects its own
+// blocks, so (b, h) pseudo-requests are flattened batch-major into the output
+// (B*num_heads rows); num_heads == 1 is the TP>=4 single-kv-head case. The page
+// index is head-encoded (head-minor) as base_page*num_heads + h -- the index into
+// an HND cache [num_pages, nkv, ps, D] reshaped to [num_pages*nkv, 1, ps, D].
+// page_table and seq_lens_out are allocated by the caller.
+template <typename SeqLenT, bool kUsePDL>
+void minimax_decode_topk_page_table(
+    tvm::ffi::TensorView score,         // [H, B, S] fp32 (H = num_kv_heads)
+    tvm::ffi::TensorView seq_lens,      // [B] int32/int64
+    tvm::ffi::TensorView req_to_token,  // [max_reqs, max_kv_len] int32
+    tvm::ffi::TensorView slot_ids,      // [B] int64 (req_pool_indices)
+    tvm::ffi::TensorView page_table,    // [B*H, max_sparse_pages] int32 (out)
+    tvm::ffi::TensorView seq_lens_out,  // [B*H] int32 (effective KV length, out)
+    int64_t block_size,
+    int64_t topk,
+    int64_t page_size) {
+  using namespace host;
+
+  SymbolicSize H = {"num_heads"};
+  SymbolicSize B = {"batch"};
+  SymbolicSize S = {"max_seqblock"};
+  SymbolicSize R = {"max_reqs"};
+  SymbolicSize KV = {"max_kv_len"};
+  SymbolicSize BH = {"batch_heads"};
+  SymbolicSize P = {"max_sparse_pages"};
+  SymbolicDevice device_;
+  device_.set_options<kDLCUDA>();
+
+  TensorMatcher({H, B, S}).with_dtype<fp32_t>().with_device(device_).verify(score);
+  TensorMatcher({B}).with_dtype<SeqLenT>().with_device(device_).verify(seq_lens);
+  TensorMatcher({R, KV}).with_dtype<int32_t>().with_device(device_).verify(req_to_token);
+  TensorMatcher({B}).with_dtype<int64_t>().with_device(device_).verify(slot_ids);
+  TensorMatcher({BH, P}).with_dtype<int32_t>().with_device(device_).verify(page_table);
+  TensorMatcher({BH}).with_dtype<int32_t>().with_device(device_).verify(seq_lens_out);
+
+  const int num_heads = static_cast<int>(H.unwrap());
+  const int batch = static_cast<int>(B.unwrap());
+  const int max_seqblock = static_cast<int>(S.unwrap());
+  const int max_kv_len = static_cast<int>(KV.unwrap());
+  const int max_sparse_pages = static_cast<int>(P.unwrap());
+  const int r2t_stride = static_cast<int>(req_to_token.stride(0));
+  const DLDevice device = device_.unwrap();
+
+  RuntimeCheck(
+      BH.unwrap() == static_cast<int64_t>(batch) * num_heads,
+      "page_table rows (",
+      BH.unwrap(),
+      ") must equal batch*num_heads (",
+      static_cast<int64_t>(batch) * num_heads,
+      ")");
+  RuntimeCheck(
+      block_size > 0 && page_size > 0 && block_size % page_size == 0,
+      "block_size must be a positive multiple of page_size");
+  RuntimeCheck(topk <= static_cast<int64_t>(TopKTrait::kMaxTopK), "topk exceeds kMaxTopK for page-table mode");
+  if (batch == 0 || num_heads == 0) return;
+
+  const dim3 grid(static_cast<unsigned>(batch), static_cast<unsigned>(num_heads));
+  LaunchKernel(grid, TopKTrait::kCTASize, device, 0)
+      .enable_pdl(kUsePDL)(
+          minimax_decode_topk_page_table_kernel<SeqLenT, kUsePDL>,
+          static_cast<const float*>(score.data_ptr()),
+          static_cast<const SeqLenT*>(seq_lens.data_ptr()),
+          static_cast<const int32_t*>(req_to_token.data_ptr()),
+          static_cast<const int64_t*>(slot_ids.data_ptr()),
+          static_cast<int32_t*>(page_table.data_ptr()),
+          static_cast<int32_t*>(seq_lens_out.data_ptr()),
+          batch,
+          num_heads,
+          max_seqblock,
+          static_cast<int>(block_size),
+          static_cast<int>(topk),
+          static_cast<int>(page_size),
+          r2t_stride,
+          max_kv_len,
+          max_sparse_pages);
+}
+
+}  // namespace

python/sglang/jit_kernel/csrc/minimax/per_token_quant_ue8m0.cuh

@@ -0,0 +1,280 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+// Per-token group quant to FP8-e4m3 with a fused UE8M0 scale. Each group of
+// kGroupSize columns gets one UE8M0 exponent byte written contiguously in
+// row-major order into ``x_sf`` (int32 [num_tokens, num_groups/4], 4 group
+// bytes per int32). This is the deep_gemm "transform_sf" pack done inline in
+// the quant, reusing the dsv4 ``cast_to_ue8m0``/``pack_fp8`` primitives -- it
+// is byte-identical to ``per_token_group_quant_fp8(scale_ue8m0=True)`` followed
+// by ``transform_sf_into_required_layout`` (both round via ceil(log2(absmax/
+// FP8_MAX))), but emits no separate transpose/pack kernel.
+struct PerTokenQuantUe8m0Params {
+  const bf16_t* __restrict__ x;  // [num_tokens, hidden]
+  fp8_e4m3_t* __restrict__ x_q;  // [num_tokens, hidden]
+  int32_t* __restrict__ x_sf;    // [num_tokens, num_groups/4]; written as bytes
+  uint32_t num_tokens;
+  uint32_t hidden;
+  uint32_t num_groups;  // hidden / kGroupSize
+};
+
+template <uint32_t kGroupSize, bool kUsePDL>
+__global__ __launch_bounds__(1024, 2) void  //
+    per_token_quant_ue8m0_kernel(const PerTokenQuantUe8m0Params __grid_constant__ params) {
+  using namespace device;
+  constexpr uint32_t kVecElems = 8;  // 8 bf16 = 16B load per thread
+  static_assert(kGroupSize % kVecElems == 0, "group_size must be a multiple of 8");
+  constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems;
+  using InputVec = AlignedVector<bf16x2_t, kVecElems / 2>;
+  using OutputVec = AlignedVector<fp8x2_e4m3_t, kVecElems / 2>;
+
+  const uint32_t token_id = blockIdx.x;
+  const uint32_t tid = threadIdx.x;
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto token_in = params.x + static_cast<uint64_t>(token_id) * params.hidden;
+  const auto token_out = params.x_q + static_cast<uint64_t>(token_id) * params.hidden;
+
+  InputVec in_vec;
+  in_vec.load(token_in, tid);
+  float local_max = 0.0f;
+  float vals[kVecElems];
+#pragma unroll
+  for (uint32_t i = 0; i < kVecElems / 2; ++i) {
+    const auto [v0, v1] = cast<fp32x2_t>(in_vec[i]);
+    vals[2 * i + 0] = v0;
+    vals[2 * i + 1] = v1;
+    local_max = fmaxf(local_max, fmaxf(fabsf(v0), fabsf(v1)));
+  }
+  // Absmax across the kThreadsPerGroup threads that cover one group.
+  local_max = warp::reduce_max<kThreadsPerGroup>(local_max);
+  const float absmax = fmaxf(local_max, 1e-10f);
+  const float raw_scale = absmax / math::FP8_E4M3_MAX;
+  const uint32_t ue8m0_exp = cast_to_ue8m0(raw_scale);
+  const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23);
+
+  OutputVec out_vec;
+#pragma unroll
+  for (uint32_t i = 0; i < kVecElems / 2; ++i) {
+    out_vec[i] = pack_fp8(vals[2 * i + 0] * inv_scale, vals[2 * i + 1] * inv_scale);
+  }
+  out_vec.store(token_out, tid);
+
+  const uint32_t group_id = tid / kThreadsPerGroup;
+  const uint32_t within_group_id = tid % kThreadsPerGroup;
+  if (within_group_id == 0 && group_id < params.num_groups) {
+    const uint32_t byte_off = token_id * params.num_groups + group_id;
+    reinterpret_cast<uint8_t*>(params.x_sf)[byte_off] = static_cast<uint8_t>(ue8m0_exp);
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// Fused quant + scatter: like per_token_quant_ue8m0_kernel, but instead of
+// writing the fp8/scale for the single source token, it scatters them straight
+// into the permuted grouped-GEMM input -- replicating each token to its ``topk``
+// destination rows -- so the separate fill_gateup_input_triton_kernel launch (and
+// the intermediate x_q/x_sf buffers) are eliminated. The fp8 value + UE8M0 scale
+// are computed exactly once per token (identical to the non-fused kernel); only
+// the stores differ.
+struct PerTokenQuantUe8m0ScatterParams {
+  const bf16_t* __restrict__ x;              // [num_tokens, hidden]
+  fp8_e4m3_t* __restrict__ gateup_input;     // [E, m_max, hidden]
+  int32_t* __restrict__ gateup_input_scale;  // [E, num_groups/4, m_max] int32 (MN-major), written as bytes
+  const int32_t* __restrict__ src2dst;       // [num_tokens, topk] -> dst row = expert*m_max + slot
+  const int32_t* __restrict__ topk_ids;      // [num_tokens, topk]; <0 = skip
+  uint32_t num_tokens;
+  uint32_t hidden;
+  uint32_t num_groups;  // hidden / kGroupSize
+  uint32_t topk;
+  uint32_t m_max;
+};
+
+template <uint32_t kGroupSize, uint32_t kTopK, bool kUsePDL>
+__global__ __launch_bounds__(1024, 2) void  //
+    per_token_quant_ue8m0_scatter_kernel(const PerTokenQuantUe8m0ScatterParams __grid_constant__ params) {
+  using namespace device;
+  constexpr uint32_t kVecElems = 8;  // 8 bf16 = 16B load per thread
+  static_assert(kGroupSize % kVecElems == 0, "group_size must be a multiple of 8");
+  static_assert(kTopK <= kWarpThreads, "kTopK must fit in a warp for the lane-parallel scale write");
+  constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems;
+  using InputVec = AlignedVector<bf16x2_t, kVecElems / 2>;
+  using OutputVec = AlignedVector<fp8x2_e4m3_t, kVecElems / 2>;
+
+  const uint32_t token_id = blockIdx.x;
+  const uint32_t tid = threadIdx.x;
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto token_in = params.x + static_cast<uint64_t>(token_id) * params.hidden;
+
+  InputVec in_vec;
+  in_vec.load(token_in, tid);
+  float local_max = 0.0f;
+  float vals[kVecElems];
+#pragma unroll
+  for (uint32_t i = 0; i < kVecElems / 2; ++i) {
+    const auto [v0, v1] = cast<fp32x2_t>(in_vec[i]);
+    vals[2 * i + 0] = v0;
+    vals[2 * i + 1] = v1;
+    local_max = fmaxf(local_max, fmaxf(fabsf(v0), fabsf(v1)));
+  }
+  // Butterfly reduce: every thread of the group ends up with the group absmax
+  // (so all of them can write scale bytes below, no broadcast needed).
+  local_max = warp::reduce_max<kThreadsPerGroup>(local_max);
+  const float absmax = fmaxf(local_max, 1e-10f);
+  const float raw_scale = absmax / math::FP8_E4M3_MAX;
+  const uint32_t ue8m0_exp = cast_to_ue8m0(raw_scale);
+  const float inv_scale = __uint_as_float((127u + 127u - ue8m0_exp) << 23);
+
+  OutputVec out_vec;
+#pragma unroll
+  for (uint32_t i = 0; i < kVecElems / 2; ++i) {
+    out_vec[i] = pack_fp8(vals[2 * i + 0] * inv_scale, vals[2 * i + 1] * inv_scale);
+  }
+
+  // Read this token's kTopK destinations once (fully unrolled).
+  const auto* src2dst_row = params.src2dst + static_cast<uint64_t>(token_id) * kTopK;
+  const auto* topk_ids_row = params.topk_ids + static_cast<uint64_t>(token_id) * kTopK;
+  int32_t dst_rows[kTopK];
+#pragma unroll
+  for (uint32_t i = 0; i < kTopK; ++i) {
+    dst_rows[i] = (topk_ids_row[i] >= 0) ? src2dst_row[i] : -1;
+  }
+
+  const uint32_t group_id = tid / kThreadsPerGroup;
+  const uint32_t within_group = tid % kThreadsPerGroup;
+  const uint32_t c = group_id / 4u;  // packed int32 index along the group axis
+  const uint32_t b = group_id % 4u;  // byte within that int32
+  const uint64_t scale_g4 = params.num_groups / 4u;
+  auto* scale_bytes = reinterpret_cast<uint8_t*>(params.gateup_input_scale);
+
+  // 1) Output fp8: every thread replicates its 16B chunk to all kTopK dst rows.
+#pragma unroll
+  for (uint32_t i = 0; i < kTopK; ++i) {
+    const int32_t dst = dst_rows[i];
+    if (dst < 0) continue;
+    out_vec.store(params.gateup_input + static_cast<uint64_t>(dst) * params.hidden, tid);
+  }
+
+  // 2) Scale bytes: distribute the kTopK experts across the group's threads
+  // (each already holds the group exponent) so they write in parallel instead
+  // of one leader looping. lane `within_group` handles experts {within_group,
+  // within_group + kThreadsPerGroup, ...}; for kTopK <= kThreadsPerGroup that is
+  // exactly one expert per lane (no loop).
+#pragma unroll
+  for (uint32_t i = within_group; i < kTopK; i += kThreadsPerGroup) {
+    const int32_t dst = dst_rows[i];
+    if (dst < 0) continue;
+    const uint32_t expert = static_cast<uint32_t>(dst) / params.m_max;
+    const uint32_t m = static_cast<uint32_t>(dst) % params.m_max;
+    // int32 element [expert, c, m] of [E, G/4, m_max], byte b within it.
+    const uint64_t int32_index =
+        static_cast<uint64_t>(expert) * scale_g4 * params.m_max + static_cast<uint64_t>(c) * params.m_max + m;
+    scale_bytes[int32_index * 4u + b] = static_cast<uint8_t>(ue8m0_exp);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kGroupSize, int64_t kTopK, bool kUsePDL>
+void per_token_quant_ue8m0_scatter(
+    tvm::ffi::TensorView x,
+    tvm::ffi::TensorView gateup_input,
+    tvm::ffi::TensorView gateup_input_scale,
+    tvm::ffi::TensorView src2dst,
+    tvm::ffi::TensorView topk_ids,
+    int64_t topk,
+    int64_t m_max) {
+  using namespace host;
+  auto device = SymbolicDevice{};
+  auto M = SymbolicSize{"num_tokens"};
+  auto H = SymbolicSize{"hidden"};
+  auto E = SymbolicSize{"num_experts"};
+  auto MM = SymbolicSize{"m_max"};
+  auto G4 = SymbolicSize{"num_groups_div_4"};
+  device.set_options<kDLCUDA>();
+  TensorMatcher({M, H}).with_dtype<bf16_t>().with_device(device).verify(x);
+  TensorMatcher({E, MM, H}).with_dtype<fp8_e4m3_t>().with_device(device).verify(gateup_input);
+  TensorMatcher({E, G4, MM}).with_dtype<int32_t>().with_device(device).verify(gateup_input_scale);
+
+  const uint32_t num_tokens = static_cast<uint32_t>(M.unwrap());
+  const uint32_t hidden = static_cast<uint32_t>(H.unwrap());
+  RuntimeCheck(hidden % kGroupSize == 0, "hidden ", hidden, " not divisible by group_size ", kGroupSize);
+  const uint32_t num_groups = hidden / static_cast<uint32_t>(kGroupSize);
+  RuntimeCheck(num_groups % 4 == 0, "num_groups must be a multiple of 4 for int32 packing");
+  RuntimeCheck(static_cast<uint32_t>(G4.unwrap()) * 4 == num_groups, "scale G/4 mismatch");
+  RuntimeCheck(static_cast<uint32_t>(MM.unwrap()) == static_cast<uint32_t>(m_max), "m_max mismatch");
+  const uint32_t threads = hidden / 8;  // kVecElems
+  RuntimeCheck(threads <= 1024, "hidden/8 must be <= 1024, got ", threads);
+  RuntimeCheck(topk == kTopK, "topk does not match compiled template");
+
+  const auto params = PerTokenQuantUe8m0ScatterParams{
+      .x = static_cast<const bf16_t*>(x.data_ptr()),
+      .gateup_input = static_cast<fp8_e4m3_t*>(gateup_input.data_ptr()),
+      .gateup_input_scale = static_cast<int32_t*>(gateup_input_scale.data_ptr()),
+      .src2dst = static_cast<const int32_t*>(src2dst.data_ptr()),
+      .topk_ids = static_cast<const int32_t*>(topk_ids.data_ptr()),
+      .num_tokens = num_tokens,
+      .hidden = hidden,
+      .num_groups = num_groups,
+      .topk = static_cast<uint32_t>(kTopK),
+      .m_max = static_cast<uint32_t>(m_max),
+  };
+  if (num_tokens == 0) return;
+  constexpr auto kernel = per_token_quant_ue8m0_scatter_kernel<kGroupSize, kTopK, kUsePDL>;
+  LaunchKernel(num_tokens, threads, device.unwrap())  //
+      .enable_pdl(kUsePDL)(kernel, params);
+}
+
+template <int64_t kGroupSize, bool kUsePDL>
+void per_token_quant_ue8m0(tvm::ffi::TensorView x, tvm::ffi::TensorView x_q, tvm::ffi::TensorView x_sf) {
+  using namespace host;
+  auto device = SymbolicDevice{};
+  auto M = SymbolicSize{"num_tokens"};
+  auto H = SymbolicSize{"hidden"};
+  auto G4 = SymbolicSize{"num_groups_div_4"};
+  device.set_options<kDLCUDA>();
+  TensorMatcher({M, H}).with_dtype<bf16_t>().with_device(device).verify(x);
+  TensorMatcher({M, H}).with_dtype<fp8_e4m3_t>().with_device(device).verify(x_q);
+  TensorMatcher({M, G4}).with_dtype<int32_t>().with_device(device).verify(x_sf);
+
+  const uint32_t num_tokens = static_cast<uint32_t>(M.unwrap());
+  const uint32_t hidden = static_cast<uint32_t>(H.unwrap());
+  RuntimeCheck(hidden % kGroupSize == 0, "hidden ", hidden, " not divisible by group_size ", kGroupSize);
+  const uint32_t num_groups = hidden / static_cast<uint32_t>(kGroupSize);
+  RuntimeCheck(static_cast<uint32_t>(G4.unwrap()) * 4 == num_groups);
+  const uint32_t threads = hidden / 8;  // kVecElems
+  RuntimeCheck(threads <= 1024, "hidden/8 must be <= 1024, got ", threads);
+
+  const auto params = PerTokenQuantUe8m0Params{
+      .x = static_cast<const bf16_t*>(x.data_ptr()),
+      .x_q = static_cast<fp8_e4m3_t*>(x_q.data_ptr()),
+      .x_sf = static_cast<int32_t*>(x_sf.data_ptr()),
+      .num_tokens = num_tokens,
+      .hidden = hidden,
+      .num_groups = num_groups,
+  };
+  if (num_tokens == 0) return;
+  constexpr auto kernel = per_token_quant_ue8m0_kernel<kGroupSize, kUsePDL>;
+  LaunchKernel(num_tokens, threads, device.unwrap())  //
+      .enable_pdl(kUsePDL)(kernel, params);
+}
+
+}  // namespace