chore: import upstream snapshot with attribution

python/sglang/jit_kernel/csrc/deepseek_v4/c128.cuh

@@ -0,0 +1,522 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.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/compress.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/object.h>
+
+#include <cstdint>
+
+namespace {
+
+using Plan128 = device::compress::PrefillPlan;
+using IndiceT = int32_t;
+
+/// \brief Each thread will handle this many elements (split along head_dim)
+constexpr int32_t kTileElements = 2;
+/// \brief Each warp will handle this many elements (split along 128)
+constexpr int32_t kElementsPerWarp = 8;
+constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
+constexpr uint32_t kBlockSize = device::kWarpThreads * kNumWarps;
+
+/// \brief Need to reduce register usage to increase occupancy
+#define C128_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
+
+struct Compress128DecodeParams {
+  /**
+   * \brief Shape: `[num_indices, 128, head_dim * 2]` \n
+   * last dimension layout:
+   * | kv current | score current |
+   */
+  void* __restrict__ kv_score_buffer;
+  /** \brief Shape: `[batch_size, head_dim * 2]` */
+  const void* __restrict__ kv_score_input;
+  /** \brief Shape: `[batch_size, head_dim]` */
+  void* __restrict__ kv_compressed_output;
+  /** \brief Shape: `[128, head_dim]` (called `ape`) */
+  const void* __restrict__ score_bias;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ indices;
+  /** \brief Shape: `[batch_size, ]` */
+  const IndiceT* __restrict__ seq_lens;
+  /** \NOTE: `batch_size` <= `num_indices` */
+  uint32_t batch_size;
+};
+
+struct Compress128PrefillParams {
+  /**
+   * \brief Shape: `[num_indices, 128, head_dim * 2]` \n
+   * last dimension layout:
+   * | kv current | score current |
+   */
+  void* __restrict__ kv_score_buffer;
+  /** \brief Shape: `[batch_size, head_dim * 2]` */
+  const void* __restrict__ kv_score_input;
+  /** \brief Shape: `[batch_size, head_dim]` */
+  void* __restrict__ kv_compressed_output;
+  /** \brief Shape: `[128, head_dim]` (called `ape`) */
+  const void* __restrict__ score_bias;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ indices;
+  /** \brief Shape: `[batch_size, ]`*/
+  const int32_t* __restrict__ load_indices;
+  /** \brief The following part is plan info. */
+  const Plan128* __restrict__ compress_plan;
+  const Plan128* __restrict__ write_plan;
+  uint32_t num_compress;
+  uint32_t num_write;
+};
+
+struct Compress128SharedBuffer {
+  using Storage = device::AlignedVector<float, kTileElements>;
+  Storage data[kNumWarps][device::kWarpThreads + 1];  // padding to avoid bank conflict
+  SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
+    return data[warp_id][lane_id];
+  }
+  SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
+    return data[warp_id][lane_id][tile_id];
+  }
+};
+
+template <typename T>
+SGL_DEVICE void c128_write(
+    T* kv_score_buf,  //
+    const T* kv_score_src,
+    const int64_t head_dim,
+    const int32_t write_pos,
+    const uint32_t lane_id) {
+  using namespace device;
+
+  using Storage = AlignedVector<T, kTileElements>;
+  const auto element_size = head_dim * 2;
+  const auto gmem = tile::Memory<Storage>{lane_id, kWarpThreads};
+  kv_score_buf += write_pos * element_size;
+
+  /// NOTE: Layout | [0] = kv | [1] = score |
+  Storage kv_score[2];
+#pragma unroll
+  for (int32_t i = 0; i < 2; ++i) {
+    kv_score[i] = gmem.load(kv_score_src + head_dim * i);
+  }
+#pragma unroll
+  for (int32_t i = 0; i < 2; ++i) {
+    gmem.store(kv_score_buf + head_dim * i, kv_score[i]);
+  }
+}
+
+template <typename InFloat, typename OutFloat>
+SGL_DEVICE void c128_forward(
+    const InFloat* kv_score_buf,
+    const InFloat* kv_score_src,
+    OutFloat* kv_out,
+    const InFloat* score_bias,
+    const int64_t head_dim,
+    const int32_t window_len,
+    const uint32_t warp_id,
+    const uint32_t lane_id) {
+  using namespace device;
+
+  const auto element_size = head_dim * 2;
+  const auto score_offset = head_dim;
+
+  /// NOTE: part 1: load kv + score
+  using StorageIn = AlignedVector<InFloat, kTileElements>;
+  const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
+  StorageIn kv[kElementsPerWarp];
+  StorageIn score[kElementsPerWarp];
+  StorageIn bias[kElementsPerWarp];
+  const int32_t warp_offset = warp_id * kElementsPerWarp;
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; ++i) {
+    const int32_t j = i + warp_offset;
+    bias[i] = gmem_in.load(score_bias + j * head_dim);
+  }
+
+#pragma unroll
+  for (int32_t i = 0; i < kElementsPerWarp; ++i) {
+    const int32_t j = i + warp_offset;
+    const InFloat* src;
+    __builtin_assume(j < 128);
+    if (j < window_len) {
+      src = kv_score_buf + j * element_size;
+    } else {
+      /// NOTE: k in [-127, 0]. We'll load from the ragged `kv_score_src`
+      const int32_t k = j - 127;
+      src = kv_score_src + k * element_size;
+    }
+    kv[i] = gmem_in.load(src);
+    score[i] = gmem_in.load(src + score_offset);
+  }
+
+  /// NOTE: part 2: safe online softmax + weighted sum
+  using TmpStorage = typename Compress128SharedBuffer::Storage;
+  __shared__ Compress128SharedBuffer s_local_val_max;
+  __shared__ Compress128SharedBuffer s_local_exp_sum;
+  __shared__ Compress128SharedBuffer s_local_product;
+
+  TmpStorage tmp_val_max;
+  TmpStorage tmp_exp_sum;
+  TmpStorage tmp_product;
+
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+    float score_fp32[kElementsPerWarp];
+
+#pragma unroll
+    for (int32_t j = 0; j < kElementsPerWarp; ++j) {
+      score_fp32[j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
+    }
+
+    float max_value = score_fp32[0];
+    float sum_exp_value = 0.0f;
+
+#pragma unroll
+    for (int32_t j = 1; j < kElementsPerWarp; ++j) {
+      const auto fp32_score = score_fp32[j];
+      max_value = fmaxf(max_value, fp32_score);
+    }
+
+    float sum_product = 0.0f;
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      const auto fp32_score = score_fp32[j];
+      const auto exp_score = expf(fp32_score - max_value);
+      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_exp_value += exp_score;
+    }
+
+    tmp_val_max[i] = max_value;
+    tmp_exp_sum[i] = sum_exp_value;
+    tmp_product[i] = sum_product;
+  }
+
+  // naturally aligned, so no bank conflict
+  s_local_val_max(warp_id, lane_id) = tmp_val_max;
+  s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
+  s_local_product(warp_id, lane_id) = tmp_product;
+
+  __syncthreads();
+
+  /// NOTE: part 3: online softmax
+  /// NOTE: We have `kTileElements * kWarpThreads * kNumWarps` values to reduce
+  /// each reduce will consume `kNumWarps` threads (use partial warp reduction)
+  constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
+  constexpr uint32_t kIteration = kReductionCount / kBlockSize;
+
+#pragma unroll
+  for (uint32_t i = 0; i < kIteration; ++i) {
+    /// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
+    const uint32_t j = i * kBlockSize + warp_id * kWarpThreads + lane_id;
+    /// NOTE: Range `[0, kNumWarps)`
+    const uint32_t local_warp_id = j % kNumWarps;
+    /// NOTE: Range `[0, kTileElements * kWarpThreads)`
+    const uint32_t local_elem_id = j / kNumWarps;
+    /// NOTE: Range `[0, kTileElements)`
+    const uint32_t local_tile_id = local_elem_id % kTileElements;
+    /// NOTE: Range `[0, kWarpThreads)`
+    const uint32_t local_lane_id = local_elem_id / kTileElements;
+    /// NOTE: each warp will access the whole tile (all `kTileElements`)
+    /// and for different lanes, the memory access only differ in `local_warp_id`
+    /// so there's no bank conflict in shared memory access.
+    static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
+    const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
+    const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
+    const auto rescale = expf(local_val_max - global_val_max);
+    const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
+    const auto final_scale = rescale / global_exp_sum;
+    const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
+    kv_out[local_elem_id] = cast<OutFloat>(global_product);
+  }
+}
+
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
+  using namespace device;
+
+  constexpr int64_t kTileDim = kTileElements * kWarpThreads;  // 64
+  constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  constexpr int64_t kElementSize = kHeadDim * 2;
+  static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
+
+  const auto& [
+    _kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
+    indices, seq_lens, batch_size // decode info
+  ] = params;
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+
+  const uint32_t global_bid = blockIdx.x / kNumSplit;  // batch id
+  const uint32_t global_sid = blockIdx.x % kNumSplit;  // split id
+  if (global_bid >= batch_size) return;
+
+  const int32_t index = indices[global_bid];
+  const int32_t seq_len = seq_lens[global_bid];
+  const int64_t split_offset = global_sid * kTileDim;
+
+  // kv score
+  const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
+  const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
+
+  // kv input
+  const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
+  const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
+
+  // kv output
+  const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
+  const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
+
+  // score bias (ape)
+  const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  /// NOTE: the write must be visible to the subsequent c128_forward,
+  /// so only the last warp can write to HBM
+  /// In addition, `position` = `seq_len - 1`. To avoid underflow, we use `seq_len + 127`
+  if (warp_id == kNumWarps - 1) {
+    c128_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 127) % 128, lane_id);
+  }
+  if (seq_len % 128 == 0) {
+    c128_forward(kv_buf, kv_src, kv_out, score_bias, kHeadDim, /*window_len=*/128, warp_id, lane_id);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// compress kernel
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kWrite, bool kUsePDL>
+C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
+  using namespace device;
+
+  constexpr int64_t kTileDim = kTileElements * kWarpThreads;  // 64
+  constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  constexpr int64_t kElementSize = kHeadDim * 2;
+  static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
+
+  const auto& [
+    _kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
+    indices, load_indices, compress_plan, write_plan, num_compress, num_write // prefill plan
+  ] = params;
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+
+  uint32_t global_id;
+  if constexpr (kWrite) {
+    // for write kernel, we use global warp_id to dispatch work
+    global_id = (blockIdx.x * blockDim.x + threadIdx.x) / kWarpThreads;
+  } else {
+    // for compress kernel, we use block id to dispatch work
+    global_id = blockIdx.x;  // block id
+  }
+  const uint32_t global_pid = global_id / kNumSplit;  // plan id
+  const uint32_t global_sid = global_id % kNumSplit;  // split id
+
+  /// NOTE: compiler can optimize this if-else at compile time
+  const auto num_plans = kWrite ? num_write : num_compress;
+  const auto plan_ptr = kWrite ? write_plan : compress_plan;
+  if (global_pid >= num_plans) return;
+
+  const auto& [ragged_id, global_bid, position, window_len] = plan_ptr[global_pid];
+  const auto indices_ptr = kWrite ? indices : load_indices;
+
+  const int64_t split_offset = global_sid * kTileDim;
+
+  // kv input
+  const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
+  const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
+
+  // kv output
+  const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
+  const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset;
+
+  // score bias (ape)
+  const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
+
+  if (ragged_id == 0xFFFFFFFF) [[unlikely]]
+    return;
+
+  const int32_t index = indices_ptr[global_bid];
+  // kv score
+  const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
+  const auto kv_buf = kv_score_buffer + index * (kElementSize * 128) + split_offset;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // only responsible for the compress part
+  if constexpr (kWrite) {
+    c128_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 128, lane_id);
+  } else {
+    c128_forward(kv_buf, kv_src, kv_out, score_bias, kHeadDim, window_len, warp_id, lane_id);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+struct FlashCompress128Kernel {
+  static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, InFloat, OutFloat, kUsePDL>;
+  template <bool kWrite>
+  static constexpr auto prefill_kernel = flash_c128_prefill<kHeadDim, InFloat, OutFloat, kWrite, kUsePDL>;
+  static constexpr auto prefill_c_kernel = prefill_kernel</*kWrite=*/false>;
+  static constexpr auto prefill_w_kernel = prefill_kernel</*kWrite=*/true>;
+  static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static constexpr uint32_t kWriteBlockSize = 128;
+  static constexpr uint32_t kWarpsPerWriteBlock = kWriteBlockSize / device::kWarpThreads;
+
+  static void run_decode(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> /* UNUSED */) {
+    using namespace host;
+
+    // this should not happen in practice
+    auto B = SymbolicSize{"batch_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, 128, kHeadDim * 2})  // kv score
+        .with_dtype<InFloat>()
+        .with_device(device)
+        .verify(kv_score_buffer);
+    TensorMatcher({B, kHeadDim * 2})  // kv score input
+        .with_dtype<InFloat>()
+        .with_device(device)
+        .verify(kv_score_input);
+    TensorMatcher({B, kHeadDim})  // kv compressed output
+        .with_dtype<OutFloat>()
+        .with_device(device)
+        .verify(kv_compressed_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<InFloat>()
+        .with_device(device)
+        .verify(ape);
+    TensorMatcher({B})  // indices
+        .with_dtype<IndiceT>()
+        .with_device(device)
+        .verify(indices);
+    TensorMatcher({B})  // seq lens
+        .with_dtype<IndiceT>()
+        .with_device(device)
+        .verify(seq_lens);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto params = Compress128DecodeParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .indices = static_cast<const IndiceT*>(indices.data_ptr()),
+        .seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
+        .batch_size = batch_size,
+    };
+
+    const uint32_t num_blocks = batch_size * kNumSplit;
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(decode_kernel, params);
+  }
+
+  static void run_prefill(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView compress_plan,
+      const tvm::ffi::TensorView write_plan,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto N = SymbolicSize{"num_q_tokens"};
+    auto X = SymbolicSize{"compress_tokens"};
+    auto Y = SymbolicSize{"write_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, 128, kHeadDim * 2})  // kv score
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(kv_score_buffer);
+    TensorMatcher({N, kHeadDim * 2})  // kv score input
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(kv_score_input);
+    TensorMatcher({N, kHeadDim})  // kv compressed output
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_compressed_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(ape);
+    TensorMatcher({B})  // indices
+        .with_dtype<IndiceT>()
+        .with_device(device_)
+        .verify(indices);
+    TensorMatcher({X, compress::kPrefillPlanDim})  // compress plan
+        .with_dtype<compress::PrefillPlanTensorDtype>()
+        .with_device(device_)
+        .verify(compress_plan);
+    TensorMatcher({Y, compress::kPrefillPlanDim})  // write plan
+        .with_dtype<compress::PrefillPlanTensorDtype>()
+        .with_device(device_)
+        .verify(write_plan);
+
+    // might be needed for prefill write
+    const auto load_indices = extra.value_or(indices);
+    TensorMatcher({B})  // [read_positions]
+        .with_dtype<IndiceT>()
+        .with_device(device_)
+        .verify(load_indices);
+
+    const auto device = device_.unwrap();
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_c = static_cast<uint32_t>(X.unwrap());
+    const auto num_w = static_cast<uint32_t>(Y.unwrap());
+    const auto params = Compress128PrefillParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .indices = static_cast<const IndiceT*>(indices.data_ptr()),
+        .load_indices = static_cast<const IndiceT*>(load_indices.data_ptr()),
+        .compress_plan = static_cast<const Plan128*>(compress_plan.data_ptr()),
+        .write_plan = static_cast<const Plan128*>(write_plan.data_ptr()),
+        .num_compress = num_c,
+        .num_write = num_w,
+    };
+    RuntimeCheck(num_q_tokens >= batch_size, "num_q_tokens must be >= batch_size");
+    RuntimeCheck(num_q_tokens >= std::max(num_c, num_w), "invalid prefill plan");
+
+    constexpr auto kBlockSize_C = kBlockSize;
+    constexpr auto kBlockSize_W = kWriteBlockSize;
+    if (const auto num_c_blocks = num_c * kNumSplit) {
+      LaunchKernel(num_c_blocks, kBlockSize_C, device)  //
+          .enable_pdl(kUsePDL)(prefill_c_kernel, params);
+    }
+    if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerWriteBlock)) {
+      LaunchKernel(num_w_blocks, kBlockSize_W, device)  //
+          .enable_pdl(kUsePDL)(prefill_w_kernel, params);
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/c128_online.cuh

@@ -0,0 +1,726 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.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/compress.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/object.h>
+
+#include <algorithm>
+#include <cfloat>
+#include <cstdint>
+
+namespace device::compress {
+
+/// \brief Plan entry for online compress 128 prefill.
+/// Each entry describes a contiguous segment of tokens that lies inside a
+/// single 128-chunk. Multiple segments can map to the same batch id when the
+/// extend tokens span chunk boundaries.
+///
+/// **Layout compatibility:** the field order/types match `PrefillPlan` so that
+/// downstream kernels (e.g. `fused_norm_rope` in `CompressExtend` mode) can
+/// consume the compress_plan tensor as-if it were a `PrefillPlan` tensor --
+/// they only read `ragged_id` and `position`, both of which carry identical
+/// semantics here (the LAST token of the segment in q-ragged and global
+/// coordinates respectively).
+///
+/// Note that `window_len` here means "number of real tokens in this segment"
+/// (1..128), which differs from `PrefillPlan::window_len`. Downstream kernels
+/// that share the tensor MUST NOT read it under that name.
+struct alignas(16) OnlinePrefillPlan {
+  /// \brief Ragged-q position of the LAST token in this segment.
+  /// Equal to `segment_start_ragged + window_len - 1`.
+  uint32_t ragged_id;
+  /// \brief Index into the `indices` / `load_indices` arrays.
+  uint32_t batch_id;
+  /// \brief Global position of the LAST token in this segment.
+  /// For compress plans, `position % 128 == 127` (chunk-closing); for write
+  /// plans, `position % 128 < 127`.
+  uint32_t position;
+  /// \brief Number of real tokens in this segment (1..128).
+  /// The first segment token sits at `position - window_len + 1` (global) and
+  /// at `ragged_id - window_len + 1` (ragged).
+  uint32_t window_len;
+};
+
+static_assert(alignof(OnlinePrefillPlan) == alignof(PrefillPlan));
+static_assert(sizeof(OnlinePrefillPlan) == sizeof(PrefillPlan));
+
+}  // namespace device::compress
+
+namespace host::compress {
+
+using device::compress::OnlinePrefillPlan;
+using OnlinePrefillPlanTensorDtype = uint8_t;
+inline constexpr int64_t kOnlinePrefillPlanDim = 16;
+
+static_assert(alignof(OnlinePrefillPlan) == sizeof(OnlinePrefillPlan));
+static_assert(sizeof(OnlinePrefillPlan) == kOnlinePrefillPlanDim * sizeof(OnlinePrefillPlanTensorDtype));
+
+}  // namespace host::compress
+
+namespace {
+
+using OnlinePlan = device::compress::OnlinePrefillPlan;
+using IndiceT = int32_t;
+
+/// \brief Need to reduce register usage to increase occupancy
+struct Compress128OnlineDecodeParams {
+  /** \brief Shape: `[num_indices, 1, head_dim * 3 (max, sum, kv) ]` \n */
+  void* __restrict__ kv_score_buffer;
+  /** \brief Shape: `[batch_size, head_dim * 2]` */
+  const void* __restrict__ kv_score_input;
+  /** \brief Shape: `[batch_size, head_dim]` */
+  void* __restrict__ kv_compressed_output;
+  /** \brief Shape: `[128, head_dim]` (called `ape`) */
+  const void* __restrict__ score_bias;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ indices;
+  /** \brief Shape: `[batch_size, ]` */
+  const IndiceT* __restrict__ seq_lens;
+  /** \NOTE: `batch_size` <= `num_indices` */
+  uint32_t batch_size;
+};
+
+/// \brief Need to reduce register usage to increase occupancy
+struct Compress128OnlinePrefillParams {
+  /** \brief Shape: `[num_indices, 1, head_dim * 3 (max, sum, kv) ]` \n */
+  void* __restrict__ kv_score_buffer;
+  /** \brief Shape: `[num_q_tokens, head_dim * 2]` */
+  const void* __restrict__ kv_score_input;
+  /** \brief Shape: `[num_q_tokens, head_dim]` */
+  void* __restrict__ kv_compressed_output;
+  /** \brief Shape: `[128, head_dim]` (called `ape`) */
+  const void* __restrict__ score_bias;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ indices;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ load_indices;
+  /// \brief Plan for segments that close a chunk (write to `kv_compressed_output`).
+  /// Shape: `[num_compress, 16]` (uint8).
+  const OnlinePlan* __restrict__ compress_plan;
+  /// \brief Plan for the trailing partial segment of each batch (write back to
+  /// `kv_score_buffer`). Shape: `[num_write, 16]` (uint8).
+  const OnlinePlan* __restrict__ write_plan;
+  uint32_t num_compress;
+  uint32_t num_write;
+};
+
+// 4 elements per thread, kHeadDim / 4 threads per block
+template <int64_t kHeadDim, bool kUsePDL>
+__global__ void flash_c128_online_decode(const __grid_constant__ Compress128OnlineDecodeParams params) {
+  using namespace device;
+  constexpr uint32_t kVecSize = 4;
+  constexpr uint32_t kBlockSize = kHeadDim / kVecSize;
+  using Vec = AlignedVector<float, kVecSize>;
+  const auto gmem = tile::Memory<Vec>::cta(kBlockSize);
+  const auto batch_id = blockIdx.x;
+  const auto index = params.indices[batch_id];
+  const auto seq_len = params.seq_lens[batch_id];
+
+  const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
+  const auto kv_buf = kv_score_buffer + index * (kHeadDim * 3);
+  const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
+  const auto kv_src = kv_score_input + batch_id * (kHeadDim * 2);
+
+  /// NOTE: kv_score_buffer layout is [max, sum, kv] (slot 0 / 1 / 2). Reads,
+  /// writes, and the prefill kernel must all agree on this order.
+  const auto max_score_vec = gmem.load(kv_buf, 0);
+  const auto sum_score_vec = gmem.load(kv_buf, 1);
+  const auto old_kv_vec = gmem.load(kv_buf, 2);
+
+  /// NOTE: kv_score_input layout is | kv | score | (head_dim each), matching
+  /// the offline c128 kernel and the online prefill kernel.
+  const auto new_kv_vec = gmem.load(kv_src, 0);
+  const auto new_score_raw_vec = gmem.load(kv_src, 1);
+
+  /// NOTE: the new token sits at global position `seq_len - 1`, so its
+  /// position inside the 128-chunk is `(seq_len - 1) % 128`. The previous
+  /// `seq_len % 128` was off by one (`bias[127]` vs `bias[0]`, etc.).
+  const auto pos_in_chunk = (seq_len - 1) % 128;
+  const auto bias_vec = gmem.load(params.score_bias, pos_in_chunk);
+
+  Vec out_kv_vec;
+  Vec out_max_vec;
+  Vec out_sum_vec;
+  if (pos_in_chunk != 0) {
+    // Mid-chunk: combine prior partial state with the new token via online softmax.
+#pragma unroll
+    for (uint32_t i = 0; i < 4; ++i) {
+      const auto old_max = max_score_vec[i];
+      const auto old_kv = old_kv_vec[i];
+      const auto new_score = new_score_raw_vec[i] + bias_vec[i];
+      const auto new_kv = new_kv_vec[i];
+      const auto new_max = fmax(old_max, new_score);
+      const auto old_sum = sum_score_vec[i] * expf(old_max - new_max);
+      const auto new_exp = expf(new_score - new_max);
+      const auto new_sum = old_sum + new_exp;
+      out_kv_vec[i] = (old_kv * old_sum + new_kv * new_exp) / new_sum;
+      out_max_vec[i] = new_max;
+      out_sum_vec[i] = new_sum;
+    }
+  } else {
+    // First token of a new 128-chunk: initialize state with this token alone.
+#pragma unroll
+    for (uint32_t i = 0; i < 4; ++i) {
+      out_kv_vec[i] = new_kv_vec[i];
+      out_max_vec[i] = new_score_raw_vec[i] + bias_vec[i];
+      out_sum_vec[i] = 1.0f;  // exp(score - max) with max == score
+    }
+  }
+
+  if (pos_in_chunk == 127) {
+    // Chunk just closed: emit the compressed kv. No need to update the buffer
+    // -- the next chunk's first token will overwrite it.
+    const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim;
+    gmem.store(kv_out, out_kv_vec);
+  } else {
+    // Otherwise persist the running [max, sum, kv] state for the next step.
+    gmem.store(kv_buf, out_max_vec, 0);
+    gmem.store(kv_buf, out_sum_vec, 1);
+    gmem.store(kv_buf, out_kv_vec, 2);
+  }
+}
+
+constexpr int32_t kTileElements = 2;  // split (along head-dim)
+/// \brief Each warp will handle this many elements (split along softmax-128)
+constexpr int32_t kElementsPerWarp = 8;
+constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
+constexpr uint32_t kPrefillBlockSize = device::kWarpThreads * kNumWarps;
+using PrefillStorage = device::AlignedVector<float, kTileElements>;
+
+struct Compress128SharedBuffer {
+  using Storage = device::AlignedVector<float, 4>;
+  Storage data[kNumWarps][device::kWarpThreads + 1];  // padding to avoid bank conflict
+  SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
+    return data[warp_id][lane_id];
+  }
+  SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
+    return data[warp_id][lane_id][tile_id];
+  }
+};
+
+template <bool kNeedData>
+SGL_DEVICE void c128_prefill_forward(
+    const PrefillStorage (&kv)[kElementsPerWarp],
+    const PrefillStorage (&score)[kElementsPerWarp],
+    float* kv_out,
+    float* max_out,
+    float* sum_out,
+    const uint32_t warp_id,
+    const uint32_t lane_id) {
+  using namespace device;
+
+  /// NOTE: part 2: safe online softmax + weighted sum
+  using TmpStorage = typename Compress128SharedBuffer::Storage;
+  __shared__ Compress128SharedBuffer s_local_val_max;
+  __shared__ Compress128SharedBuffer s_local_exp_sum;
+  __shared__ Compress128SharedBuffer s_local_product;
+
+  TmpStorage tmp_val_max;
+  TmpStorage tmp_exp_sum;
+  TmpStorage tmp_product;
+
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+    float score_fp32[kElementsPerWarp];
+
+#pragma unroll
+    for (int32_t j = 0; j < kElementsPerWarp; ++j) {
+      score_fp32[j] = score[j][i];
+    }
+
+    float max_value = score_fp32[0];
+    float sum_exp_value = 0.0f;
+
+#pragma unroll
+    for (int32_t j = 1; j < kElementsPerWarp; ++j) {
+      const auto fp32_score = score_fp32[j];
+      max_value = fmaxf(max_value, fp32_score);
+    }
+
+    float sum_product = 0.0f;
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      const auto fp32_score = score_fp32[j];
+      const auto exp_score = expf(fp32_score - max_value);
+      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_exp_value += exp_score;
+    }
+
+    tmp_val_max[i] = max_value;
+    tmp_exp_sum[i] = sum_exp_value;
+    tmp_product[i] = sum_product;
+  }
+
+  // naturally aligned, so no bank conflict
+  s_local_val_max(warp_id, lane_id) = tmp_val_max;
+  s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
+  s_local_product(warp_id, lane_id) = tmp_product;
+
+  __syncthreads();
+
+  /// NOTE: part 3: online softmax
+  /// NOTE: We have `kTileElements * kWarpThreads * kNumWarps` values to reduce
+  /// each reduce will consume `kNumWarps` threads (use partial warp reduction)
+  constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
+  constexpr uint32_t kIteration = kReductionCount / kPrefillBlockSize;
+
+#pragma unroll
+  for (uint32_t i = 0; i < kIteration; ++i) {
+    /// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
+    const uint32_t j = i * kPrefillBlockSize + warp_id * kWarpThreads + lane_id;
+    /// NOTE: Range `[0, kNumWarps)`
+    const uint32_t local_warp_id = j % kNumWarps;
+    /// NOTE: Range `[0, kTileElements * kWarpThreads)`
+    const uint32_t local_elem_id = j / kNumWarps;
+    /// NOTE: Range `[0, kTileElements)`
+    const uint32_t local_tile_id = local_elem_id % kTileElements;
+    /// NOTE: Range `[0, kWarpThreads)`
+    const uint32_t local_lane_id = local_elem_id / kTileElements;
+    /// NOTE: each warp will access the whole tile (all `kTileElements`)
+    /// and for different lanes, the memory access only differ in `local_warp_id`
+    /// so there's no bank conflict in shared memory access.
+    static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
+    const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
+    const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
+    const auto rescale = expf(local_val_max - global_val_max);
+    const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
+    const auto final_scale = rescale / global_exp_sum;
+    const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
+    kv_out[local_elem_id] = global_product;
+    if constexpr (kNeedData) {
+      max_out[local_elem_id] = global_val_max;
+      sum_out[local_elem_id] = global_exp_sum;
+    }
+  }
+  if constexpr (kNeedData) __syncthreads();
+}
+
+/// \brief Sentinel score for padded positions in a 128-segment.
+/// Must be finite so that `score - max` never produces NaN even when an
+/// entire warp has only padded positions.
+constexpr float kPadScore = -FLT_MAX;
+
+/// \brief Online compress 128 prefill. Two passes share this body:
+/// - `kWrite=false` (compress pass): handles segments that close a chunk.
+///   May load prior partial state from the buffer, but never writes to it,
+///   so concurrent blocks can read the same slot without racing.
+/// - `kWrite=true` (write pass): handles the trailing partial segment of each
+///   batch. Each batch contributes at most one such plan, so concurrent blocks
+///   touch disjoint buffer slots.
+///
+/// The two passes MUST run as separate kernel launches (in stream order) so
+/// that all reads in pass 1 finish before any writes in pass 2 start.
+template <int64_t kHeadDim, bool kWrite, bool kUsePDL>
+__global__ __launch_bounds__(kPrefillBlockSize, 2)  //
+    void flash_c128_online_prefill(const __grid_constant__ Compress128OnlinePrefillParams params) {
+  using namespace device;
+
+  constexpr int64_t kTileDim = kTileElements * kWarpThreads;  // 64
+  constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
+
+  /// NOTE: the compiler folds the if-else at compile time.
+  const auto num_plans = kWrite ? params.num_write : params.num_compress;
+  const auto plan_ptr = kWrite ? params.write_plan : params.compress_plan;
+  const uint32_t global_id = blockIdx.x;
+  const uint32_t global_pid = global_id / kNumSplit;  // plan id
+  const uint32_t global_sid = global_id % kNumSplit;  // split id
+  if (global_pid >= num_plans) return;
+  const auto [ragged_id, batch_id, position, window_len] = plan_ptr[global_pid];
+  if (ragged_id == 0xFFFFFFFFu) [[unlikely]]
+    return;
+
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+  const int32_t split_offset = global_sid * kTileDim;  // int32 is enough
+
+  const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
+  const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
+  const auto kv_compressed_output = static_cast<float*>(params.kv_compressed_output);
+  const auto score_bias_base = static_cast<const float*>(params.score_bias);
+
+  constexpr int64_t kElementSize = kHeadDim * 2;  // | kv | score |
+  const uint32_t chunk_offset = (position % 128u) + 1u - window_len;
+  const uint32_t window_end = chunk_offset + window_len;        // exclusive, in [1, 128]
+  const int32_t segment_start = ragged_id - (position % 128u);  // can be negative, but safe
+  const int32_t load_index = chunk_offset != 0 ? params.load_indices[batch_id] : -1;
+  const int32_t store_index = kWrite ? params.indices[batch_id] : -1;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // 2 * 8 = 16 register per elem. in theory we should consume 48 register here
+  PrefillStorage kv[kElementsPerWarp];
+  PrefillStorage score[kElementsPerWarp];
+  PrefillStorage bias[kElementsPerWarp];
+  const auto warp_offset = warp_id * kElementsPerWarp;
+
+#pragma unroll
+  for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
+    const uint32_t j = i + warp_offset;
+    if (j >= chunk_offset && j < window_end) {
+      const auto kv_src_ptr = kv_score_input + (segment_start + j) * kElementSize + split_offset;
+      const auto score_src_ptr = kv_src_ptr + kHeadDim;
+      const auto bias_src_ptr = score_bias_base + j * kHeadDim + split_offset;
+      kv[i].load(kv_src_ptr, lane_id);
+      score[i].load(score_src_ptr, lane_id);
+      bias[i].load(bias_src_ptr, lane_id);
+    }
+  }
+
+#pragma unroll
+  for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
+    const uint32_t j = i + warp_offset;
+    const bool is_valid = (j >= chunk_offset && j < window_end);
+#pragma unroll
+    for (uint32_t ii = 0; ii < kTileElements; ++ii) {
+      score[i][ii] = is_valid ? score[i][ii] + bias[i][ii] : kPadScore;
+      /// NOTE: must zero out kv on padded slots -- `c128_prefill_forward`
+      /// computes `kv * exp_score` where `exp_score = expf(-FLT_MAX - max) ??? 0`,
+      /// and IEEE-754 makes `NaN * 0 = NaN` / `+-inf * 0 = NaN`. An
+      /// uninitialized register can hold a NaN/inf bit pattern, so without
+      /// this reset a single padded warp can poison the whole softmax.
+      kv[i][ii] = is_valid ? kv[i][ii] : 0.0f;
+    }
+  }
+
+  __shared__ alignas(16) float seg_kv[kTileDim];
+  __shared__ alignas(16) float seg_max[kTileDim];
+  __shared__ alignas(16) float seg_sum[kTileDim];
+
+  c128_prefill_forward<true>(kv, score, seg_kv, seg_max, seg_sum, warp_id, lane_id);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  if (warp_id == 0) {
+    PrefillStorage out_kv_vec, out_max_vec, out_sum_vec;
+    out_kv_vec.load(seg_kv, lane_id);
+    out_max_vec.load(seg_max, lane_id);
+    out_sum_vec.load(seg_sum, lane_id);
+    if (chunk_offset != 0) {
+      /// NOTE: load (max, sum, kv) of the in-progress chunk for this index.
+      /// `load_indices` may differ from `indices` when the prior partial state
+      /// lives on a different slot than the slot we ultimately write to.
+      const auto buf_load = kv_score_buffer + load_index * (kHeadDim * 3) + split_offset;
+      PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec;
+      buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id);
+      buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id);
+      buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id);
+#pragma unroll
+      for (uint32_t ii = 0; ii < kTileElements; ++ii) {
+        const float m1 = buf_max_vec[ii];
+        const float s1 = buf_sum_vec[ii];
+        const float k1 = buf_kv_vec[ii];
+        const float m2 = out_max_vec[ii];
+        const float s2 = out_sum_vec[ii];
+        const float k2 = out_kv_vec[ii];
+        const float new_max = fmaxf(m1, m2);
+        const float new_s1 = s1 * expf(m1 - new_max);
+        const float new_s2 = s2 * expf(m2 - new_max);
+        const float new_sum = new_s1 + new_s2;
+        const float new_kv = (k1 * new_s1 + k2 * new_s2) / new_sum;
+        out_max_vec[ii] = new_max;
+        out_sum_vec[ii] = new_sum;
+        out_kv_vec[ii] = new_kv;
+      }
+    }
+
+    if constexpr (kWrite) {
+      const auto buf_store = kv_score_buffer + store_index * (kHeadDim * 3) + split_offset;
+      reinterpret_cast<PrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec;
+      reinterpret_cast<PrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec;
+      reinterpret_cast<PrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec;
+    } else {
+      const auto out_ptr = kv_compressed_output + ragged_id * kHeadDim + split_offset;
+      reinterpret_cast<PrefillStorage*>(out_ptr)[lane_id] = out_kv_vec;
+    }
+  }
+}
+
+template <int64_t kHeadDim, bool kUsePDL>
+struct FlashCompress128OnlineKernel {
+  static constexpr auto decode_kernel = flash_c128_online_decode<kHeadDim, kUsePDL>;
+  template <bool kWrite>
+  static constexpr auto prefill_kernel = flash_c128_online_prefill<kHeadDim, kWrite, kUsePDL>;
+  static constexpr auto prefill_c_kernel = prefill_kernel</*kWrite=*/false>;
+  static constexpr auto prefill_w_kernel = prefill_kernel</*kWrite=*/true>;
+  static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4;
+
+  static void run_decode(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> /* UNUSED */) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, 1, kHeadDim * 3})  // kv score buffer (max, sum, kv)
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(kv_score_buffer);
+    TensorMatcher({B, kHeadDim * 2})  // kv score input
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(kv_score_input);
+    TensorMatcher({B, kHeadDim})  // kv compressed output
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(kv_compressed_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(ape);
+    TensorMatcher({B}).with_dtype<IndiceT>().with_device(device).verify(indices);
+    TensorMatcher({B}).with_dtype<IndiceT>().with_device(device).verify(seq_lens);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto params = Compress128OnlineDecodeParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .indices = static_cast<const IndiceT*>(indices.data_ptr()),
+        .seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
+        .batch_size = batch_size,
+    };
+    LaunchKernel(batch_size, kDecodeBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(decode_kernel, params);
+  }
+
+  static void run_prefill(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView compress_plan,
+      const tvm::ffi::TensorView write_plan,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
+    using namespace host;
+    using host::compress::kOnlinePrefillPlanDim;
+    using host::compress::OnlinePrefillPlanTensorDtype;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto N = SymbolicSize{"num_q_tokens"};
+    auto X = SymbolicSize{"compress_tokens"};
+    auto Y = SymbolicSize{"write_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, 1, kHeadDim * 3})  // kv score buffer (max, sum, kv) ??? 2D
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_score_buffer);
+    TensorMatcher({N, kHeadDim * 2})  // kv score input
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_score_input);
+    TensorMatcher({N, kHeadDim})  // kv compressed output
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_compressed_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(ape);
+    TensorMatcher({B})  // indices
+        .with_dtype<IndiceT>()
+        .with_device(device_)
+        .verify(indices);
+    TensorMatcher({X, kOnlinePrefillPlanDim})  // compress plan
+        .with_dtype<OnlinePrefillPlanTensorDtype>()
+        .with_device(device_)
+        .verify(compress_plan);
+    TensorMatcher({Y, kOnlinePrefillPlanDim})  // write plan
+        .with_dtype<OnlinePrefillPlanTensorDtype>()
+        .with_device(device_)
+        .verify(write_plan);
+
+    /// NOTE: `extra` is `load_indices`. When the previous partial state lives
+    /// on a slot different from the destination slot (e.g. paged buffers), the
+    /// caller must supply this; otherwise it defaults to `indices`.
+    const auto load_indices = extra.value_or(indices);
+    TensorMatcher({B}).with_dtype<IndiceT>().with_device(device_).verify(load_indices);
+
+    const auto device = device_.unwrap();
+    const auto num_c = static_cast<uint32_t>(X.unwrap());
+    const auto num_w = static_cast<uint32_t>(Y.unwrap());
+    const auto params = Compress128OnlinePrefillParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .indices = static_cast<const IndiceT*>(indices.data_ptr()),
+        .load_indices = static_cast<const IndiceT*>(load_indices.data_ptr()),
+        .compress_plan = static_cast<const OnlinePlan*>(compress_plan.data_ptr()),
+        .write_plan = static_cast<const OnlinePlan*>(write_plan.data_ptr()),
+        .num_compress = num_c,
+        .num_write = num_w,
+    };
+
+    /// NOTE: pass 1 reads the buffer (for the first segment of each batch
+    /// that started mid-chunk) and writes only to `kv_compressed_output`.
+    /// Pass 2 then writes the trailing partial state of each batch back to
+    /// the buffer. Stream serialization between the two launches enforces
+    /// read-before-write on shared buffer slots.
+    if (const auto num_c_blocks = num_c * kNumSplit) {
+      LaunchKernel(num_c_blocks, kPrefillBlockSize, device)  //
+          .enable_pdl(kUsePDL)(prefill_c_kernel, params);
+    }
+    if (const auto num_w_blocks = num_w * kNumSplit) {
+      LaunchKernel(num_w_blocks, kPrefillBlockSize, device)  //
+          .enable_pdl(kUsePDL)(prefill_w_kernel, params);
+    }
+  }
+};
+
+}  // namespace
+
+namespace host::compress {
+
+using OnlinePlanResult = tvm::ffi::Tuple<uint32_t, uint32_t>;
+
+struct OnlinePrefillCompressParams {
+  OnlinePrefillPlan* __restrict__ compress_plan;
+  OnlinePrefillPlan* __restrict__ write_plan;
+  const int64_t* __restrict__ seq_lens;
+  const int64_t* __restrict__ extend_lens;
+  uint32_t batch_size;
+  uint32_t num_tokens;
+};
+
+/// \brief Build the compress + write plans for online compress 128 prefill.
+///
+/// Each batch's `[prefix_len, prefix_len + extend_len)` range is split at
+/// 128-aligned boundaries. Every resulting segment falls into one of:
+/// - **compress**: closes a 128-chunk (`chunk_offset + window_len == 128`).
+///   These plans only read the buffer (when starting mid-chunk) and write the
+///   compressed kv to `kv_compressed_output`.
+/// - **write**: trailing partial of the batch (`chunk_offset + window_len < 128`).
+///   May read the buffer and always writes the new partial state back to it.
+///   Each batch produces at most one such plan.
+///
+/// The two plans MUST be dispatched as separate kernel launches in stream
+/// order so that pass-1 reads of a buffer slot complete before any pass-2
+/// write of the same slot.
+inline OnlinePlanResult plan_online_prefill_host(const OnlinePrefillCompressParams& params, const bool use_cuda_graph) {
+  const auto& [compress_plan, write_plan, seq_lens, extend_lens, batch_size, num_tokens] = params;
+
+  uint32_t counter = 0;
+  uint32_t compress_count = 0;
+  uint32_t write_count = 0;
+  for (const auto i : irange(batch_size)) {
+    const uint32_t seq_len = static_cast<uint32_t>(seq_lens[i]);
+    const uint32_t extend_len = static_cast<uint32_t>(extend_lens[i]);
+    RuntimeCheck(0 < extend_len && extend_len <= seq_len);
+    const uint32_t prefix_len = seq_len - extend_len;
+    const uint32_t end_pos = prefix_len + extend_len;
+    /// NOTE: split the extend range into per-128-chunk segments. Each segment
+    /// stays inside one chunk, so the kernel can decide load/store from
+    /// `chunk_offset` and `window_len` alone.
+    uint32_t pos = prefix_len;
+    while (pos < end_pos) {
+      const uint32_t chunk_start = (pos / 128u) * 128u;
+      const uint32_t seg_end = std::min(end_pos, chunk_start + 128u);  // exclusive
+      const uint32_t seg_len = seg_end - pos;
+      const uint32_t chunk_off = pos - chunk_start;
+      /// NOTE: store last-token coordinates so that downstream consumers
+      /// (e.g. `fused_norm_rope`) can read `ragged_id` and `position` with the
+      /// same semantics as `PrefillPlan`. The segment start is recoverable as
+      /// `ragged_id - window_len + 1` and `position - window_len + 1`.
+      const uint32_t last_pos = seg_end - 1;
+      const uint32_t last_ragged = counter + (last_pos - prefix_len);
+      const auto plan = OnlinePrefillPlan{
+          .ragged_id = last_ragged,
+          .batch_id = i,
+          .position = last_pos,
+          .window_len = seg_len,
+      };
+      if (chunk_off + seg_len == 128u) {
+        // full chunk, must be complete, maybe read the buffer, no write
+        RuntimeCheck(compress_count < num_tokens);
+        compress_plan[compress_count++] = plan;
+      } else {
+        // last chunk, must be incomplete, maybe read the buffer, must write
+        RuntimeCheck(write_count < num_tokens);
+        write_plan[write_count++] = plan;
+      }
+      pos = seg_end;
+    }
+    counter += extend_len;
+  }
+  RuntimeCheck(counter == num_tokens, "input size ", counter, " != num_q_tokens ", num_tokens);
+  if (!use_cuda_graph) return OnlinePlanResult{compress_count, write_count};
+  /// NOTE: pad both plans with sentinel entries so cuda-graph runs always see
+  /// the same number of blocks. The kernel skips plans whose `ragged_id` is -1.
+  constexpr auto kInvalid = static_cast<uint32_t>(-1);
+  constexpr auto kInvalidPlan = OnlinePrefillPlan{kInvalid, kInvalid, kInvalid, kInvalid};
+  for (const auto i : irange(compress_count, num_tokens)) {
+    compress_plan[i] = kInvalidPlan;
+  }
+  for (const auto i : irange(write_count, num_tokens)) {
+    write_plan[i] = kInvalidPlan;
+  }
+  return OnlinePlanResult{num_tokens, num_tokens};
+}
+
+inline OnlinePlanResult plan_online_prefill(
+    const tvm::ffi::TensorView extend_lens,
+    const tvm::ffi::TensorView seq_lens,
+    const tvm::ffi::TensorView compress_plan,
+    const tvm::ffi::TensorView write_plan,
+    const bool use_cuda_graph) {
+  auto N = SymbolicSize{"batch_size"};
+  auto M = SymbolicSize{"num_tokens"};
+  auto device = SymbolicDevice{};
+  /// NOTE: only host (CPU/cuda-host) planning is implemented for now. The
+  device.set_options<kDLCPU, kDLCUDAHost>();
+  TensorMatcher({N})  //
+      .with_dtype<int64_t>()
+      .with_device(device)
+      .verify(extend_lens)
+      .verify(seq_lens);
+  TensorMatcher({M, kOnlinePrefillPlanDim})  //
+      .with_dtype<OnlinePrefillPlanTensorDtype>()
+      .with_device(device)
+      .verify(compress_plan)
+      .verify(write_plan);
+  const auto params = OnlinePrefillCompressParams{
+      .compress_plan = static_cast<OnlinePrefillPlan*>(compress_plan.data_ptr()),
+      .write_plan = static_cast<OnlinePrefillPlan*>(write_plan.data_ptr()),
+      .seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
+      .extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
+      .batch_size = static_cast<uint32_t>(N.unwrap()),
+      .num_tokens = static_cast<uint32_t>(M.unwrap()),
+  };
+  return plan_online_prefill_host(params, use_cuda_graph);
+}
+
+}  // namespace host::compress
+
+namespace {
+
+[[maybe_unused]]
+constexpr auto& plan_compress_online_prefill = host::compress::plan_online_prefill;
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/c128_online_v2.cuh

@@ -0,0 +1,876 @@
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.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/compress_v2.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/container/tuple.h>
+
+#include <algorithm>
+#include <cfloat>
+#include <cstdint>
+#include <cstdlib>
+#include <cstring>
+
+namespace {
+
+using PlanD = device::compress::DecodePlan;
+using PlanC = device::compress::CompressPlan;
+
+// ---------------------------------------------------------------------------
+// Decode kernel: 1 token / batch. Each block handles one batch.
+// 4 elements per thread -> kBlockSize = head_dim / 4.
+// ---------------------------------------------------------------------------
+
+struct Compress128OnlineDecodeParams {
+  void* __restrict__ kv_score_buffer;       // [num_slots, 1, head_dim * 3]
+  const void* __restrict__ kv_score_input;  // [batch_size, head_dim * 2]
+  void* __restrict__ kv_compressed_output;  // [batch_size, head_dim]
+  const void* __restrict__ score_bias;      // [128, head_dim]
+  const PlanD* __restrict__ plan_d;
+  uint32_t batch_size;
+};
+
+template <int64_t kHeadDim, bool kUsePDL>
+__global__ void flash_c128_online_decode_v2(const __grid_constant__ Compress128OnlineDecodeParams params) {
+  using namespace device;
+  constexpr uint32_t kVecSize = 4;
+  constexpr uint32_t kBlockSize = kHeadDim / kVecSize;
+  using Vec = AlignedVector<float, kVecSize>;
+  const auto gmem = tile::Memory<Vec>::cta(kBlockSize);
+  const auto batch_id = blockIdx.x;
+  if (batch_id >= params.batch_size) return;
+
+  // Wait for the plan-finalize kernel to publish `plan.read_page_0 / write_loc`
+  // before reading the plan. The plan kernel runs on the same stream and does
+  // NOT issue a PDL trigger, so launching this kernel with PDL means our
+  // pre-wait global reads can race with the plan kernel's writes.
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto plan = params.plan_d[batch_id];
+  const auto pos_in_chunk = (plan.seq_len - 1) % 128;
+
+  const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
+  const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
+  const auto kv_load_buf = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3);
+  const auto kv_store_buf = kv_score_buffer + plan.write_loc * (kHeadDim * 3);
+  const auto kv_src = kv_score_input + batch_id * (kHeadDim * 2);
+
+  // Buffer layout: [max | sum | kv] (slot 0 / 1 / 2 of the head_dim*3 row).
+  const auto new_kv_vec = gmem.load(kv_src, 0);
+  const auto new_score_raw_vec = gmem.load(kv_src, 1);
+  const auto bias_vec = gmem.load(params.score_bias, pos_in_chunk);
+
+  Vec out_kv_vec;
+  Vec out_max_vec;
+  Vec out_sum_vec;
+  if (pos_in_chunk != 0) {
+    // Mid-chunk: combine prior partial state with the new token.
+    const auto max_score_vec = gmem.load(kv_load_buf, 0);
+    const auto sum_score_vec = gmem.load(kv_load_buf, 1);
+    const auto old_kv_vec = gmem.load(kv_load_buf, 2);
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      const auto old_max = max_score_vec[i];
+      const auto old_kv = old_kv_vec[i];
+      const auto new_score = new_score_raw_vec[i] + bias_vec[i];
+      const auto new_kv = new_kv_vec[i];
+      const auto new_max = fmaxf(old_max, new_score);
+      const auto old_sum = sum_score_vec[i] * expf(old_max - new_max);
+      const auto new_exp = expf(new_score - new_max);
+      const auto new_sum = old_sum + new_exp;
+      out_kv_vec[i] = (old_kv * old_sum + new_kv * new_exp) / new_sum;
+      out_max_vec[i] = new_max;
+      out_sum_vec[i] = new_sum;
+    }
+  } else {
+    // First token of a new chunk: state == this token alone.
+#pragma unroll
+    for (uint32_t i = 0; i < kVecSize; ++i) {
+      out_kv_vec[i] = new_kv_vec[i];
+      out_max_vec[i] = new_score_raw_vec[i] + bias_vec[i];
+      out_sum_vec[i] = 1.0f;
+    }
+  }
+
+  if (pos_in_chunk == 127) {
+    // Chunk just closed: emit compressed kv, no buffer update.
+    const auto kv_out = static_cast<float*>(params.kv_compressed_output) + batch_id * kHeadDim;
+    gmem.store(kv_out, out_kv_vec);
+  } else {
+    gmem.store(kv_store_buf, out_max_vec, 0);
+    gmem.store(kv_store_buf, out_sum_vec, 1);
+    gmem.store(kv_store_buf, out_kv_vec, 2);
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Prefill kernel: 1 segment / block. Two passes (compress + write) share the
+// kernel template, parameterized by `kWrite`.
+// 16 warps per block; each warp handles 8 of the 128 chunk positions.
+// ---------------------------------------------------------------------------
+
+constexpr int32_t kTileElements = 2;     // split along head-dim
+constexpr int32_t kElementsPerWarp = 8;  // split along the 128-chunk
+constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
+constexpr uint32_t kPrefillBlockSize = device::kWarpThreads * kNumWarps;
+using PrefillStorage = device::AlignedVector<float, kTileElements>;
+
+struct Compress128OnlinePrefillParams {
+  void* __restrict__ kv_score_buffer;       // [num_slots, 1, head_dim * 3]
+  const void* __restrict__ kv_score_input;  // [num_q_tokens, head_dim * 2]
+  void* __restrict__ kv_compressed_output;  // [num_compress, head_dim]
+  const void* __restrict__ score_bias;      // [128, head_dim]
+  const PlanC* __restrict__ plan_c;         // close-chunk segments
+  const PlanC* __restrict__ plan_w;         // trailing partial segments
+  uint32_t num_compress;
+  uint32_t num_write;
+};
+
+struct Compress128SharedBuffer {
+  using Storage = device::AlignedVector<float, 4>;
+  Storage data[kNumWarps][device::kWarpThreads + 1];  // +1 to avoid bank conflict
+  SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
+    return data[warp_id][lane_id];
+  }
+  SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
+    return data[warp_id][lane_id][tile_id];
+  }
+};
+
+/// \brief Sentinel score for padded positions in a 128-segment.
+constexpr float kPadScore = -FLT_MAX;
+
+[[maybe_unused]]
+SGL_DEVICE void c128_prefill_segment_softmax(
+    const PrefillStorage (&kv)[kElementsPerWarp],
+    const PrefillStorage (&score)[kElementsPerWarp],
+    float* seg_kv,
+    float* seg_max,
+    float* seg_sum,
+    const uint32_t warp_id,
+    const uint32_t lane_id) {
+  using namespace device;
+
+  // Per-warp running state (max, sum, kv) for kTileElements head-dim slots.
+  using TmpStorage = typename Compress128SharedBuffer::Storage;
+  __shared__ Compress128SharedBuffer s_local_val_max;
+  __shared__ Compress128SharedBuffer s_local_exp_sum;
+  __shared__ Compress128SharedBuffer s_local_product;
+
+  TmpStorage tmp_val_max;
+  TmpStorage tmp_exp_sum;
+  TmpStorage tmp_product;
+
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+    float score_fp32[kElementsPerWarp];
+#pragma unroll
+    for (int32_t j = 0; j < kElementsPerWarp; ++j) {
+      score_fp32[j] = score[j][i];
+    }
+    float max_value = score_fp32[0];
+#pragma unroll
+    for (int32_t j = 1; j < kElementsPerWarp; ++j) {
+      max_value = fmaxf(max_value, score_fp32[j]);
+    }
+    float sum_exp_value = 0.0f;
+    float sum_product = 0.0f;
+#pragma unroll
+    for (int32_t j = 0; j < kElementsPerWarp; ++j) {
+      const auto exp_score = expf(score_fp32[j] - max_value);
+      sum_product += kv[j][i] * exp_score;
+      sum_exp_value += exp_score;
+    }
+    tmp_val_max[i] = max_value;
+    tmp_exp_sum[i] = sum_exp_value;
+    tmp_product[i] = sum_product;
+  }
+
+  // Aligned writes (no bank conflict thanks to `+1` padding).
+  s_local_val_max(warp_id, lane_id) = tmp_val_max;
+  s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
+  s_local_product(warp_id, lane_id) = tmp_product;
+
+  __syncthreads();
+
+  // Cross-warp reduction. Same recipe as c128_online.cuh: each block-thread
+  // pair reduces a (tile_id, lane_id) slot using a kNumWarps-wide warp shuffle.
+  constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
+  constexpr uint32_t kIteration = kReductionCount / kPrefillBlockSize;
+  static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
+
+#pragma unroll
+  for (uint32_t i = 0; i < kIteration; ++i) {
+    const uint32_t j = i * kPrefillBlockSize + warp_id * kWarpThreads + lane_id;
+    const uint32_t local_warp_id = j % kNumWarps;
+    const uint32_t local_elem_id = j / kNumWarps;
+    const uint32_t local_tile_id = local_elem_id % kTileElements;
+    const uint32_t local_lane_id = local_elem_id / kTileElements;
+    const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
+    const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
+    const auto rescale = expf(local_val_max - global_val_max);
+    const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
+    const auto final_scale = rescale / global_exp_sum;
+    const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
+    seg_kv[local_elem_id] = global_product;
+    seg_max[local_elem_id] = global_val_max;
+    seg_sum[local_elem_id] = global_exp_sum;
+  }
+  __syncthreads();
+}
+
+/// \brief Online compress 128 prefill v2.
+///
+/// `kWrite=false` (compress pass): handles segments that close a 128-chunk.
+///   Reads optional prior state from `read_page_0` (-1 = none), emits compressed
+///   kv to `kv_compressed_output[plan_id]` (compact).
+/// `kWrite=true`  (write pass)   : handles trailing partial segments.
+///   Reads optional prior state from `read_page_1` (-1 = fallback to
+///   `read_page_0`), writes new running state to `read_page_0`.
+template <int64_t kHeadDim, bool kWrite, bool kUsePDL>
+__global__ __launch_bounds__(kPrefillBlockSize, 2)  //
+    void flash_c128_online_prefill_v2(const __grid_constant__ Compress128OnlinePrefillParams params) {
+  using namespace device;
+
+  constexpr int64_t kTileDim = kTileElements * kWarpThreads;  // 64
+  constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static_assert(kHeadDim % kTileDim == 0);
+
+  // Compile-time fold to the right plan list.
+  const auto num_plans = kWrite ? params.num_write : params.num_compress;
+  const auto plan_ptr = kWrite ? params.plan_w : params.plan_c;
+  const uint32_t global_id = blockIdx.x;
+  const uint32_t global_pid = global_id / kNumSplit;
+  const uint32_t global_sid = global_id % kNumSplit;
+  if (global_pid >= num_plans) return;
+
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t lane_id = threadIdx.x % kWarpThreads;
+  const int32_t split_offset = global_sid * kTileDim;
+
+  // The previous kernel (plan-finalize stage 1) does NOT issue a PDL trigger,
+  // so PDLWaitPrimary effectively waits for stage 1 to complete. Read the plan
+  // AFTER the wait so the freshly-written `read_page_0` (= state-pool slot) is
+  // visible. Reading it before the wait is a real race -- with PDL enabled the
+  // kernel can begin executing before stage 1's stores propagate, and we'd see
+  // the stage-0 batch_id placeholder in `read_page_0` instead of the slot.
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto plan = plan_ptr[global_pid];
+  if (plan.is_invalid()) [[unlikely]]
+    return;
+
+  const auto kv_score_buffer = static_cast<float*>(params.kv_score_buffer);
+  const auto kv_score_input = static_cast<const float*>(params.kv_score_input);
+  const auto kv_compressed_output = static_cast<float*>(params.kv_compressed_output);
+  const auto score_bias_base = static_cast<const float*>(params.score_bias);
+
+  constexpr int64_t kElementSize = kHeadDim * 2;  // | kv | score |
+
+  // `j` below is a chunk-local offset. Convert it to the ragged-input row by
+  // anchoring on the last token in this segment: ragged_id - pos_in_chunk_end + 1 + j.
+  const uint32_t window_len = plan.buffer_len;
+  const uint32_t position = plan.seq_len - 1;
+  const uint32_t pos_in_chunk_end = (position % 128u) + 1u;     // exclusive, in [1, 128]
+  const uint32_t chunk_offset = pos_in_chunk_end - window_len;  // in [0, 127]
+  const int32_t chunk_start_ragged = static_cast<int32_t>(plan.ragged_id) - static_cast<int32_t>(pos_in_chunk_end) + 1;
+
+  // --- Stage 1: load kv / score / bias for this warp's 8 chunk positions.
+  PrefillStorage kv[kElementsPerWarp];
+  PrefillStorage score[kElementsPerWarp];
+  PrefillStorage bias[kElementsPerWarp];
+  const uint32_t warp_offset = warp_id * kElementsPerWarp;
+
+#pragma unroll
+  for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
+    const uint32_t j = i + warp_offset;
+    if (j >= chunk_offset && j < pos_in_chunk_end) {
+      const int32_t ragged_id = chunk_start_ragged + static_cast<int32_t>(j);
+      const auto kv_src_ptr = kv_score_input + ragged_id * kElementSize + split_offset;
+      const auto score_src_ptr = kv_src_ptr + kHeadDim;
+      const auto bias_src_ptr = score_bias_base + j * kHeadDim + split_offset;
+      kv[i].load(kv_src_ptr, lane_id);
+      score[i].load(score_src_ptr, lane_id);
+      bias[i].load(bias_src_ptr, lane_id);
+    }
+  }
+
+  // --- Stage 2: pad invalid positions. score = -FLT_MAX, kv = 0 (so that
+  // kv * exp(score-max) ??? 0 / 0 cleanly without producing NaN/inf).
+#pragma unroll
+  for (uint32_t i = 0; i < kElementsPerWarp; ++i) {
+    const uint32_t j = i + warp_offset;
+    const bool is_valid = (j >= chunk_offset && j < pos_in_chunk_end);
+#pragma unroll
+    for (uint32_t ii = 0; ii < kTileElements; ++ii) {
+      score[i][ii] = is_valid ? score[i][ii] + bias[i][ii] : kPadScore;
+      kv[i][ii] = is_valid ? kv[i][ii] : 0.0f;
+    }
+  }
+
+  // --- Stage 3: warp-tile online softmax over the 128-position chunk.
+  __shared__ alignas(16) float seg_kv[kTileDim];
+  __shared__ alignas(16) float seg_max[kTileDim];
+  __shared__ alignas(16) float seg_sum[kTileDim];
+  c128_prefill_segment_softmax(kv, score, seg_kv, seg_max, seg_sum, warp_id, lane_id);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // --- Stage 4: warp 0 folds with prior partial state (if any) and writes.
+  if (warp_id == 0) {
+    PrefillStorage out_kv_vec, out_max_vec, out_sum_vec;
+    out_kv_vec.load(seg_kv, lane_id);
+    out_max_vec.load(seg_max, lane_id);
+    out_sum_vec.load(seg_sum, lane_id);
+
+    const int32_t read_page = plan.read_page_1 >= 0 ? plan.read_page_1 : plan.read_page_0;
+    if (chunk_offset != 0 && read_page >= 0) {
+      // Combine with prior partial state for this slot.
+      const auto buf_load = kv_score_buffer + read_page * (kHeadDim * 3) + split_offset;
+      PrefillStorage buf_max_vec, buf_sum_vec, buf_kv_vec;
+      buf_max_vec.load(buf_load + 0 * kHeadDim, lane_id);
+      buf_sum_vec.load(buf_load + 1 * kHeadDim, lane_id);
+      buf_kv_vec.load(buf_load + 2 * kHeadDim, lane_id);
+#pragma unroll
+      for (uint32_t ii = 0; ii < kTileElements; ++ii) {
+        const float m1 = buf_max_vec[ii];
+        const float s1 = buf_sum_vec[ii];
+        const float k1 = buf_kv_vec[ii];
+        const float m2 = out_max_vec[ii];
+        const float s2 = out_sum_vec[ii];
+        const float k2 = out_kv_vec[ii];
+        const float new_max = fmaxf(m1, m2);
+        const float new_s1 = s1 * expf(m1 - new_max);
+        const float new_s2 = s2 * expf(m2 - new_max);
+        const float new_sum = new_s1 + new_s2;
+        const float new_kv = (k1 * new_s1 + k2 * new_s2) / new_sum;
+        out_max_vec[ii] = new_max;
+        out_sum_vec[ii] = new_sum;
+        out_kv_vec[ii] = new_kv;
+      }
+    }
+
+    if constexpr (kWrite) {
+      // For trailing-partial segments the load and store slots collapse to the
+      // segment's own chunk slot (the request keeps a single in-progress
+      // chunk's running state at any time), so we reuse `read_page_0`.
+      const auto buf_store = kv_score_buffer + plan.read_page_0 * (kHeadDim * 3) + split_offset;
+      reinterpret_cast<PrefillStorage*>(buf_store + 0 * kHeadDim)[lane_id] = out_max_vec;
+      reinterpret_cast<PrefillStorage*>(buf_store + 1 * kHeadDim)[lane_id] = out_sum_vec;
+      reinterpret_cast<PrefillStorage*>(buf_store + 2 * kHeadDim)[lane_id] = out_kv_vec;
+    } else {
+      // Compact output: one row per compress plan, indexed by `global_pid`.
+      const auto out_ptr = kv_compressed_output + global_pid * kHeadDim + split_offset;
+      reinterpret_cast<PrefillStorage*>(out_ptr)[lane_id] = out_kv_vec;
+    }
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Host wrapper: matches the c128_v2 / c4_v2 host API style (run_decode /
+// run_prefill methods on a kernel-class template). We only expose `kHeadDim`
+// + `kUsePDL`; the dtype is fixed to fp32 for the online state pool.
+// ---------------------------------------------------------------------------
+
+template <int64_t kHeadDim, bool kUsePDL>
+struct FlashCompress128OnlineKernel {
+  static constexpr auto decode_kernel = flash_c128_online_decode_v2<kHeadDim, kUsePDL>;
+  template <bool kWrite>
+  static constexpr auto prefill_kernel = flash_c128_online_prefill_v2<kHeadDim, kWrite, kUsePDL>;
+  static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static constexpr uint32_t kDecodeBlockSize = kHeadDim / 4;
+
+  static void run_decode(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView plan_d_) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, 1, kHeadDim * 3})  // kv score buffer (max, sum, kv)
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_score_buffer);
+    TensorMatcher({B, kHeadDim * 2})  // kv score input
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_score_input);
+    TensorMatcher({B, kHeadDim})  // kv compressed output (sparse by batch_id)
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_compressed_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(ape);
+
+    const auto plan_d = compress::verify_plan_d(plan_d_, B, device_);
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    if (batch_size == 0) return;
+    const auto params = Compress128OnlineDecodeParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .plan_d = plan_d,
+        .batch_size = batch_size,
+    };
+    LaunchKernel(batch_size, kDecodeBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(decode_kernel, params);
+  }
+
+  static void run_prefill(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView plan_c_,
+      const tvm::ffi::TensorView plan_w_) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_q_tokens"};
+    auto C = SymbolicSize{"num_c_plans"};
+    auto W = SymbolicSize{"num_w_plans"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, 1, kHeadDim * 3})  // kv score buffer
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_score_buffer);
+    TensorMatcher({N, kHeadDim * 2})  // kv score input (ragged)
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_score_input);
+    TensorMatcher({C, kHeadDim})  // kv compressed output (compact, by plan_c index)
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(kv_compressed_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(ape);
+
+    // Both compress and write segments use PlanC layout. Stage 1 stores the
+    // committed-bank load slot in read_page_1 and the write slot in read_page_0.
+    const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
+    const auto plan_w = compress::verify_plan_c(plan_w_, W, device_);
+    const auto device = device_.unwrap();
+    const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_c = static_cast<uint32_t>(C.unwrap());
+    const auto num_w = static_cast<uint32_t>(W.unwrap());
+    RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
+    const auto params = Compress128OnlinePrefillParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .plan_c = plan_c,
+        .plan_w = plan_w,
+        .num_compress = num_c,
+        .num_write = num_w,
+    };
+
+    // The two passes MUST be serialized in stream order: pass 1 reads slots
+    // that pass 2 may write to; running them in parallel would race.
+    if (const auto num_c_blocks = num_c * kNumSplit) {
+      LaunchKernel(num_c_blocks, kPrefillBlockSize, device)  //
+          .enable_pdl(kUsePDL)(prefill_kernel</*kWrite=*/false>, params);
+    }
+    if (const auto num_w_blocks = num_w * kNumSplit) {
+      LaunchKernel(num_w_blocks, kPrefillBlockSize, device)  //
+          .enable_pdl(kUsePDL)(prefill_kernel</*kWrite=*/true>, params);
+    }
+  }
+};
+
+}  // namespace
+
+// ===========================================================================
+// Plan builders. Mirrors the offline v2 pattern (`c_plan.cuh`):
+//   - Decode: a single GPU kernel reads seq_lens / req_to_token /
+//     req_pool_indices on device and emits the final PlanD tensor in one go.
+//   - Prefill: stage 0 (host, on CPU pinned memory) splits each batch's
+//     extend range into per-chunk segments and emits PlanC entries with the
+//     batch_id stashed in `read_page_0` as a placeholder. Stage 1 is a tiny
+//     GPU kernel that finalizes `read_page_0` to `req_to_token[rid][chunk_start]`,
+//     so the slot tensors never leave GPU memory. The online state pool keeps
+//     a single in-progress chunk per request, so each segment's load and
+//     store slot collapse to one value (the slot for the segment's own chunk).
+//     For online-c128 MTP, stage 1 keeps that write slot in `read_page_0` and
+//     stores the committed-bank load slot in `read_page_1`.
+// ===========================================================================
+
+namespace host::compress {
+
+using device::compress::CompressPlan;
+using device::compress::DecodePlan;
+
+// ---------------------------------------------------------------------------
+// Decode plan builder.
+// ---------------------------------------------------------------------------
+
+struct OnlineDecodePlanParams {
+  DecodePlan* __restrict__ plan_d;
+  const int64_t* __restrict__ seq_lens;
+  const int64_t* __restrict__ req_pool_indices;
+  const int32_t* __restrict__ req_to_token;
+  int64_t stride_r2t;
+  int32_t state_slot_offset;
+  uint32_t batch_size;
+};
+
+__global__ void plan_c128_online_decode_kernel(const OnlineDecodePlanParams params) {
+  const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= params.batch_size) return;
+  const auto seq_len = static_cast<uint32_t>(params.seq_lens[idx]);
+  const auto rid = params.req_pool_indices[idx];
+  const int32_t slot = static_cast<int32_t>(rid) + params.state_slot_offset;
+  params.plan_d[idx] = DecodePlan{
+      .seq_len = seq_len,
+      .write_loc = slot,
+      .read_page_0 = slot,
+      .read_page_1 = -1,
+  };
+}
+
+/// \brief Build the decode plan tensor. Caller (Python) pre-allocates
+/// `plan_d_dev` as a `(batch_size, 16)` device uint8 tensor; this routine
+/// only fills it. See `plan_online_prefill` for the rationale (avoid
+/// `ffi::empty` + dlpack roundtrip / PyTorch caching-allocator stream
+/// tracking issue that surfaces as IMA in unrelated downstream kernels).
+inline void plan_online_decode(
+    const tvm::ffi::TensorView seq_lens,
+    const tvm::ffi::TensorView req_pool_indices,
+    const tvm::ffi::TensorView req_to_token,
+    const tvm::ffi::TensorView plan_d_dev_,
+    const int32_t state_slot_offset) {
+  auto B = SymbolicSize{"batch_size"};
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLCUDA>();
+
+  auto seq_dtype = SymbolicDType{};
+  TensorMatcher({B})  //
+      .with_dtype<int64_t>(seq_dtype)
+      .with_device(device_)
+      .verify(seq_lens);
+  TensorMatcher({B})  //
+      .with_dtype<int64_t>()
+      .with_device(device_)
+      .verify(req_pool_indices);
+  TensorMatcher({-1, -1})  //
+      .with_dtype<int32_t>()
+      .with_device(device_)
+      .verify(req_to_token);
+  TensorMatcher({B, sizeof(DecodePlan)})  //
+      .with_dtype<uint8_t>()
+      .with_device(device_)
+      .verify(plan_d_dev_);
+  RuntimeCheck(state_slot_offset >= 0);
+
+  const auto batch_size = static_cast<uint32_t>(B.unwrap());
+  if (batch_size == 0) return;
+
+  const auto device = device_.unwrap();
+  constexpr uint32_t kBlockSize = 256;
+  const uint32_t num_blocks = host::div_ceil(batch_size, kBlockSize);
+  const auto stride_r2t = req_to_token.stride(0);
+  const auto params = OnlineDecodePlanParams{
+      .plan_d = static_cast<DecodePlan*>(plan_d_dev_.data_ptr()),
+      .seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
+      .req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
+      .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
+      .stride_r2t = stride_r2t,
+      .state_slot_offset = state_slot_offset,
+      .batch_size = batch_size,
+  };
+  LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_decode_kernel, params);
+}
+
+// ---------------------------------------------------------------------------
+// Prefill plan builder: host stage 0 + GPU stage 1.
+// ---------------------------------------------------------------------------
+
+struct OnlinePrefillStage0Params {
+  CompressPlan* __restrict__ plan_c;
+  CompressPlan* __restrict__ plan_w;
+  const int64_t* __restrict__ seq_lens;
+  const int64_t* __restrict__ extend_lens;
+  uint32_t batch_size;
+  uint32_t num_q_tokens;
+};
+
+inline std::tuple<uint32_t, uint32_t> _plan_prefill_partial(const OnlinePrefillStage0Params& p) {
+  uint32_t counter = 0;
+  uint32_t compress_count = 0;
+  uint32_t write_count = 0;
+  for (const auto i : irange(p.batch_size)) {
+    const uint32_t seq_len = static_cast<uint32_t>(p.seq_lens[i]);
+    const uint32_t extend_len = static_cast<uint32_t>(p.extend_lens[i]);
+    RuntimeCheck(0 < extend_len && extend_len <= seq_len);
+    const uint32_t prefix_len = seq_len - extend_len;
+    const uint32_t end_pos = prefix_len + extend_len;
+
+    uint32_t pos = prefix_len;
+    while (pos < end_pos) {
+      const uint32_t chunk_start = (pos / 128u) * 128u;
+      const uint32_t seg_end = std::min(end_pos, chunk_start + 128u);  // exclusive
+      const uint32_t seg_len = seg_end - pos;
+      const uint32_t chunk_off = pos - chunk_start;
+      const uint32_t last_pos = seg_end - 1;
+      const uint32_t last_ragged = counter + (last_pos - prefix_len);
+      RuntimeCheck(last_ragged < (1u << 16), "PlanC.ragged_id is uint16; ragged ", last_ragged, " overflows");
+      RuntimeCheck(seg_len <= 128u);
+      // Stash batch_id in `read_page_0` for stage 1 to translate. A
+      // chunk-aligned segment never loads, so we still need stage 1 to fill
+      // a slot in -- the kernel keys the load on `chunk_offset != 0`.
+      const auto plan = CompressPlan{
+          .seq_len = last_pos + 1u,
+          .ragged_id = static_cast<uint16_t>(last_ragged),
+          .buffer_len = static_cast<uint16_t>(seg_len),
+          .read_page_0 = static_cast<int32_t>(i),  // batch_id placeholder
+          .read_page_1 = -1,                       // filled by stage 1 with committed-bank slot
+      };
+      if (chunk_off + seg_len == 128u) {
+        // close-chunk segment
+        RuntimeCheck(compress_count < p.num_q_tokens);
+        p.plan_c[compress_count++] = plan;
+      } else {
+        // trailing partial segment
+        RuntimeCheck(write_count < p.num_q_tokens);
+        p.plan_w[write_count++] = plan;
+      }
+      pos = seg_end;
+    }
+    counter += extend_len;
+  }
+  RuntimeCheck(counter == p.num_q_tokens, "input size ", counter, " != num_q_tokens ", p.num_q_tokens);
+  return std::tuple<uint32_t, uint32_t>{compress_count, write_count};
+}
+
+struct OnlinePrefillStage1Params {
+  CompressPlan* __restrict__ plan_c;
+  CompressPlan* __restrict__ plan_w;
+  const int64_t* __restrict__ req_pool_indices;  // (batch_size,)
+  const int32_t* __restrict__ req_to_token;      // (num_reqs, max_tokens)
+  int64_t stride_r2t;
+  int32_t state_slot_offset;
+  uint32_t num_c;
+  uint32_t num_w;
+};
+
+__global__ void plan_c128_online_prefill_kernel(const OnlinePrefillStage1Params params) {
+  const uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t total = params.num_c + params.num_w;
+  if (idx >= total) return;
+
+  const bool is_compress = idx < params.num_c;
+  CompressPlan* const plan_ptr = is_compress ? &params.plan_c[idx] : &params.plan_w[idx - params.num_c];
+  auto plan = *plan_ptr;
+  if (plan.is_invalid()) return;
+  const auto batch_id = plan.read_page_0;
+  const auto rid = params.req_pool_indices[batch_id];
+  const int32_t main_slot = static_cast<int32_t>(rid);
+  plan.read_page_0 = main_slot + params.state_slot_offset;
+  plan.read_page_1 = main_slot;
+  *plan_ptr = plan;
+}
+
+using OnlinePrefillPlan = tvm::ffi::Tuple<uint32_t, uint32_t>;
+
+inline OnlinePrefillPlan plan_online_prefill(
+    const tvm::ffi::TensorView seq_lens,
+    const tvm::ffi::TensorView extend_lens,
+    const tvm::ffi::TensorView req_pool_indices,
+    const tvm::ffi::TensorView req_to_token,
+    const tvm::ffi::TensorView plan_c_pin,
+    const tvm::ffi::TensorView plan_w_pin,
+    const tvm::ffi::TensorView plan_c_dev_,
+    const tvm::ffi::TensorView plan_w_dev_,
+    const int32_t state_slot_offset,
+    const bool use_cuda_graph) {
+  auto B = SymbolicSize{"batch_size"};
+  auto N = SymbolicSize{"num_q_tokens"};
+  auto cpu = SymbolicDevice{};
+  auto device_ = SymbolicDevice{};
+  cpu.set_options<kDLCPU, kDLCUDAHost>();
+  device_.set_options<kDLCUDA>();
+
+  TensorMatcher({B})  //
+      .with_dtype<int64_t>()
+      .with_device(cpu)
+      .verify(seq_lens)
+      .verify(extend_lens);
+  TensorMatcher({B})  //
+      .with_dtype<int64_t>()
+      .with_device(device_)
+      .verify(req_pool_indices);
+  TensorMatcher({-1, -1})  //
+      .with_dtype<int32_t>()
+      .with_device(device_)
+      .verify(req_to_token);
+  TensorMatcher({N, sizeof(CompressPlan)})  //
+      .with_dtype<uint8_t>()
+      .with_device(cpu)
+      .verify(plan_c_pin)
+      .verify(plan_w_pin);
+  TensorMatcher({N, sizeof(CompressPlan)})  //
+      .with_dtype<uint8_t>()
+      .with_device(device_)
+      .verify(plan_c_dev_)
+      .verify(plan_w_dev_);
+  RuntimeCheck(state_slot_offset >= 0);
+
+  const auto stage0_params = OnlinePrefillStage0Params{
+      .plan_c = static_cast<CompressPlan*>(plan_c_pin.data_ptr()),
+      .plan_w = static_cast<CompressPlan*>(plan_w_pin.data_ptr()),
+      .seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
+      .extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
+      .batch_size = static_cast<uint32_t>(B.unwrap()),
+      .num_q_tokens = static_cast<uint32_t>(N.unwrap()),
+  };
+
+  // Debug instrumentation: SGLANG_DEBUG_C128_ONLINE_GUARD=1 wraps stage 0
+  // with redzone + post-write magic-check on the pin buffers, plus a strict
+  // upper-bound check on `batch_size` and `num_q_tokens`. If stage 0 has a
+  // CPU OOB this trips a clear panic at the offending byte instead of a
+  // delayed CUDA IMA from corrupted heap memory.
+  static const bool kGuard = []() {
+    const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_GUARD");
+    return v != nullptr && v[0] == '1';
+  }();
+  if (kGuard) {
+    RuntimeCheck(stage0_params.batch_size <= 65536u, "batch_size out of bound: ", stage0_params.batch_size);
+    RuntimeCheck(stage0_params.num_q_tokens <= 65536u, "num_q_tokens out of bound: ", stage0_params.num_q_tokens);
+    // Stamp the pin buffers with 0xAB so we can detect any byte still 0xAB
+    // beyond what stage 0 should have written (= OOB never reached, that's fine)
+    // or any byte BEYOND num_q_tokens*16 written to (= true OOB into
+    // adjacent allocation).
+    auto* pc = static_cast<uint8_t*>(plan_c_pin.data_ptr());
+    auto* pw = static_cast<uint8_t*>(plan_w_pin.data_ptr());
+    const auto bytes = static_cast<size_t>(N.unwrap()) * sizeof(CompressPlan);
+    std::memset(pc, 0xAB, bytes);
+    std::memset(pw, 0xAB, bytes);
+  }
+
+  const auto [num_c, num_w] = _plan_prefill_partial(stage0_params);
+  const auto num_c_padded = use_cuda_graph ? static_cast<uint32_t>(N.unwrap()) : num_c;
+  const auto num_w_padded = use_cuda_graph ? static_cast<uint32_t>(N.unwrap()) : num_w;
+
+  if (kGuard) {
+    // Verify stage 0 wrote ONLY to the [0, num_c*16) and [0, num_w*16) prefix.
+    auto* pc = static_cast<const uint8_t*>(plan_c_pin.data_ptr());
+    auto* pw = static_cast<const uint8_t*>(plan_w_pin.data_ptr());
+    const auto end_c = static_cast<size_t>(num_c) * sizeof(CompressPlan);
+    const auto end_w = static_cast<size_t>(num_w) * sizeof(CompressPlan);
+    const auto pin_bytes = static_cast<size_t>(N.unwrap()) * sizeof(CompressPlan);
+    for (size_t k = end_c; k < pin_bytes; ++k) {
+      RuntimeCheck(
+          pc[k] == 0xAB,
+          "GUARD: plan_c_pin OOB write at byte ",
+          k,
+          " (num_c=",
+          num_c,
+          ", num_q_tokens=",
+          N.unwrap(),
+          ")");
+    }
+    for (size_t k = end_w; k < pin_bytes; ++k) {
+      RuntimeCheck(
+          pw[k] == 0xAB,
+          "GUARD: plan_w_pin OOB write at byte ",
+          k,
+          " (num_w=",
+          num_w,
+          ", num_q_tokens=",
+          N.unwrap(),
+          ")");
+    }
+  }
+
+  const auto device = device_.unwrap();
+  // Out-params pre-allocated by Python. Cast to typed pointers for use.
+  auto* const plan_c_dev_ptr = static_cast<CompressPlan*>(plan_c_dev_.data_ptr());
+  auto* const plan_w_dev_ptr = static_cast<CompressPlan*>(plan_w_dev_.data_ptr());
+
+  if (use_cuda_graph) {
+    const auto kInvalidPlan = CompressPlan::invalid();
+    auto* const plan_c_pin_ptr = static_cast<CompressPlan*>(plan_c_pin.data_ptr());
+    auto* const plan_w_pin_ptr = static_cast<CompressPlan*>(plan_w_pin.data_ptr());
+    for (const auto i : irange(num_c, num_c_padded)) {
+      plan_c_pin_ptr[i] = kInvalidPlan;
+    }
+    for (const auto i : irange(num_w, num_w_padded)) {
+      plan_w_pin_ptr[i] = kInvalidPlan;
+    }
+  }
+
+  if (const auto total = num_c_padded + num_w_padded) {
+    const auto stream = LaunchKernel::resolve_device(device);
+    // SGLANG_DEBUG_C128_ONLINE_SYNC_H2D=1 forces a synchronous H2D copy.
+    static const bool kSyncH2D = []() {
+      const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_SYNC_H2D");
+      return v != nullptr && v[0] == '1';
+    }();
+    // SGLANG_DEBUG_C128_ONLINE_NO_H2D=1 skips the H2D copy entirely (debug only).
+    static const bool kNoH2D = []() {
+      const char* v = std::getenv("SGLANG_DEBUG_C128_ONLINE_NO_H2D");
+      return v != nullptr && v[0] == '1';
+    }();
+    const auto copy_to_device = [stream](void* dst, void* src, int64_t count) {
+      if (kNoH2D) return;
+      const auto bytes = count * sizeof(CompressPlan);
+      if (kSyncH2D) {
+        RuntimeDeviceCheck(::cudaMemcpy(dst, src, bytes, ::cudaMemcpyHostToDevice));
+      } else {
+        RuntimeDeviceCheck(::cudaMemcpyAsync(dst, src, bytes, ::cudaMemcpyHostToDevice, stream));
+      }
+    };
+    if (num_c_padded) copy_to_device(plan_c_dev_ptr, plan_c_pin.data_ptr(), num_c_padded);
+    if (num_w_padded) copy_to_device(plan_w_dev_ptr, plan_w_pin.data_ptr(), num_w_padded);
+
+    const auto stage1_params = OnlinePrefillStage1Params{
+        .plan_c = plan_c_dev_ptr,
+        .plan_w = plan_w_dev_ptr,
+        .req_pool_indices = static_cast<const int64_t*>(req_pool_indices.data_ptr()),
+        .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
+        .stride_r2t = req_to_token.stride(0),
+        .state_slot_offset = state_slot_offset,
+        .num_c = num_c_padded,
+        .num_w = num_w_padded,
+    };
+    constexpr uint32_t kBlockSize = 128;
+    const auto num_blocks = host::div_ceil(total, kBlockSize);
+    LaunchKernel(num_blocks, kBlockSize, device)(plan_c128_online_prefill_kernel, stage1_params);
+  }
+  return OnlinePrefillPlan{num_c_padded, num_w_padded};
+}
+
+}  // namespace host::compress
+
+namespace {
+
+[[maybe_unused]]
+constexpr auto& plan_compress_128_online_decode = host::compress::plan_online_decode;
+[[maybe_unused]]
+constexpr auto& plan_compress_128_online_prefill = host::compress::plan_online_prefill;
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/c128_v2.cuh

@@ -0,0 +1,512 @@
+/**
+ * \brief Here's some dimension info for the main buffer used in C128 prefill and decode.
+ *
+ * kv_buffer: [num_indices, 128, head_dim * 2]
+ * - last dimension layout: | kv | score |
+ * kv_input: [batch_size, head_dim * 2]
+ * kv_output: [batch_size, head_dim]
+ * score_bias (ape): [128, head_dim]
+ * plan_c/plan_w: [variable length]
+ *
+ * For prefill, batch_size = num_q_tokens
+ */
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.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/compress_v2.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/object.h>
+
+#include <cfloat>
+#include <cstdint>
+#include <type_traits>
+
+namespace {
+
+using PlanD = device::compress::DecodePlan;
+using PlanC = device::compress::CompressPlan;
+using PlanW = device::compress::WritePlan;
+
+/// \brief Each thread will handle this many elements (split along head_dim)
+constexpr int32_t kTileElements = 2;
+/// \brief Each warp will handle this many elements (split along 128)
+constexpr int32_t kElementsPerWarp = 8;
+constexpr uint32_t kNumWarps = 128 / kElementsPerWarp;
+constexpr uint32_t kBlockSize = device::kWarpThreads * kNumWarps;
+constexpr uint32_t kWriteBlockSize = 128;  // one warp per write
+
+/// \brief Need to reduce register usage to increase occupancy
+#define C128_KERNEL __global__ __launch_bounds__(kBlockSize, 2)
+#define WRITE_KERNEL __global__ __launch_bounds__(kWriteBlockSize, 16)
+
+struct Compress128DecodeParams {
+  void* __restrict__ kv_buffer;
+  const void* __restrict__ kv_input;
+  void* __restrict__ kv_output;
+  const void* __restrict__ score_bias;
+  const PlanD* __restrict__ plan_d;
+  uint32_t batch_size;
+};
+
+struct Compress128PrefillParams {
+  void* __restrict__ kv_buffer;
+  const void* __restrict__ kv_input;
+  void* __restrict__ kv_output;
+  const void* __restrict__ score_bias;
+  const PlanC* __restrict__ plan_c;
+  const PlanW* __restrict__ plan_w;
+  uint32_t num_compress;
+  uint32_t num_write;
+};
+
+struct Compress128SharedBuffer {
+  using Storage = device::AlignedVector<float, kTileElements>;
+  Storage data[kNumWarps][device::kWarpThreads + 1];  // padding to avoid bank conflict
+  SGL_DEVICE Storage& operator()(uint32_t warp_id, uint32_t lane_id) {
+    return data[warp_id][lane_id];
+  }
+  SGL_DEVICE float& operator()(uint32_t warp_id, uint32_t lane_id, uint32_t tile_id) {
+    return data[warp_id][lane_id][tile_id];
+  }
+};
+
+template <int64_t kHeadDim_>
+struct C128Trait {
+  static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
+  static constexpr int64_t kHeadDim = kHeadDim_;
+  static constexpr int64_t kScoreOffset = kHeadDim;
+  static constexpr int64_t kElementSize = kHeadDim * 2;
+  static constexpr int64_t kPageElementSize = 128 * kElementSize;  // page size = 128
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static_assert(kHeadDim % kTileDim == 0);
+};
+
+template <typename Trait, bool kUsePDL, typename BufferFloat, typename InputFloat, typename OutFloat>
+SGL_DEVICE void c128_forward(
+    const BufferFloat* kv_buf,  // [128n, 128n + 127]
+    const InputFloat* kv_src,   // ragged pointer at position = 128n + 127
+    OutFloat* kv_out,
+    const InputFloat* score_bias,
+    const int32_t buffer_len) {
+  using namespace device;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+
+  /// NOTE: part 1: load kv + score
+  using StorageIn = AlignedVector<InputFloat, kTileElements>;
+  const auto gmem_in = tile::Memory<StorageIn>{lane_id, kWarpThreads};
+  StorageIn kv[kElementsPerWarp];
+  StorageIn score[kElementsPerWarp];
+  StorageIn bias[kElementsPerWarp];
+  const int32_t warp_offset = warp_id * kElementsPerWarp;
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; ++i) {
+    const int32_t j = i + warp_offset;
+    bias[i] = gmem_in.load(score_bias + j * Trait::kHeadDim);
+  }
+
+  const auto kv_start = kv_src - 127 * Trait::kElementSize;  // point to start
+
+  if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
+#pragma unroll
+    for (int32_t i = 0; i < kElementsPerWarp; ++i) {
+      const int32_t j = i + warp_offset;
+      __builtin_assume(j < 128);
+      const auto src = j < buffer_len ? kv_buf : kv_start;
+      kv[i] = gmem_in.load(src + j * Trait::kElementSize);
+      score[i] = gmem_in.load(src + j * Trait::kElementSize + Trait::kScoreOffset);
+    }
+  } else {  // mixed dtype
+    using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
+    const auto gmem_buffer = tile::Memory<StorageBuffer>{lane_id, kWarpThreads};
+
+#pragma unroll
+    for (int32_t i = 0; i < kElementsPerWarp; ++i) {
+      const int32_t j = i + warp_offset;
+      __builtin_assume(j < 128);
+      if (j < buffer_len) {
+        const auto src = kv_buf + j * Trait::kElementSize;
+        const auto kv_tmp = gmem_buffer.load(src);
+        const auto score_tmp = gmem_buffer.load(src + Trait::kScoreOffset);
+#pragma unroll
+        for (int32_t k = 0; k < kTileElements; ++k) {
+          kv[i][k] = cast<InputFloat>(kv_tmp[k]);
+          score[i][k] = cast<InputFloat>(score_tmp[k]);
+        }
+      } else {
+        const auto src = kv_start + j * Trait::kElementSize;
+        kv[i] = gmem_in.load(src);
+        score[i] = gmem_in.load(src + Trait::kScoreOffset);
+      }
+    }
+  }
+
+  /// NOTE: part 2: safe online softmax + weighted sum
+  using TmpStorage = typename Compress128SharedBuffer::Storage;
+  __shared__ Compress128SharedBuffer s_local_val_max;
+  __shared__ Compress128SharedBuffer s_local_exp_sum;
+  __shared__ Compress128SharedBuffer s_local_product;
+
+  TmpStorage tmp_val_max;
+  TmpStorage tmp_exp_sum;
+  TmpStorage tmp_product;
+
+  float score_fp32[kTileElements][kElementsPerWarp];
+
+  // convert to fp32 and apply bias first
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+#pragma unroll
+    for (int32_t j = 0; j < kElementsPerWarp; ++j) {
+      score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
+    }
+  }
+
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+    const auto& score = score_fp32[i];
+    float max_value = score[0];
+    float sum_exp_value = 0.0f;
+
+#pragma unroll
+    for (int32_t j = 1; j < kElementsPerWarp; ++j) {
+      const auto fp32_score = score[j];
+      max_value = fmaxf(max_value, fp32_score);
+    }
+
+    float sum_product = 0.0f;
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      const auto fp32_score = score[j];
+      const auto exp_score = expf(fp32_score - max_value);
+      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_exp_value += exp_score;
+    }
+
+    tmp_val_max[i] = max_value;
+    tmp_exp_sum[i] = sum_exp_value;
+    tmp_product[i] = sum_product;
+  }
+
+  // naturally aligned, so no bank conflict
+  s_local_val_max(warp_id, lane_id) = tmp_val_max;
+  s_local_exp_sum(warp_id, lane_id) = tmp_exp_sum;
+  s_local_product(warp_id, lane_id) = tmp_product;
+
+  __syncthreads();
+
+  /// NOTE: part 3: online softmax
+  /// NOTE: We have `kTileElements * kWarpThreads * kNumWarps` values to reduce
+  /// each reduce will consume `kNumWarps` threads (use partial warp reduction)
+  constexpr uint32_t kReductionCount = kTileElements * kWarpThreads * kNumWarps;
+  constexpr uint32_t kIteration = kReductionCount / kBlockSize;
+
+  PDLTriggerSecondary<kUsePDL>();
+
+#pragma unroll
+  for (uint32_t i = 0; i < kIteration; ++i) {
+    /// NOTE: Range `[0, kTileElements * kWarpThreads * kNumWarps)`
+    const uint32_t j = i * kBlockSize + warp_id * kWarpThreads + lane_id;
+    /// NOTE: Range `[0, kNumWarps)`
+    const uint32_t local_warp_id = j % kNumWarps;
+    /// NOTE: Range `[0, kTileElements * kWarpThreads)`
+    const uint32_t local_elem_id = j / kNumWarps;
+    /// NOTE: Range `[0, kTileElements)`
+    const uint32_t local_tile_id = local_elem_id % kTileElements;
+    /// NOTE: Range `[0, kWarpThreads)`
+    const uint32_t local_lane_id = local_elem_id / kTileElements;
+    /// NOTE: each warp will access the whole tile (all `kTileElements`)
+    /// and for different lanes, the memory access only differ in `local_warp_id`
+    /// so there's no bank conflict in shared memory access.
+    static_assert(kTileElements * kNumWarps == kWarpThreads, "TODO: support other configs");
+    const auto local_val_max = s_local_val_max(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_exp_sum = s_local_exp_sum(local_warp_id, local_lane_id, local_tile_id);
+    const auto local_product = s_local_product(local_warp_id, local_lane_id, local_tile_id);
+    const auto global_val_max = warp::reduce_max<kNumWarps>(local_val_max);
+    const auto rescale = expf(local_val_max - global_val_max);
+    const auto global_exp_sum = warp::reduce_sum<kNumWarps>(local_exp_sum * rescale);
+    const auto final_scale = rescale / global_exp_sum;
+    const auto global_product = warp::reduce_sum<kNumWarps>(local_product * final_scale);
+    kv_out[local_elem_id] = cast<OutFloat>(global_product);
+  }
+}
+
+template <typename Trait, typename BufferFloat, typename InputFloat>
+SGL_DEVICE void c128_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) {
+  using namespace device;
+
+  using StorageInput = AlignedVector<InputFloat, kTileElements>;
+  const auto gmem_input = tile::Memory<StorageInput>::warp();
+
+  StorageInput data[2];
+#pragma unroll
+  for (int32_t i = 0; i < 2; ++i) {
+    data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i);
+  }
+
+  if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
+#pragma unroll
+    for (int32_t i = 0; i < 2; ++i) {
+      gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]);
+    }
+  } else {
+    using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
+    const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
+
+    StorageBuffer data_cast[2];
+#pragma unroll
+    for (int32_t i = 0; i < 2; ++i) {
+#pragma unroll
+      for (int32_t j = 0; j < kTileElements; ++j) {
+        data_cast[i][j] = cast<BufferFloat>(data[i][j]);
+      }
+      gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]);
+    }
+  }
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+C128_KERNEL void flash_c128_decode(const __grid_constant__ Compress128DecodeParams params) {
+  using namespace device;
+  using Trait = C128Trait<kHeadDim>;
+
+  const uint32_t warp_id = threadIdx.x / kWarpThreads;
+  const uint32_t global_bid = blockIdx.x / Trait::kNumSplit;  // batch id
+  const uint32_t global_sid = blockIdx.x % Trait::kNumSplit;  // split id
+  const int64_t split_offset = global_sid * Trait::kTileDim;
+  if (global_bid >= params.batch_size) return;
+
+  const auto plan = params.plan_d[global_bid];
+  const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
+  const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
+  const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
+  const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
+
+  const auto kv_src = kv_input + global_bid * Trait::kElementSize;
+  const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
+  const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
+  const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
+
+  PDLWaitPrimary<kUsePDL>();
+  // the write warp must match the load warp in the following `c128_forward`
+  if (warp_id == kNumWarps - 1) {
+    c128_write_decode<Trait, BufferFloat, InputFloat>(kv_dst, kv_src);
+  }
+  if (plan.write_loc % 128 == 127) {
+    c128_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, 128);
+  }
+}
+
+// compress kernel
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+C128_KERNEL void flash_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
+  using namespace device;
+  using Trait = C128Trait<kHeadDim>;
+
+  const uint32_t global_pid = blockIdx.x / Trait::kNumSplit;  // plan id
+  const uint32_t global_sid = blockIdx.x % Trait::kNumSplit;  // split id
+  const int64_t split_offset = global_sid * Trait::kTileDim;
+  if (global_pid >= params.num_compress) return;
+
+  const auto plan = params.plan_c[global_pid];
+  const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
+  const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
+  const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
+  const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
+  if (plan.is_invalid()) return;
+
+  const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
+  // Compact output: one row per compress plan, indexed by `global_pid`.
+  const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
+  const auto kv_buf = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
+  PDLWaitPrimary<kUsePDL>();
+  c128_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(kv_buf, kv_src, kv_out, score_bias, plan.buffer_len);
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+WRITE_KERNEL void write_c128_prefill(const __grid_constant__ Compress128PrefillParams params) {
+  using namespace device;
+  using Trait = C128Trait<kHeadDim>;
+  using StorageInput = AlignedVector<InputFloat, kTileElements>;
+
+  const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t global_wid = global_tid / kWarpThreads;      // warp id
+  const uint32_t global_pid = global_wid / Trait::kNumSplit;  // plan id
+  const uint32_t global_sid = global_wid % Trait::kNumSplit;  // split id
+  // split the contiguous `kHeadDim * 2` into `kNumSplit` tiles
+  // each warp handles 1 contiguous tile (in contrast, decode handle the strided head_dim)
+  const int64_t split_offset = global_sid * (Trait::kTileDim * 2);
+  if (global_pid >= params.num_write) return;
+
+  const auto plan = params.plan_w[global_pid];
+  const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
+  const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
+  if (plan.is_invalid()) return;
+
+  // each warp will handle a contiguous region
+  const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
+  const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
+  const auto gmem_input = tile::Memory<StorageInput>::warp();
+
+  PDLWaitPrimary<kUsePDL>();
+  StorageInput data[2];
+#pragma unroll
+  for (int32_t i = 0; i < 2; ++i) {
+    data[i] = gmem_input.load(kv_src, i);
+  }
+
+  if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
+    PDLTriggerSecondary<kUsePDL>();
+#pragma unroll
+    for (int32_t i = 0; i < 2; ++i) {
+      gmem_input.store(kv_buf, data[i], i);
+    }
+  } else {
+    using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
+    const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
+
+    StorageBuffer data_cast[2];
+#pragma unroll
+    for (int32_t i = 0; i < 2; ++i) {
+#pragma unroll
+      for (int32_t j = 0; j < kTileElements; ++j) {
+        data_cast[i][j] = cast<BufferFloat>(data[i][j]);
+      }
+    }
+    PDLTriggerSecondary<kUsePDL>();
+#pragma unroll
+    for (int32_t i = 0; i < 2; ++i) {
+      gmem_buffer.store(kv_buf, data_cast[i], i);
+    }
+  }
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+struct FlashCompress128Kernel {
+  static constexpr auto decode_kernel = flash_c128_decode<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_c_kernel = flash_c128_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_w_kernel = write_c128_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
+  static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 64
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  using Trait = C128Trait<kHeadDim>;
+
+  static void run_decode(
+      const tvm::ffi::TensorView kv_buffer,
+      const tvm::ffi::TensorView kv_input,
+      const tvm::ffi::TensorView kv_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView plan_d_) {
+    using namespace host;
+
+    auto N = SymbolicSize{"batch_size"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({-1, 128, Trait::kElementSize})  // kv score
+        .with_dtype<BufferFloat>()
+        .with_device(device_)
+        .verify(kv_buffer);
+    TensorMatcher({N, Trait::kElementSize})  // kv score input
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(kv_input);
+    TensorMatcher({N, kHeadDim})  // kv compressed output
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(ape);
+
+    const auto plan_d = compress::verify_plan_d(plan_d_, N, device_);
+    const auto batch_size = static_cast<uint32_t>(N.unwrap());
+    const auto params = Compress128DecodeParams{
+        .kv_buffer = kv_buffer.data_ptr(),
+        .kv_input = kv_input.data_ptr(),
+        .kv_output = kv_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .plan_d = plan_d,
+        .batch_size = batch_size,
+    };
+    const uint32_t num_blocks = batch_size * kNumSplit;
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(decode_kernel, params);
+  }
+
+  static void run_prefill(
+      const tvm::ffi::TensorView kv_buffer,
+      const tvm::ffi::TensorView kv_input,
+      const tvm::ffi::TensorView kv_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView plan_c_,
+      const tvm::ffi::TensorView plan_w_) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_q_tokens"};
+    auto C = SymbolicSize{"num_c_plans"};
+    auto W = SymbolicSize{"num_w_plans"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({-1, 128, Trait::kElementSize})  // kv score
+        .with_dtype<BufferFloat>()
+        .with_device(device_)
+        .verify(kv_buffer);
+    TensorMatcher({N, Trait::kElementSize})  // kv score input (ragged)
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(kv_input);
+    TensorMatcher({C, kHeadDim})  // kv compressed output (compact)
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_output);
+    TensorMatcher({128, kHeadDim})  // ape
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(ape);
+
+    const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
+    const auto plan_w = compress::verify_plan_w(plan_w_, W, device_);
+    const auto device = device_.unwrap();
+    const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_c = static_cast<uint32_t>(C.unwrap());
+    const auto num_w = static_cast<uint32_t>(W.unwrap());
+    const auto params = Compress128PrefillParams{
+        .kv_buffer = kv_buffer.data_ptr(),
+        .kv_input = kv_input.data_ptr(),
+        .kv_output = kv_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .plan_c = plan_c,
+        .plan_w = plan_w,
+        .num_compress = num_c,
+        .num_write = num_w,
+    };
+    RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
+    if (const auto num_c_blocks = num_c * kNumSplit) {
+      constexpr auto kBlockSize_C = kBlockSize;
+      LaunchKernel(num_c_blocks, kBlockSize_C, device)  //
+          .enable_pdl(kUsePDL)(prefill_c_kernel, params);
+    }
+    constexpr uint32_t kWarpsPerWriteBlock = kWriteBlockSize / device::kWarpThreads;
+    if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerWriteBlock)) {
+      constexpr auto kBlockSize_W = kWriteBlockSize;
+      LaunchKernel(num_w_blocks, kBlockSize_W, device)  //
+          .enable_pdl(kUsePDL)(prefill_w_kernel, params);
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/c4.cuh

@@ -0,0 +1,549 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.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/compress.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/object.h>
+
+#include <cstdint>
+
+namespace {
+
+using Plan4 = device::compress::PrefillPlan;
+using IndiceT = int32_t;
+
+/// \brief Each thread will handle this many elements (split along head_dim)
+constexpr int kTileElements = 4;
+
+/// \brief Need to improve register usage to reduce latency
+#define C4_KERNEL __global__ __launch_bounds__(128, 4)
+
+enum class PageMode {
+  RingBuffer = 8,
+  Page4Align = 4,
+};
+
+struct alignas(16) C4IndexBundle {
+  int32_t load_first_page;
+  int32_t load_second_page;
+  int32_t write_first_page;
+  int32_t last_position;
+};
+
+struct Compress4DecodeParams {
+  /**
+   * \brief Shape: `[num_indices, 8, head_dim * 4]` \n
+   * last dimension layout:
+   * | kv overlap | kv | score overlap | score |
+   */
+  void* __restrict__ kv_score_buffer;
+  /** \brief Shape: `[batch_size, head_dim * 4]` */
+  const void* __restrict__ kv_score_input;
+  /** \brief Shape: `[batch_size, head_dim]` */
+  void* __restrict__ kv_compressed_output;
+  /** \brief Shape: `[8, head_dim]` (called `ape`) */
+  const void* __restrict__ score_bias;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ indices;
+  /** \brief Shape: `[batch_size, ]` */
+  const IndiceT* __restrict__ seq_lens;
+  /** \brief Shape: `[batch_size, 1]` */
+  const int32_t* __restrict__ extra;
+  /** \NOTE: `batch_size` <= `num_indices` */
+  uint32_t batch_size;
+};
+
+struct Compress4PrefillParams {
+  /**
+   * \brief Shape: `[num_indices, 8, head_dim * 4]` \n
+   * last dimension layout:
+   * | kv overlap | kv | score overlap | score |
+   */
+  void* __restrict__ kv_score_buffer;
+  /** \brief Shape: `[num_q_tokens, head_dim * 4]` */
+  const void* __restrict__ kv_score_input;
+  /** \brief Shape: `[num_q_tokens, head_dim]` */
+  void* __restrict__ kv_compressed_output;
+  /** \brief Shape: `[8, head_dim]` (called `ape`) */
+  const void* __restrict__ score_bias;
+  /** \brief Shape: `[batch_size, ]`*/
+  const IndiceT* __restrict__ indices;
+  /** \brief Shape: `[batch_size, 4]` */
+  const C4IndexBundle* __restrict__ extra;
+  /** \brief The following part is plan info. */
+
+  const Plan4* __restrict__ compress_plan;
+  const Plan4* __restrict__ write_plan;
+  uint32_t num_compress;
+  uint32_t num_write;
+};
+
+template <typename T>
+SGL_DEVICE void c4_write(
+    T* kv_score_buf,  //
+    const T* kv_score_src,
+    const int64_t head_dim,
+    const int32_t write_pos) {
+  using namespace device;
+
+  using Storage = AlignedVector<T, kTileElements>;
+  const auto element_size = head_dim * 4;
+  const auto gmem = tile::Memory<Storage>::warp();
+  kv_score_buf += write_pos * element_size;
+
+  /// NOTE: Layout | [0] = kv overlap | [1] = kv | [2] = score overlap | [3] = score |
+  Storage kv_score[4];
+#pragma unroll
+  for (int32_t i = 0; i < 4; ++i) {
+    kv_score[i] = gmem.load(kv_score_src + head_dim * i);
+  }
+#pragma unroll
+  for (int32_t i = 0; i < 4; ++i) {
+    gmem.store(kv_score_buf + head_dim * i, kv_score[i]);
+  }
+}
+
+template <bool kPaged, typename InFloat, typename OutFloat>
+SGL_DEVICE void c4_forward(
+    const InFloat* kv_score_buf,
+    const InFloat* kv_score_src,
+    OutFloat* kv_out,
+    const InFloat* score_bias,
+    const int64_t head_dim,
+    const int32_t seq_len,
+    const int32_t window_len,
+    [[maybe_unused]] const InFloat* kv_score_overlap_buf = nullptr) {
+  using namespace device;
+
+  const auto element_size = head_dim * 4;
+  const auto score_offset = head_dim * 2;
+  const auto overlap_stride = head_dim;
+
+  /// NOTE: part 1: load kv + score
+  using StorageIn = AlignedVector<InFloat, kTileElements>;
+  const auto gmem_in = tile::Memory<StorageIn>::warp();
+  StorageIn kv[8];
+  StorageIn score[8];
+  StorageIn bias[8];
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; ++i) {
+    bias[i] = gmem_in.load(score_bias + i * head_dim);
+  }
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; ++i) {
+    const bool is_overlap = i < 4;
+    const InFloat* src;
+    if (i < window_len) {
+      /// NOTE: `seq_len` must be a multiple of 4 here
+      if constexpr (kPaged) {
+        const auto kv_score_ptr = is_overlap ? kv_score_overlap_buf : kv_score_buf;
+        const int32_t k = i % 4;
+        src = kv_score_ptr + k * element_size;
+      } else {
+        const int32_t k = (seq_len + i) % 8;
+        src = kv_score_buf + k * element_size;
+      }
+    } else {
+      /// NOTE: k in [-7, 0]. We'll load from the ragged `kv_score_src`
+      const int32_t k = i - 7;
+      src = kv_score_src + k * element_size;
+    }
+    src += (is_overlap ? 0 : overlap_stride);
+    kv[i] = gmem_in.load(src);
+    score[i] = gmem_in.load(src + score_offset);
+  }
+
+  if (seq_len == 4) {
+    [[unlikely]];
+    constexpr float kFloatNegInf = -1e9f;
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      kv[i].fill(cast<InFloat>(0.0f));
+      score[i].fill(cast<InFloat>(kFloatNegInf));
+    }
+  }
+
+  /// NOTE: part 2: safe online softmax + weighted sum
+  using StorageOut = AlignedVector<OutFloat, kTileElements>;
+  const auto gmem_out = tile::Memory<StorageOut>::warp();
+  StorageOut result;
+
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+    float score_fp32[8];
+
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      score_fp32[j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
+    }
+
+    float max_value = score_fp32[0];
+    float sum_exp_value = 0.0f;
+
+#pragma unroll
+    for (int32_t j = 1; j < 8; ++j) {
+      const auto fp32_score = score_fp32[j];
+      max_value = fmaxf(max_value, fp32_score);
+    }
+
+    float sum_product = 0.0f;
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      const auto fp32_score = score_fp32[j];
+      const auto exp_score = expf(fp32_score - max_value);
+      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_exp_value += exp_score;
+    }
+
+    result[i] = cast<OutFloat>(sum_product / sum_exp_value);
+  }
+
+  gmem_out.store(kv_out, result);
+}
+
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, PageMode kMode, bool kUsePDL>
+C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
+  using namespace device;
+
+  constexpr int64_t kTileDim = kTileElements * kWarpThreads;  // 128
+  constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  constexpr int64_t kElementSize = kHeadDim * 4;  // `* 4` due to overlap transform + score
+  static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
+
+  const auto& [
+    _kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
+    indices, seq_lens, extra, batch_size // decode info
+  ] = params;
+  const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t global_wid = global_tid / kWarpThreads;  // warp id
+  const uint32_t global_bid = global_wid / kNumSplit;     // batch id
+  const uint32_t global_sid = global_wid % kNumSplit;     // split id
+
+  if (global_bid >= batch_size) return;
+
+  const int32_t index = indices[global_bid];
+  const int32_t seq_len = seq_lens[global_bid];
+  const int64_t split_offset = global_sid * kTileDim;
+
+  // kv score
+  const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
+
+  // kv input
+  const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
+  const auto kv_src = kv_score_input + global_bid * kElementSize + split_offset;
+
+  // kv output
+  const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
+  const auto kv_out = kv_compressed_output + global_bid * kHeadDim + split_offset;
+
+  // score bias (ape)
+  const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  /// NOTE: `position` = `seq_len - 1`. To avoid underflow, we use `seq_len + page_size - 1`
+  if constexpr (kMode == PageMode::Page4Align) {
+    const auto index_prev = extra[global_bid];
+    const auto kv_buf = kv_score_buffer + index * (kElementSize * 4) + split_offset;
+    c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 3) % 4);
+    if (seq_len % 4 == 0) {
+      const auto kv_overlap = kv_buf + (index_prev - index) * (kElementSize * 4);
+      c4_forward<true>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, 8, kv_overlap);
+    }
+  } else {
+    static_assert(kMode == PageMode::RingBuffer, "Unsupported PageMode");
+    const auto kv_buf = kv_score_buffer + index * (kElementSize * 8) + split_offset;
+    c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/(seq_len + 7) % 8);
+    if (seq_len % 4 == 0) {
+      c4_forward<false>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, /*window_size=*/8);
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, PageMode kMode, bool kWrite, bool kUsePDL>
+C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
+  using namespace device;
+
+  constexpr int64_t kTileDim = kTileElements * kWarpThreads;  // 128
+  constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  constexpr int64_t kElementSize = kHeadDim * 4;  // `* 4` due to overlap transform + score
+  static_assert(kHeadDim % kTileDim == 0, "Head dim must be multiple of tile dim");
+
+  const auto& [
+    _kv_score_buffer, _kv_score_input, _kv_compressed_output, _score_bias, // kv score
+    indices, extra, compress_plan, write_plan, num_compress, num_write // prefill plan
+  ] = params;
+
+  const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t global_wid = global_tid / kWarpThreads;  // warp id
+  const uint32_t global_pid = global_wid / kNumSplit;     // plan id
+  const uint32_t global_sid = global_wid % kNumSplit;     // split id
+
+  /// NOTE: compiler can optimize this if-else at compile time
+  const auto num_plans = kWrite ? num_write : num_compress;
+  const auto plan_ptr = kWrite ? write_plan : compress_plan;
+  if (global_pid >= num_plans) return;
+
+  const auto& [ragged_id, global_bid, position, window_len] = plan_ptr[global_pid];
+  const int64_t split_offset = global_sid * kTileDim;
+
+  // kv score
+  const auto kv_score_buffer = static_cast<InFloat*>(_kv_score_buffer);
+
+  // kv input
+  const auto kv_score_input = static_cast<const InFloat*>(_kv_score_input);
+  const auto kv_src = kv_score_input + ragged_id * kElementSize + split_offset;
+
+  // kv output
+  const auto kv_compressed_output = static_cast<OutFloat*>(_kv_compressed_output);
+  const auto kv_out = kv_compressed_output + ragged_id * kHeadDim + split_offset;
+
+  if (ragged_id == 0xFFFFFFFF) [[unlikely]]
+    return;
+
+  // score bias (ape)
+  const auto score_bias = static_cast<const InFloat*>(_score_bias) + split_offset;
+  const auto seq_len = position + 1;
+  const int32_t index = indices[global_bid];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  if constexpr (kMode == PageMode::Page4Align) {
+    const auto write_second_page = index;
+    const auto [load_first_page, load_second_page, write_first_page, last_pos] = extra[global_bid];
+    if constexpr (kWrite) {
+      int32_t index;
+      if (position < static_cast<uint32_t>(last_pos)) {
+        index = write_first_page;
+      } else {
+        index = write_second_page;
+      }
+      const auto kv_buf = kv_score_buffer + index * (kElementSize * 4) + split_offset;
+      c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 4);
+    } else {
+      int32_t index_overlap, index_normal;
+      if (window_len <= 4) {
+        index_overlap = load_second_page;
+        index_normal = load_second_page;  // not used
+      } else {
+        index_overlap = load_first_page;
+        index_normal = load_second_page;
+      }
+      const auto kv_buf = kv_score_buffer + index_normal * (kElementSize * 4) + split_offset;
+      const auto kv_overlap = kv_score_buffer + index_overlap * (kElementSize * 4) + split_offset;
+      c4_forward<true>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, window_len, kv_overlap);
+    }
+  } else {
+    static_assert(kMode == PageMode::RingBuffer, "Unsupported PageMode");
+    const auto kv_buf = kv_score_buffer + index * (kElementSize * 8) + split_offset;
+    if constexpr (kWrite) {
+      c4_write(kv_buf, kv_src, kHeadDim, /*write_pos=*/position % 8);
+    } else {
+      c4_forward<false>(kv_buf, kv_src, kv_out, score_bias, kHeadDim, seq_len, window_len);
+    }
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <int64_t kHeadDim, typename InFloat, typename OutFloat, bool kUsePDL>
+struct FlashCompress4Kernel {
+  template <PageMode kMode>
+  static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, InFloat, OutFloat, kMode, kUsePDL>;
+  template <PageMode kMode, bool kWrite>
+  static constexpr auto prefill_kernel = flash_c4_prefill<kHeadDim, InFloat, OutFloat, kMode, kWrite, kUsePDL>;
+  template <PageMode kMode>
+  static constexpr auto prefill_c_kernel = prefill_kernel<kMode, /*kWrite=*/false>;
+  template <PageMode kMode>
+  static constexpr auto prefill_w_kernel = prefill_kernel<kMode, /*kWrite=*/true>;
+  static constexpr uint32_t kBlockSize = 128;
+  static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static constexpr uint32_t kWarpsPerBlock = kBlockSize / device::kWarpThreads;
+
+  using Self = FlashCompress4Kernel;
+
+  static void run_decode(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
+    using namespace host;
+
+    // this should not happen in practice
+    auto B = SymbolicSize{"batch_size"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    const auto extra_ptr = _get_extra_pointer(B, device_, extra);
+    const auto page_size = extra_ptr != nullptr ? 4 : 8;
+
+    TensorMatcher({-1, page_size, kHeadDim * 4})  // kv score
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(kv_score_buffer);
+    TensorMatcher({B, kHeadDim * 4})  // kv score input
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(kv_score_input);
+    TensorMatcher({B, kHeadDim})  // kv compressed output
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_compressed_output);
+    TensorMatcher({8, kHeadDim})  // ape
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(ape);
+    TensorMatcher({B})  // indices
+        .with_dtype<IndiceT>()
+        .with_device(device_)
+        .verify(indices);
+    TensorMatcher({B})  // seq lens
+        .with_dtype<IndiceT>()
+        .with_device(device_)
+        .verify(seq_lens);
+
+    const auto device = device_.unwrap();
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto params = Compress4DecodeParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .indices = static_cast<const IndiceT*>(indices.data_ptr()),
+        .seq_lens = static_cast<const IndiceT*>(seq_lens.data_ptr()),
+        .extra = static_cast<const int32_t*>(extra_ptr),
+        .batch_size = batch_size,
+    };
+    const auto kernel = extra_ptr != nullptr ? decode_kernel<PageMode::Page4Align>  //
+                                             : decode_kernel<PageMode::RingBuffer>;
+    const uint32_t num_blocks = div_ceil(batch_size * kNumSplit, kWarpsPerBlock);
+    LaunchKernel(num_blocks, kBlockSize, device)  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+
+  static void run_prefill(
+      const tvm::ffi::TensorView kv_score_buffer,
+      const tvm::ffi::TensorView kv_score_input,
+      const tvm::ffi::TensorView kv_compressed_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView indices,
+      const tvm::ffi::TensorView compress_plan,
+      const tvm::ffi::TensorView write_plan,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> extra) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto N = SymbolicSize{"num_q_tokens"};
+    auto X = SymbolicSize{"compress_tokens"};
+    auto Y = SymbolicSize{"write_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    const auto extra_ptr = _get_extra_pointer(B, device_, extra, /*is_prefill=*/true);
+    const auto page_size = extra_ptr != nullptr ? 4 : 8;
+
+    TensorMatcher({-1, page_size, kHeadDim * 4})  // kv score
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(kv_score_buffer);
+    TensorMatcher({N, kHeadDim * 4})  // kv score input
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(kv_score_input);
+    TensorMatcher({N, kHeadDim})  // kv compressed output
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_compressed_output);
+    TensorMatcher({8, kHeadDim})  // ape
+        .with_dtype<InFloat>()
+        .with_device(device_)
+        .verify(ape);
+    TensorMatcher({B})  // indices
+        .with_dtype<IndiceT>()
+        .with_device(device_)
+        .verify(indices);
+    TensorMatcher({X, compress::kPrefillPlanDim})  // compress plan
+        .with_dtype<compress::PrefillPlanTensorDtype>()
+        .with_device(device_)
+        .verify(compress_plan);
+    TensorMatcher({Y, compress::kPrefillPlanDim})  // write plan
+        .with_dtype<compress::PrefillPlanTensorDtype>()
+        .with_device(device_)
+        .verify(write_plan);
+
+    const auto device = device_.unwrap();
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_c = static_cast<uint32_t>(X.unwrap());
+    const auto num_w = static_cast<uint32_t>(Y.unwrap());
+    const auto params = Compress4PrefillParams{
+        .kv_score_buffer = kv_score_buffer.data_ptr(),
+        .kv_score_input = kv_score_input.data_ptr(),
+        .kv_compressed_output = kv_compressed_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .indices = static_cast<const IndiceT*>(indices.data_ptr()),
+        .extra = static_cast<const C4IndexBundle*>(extra_ptr),
+        .compress_plan = static_cast<const Plan4*>(compress_plan.data_ptr()),
+        .write_plan = static_cast<const Plan4*>(write_plan.data_ptr()),
+        .num_compress = num_c,
+        .num_write = num_w,
+    };
+    RuntimeCheck(num_q_tokens >= batch_size, "num_q_tokens must be >= batch_size");
+    RuntimeCheck(num_q_tokens >= std::max(num_c, num_w), "invalid prefill plan");
+    if (const auto num_c_blocks = div_ceil(num_c * kNumSplit, kWarpsPerBlock)) {
+      const auto c_kernel = extra_ptr != nullptr ? prefill_c_kernel<PageMode::Page4Align>  //
+                                                 : prefill_c_kernel<PageMode::RingBuffer>;
+      LaunchKernel(num_c_blocks, kBlockSize, device)  //
+          .enable_pdl(kUsePDL)(c_kernel, params);
+    }
+    if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerBlock)) {
+      const auto w_kernel = extra_ptr != nullptr ? prefill_w_kernel<PageMode::Page4Align>  //
+                                                 : prefill_w_kernel<PageMode::RingBuffer>;
+      LaunchKernel(num_w_blocks, kBlockSize, device)  //
+          .enable_pdl(kUsePDL)(w_kernel, params);
+    }
+  }
+
+  // some auxiliary functions
+ private:
+  static const void* _get_extra_pointer(
+      host::SymbolicSize& B,  // batch_size
+      host::SymbolicDevice& device,
+      const tvm::ffi::Optional<tvm::ffi::TensorView>& extra,
+      bool is_prefill = false) {
+    // only have value when using page-aligned mode
+    if (!extra.has_value()) return nullptr;
+    const auto& extra_tensor = extra.value();
+    /// NOTE: the metadata layout is different for prefill and decode:
+    /// for prefill, last 4 are:
+    /// load overlap | load normal | write overlap | last written page
+    /// for decode, last 1 is the write (also load) overlap
+    host::TensorMatcher({B, is_prefill ? 4 : 1})  // extra tensor
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(extra_tensor);
+    const auto data_ptr = extra_tensor.data_ptr();
+    host::RuntimeCheck(data_ptr != nullptr, "extra tensor data ptr is null");
+    if (is_prefill) {
+      static_assert(alignof(C4IndexBundle) == 16);
+      host::RuntimeCheck(std::bit_cast<uintptr_t>(data_ptr) % 16 == 0, "extra tensor is not properly aligned");
+    }
+    return data_ptr;
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/c4_v2.cuh

@@ -0,0 +1,491 @@
+/**
+ * \brief Here's some dimension info for the main buffer used in C4 prefill and decode.
+ *
+ * kv_buffer: [num_indices, 8, head_dim * 4]
+ * - last dimension layout: | kv overlap | kv | score overlap | score |
+ * kv_input: [batch_size, head_dim * 4]
+ * kv_output: [batch_size, head_dim]
+ * score_bias (ape): [8, head_dim]
+ * plan_c/plan_w: [variable length]
+ *
+ * For prefill, batch_size = num_q_tokens
+ */
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.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/compress_v2.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/object.h>
+
+#include <cfloat>
+#include <cstdint>
+#include <type_traits>
+
+namespace {
+
+using PlanD = device::compress::DecodePlan;
+using PlanC = device::compress::CompressPlan;
+using PlanW = device::compress::WritePlan;
+
+/// \brief Each thread will handle this many elements (split along head_dim)
+constexpr int32_t kTileElements = 4;
+
+/// \brief Need to improve register usage to reduce latency
+#define C4_KERNEL __global__ __launch_bounds__(128, 4)
+#define WRITE_KERNEL __global__ __launch_bounds__(128, 16)
+
+struct Compress4DecodeParams {
+  void* __restrict__ kv_buffer;
+  const void* __restrict__ kv_input;
+  void* __restrict__ kv_output;
+  const void* __restrict__ score_bias;
+  const PlanD* __restrict__ plan_d;
+  uint32_t batch_size;
+};
+
+struct Compress4PrefillParams {
+  void* __restrict__ kv_buffer;
+  const void* __restrict__ kv_input;
+  void* __restrict__ kv_output;
+  const void* __restrict__ score_bias;
+  const PlanC* __restrict__ plan_c;
+  const PlanW* __restrict__ plan_w;
+  uint32_t num_compress;
+  uint32_t num_write;
+};
+
+template <int64_t kHeadDim_>
+struct C4Trait {
+  static constexpr int64_t kTileDim = kTileElements * device::kWarpThreads;  // 128
+  static constexpr int64_t kHeadDim = kHeadDim_;
+  static constexpr int64_t kOverlapOffset = kHeadDim;
+  static constexpr int64_t kScoreOffset = kHeadDim * 2;
+  static constexpr int64_t kElementSize = kHeadDim * 4;
+  static constexpr int64_t kPageElementSize = 4 * kElementSize;  // page size = 4
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static_assert(kHeadDim % kTileDim == 0);
+};
+
+template <typename Trait, bool kUsePDL, typename BufferFloat, typename InputFloat, typename OutFloat>
+SGL_DEVICE void c4_forward(
+    const BufferFloat* kv_buf_0,  // overlap [4n - 4, 4n - 1]
+    const BufferFloat* kv_buf_1,  // normal [4n + 0, 4n + 3]
+    const InputFloat* kv_src,     // ragged pointer at position = 4n + 3
+    OutFloat* kv_out,
+    const InputFloat* score_bias,
+    const bool should_overlap,
+    const int32_t buffer_len) {
+  using namespace device;
+
+  using StorageIn = AlignedVector<InputFloat, kTileElements>;
+  const auto gmem_in = tile::Memory<StorageIn>::warp();
+  StorageIn kv[8];
+  StorageIn score[8];
+  StorageIn bias[8];
+
+#pragma unroll
+  for (int32_t i = 0; i < 8; ++i) {
+    bias[i] = gmem_in.load(score_bias + i * Trait::kHeadDim);
+  }
+
+  if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
+    if (should_overlap) {
+      const auto kv_start = kv_src - 7 * Trait::kElementSize;  // point to start
+#pragma unroll
+      for (int32_t i = 0; i < 4; ++i) {
+        const auto src = i < buffer_len ? kv_buf_0 : kv_start;
+        const auto base = src + i * Trait::kElementSize;
+        kv[i] = gmem_in.load(base);
+        score[i] = gmem_in.load(base + Trait::kScoreOffset);
+      }
+    } else {
+      [[unlikely]];
+      constexpr float kFloatNegInf = -FLT_MAX;
+#pragma unroll
+      for (int32_t i = 0; i < 4; ++i) {
+        kv[i].fill(cast<InputFloat>(0.0f));
+        score[i].fill(cast<InputFloat>(kFloatNegInf));
+      }
+    }
+
+    const auto kv_start = kv_src - 3 * Trait::kElementSize;  // point to start
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      const auto src = i + 4 < buffer_len ? kv_buf_1 : kv_start;
+      const auto base = src + i * Trait::kElementSize + Trait::kOverlapOffset;
+      kv[i + 4] = gmem_in.load(base);
+      score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
+    }
+  } else {  // mixed dtype
+    using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
+    const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
+    const auto kv_start_0 = kv_src - 7 * Trait::kElementSize;  // point to start
+
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      if (should_overlap && i < buffer_len) {
+        const auto base = kv_buf_0 + i * Trait::kElementSize;
+        const auto kv_tmp = gmem_buffer.load(base);
+        const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset);
+#pragma unroll
+        for (int32_t j = 0; j < kTileElements; ++j) {
+          kv[i][j] = cast<InputFloat>(kv_tmp[j]);
+          score[i][j] = cast<InputFloat>(score_tmp[j]);
+        }
+      } else if (should_overlap) {
+        const auto base = kv_start_0 + i * Trait::kElementSize;
+        kv[i] = gmem_in.load(base);
+        score[i] = gmem_in.load(base + Trait::kScoreOffset);
+      } else {
+        [[unlikely]];
+        constexpr float kFloatNegInf = -FLT_MAX;
+        kv[i].fill(cast<InputFloat>(0.0f));
+        score[i].fill(cast<InputFloat>(kFloatNegInf));
+      }
+    }
+
+    const auto kv_start = kv_src - 3 * Trait::kElementSize;  // point to start
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      if (i + 4 < buffer_len) {
+        const auto base = kv_buf_1 + i * Trait::kElementSize + Trait::kOverlapOffset;
+        const auto kv_tmp = gmem_buffer.load(base);
+        const auto score_tmp = gmem_buffer.load(base + Trait::kScoreOffset);
+#pragma unroll
+        for (int32_t j = 0; j < kTileElements; ++j) {
+          kv[i + 4][j] = cast<InputFloat>(kv_tmp[j]);
+          score[i + 4][j] = cast<InputFloat>(score_tmp[j]);
+        }
+      } else {
+        const auto base = kv_start + i * Trait::kElementSize + Trait::kOverlapOffset;
+        kv[i + 4] = gmem_in.load(base);
+        score[i + 4] = gmem_in.load(base + Trait::kScoreOffset);
+      }
+    }
+  }
+
+  /// NOTE: part 2: safe online softmax + weighted sum
+  using StorageOut = AlignedVector<OutFloat, kTileElements>;
+  const auto gmem_out = tile::Memory<StorageOut>::warp();
+  StorageOut result;
+
+  // consume 32 fp registers
+  float score_fp32[kTileElements][8];
+
+  // convert to fp32 and apply bias first
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      score_fp32[i][j] = cast<float>(score[j][i]) + cast<float>(bias[j][i]);
+    }
+  }
+
+#pragma unroll
+  for (int32_t i = 0; i < kTileElements; ++i) {
+    const auto& score = score_fp32[i];
+    float max_value = score[0];
+    float sum_exp_value = 0.0f;
+
+#pragma unroll
+    for (int32_t j = 1; j < 8; ++j) {
+      const auto fp32_score = score[j];
+      max_value = fmaxf(max_value, fp32_score);
+    }
+
+    float sum_product = 0.0f;
+#pragma unroll
+    for (int32_t j = 0; j < 8; ++j) {
+      const auto fp32_score = score[j];
+      const auto exp_score = expf(fp32_score - max_value);
+      sum_product += cast<float>(kv[j][i]) * exp_score;
+      sum_exp_value += exp_score;
+    }
+
+    result[i] = cast<OutFloat>(sum_product / sum_exp_value);
+  }
+
+  // overlap the store with the next iteration's load
+  PDLTriggerSecondary<kUsePDL>();
+  gmem_out.store(kv_out, result);
+}
+
+template <typename Trait, typename BufferFloat, typename InputFloat>
+SGL_DEVICE void c4_write_decode(BufferFloat* kv_buf, const InputFloat* kv_src) {
+  using namespace device;
+
+  using StorageInput = AlignedVector<InputFloat, kTileElements>;
+  const auto gmem_input = tile::Memory<StorageInput>::warp();
+
+  StorageInput data[4];
+#pragma unroll
+  for (int32_t i = 0; i < 4; ++i) {
+    data[i] = gmem_input.load(kv_src + Trait::kHeadDim * i);
+  }
+
+  if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      gmem_input.store(kv_buf + Trait::kHeadDim * i, data[i]);
+    }
+  } else {
+    using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
+    const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
+
+    StorageBuffer data_cast[4];
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+#pragma unroll
+      for (int32_t j = 0; j < kTileElements; ++j) {
+        data_cast[i][j] = cast<BufferFloat>(data[i][j]);
+      }
+      gmem_buffer.store(kv_buf + Trait::kHeadDim * i, data_cast[i]);
+    }
+  }
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+C4_KERNEL void flash_c4_decode(const __grid_constant__ Compress4DecodeParams params) {
+  using namespace device;
+  using Trait = C4Trait<kHeadDim>;
+
+  const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t global_wid = global_tid / kWarpThreads;      // warp id
+  const uint32_t global_bid = global_wid / Trait::kNumSplit;  // batch id
+  const uint32_t global_sid = global_wid % Trait::kNumSplit;  // split id
+  const int64_t split_offset = global_sid * Trait::kTileDim;
+  if (global_bid >= params.batch_size) return;
+
+  const auto plan = params.plan_d[global_bid];
+  const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
+  const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
+  const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
+  const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
+
+  const auto kv_src = kv_input + global_bid * Trait::kElementSize;
+  const auto kv_out = kv_output + global_bid * Trait::kHeadDim;
+  const auto kv_buf_0 = kv_buffer + plan.read_page_0 * Trait::kPageElementSize;
+  const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
+  const auto kv_dst = kv_buffer + plan.write_loc * Trait::kElementSize;
+
+  PDLWaitPrimary<kUsePDL>();
+  c4_write_decode<Trait, BufferFloat, InputFloat>(kv_dst, kv_src);
+  if (plan.seq_len % 4 == 0) {
+    const auto need_overlap = plan.seq_len > 4;
+    c4_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(
+        kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, 8);
+  }
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+C4_KERNEL void flash_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
+  using namespace device;
+  using Trait = C4Trait<kHeadDim>;
+
+  const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t global_wid = global_tid / kWarpThreads;      // warp id
+  const uint32_t global_pid = global_wid / Trait::kNumSplit;  // plan id
+  const uint32_t global_sid = global_wid % Trait::kNumSplit;  // split id
+  const int64_t split_offset = global_sid * Trait::kTileDim;
+  if (global_pid >= params.num_compress) return;
+
+  const auto plan = params.plan_c[global_pid];
+  const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
+  const auto kv_output = static_cast<OutFloat*>(params.kv_output) + split_offset;
+  const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
+  const auto score_bias = static_cast<const InputFloat*>(params.score_bias) + split_offset;
+  if (plan.is_invalid()) return;
+
+  const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
+  // Compact output: one row per compress plan, indexed by `global_pid`.
+  const auto kv_out = kv_output + global_pid * Trait::kHeadDim;
+  const auto kv_buf_0 = kv_buffer + plan.read_page_0 * Trait::kPageElementSize;
+  const auto kv_buf_1 = kv_buffer + plan.read_page_1 * Trait::kPageElementSize;
+  const bool need_overlap = plan.seq_len > 4;
+  PDLWaitPrimary<kUsePDL>();
+  c4_forward<Trait, kUsePDL, BufferFloat, InputFloat, OutFloat>(
+      kv_buf_0, kv_buf_1, kv_src, kv_out, score_bias, need_overlap, plan.buffer_len);
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+WRITE_KERNEL void write_c4_prefill(const __grid_constant__ Compress4PrefillParams params) {
+  using namespace device;
+  using Trait = C4Trait<kHeadDim>;
+  using StorageInput = AlignedVector<InputFloat, kTileElements>;
+
+  const uint32_t global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t global_wid = global_tid / kWarpThreads;      // warp id
+  const uint32_t global_pid = global_wid / Trait::kNumSplit;  // plan id
+  const uint32_t global_sid = global_wid % Trait::kNumSplit;  // split id
+  // split the contiguous `kHeadDim * 4` into `kNumSplit` tiles
+  // each warp handles 1 contiguous tile (in contrast, decode handle the strided head_dim)
+  const int64_t split_offset = global_sid * (Trait::kTileDim * 4);
+  if (global_pid >= params.num_write) return;
+
+  const auto plan = params.plan_w[global_pid];
+  const auto kv_input = static_cast<const InputFloat*>(params.kv_input) + split_offset;
+  const auto kv_buffer = static_cast<BufferFloat*>(params.kv_buffer) + split_offset;
+  if (plan.is_invalid()) return;
+
+  // each warp will handle a contiguous region
+  const auto kv_src = kv_input + plan.ragged_id * Trait::kElementSize;
+  const auto kv_buf = kv_buffer + plan.write_loc * Trait::kElementSize;
+  const auto gmem_input = tile::Memory<StorageInput>::warp();
+
+  PDLWaitPrimary<kUsePDL>();
+  StorageInput data[4];
+#pragma unroll
+  for (int32_t i = 0; i < 4; ++i) {
+    data[i] = gmem_input.load(kv_src, i);
+  }
+
+  if constexpr (std::is_same_v<BufferFloat, InputFloat>) {
+    PDLTriggerSecondary<kUsePDL>();
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      gmem_input.store(kv_buf, data[i], i);
+    }
+  } else {
+    using StorageBuffer = AlignedVector<BufferFloat, kTileElements>;
+    const auto gmem_buffer = tile::Memory<StorageBuffer>::warp();
+
+    StorageBuffer data_cast[4];
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+#pragma unroll
+      for (int32_t j = 0; j < kTileElements; ++j) {
+        data_cast[i][j] = cast<BufferFloat>(data[i][j]);
+      }
+    }
+    PDLTriggerSecondary<kUsePDL>();
+#pragma unroll
+    for (int32_t i = 0; i < 4; ++i) {
+      gmem_buffer.store(kv_buf, data_cast[i], i);
+    }
+  }
+}
+
+template <int64_t kHeadDim, typename BufferFloat, typename InputFloat, typename OutFloat, bool kUsePDL>
+struct FlashCompress4Kernel {
+  static constexpr auto decode_kernel = flash_c4_decode<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_c_kernel = flash_c4_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
+  static constexpr auto prefill_w_kernel = write_c4_prefill<kHeadDim, BufferFloat, InputFloat, OutFloat, kUsePDL>;
+  static constexpr uint32_t kBlockSize = 128;
+  static constexpr uint32_t kTileDim = kTileElements * device::kWarpThreads;
+  static constexpr uint32_t kNumSplit = kHeadDim / kTileDim;
+  static constexpr uint32_t kWarpsPerBlock = kBlockSize / device::kWarpThreads;
+  using Trait = C4Trait<kHeadDim>;
+
+  static void run_decode(
+      const tvm::ffi::TensorView kv_buffer,
+      const tvm::ffi::TensorView kv_input,
+      const tvm::ffi::TensorView kv_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView plan_d_) {
+    using namespace host;
+
+    auto N = SymbolicSize{"batch_size"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({-1, 4, Trait::kElementSize})  // kv score
+        .with_dtype<BufferFloat>()
+        .with_device(device_)
+        .verify(kv_buffer);
+    TensorMatcher({N, Trait::kElementSize})  // kv score input
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(kv_input);
+    TensorMatcher({N, kHeadDim})  // kv compressed output
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_output);
+    TensorMatcher({8, kHeadDim})  // ape
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(ape);
+
+    const auto plan_d = compress::verify_plan_d(plan_d_, N, device_);
+    const auto batch_size = static_cast<uint32_t>(N.unwrap());
+    const auto params = Compress4DecodeParams{
+        .kv_buffer = kv_buffer.data_ptr(),
+        .kv_input = kv_input.data_ptr(),
+        .kv_output = kv_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .plan_d = plan_d,
+        .batch_size = batch_size,
+    };
+    const uint32_t num_blocks = div_ceil(batch_size * kNumSplit, kWarpsPerBlock);
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(decode_kernel, params);
+  }
+
+  static void run_prefill(
+      const tvm::ffi::TensorView kv_buffer,
+      const tvm::ffi::TensorView kv_input,
+      const tvm::ffi::TensorView kv_output,
+      const tvm::ffi::TensorView ape,
+      const tvm::ffi::TensorView plan_c_,
+      const tvm::ffi::TensorView plan_w_) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_q_tokens"};
+    auto C = SymbolicSize{"num_c_plans"};
+    auto W = SymbolicSize{"num_w_plans"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({-1, 4, Trait::kElementSize})  // kv score
+        .with_dtype<BufferFloat>()
+        .with_device(device_)
+        .verify(kv_buffer);
+    TensorMatcher({N, Trait::kElementSize})  // kv score input (ragged)
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(kv_input);
+    TensorMatcher({C, kHeadDim})  // kv compressed output (compact)
+        .with_dtype<OutFloat>()
+        .with_device(device_)
+        .verify(kv_output);
+    TensorMatcher({8, kHeadDim})  // ape
+        .with_dtype<InputFloat>()
+        .with_device(device_)
+        .verify(ape);
+    const auto plan_c = compress::verify_plan_c(plan_c_, C, device_);
+    const auto plan_w = compress::verify_plan_w(plan_w_, W, device_);
+    const auto device = device_.unwrap();
+    const auto num_q_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto num_c = static_cast<uint32_t>(C.unwrap());
+    const auto num_w = static_cast<uint32_t>(W.unwrap());
+    const auto params = Compress4PrefillParams{
+        .kv_buffer = kv_buffer.data_ptr(),
+        .kv_input = kv_input.data_ptr(),
+        .kv_output = kv_output.data_ptr(),
+        .score_bias = ape.data_ptr(),
+        .plan_c = plan_c,
+        .plan_w = plan_w,
+        .num_compress = num_c,
+        .num_write = num_w,
+    };
+    RuntimeCheck(num_q_tokens >= num_w, "invalid prefill plan: num_q < num_w");
+    if (const auto num_c_blocks = div_ceil(num_c * kNumSplit, kWarpsPerBlock)) {
+      LaunchKernel(num_c_blocks, kBlockSize, device)  //
+          .enable_pdl(kUsePDL)(prefill_c_kernel, params);
+    }
+    if (const auto num_w_blocks = div_ceil(num_w * kNumSplit, kWarpsPerBlock)) {
+      LaunchKernel(num_w_blocks, kBlockSize, device)  //
+          .enable_pdl(kUsePDL)(prefill_w_kernel, params);
+    }
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/c_plan.cuh

@@ -0,0 +1,868 @@
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <sgl_kernel/deepseek_v4/compress_v2.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tuple.h>
+
+#include <cstdint>
+#include <limits>
+
+namespace host::compress {
+
+constexpr auto kDLUInt8 = DLDataType{.code = kDLUInt, .bits = 8, .lanes = 1};
+
+using PlanC = CompressPlan;
+using PlanW = WritePlan;
+using PlanD = DecodePlan;
+
+using RID_T = int64_t;
+using R2T_T = int32_t;
+using F2S_T = int64_t;
+using IDX_T = int64_t;
+
+/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536
+SGL_DEVICE __host__ PlanW pack_w(uint32_t ragged_id, uint32_t batch_id, int32_t seq_len) {
+  return {static_cast<uint32_t>(ragged_id | batch_id << 16), seq_len};
+}
+
+/// NOTE: for the internal use, we pack the ragged and batch id, since both not exceed 65536
+SGL_DEVICE uint2 unpack_w(PlanW plan) {
+  return {static_cast<uint16_t>(plan.ragged_id), static_cast<uint16_t>(plan.ragged_id >> 16)};
+}
+
+struct Prefill0Params {
+  PlanC* plan_c;
+  PlanW* plan_w;
+  const IDX_T* seq_lens_ptr;     // [batch_size]
+  const IDX_T* extend_lens_ptr;  // [batch_size]
+  uint32_t batch_size;
+  uint32_t num_q_tokens;
+  int32_t compress_ratio;
+  int32_t swa_page_size;
+  int32_t mtp_pad;
+};
+
+struct Prefill1Params {
+  PlanC* plan_c;
+  PlanW* plan_w;
+  const RID_T* rid_ptr;  // [batch_size]
+  const R2T_T* r2t_ptr;  // [num_reqs, stride_r2t]
+  const F2S_T* f2s_ptr;  // [num_full_slots], full_loc -> swa_loc
+  int64_t stride_r2t;
+  uint32_t num_c;
+  uint32_t num_w;
+  uint32_t num_c_padded;
+  uint32_t num_w_padded;
+  uint32_t num_work;
+  int32_t swa_page_size;
+  int32_t ring_size;
+  int32_t compress_ratio;
+};
+
+struct DecodeParams {
+  PlanD* plan_d;
+  const RID_T* rid_ptr;  // [batch_size]
+  const R2T_T* r2t_ptr;  // [num_reqs, stride_r2t]
+  const F2S_T* f2s_ptr;  // [num_full_slots], full_loc -> swa_loc
+  const IDX_T* seq_ptr;  // [batch_size]
+  int64_t stride_r2t;
+  uint32_t batch_size;
+  int32_t swa_page_size;
+  int32_t ring_size;
+  int32_t compress_ratio;
+};
+
+struct Prefill1ParamsLegacy {
+  PlanC* plan_c;
+  PlanW* plan_w;
+  const RID_T* rid_ptr;  // [batch_size]
+  uint32_t num_c;
+  uint32_t num_w;
+  uint32_t num_c_padded;
+  uint32_t num_w_padded;
+  uint32_t num_work;
+  int32_t compress_ratio;
+};
+
+struct DecodeParamsLegacy {
+  PlanD* plan_d;
+  const RID_T* rid_ptr;  // [batch_size]
+  const IDX_T* seq_ptr;  // [batch_size]
+  uint32_t batch_size;
+  int32_t compress_ratio;
+};
+
+inline constexpr uint32_t kMaxPrefillBatchSize = 1024;
+
+SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
+  static_assert(device::kWarpThreads == 32);
+#pragma unroll
+  for (uint32_t offset = 1; offset < 32; offset *= 2) {
+#ifndef USE_ROCM
+    uint32_t n = __shfl_up_sync(device::kFullMask, val, offset);
+#else
+    uint32_t n = __shfl_up(val, offset, 32);
+#endif
+    if (lane_id >= offset) val += n;
+  }
+  return val;
+}
+
+/// Warp-wide max/min for integer types. `device::warp::reduce_max` routes through
+/// `dtype_trait<T>::max` which is only specialized for FP types.
+SGL_DEVICE uint32_t warp_reduce_max_u32(uint32_t val) {
+#pragma unroll
+  for (uint32_t mask = 16; mask > 0; mask >>= 1) {
+#ifndef USE_ROCM
+    val = max(val, __shfl_xor_sync(device::kFullMask, val, mask, 32));
+#else
+    val = max(val, __shfl_xor(val, mask, 32));
+#endif
+  }
+  return val;
+}
+
+SGL_DEVICE uint32_t warp_reduce_min_u32(uint32_t val) {
+#pragma unroll
+  for (uint32_t mask = 16; mask > 0; mask >>= 1) {
+#ifndef USE_ROCM
+    val = min(val, __shfl_xor_sync(device::kFullMask, val, mask, 32));
+#else
+    val = min(val, __shfl_xor(val, mask, 32));
+#endif
+  }
+  return val;
+}
+
+__global__ __launch_bounds__(1024, 1)  //
+    void plan_compress_prefill_kernel0(const Prefill0Params params) {
+  using namespace device;
+  const auto tx = threadIdx.x;
+  const auto block_size = kMaxPrefillBatchSize;
+  constexpr auto kNumWarps = kMaxPrefillBatchSize / kWarpThreads;
+  const auto cr = params.compress_ratio;
+  const auto sps = params.swa_page_size;
+  const bool is_overlap = (cr == 4);
+  const int32_t window_size = cr * (is_overlap ? 2 : 1);
+
+  alignas(128) __shared__ uint32_t counter_c;
+  alignas(128) __shared__ uint32_t counter_w;
+  __shared__ int32_t s_seq_len[kMaxPrefillBatchSize];
+  __shared__ int32_t s_prefix_len[kMaxPrefillBatchSize];
+  __shared__ uint32_t warp_max[kNumWarps];
+  __shared__ uint32_t warp_min[kNumWarps];
+  __shared__ uint32_t s_max_extend;
+  __shared__ uint32_t s_min_extend;
+
+  const auto lane_id = tx % kWarpThreads;
+  const auto warp_id = tx / kWarpThreads;
+
+  // === Stage A: load per-batch fields, init shared scratch ===
+  int32_t seq_len = 0, extend_len = 0, prefix_len = 0;
+  if (tx < params.batch_size) {
+    seq_len = static_cast<int32_t>(params.seq_lens_ptr[tx]);
+    extend_len = static_cast<int32_t>(params.extend_lens_ptr[tx]);
+    prefix_len = seq_len - extend_len;
+    s_seq_len[tx] = seq_len;
+    s_prefix_len[tx] = prefix_len;
+  }
+  if (tx == 0) {
+    counter_c = 0;
+    counter_w = 0;
+  }
+  if (tx < kNumWarps) {
+    warp_max[tx] = 0;
+    warp_min[tx] = 0xFFFFFFFFu;
+  }
+
+  // === Stage B: min/max(extend_len) for MTP-uniform detection ===
+  // For min, treat threads outside `batch_size` as +inf so they don't pull the min down.
+  const uint32_t e_for_max = static_cast<uint32_t>(extend_len);
+  const uint32_t e_for_min = (tx < params.batch_size) ? e_for_max : 0xFFFFFFFFu;
+  warp_max[warp_id] = warp_reduce_max_u32(e_for_max);
+  warp_min[warp_id] = warp_reduce_min_u32(e_for_min);
+  __syncthreads();
+  if (warp_id == 0) {
+    s_max_extend = warp_reduce_max_u32(warp_max[lane_id]);
+    s_min_extend = warp_reduce_min_u32(warp_min[lane_id]);
+  }
+  __syncthreads();
+
+  const auto num_q = params.num_q_tokens;
+  // MTP-uniform: every batch shares the same small extend_len `E`, so we can decompose
+  // a global token id `k` into (batch_id, j) = (k / E, k % E) and skip the per-batch loop.
+  const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32);
+
+  // === Stage C: emit valid plans, slot allocation via shared-mem atomicAdd ===
+  if (is_mtp_extend) {
+    // Path 1: token-driven. Each global token id maps to exactly one (batch_id, j).
+    const uint32_t E = s_max_extend;
+    // num_q is the padded buffer size (graph bucket), not the work size: cap the
+    // loop at the real token count so batch_id = k / E stays < batch_size on an
+    // underfilled replay; Stage D pads [counter, num_q) with invalid.
+    const uint32_t num_real_q = params.batch_size * E;
+    for (uint32_t k = tx; k < num_real_q; k += block_size) {
+      const uint32_t batch_id = k / E;
+      const uint32_t j = k % E;
+      const int32_t pl = s_prefix_len[batch_id];
+      const int32_t sl = s_seq_len[batch_id];
+      const int32_t position = pl + static_cast<int32_t>(j);
+      const uint32_t ragged_id = k;
+
+      if ((position + 1) % cr == 0) {
+        const int32_t buffer_len = window_size - min(static_cast<int32_t>(j) + 1, window_size);
+        const uint32_t out_idx = atomicAdd(&counter_c, 1u);
+        params.plan_c[out_idx] = {
+            .seq_len = static_cast<uint32_t>(position + 1),
+            .ragged_id = static_cast<uint16_t>(ragged_id),
+            .buffer_len = static_cast<uint16_t>(buffer_len),
+            .read_page_0 = -1,
+            .read_page_1 = static_cast<int32_t>(batch_id),
+        };
+      }
+
+      const int32_t last_c_pos = (sl / cr) * cr;
+      const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
+      bool do_write = position >= first_w_pos;
+      if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
+      if (do_write) {
+        const uint32_t out_idx = atomicAdd(&counter_w, 1u);
+        params.plan_w[out_idx] = pack_w(ragged_id, batch_id, position + 1);
+      }
+    }
+  } else {
+    // Path 2: general prefill (long extend_len). Iterate batches in an outer loop;
+    // the whole block sweeps each batch's tokens in parallel.
+    uint32_t base_e = 0;
+    for (uint32_t batch_id = 0; batch_id < params.batch_size; ++batch_id) {
+      const int32_t pl = s_prefix_len[batch_id];
+      const int32_t sl = s_seq_len[batch_id];
+      const int32_t el = sl - pl;
+      const int32_t last_c_pos = (sl / cr) * cr;
+      const int32_t first_w_pos = min(last_c_pos - (is_overlap ? cr : 0), sl - params.mtp_pad);
+      for (int32_t j = static_cast<int32_t>(tx); j < el; j += static_cast<int32_t>(block_size)) {
+        const int32_t position = pl + j;
+        const uint32_t ragged_id = base_e + static_cast<uint32_t>(j);
+
+        if ((position + 1) % cr == 0) {
+          const int32_t buffer_len = window_size - min(j + 1, window_size);
+          const uint32_t out_idx = atomicAdd(&counter_c, 1u);
+          params.plan_c[out_idx] = {
+              .seq_len = static_cast<uint32_t>(position + 1),
+              .ragged_id = static_cast<uint16_t>(ragged_id),
+              .buffer_len = static_cast<uint16_t>(buffer_len),
+              .read_page_0 = -1,
+              .read_page_1 = static_cast<int32_t>(batch_id),
+          };
+        }
+
+        bool do_write = position >= first_w_pos;
+        if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr);
+        if (do_write) {
+          const uint32_t out_idx = atomicAdd(&counter_w, 1u);
+          params.plan_w[out_idx] = pack_w(ragged_id, static_cast<uint32_t>(batch_id), position + 1);
+        }
+      }
+      base_e += static_cast<uint32_t>(el);
+    }
+  }
+  __syncthreads();
+
+  // === Stage D: pad [counter_c, num_q) / [counter_w, num_q) with invalid ===
+  const auto total_c = counter_c;
+  const auto total_w = counter_w;
+  for (uint32_t k = total_c + tx; k < num_q; k += block_size) {
+    params.plan_c[k] = PlanC::invalid();
+  }
+  for (uint32_t k = total_w + tx; k < num_q; k += block_size) {
+    params.plan_w[k] = PlanW::invalid();
+  }
+}
+
+/// NOTE: stage 1
+__global__ void plan_compress_prefill_kernel_1(const Prefill1Params params) {
+  const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= params.num_work) return;
+  auto plan_c = idx < params.num_c ? params.plan_c[idx] : PlanC::invalid();
+  auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid();
+
+  const auto compute_loc = [&](int32_t swa_loc) {
+    const auto swa_page = swa_loc / params.swa_page_size;
+    const auto ring_offset = swa_loc % params.ring_size;
+    return swa_page * params.ring_size + ring_offset;
+  };
+  const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
+    return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
+  };
+
+  if (!plan_c.is_invalid()) {  // 1. in bound. 2. not masked
+    if (plan_c.buffer_len > 0) {
+      const auto batch_id = plan_c.read_page_1;
+      const auto rid = params.rid_ptr[batch_id];
+      const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
+      // `seq_len` should be ratio-aligned here
+      const auto position_1 = static_cast<int32_t>(plan_c.seq_len - 1);
+      // only used for c4, harmless for c128
+      const auto position_0 = max(position_1 - params.compress_ratio, 0);
+      if (params.compress_ratio == 128) {
+        plan_c.read_page_0 = compute_c128_loc(rid, position_0) / 128;
+        plan_c.read_page_1 = compute_c128_loc(rid, position_1) / 128;
+      } else {
+        const auto raw_loc_0 = mapping[position_0];
+        const auto raw_loc_1 = mapping[position_1];
+        const auto state_loc_0 = params.f2s_ptr[raw_loc_0];
+        const auto state_loc_1 = params.f2s_ptr[raw_loc_1];
+        plan_c.read_page_0 = compute_loc(state_loc_0) / params.compress_ratio;
+        plan_c.read_page_1 = compute_loc(state_loc_1) / params.compress_ratio;
+      }
+      params.plan_c[idx] = plan_c;
+    }
+  } else if (idx < params.num_c_padded) {
+    params.plan_c[idx] = PlanC::invalid();
+  }
+
+  if (!plan_w.is_invalid()) {  // 1. in bound. 2. not masked
+    const auto [ragged_id, batch_id] = unpack_w(plan_w);
+    const auto rid = params.rid_ptr[batch_id];
+    const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
+    // `seq_len` (`write_loc`) may not be aligned here
+    const auto position = static_cast<int32_t>(plan_w.write_loc - 1);
+    plan_w.ragged_id = ragged_id;
+    if (params.compress_ratio == 128) {
+      plan_w.write_loc = compute_c128_loc(rid, position);
+    } else {
+      const auto raw_loc = mapping[position];
+      plan_w.write_loc = compute_loc(params.f2s_ptr[raw_loc]);
+    }
+    params.plan_w[idx] = plan_w;
+  } else if (idx < params.num_w_padded) {
+    params.plan_w[idx] = PlanW::invalid();
+  }
+}
+
+__global__ void plan_compress_decode_kernel(const DecodeParams params) {
+  const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= params.batch_size) return;
+  const auto rid = params.rid_ptr[idx];
+  const auto mapping = params.r2t_ptr + rid * params.stride_r2t;
+  const auto compute_loc = [&](int32_t swa_loc) {
+    const auto swa_page = swa_loc / params.swa_page_size;
+    const auto ring_offset = swa_loc % params.ring_size;
+    return swa_page * params.ring_size + ring_offset;
+  };
+  const auto compute_c128_loc = [&](int64_t rid, int32_t position) {
+    return static_cast<int32_t>(rid * params.ring_size + position % params.ring_size);
+  };
+  const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
+  const auto position_1 = static_cast<int32_t>(seq_len - 1);
+  const auto position_0 = max(position_1 - params.compress_ratio, 0);
+  int32_t write_loc;
+  int32_t read_page_0;
+  int32_t read_page_1;
+  if (params.compress_ratio == 128) {
+    write_loc = compute_c128_loc(rid, position_1);
+    read_page_0 = compute_c128_loc(rid, position_0) / 128;
+    read_page_1 = compute_c128_loc(rid, position_1) / 128;
+  } else {
+    const auto raw_loc_0 = mapping[position_0];
+    const auto raw_loc_1 = mapping[position_1];
+    const auto state_loc_0 = params.f2s_ptr[raw_loc_0];
+    const auto state_loc_1 = params.f2s_ptr[raw_loc_1];
+    write_loc = static_cast<int32_t>(compute_loc(state_loc_1));
+    read_page_0 = static_cast<int32_t>(compute_loc(state_loc_0) / params.compress_ratio);
+    read_page_1 = static_cast<int32_t>(write_loc / params.compress_ratio);
+  }
+  params.plan_d[idx] = {
+      .seq_len = static_cast<uint32_t>(seq_len),
+      .write_loc = write_loc,
+      .read_page_0 = read_page_0,
+      .read_page_1 = read_page_1,
+  };
+}
+
+__global__ void plan_compress_prefill_legacy_kernel(const Prefill1ParamsLegacy params) {
+  const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= params.num_work) return;
+  auto plan_c = idx < params.num_c ? params.plan_c[idx] : PlanC::invalid();
+  auto plan_w = idx < params.num_w ? params.plan_w[idx] : PlanW::invalid();
+
+  /// Per-request ring buffer slot translation:
+  /// - c4:   page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4
+  /// - c128: page = rid;                          slot = rid * 128 + position % 128
+  const auto legacy_compute_page = [&](int32_t rid, int32_t position) {
+    if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1);
+    return rid;  // c128
+  };
+  const auto legacy_compute_loc = [&](int32_t rid, int32_t position) {
+    const auto remainder = position % params.compress_ratio;
+    return legacy_compute_page(rid, position) * params.compress_ratio + remainder;
+  };
+
+  if (!plan_c.is_invalid()) {
+    const auto batch_id = plan_c.read_page_1;
+    const auto rid = static_cast<int32_t>(params.rid_ptr[batch_id]);
+    // `seq_len` is ratio-aligned for compress events
+    const auto position_1 = static_cast<int32_t>(plan_c.seq_len) - 1;
+    const auto position_0 = max(position_1 - params.compress_ratio, 0);
+    plan_c.read_page_0 = legacy_compute_page(rid, position_0);
+    plan_c.read_page_1 = legacy_compute_page(rid, position_1);
+    params.plan_c[idx] = plan_c;
+  } else if (idx < params.num_c_padded) {
+    params.plan_c[idx] = PlanC::invalid();
+  }
+
+  if (!plan_w.is_invalid()) {
+    const auto [ragged_id, batch_id] = unpack_w(plan_w);
+    const auto rid = static_cast<int32_t>(params.rid_ptr[batch_id]);
+    // `write_loc` carries (position + 1) at this stage; may not be ratio-aligned
+    const auto position = static_cast<int32_t>(plan_w.write_loc) - 1;
+    plan_w.ragged_id = ragged_id;
+    plan_w.write_loc = legacy_compute_loc(rid, position);
+    params.plan_w[idx] = plan_w;
+  } else if (idx < params.num_w_padded) {
+    params.plan_w[idx] = PlanW::invalid();
+  }
+}
+
+__global__ void plan_compress_decode_legacy_kernel(const DecodeParamsLegacy params) {
+  const auto idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= params.batch_size) return;
+  /// Per-request ring buffer slot translation:
+  /// - c4:   page = rid * 2 + (position / 4) % 2; slot = page * 4 + position % 4
+  /// - c128: page = rid;                          slot = rid * 128 + position % 128
+  const auto legacy_compute_page = [&](int32_t rid, int32_t position) {
+    if (params.compress_ratio == 4) return rid * 2 + ((position / 4) & 1);
+    return rid;  // c128
+  };
+  const auto legacy_compute_loc = [&](int32_t rid, int32_t position) {
+    const auto remainder = position % params.compress_ratio;
+    return legacy_compute_page(rid, position) * params.compress_ratio + remainder;
+  };
+  const auto rid = static_cast<int32_t>(params.rid_ptr[idx]);
+  const auto seq_len = static_cast<int32_t>(params.seq_ptr[idx]);
+  const auto position_1 = seq_len - 1;
+  const auto position_0 = max(position_1 - params.compress_ratio, 0);
+  const int32_t write_loc = legacy_compute_loc(rid, position_1);
+  const int32_t read_page_0 = legacy_compute_page(rid, position_0);
+  const int32_t read_page_1 = legacy_compute_page(rid, position_1);
+  params.plan_d[idx] = {
+      .seq_len = static_cast<uint32_t>(seq_len),
+      .write_loc = write_loc,
+      .read_page_0 = read_page_0,
+      .read_page_1 = read_page_1,
+  };
+}
+
+using PrefillPlan = tvm::ffi::Tuple<tvm::ffi::Tensor, tvm::ffi::Tensor>;
+
+/**
+ * \brief Build c4/c128 prefill plan tensors. CPU-resident.
+ * Inputs (all CPU-resident):
+ * @param req_pool_indices  `[batch_size]` int64_t
+ * @param req_to_token      `[num_reqs, max_tokens_per_req]` int64_t
+ * @param full_to_state     `[full_cache_size]` int64_t. For c4 this maps
+ *                          full loc -> SWA loc; ignored for c128, whose
+ *                          state slot is request-scoped.
+ * @param seq_lens          `[batch_size]` int64
+ * @param extend_lens       `[batch_size]` int64
+ * @param compress_plan     `[num_q_tokens, 16]` uint8 (output)
+ * @param write_plan        `[num_q_tokens,  8]` uint8 (output)
+ * @param compress_ratio 4 for c4, 128 for c128
+ * @param use_cuda_graph Whether the plans will be used with cuda graph (affects padding)
+ * @return (compress plan tensor, write plan tensor)
+ */
+inline PrefillPlan plan_compress_prefill(
+    const tvm::ffi::TensorView req_pool_indices,  // GPU
+    const tvm::ffi::TensorView req_to_token,      // GPU
+    const tvm::ffi::TensorView full_to_state,     // GPU
+    const tvm::ffi::TensorView seq_lens,          // CPU/GPU
+    const tvm::ffi::TensorView extend_lens,       // CPU/GPU
+    const tvm::ffi::TensorView pin_buffer,        // CPU
+    const uint32_t num_q_tokens,
+    const int32_t compress_ratio,
+    const int32_t swa_page_size,
+    const int32_t ring_size,
+    const bool use_cuda_graph) {
+  auto B = SymbolicSize{"batch_size"};
+  auto N = SymbolicSize{"num_q_tokens"};
+  auto cpu_or_gpu = SymbolicDevice{};
+  auto device_ = SymbolicDevice{};
+  cpu_or_gpu.set_options<kDLCPU, kDLGPU>();
+  device_.set_options<kDLGPU>();
+
+  TensorMatcher({B})  //
+      .with_dtype<RID_T>()
+      .with_device(device_)
+      .verify(req_pool_indices);
+  TensorMatcher({-1, -1})  //
+      .with_dtype<R2T_T>()
+      .with_device(device_)
+      .verify(req_to_token);
+  TensorMatcher({-1})  //
+      .with_dtype<F2S_T>()
+      .with_device(device_)
+      .verify(full_to_state);
+  TensorMatcher({B})  //
+      .with_dtype<IDX_T>()
+      .with_device(cpu_or_gpu)
+      .verify(seq_lens)
+      .verify(extend_lens);
+  TensorMatcher({-1})  //
+      .with_dtype<uint8_t>()
+      .with_device<kDLCPU>()
+      .verify(pin_buffer);
+
+  const bool is_overlap = (compress_ratio == 4);
+  const int32_t window_size = compress_ratio * (is_overlap ? 2 : 1);
+
+  const auto seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr());
+  const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
+  const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
+  const auto r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr());
+  const auto f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr());
+
+  const auto batch_size = static_cast<uint32_t>(B.unwrap());
+  constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
+  RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
+  RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
+  // `swa_page_size` >= `ring_size` >= `compress_ratio`
+  RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
+
+  const auto device = device_.unwrap();
+  const auto stream = LaunchKernel::resolve_device(device);
+
+  constexpr int32_t kMaxMTPDraftTokens = 4;
+  const auto mtp_pad = std::min(ring_size - compress_ratio, kMaxMTPDraftTokens);
+
+  if (cpu_or_gpu.unwrap().device_type == kDLGPU) {
+    // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly
+    // on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the
+    // SWA-translated read/write locations. Used for MTP / cuda-graph capture where
+    // a host sync would be expensive.
+    RuntimeCheck(batch_size <= kMaxPrefillBatchSize, "GPU plan only support batch size up to ", kMaxPrefillBatchSize);
+    auto C = ffi::empty({num_q_tokens, sizeof(PlanC)}, kDLUInt8, device);
+    auto W = ffi::empty({num_q_tokens, sizeof(PlanW)}, kDLUInt8, device);
+    const auto params0 = Prefill0Params{
+        .plan_c = static_cast<PlanC*>(C.data_ptr()),
+        .plan_w = static_cast<PlanW*>(W.data_ptr()),
+        .seq_lens_ptr = seq_ptr,
+        .extend_lens_ptr = ext_ptr,
+        .batch_size = batch_size,
+        .num_q_tokens = num_q_tokens,
+        .compress_ratio = compress_ratio,
+        .swa_page_size = swa_page_size,
+        .mtp_pad = mtp_pad,
+    };
+    LaunchKernel(1, kMaxPrefillBatchSize, device)(plan_compress_prefill_kernel0, params0);
+    // kernel_1 sees the already-padded buffers, so num_c == num_w == num_padded == num_q_tokens.
+    const auto params1 = Prefill1Params{
+        .plan_c = static_cast<PlanC*>(C.data_ptr()),
+        .plan_w = static_cast<PlanW*>(W.data_ptr()),
+        .rid_ptr = rid_ptr,
+        .r2t_ptr = r2t_ptr,
+        .f2s_ptr = f2s_ptr,
+        .stride_r2t = req_to_token.stride(0),
+        .num_c = num_q_tokens,
+        .num_w = num_q_tokens,
+        .num_c_padded = num_q_tokens,
+        .num_w_padded = num_q_tokens,
+        .num_work = num_q_tokens,
+        .swa_page_size = swa_page_size,
+        .ring_size = ring_size,
+        .compress_ratio = compress_ratio,
+    };
+    const auto block_size_1 = 256;
+    const auto num_blocks_1 = div_ceil(params1.num_work, block_size_1);
+    LaunchKernel(num_blocks_1, block_size_1, device)(plan_compress_prefill_kernel_1, params1);
+    return PrefillPlan{std::move(C), std::move(W)};
+  }
+
+  // CPU input path: only here do we need the pinned scratch buffer.
+  const auto pin_buffer_bytes = static_cast<size_t>(pin_buffer.numel()) * sizeof(uint8_t);
+  RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW)));
+  const auto plan_c_ptr = reinterpret_cast<PlanC*>(pin_buffer.data_ptr());
+  const auto plan_w_ptr = reinterpret_cast<PlanW*>(plan_c_ptr + num_q_tokens);
+
+  uint32_t counter = 0;
+  uint32_t counter_c = 0;
+  uint32_t counter_w = 0;
+
+  const auto should_compress = [=](int32_t position) { return (position + 1) % compress_ratio == 0; };
+  for (const auto i : irange(batch_size)) {
+    const int32_t seq_len = seq_ptr[i];
+    const int32_t extend_len = ext_ptr[i];
+    const int32_t prefix_len = seq_len - extend_len;
+    const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
+    const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
+    RuntimeCheck(0 < extend_len && extend_len <= seq_len);
+    const auto should_write = [=](int32_t position) {
+      if (position >= first_w_pos) return true;
+      return is_overlap && position % swa_page_size >= (swa_page_size - compress_ratio);
+    };
+    for (const auto j : irange(extend_len)) {
+      const int32_t position = prefix_len + j;
+      const int32_t ragged_id = counter + j;
+      if (should_compress(position)) {
+        const auto buffer_len = window_size - std::min(j + 1, window_size);
+        plan_c_ptr[counter_c++] = {
+            .seq_len = static_cast<uint32_t>(position + 1),
+            .ragged_id = static_cast<uint16_t>(ragged_id),
+            .buffer_len = static_cast<uint16_t>(buffer_len),
+            // to be filled by kernel
+            .read_page_0 = -1,
+            .read_page_1 = static_cast<int32_t>(i),
+        };
+      }
+      if (should_write(position)) {
+        plan_w_ptr[counter_w++] = pack_w(ragged_id, i, position + 1);
+      }
+    }
+    counter += extend_len;
+  }
+  RuntimeCheck(counter == num_q_tokens);
+
+  const auto copy_to_device = [stream](void* cuda_ptr, auto* host_ptr, size_t count) {
+    const auto size_bytes = count * sizeof(*host_ptr);
+    RuntimeDeviceCheck(cudaMemcpyAsync(cuda_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice, stream));
+  };
+  const auto num_c_padded = use_cuda_graph ? num_q_tokens : counter_c;
+  const auto num_w_padded = use_cuda_graph ? num_q_tokens : counter_w;
+  auto C = ffi::empty({num_c_padded, sizeof(PlanC)}, kDLUInt8, device);
+  auto W = ffi::empty({num_w_padded, sizeof(PlanW)}, kDLUInt8, device);
+  copy_to_device(C.data_ptr(), plan_c_ptr, counter_c);
+  copy_to_device(W.data_ptr(), plan_w_ptr, counter_w);
+  const auto params = Prefill1Params{
+      .plan_c = static_cast<PlanC*>(C.data_ptr()),
+      .plan_w = static_cast<PlanW*>(W.data_ptr()),
+      .rid_ptr = rid_ptr,
+      .r2t_ptr = r2t_ptr,
+      .f2s_ptr = f2s_ptr,
+      .stride_r2t = req_to_token.size(1),
+      .num_c = counter_c,
+      .num_w = counter_w,
+      .num_c_padded = num_c_padded,
+      .num_w_padded = num_w_padded,
+      .num_work = std::max(num_c_padded, num_w_padded),
+      .swa_page_size = swa_page_size,
+      .ring_size = ring_size,
+      .compress_ratio = compress_ratio,
+  };
+  const auto block_size = 256;
+  const auto num_blocks = div_ceil(params.num_work, block_size);
+  LaunchKernel(num_blocks, block_size, device)(plan_compress_prefill_kernel_1, params);
+  return PrefillPlan{std::move(C), std::move(W)};
+}
+
+inline tvm::ffi::Tensor plan_compress_decode(
+    const tvm::ffi::TensorView req_pool_indices,  // GPU
+    const tvm::ffi::TensorView req_to_token,      // GPU
+    const tvm::ffi::TensorView full_to_state,     // GPU
+    const tvm::ffi::TensorView seq_lens,          // CPU/GPU
+    const int32_t compress_ratio,
+    const int32_t swa_page_size,
+    const int32_t ring_size) {
+  auto B = SymbolicSize{"batch_size"};
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLGPU>();
+
+  TensorMatcher({B})  //
+      .with_dtype<RID_T>()
+      .with_device(device_)
+      .verify(req_pool_indices);
+  TensorMatcher({-1, -1})  //
+      .with_dtype<R2T_T>()
+      .with_device(device_)
+      .verify(req_to_token);
+  TensorMatcher({-1})  //
+      .with_dtype<F2S_T>()
+      .with_device(device_)
+      .verify(full_to_state);
+  TensorMatcher({B})  //
+      .with_dtype<IDX_T>()
+      .with_device(device_)
+      .verify(seq_lens);
+
+  const auto batch_size = static_cast<uint32_t>(B.unwrap());
+  const auto device = device_.unwrap();
+  auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
+  const auto params = DecodeParams{
+      .plan_d = static_cast<PlanD*>(D.data_ptr()),
+      .rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
+      .r2t_ptr = static_cast<const R2T_T*>(req_to_token.data_ptr()),
+      .f2s_ptr = static_cast<const F2S_T*>(full_to_state.data_ptr()),
+      .seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
+      .stride_r2t = req_to_token.size(1),
+      .batch_size = batch_size,
+      .swa_page_size = swa_page_size,
+      .ring_size = ring_size,
+      .compress_ratio = compress_ratio,
+  };
+  const auto block_size = 256;
+  const auto num_blocks = div_ceil(batch_size, block_size);
+  LaunchKernel(num_blocks, block_size, device)(plan_compress_decode_kernel, params);
+  return D;
+}
+
+/**
+ * \brief Build c4/c128 prefill plan tensors for the legacy non-paged ring
+ * buffer. Uses only `req_pool_indices` to derive ring slots:
+ *   - c4 (overlap):  each request occupies 2 contiguous pages (8 token slots)
+ *   - c128:          each request occupies 1 page (128 token slots)
+ *
+ * Inputs:
+ * @param req_pool_indices  `[batch_size]` int64 (GPU)
+ * @param seq_lens          `[batch_size]` int64 (CPU)
+ * @param extend_lens       `[batch_size]` int64 (CPU)
+ * @param pin_buffer        pinned scratch (CPU uint8)
+ * @return (compress plan tensor, write plan tensor)
+ */
+inline PrefillPlan plan_compress_prefill_legacy(
+    const tvm::ffi::TensorView req_pool_indices,  // GPU
+    const tvm::ffi::TensorView seq_lens,          // CPU
+    const tvm::ffi::TensorView extend_lens,       // CPU
+    const tvm::ffi::TensorView pin_buffer,        // CPU
+    const uint32_t num_q_tokens,
+    const int32_t compress_ratio,
+    const bool use_cuda_graph) {
+  auto B = SymbolicSize{"batch_size"};
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLGPU>();
+
+  TensorMatcher({B})  //
+      .with_dtype<RID_T>()
+      .with_device(device_)
+      .verify(req_pool_indices);
+  TensorMatcher({B})  //
+      .with_dtype<IDX_T>()
+      .with_device<kDLCPU>()
+      .verify(seq_lens)
+      .verify(extend_lens);
+  TensorMatcher({-1})  //
+      .with_dtype<uint8_t>()
+      .with_device<kDLCPU>()
+      .verify(pin_buffer);
+
+  const auto pin_buffer_bytes = static_cast<size_t>(pin_buffer.numel()) * sizeof(uint8_t);
+  RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW)));
+  const auto plan_c_ptr = reinterpret_cast<PlanC*>(pin_buffer.data_ptr());
+  const auto plan_w_ptr = reinterpret_cast<PlanW*>(plan_c_ptr + num_q_tokens);
+
+  const bool is_overlap = (compress_ratio == 4);
+  const auto seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr());
+  const auto ext_ptr = static_cast<const IDX_T*>(extend_lens.data_ptr());
+  const auto rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr());
+
+  const auto window_size = compress_ratio * (is_overlap ? 2 : 1);
+  const auto batch_size = static_cast<uint32_t>(B.unwrap());
+  constexpr auto kMaxTokens = static_cast<uint32_t>(std::numeric_limits<uint16_t>::max());
+  RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
+  RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
+
+  uint32_t counter = 0;
+  uint32_t counter_c = 0;
+  uint32_t counter_w = 0;
+  const auto should_compress = [=](int32_t position) { return (position + 1) % compress_ratio == 0; };
+  for (const auto i : irange(batch_size)) {
+    const int32_t seq_len = seq_ptr[i];
+    const int32_t extend_len = ext_ptr[i];
+    const int32_t prefix_len = seq_len - extend_len;
+    const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
+    const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
+    RuntimeCheck(0 < extend_len && extend_len <= seq_len);
+    const auto should_write = [=](int32_t position) { return position >= first_w_pos; };
+    for (const auto j : irange(extend_len)) {
+      const int32_t position = prefix_len + j;
+      const int32_t ragged_id = counter + j;
+      if (should_compress(position)) {
+        const auto buffer_len = window_size - std::min(j + 1, window_size);
+        plan_c_ptr[counter_c++] = {
+            .seq_len = static_cast<uint32_t>(position + 1),
+            .ragged_id = static_cast<uint16_t>(ragged_id),
+            .buffer_len = static_cast<uint16_t>(buffer_len),
+            // to be filled by kernel
+            .read_page_0 = -1,
+            .read_page_1 = static_cast<int32_t>(i),
+        };
+      }
+      if (should_write(position)) {
+        plan_w_ptr[counter_w++] = pack_w(ragged_id, i, position + 1);
+      }
+    }
+    counter += extend_len;
+  }
+  RuntimeCheck(counter == num_q_tokens);
+
+  const auto device = device_.unwrap();
+  const auto stream = LaunchKernel::resolve_device(device);
+  const auto copy_to_device = [stream](void* cuda_ptr, auto* host_ptr, size_t count) {
+    const auto size_bytes = count * sizeof(*host_ptr);
+    RuntimeDeviceCheck(cudaMemcpyAsync(cuda_ptr, host_ptr, size_bytes, cudaMemcpyHostToDevice, stream));
+  };
+  const auto num_c_padded = use_cuda_graph ? num_q_tokens : counter_c;
+  const auto num_w_padded = use_cuda_graph ? num_q_tokens : counter_w;
+  auto C = ffi::empty({num_c_padded, sizeof(PlanC)}, kDLUInt8, device);
+  auto W = ffi::empty({num_w_padded, sizeof(PlanW)}, kDLUInt8, device);
+  copy_to_device(C.data_ptr(), plan_c_ptr, counter_c);
+  copy_to_device(W.data_ptr(), plan_w_ptr, counter_w);
+  const auto params = Prefill1ParamsLegacy{
+      .plan_c = static_cast<PlanC*>(C.data_ptr()),
+      .plan_w = static_cast<PlanW*>(W.data_ptr()),
+      .rid_ptr = rid_ptr,
+      .num_c = counter_c,
+      .num_w = counter_w,
+      .num_c_padded = num_c_padded,
+      .num_w_padded = num_w_padded,
+      .num_work = std::max(num_c_padded, num_w_padded),
+      .compress_ratio = compress_ratio,
+  };
+  const auto block_size = 256;
+  const auto num_blocks = div_ceil(params.num_work, block_size);
+  if (num_blocks > 0) {
+    LaunchKernel(num_blocks, block_size, device)(plan_compress_prefill_legacy_kernel, params);
+  }
+  return PrefillPlan{std::move(C), std::move(W)};
+}
+
+inline tvm::ffi::Tensor plan_compress_decode_legacy(
+    const tvm::ffi::TensorView req_pool_indices,  // GPU
+    const tvm::ffi::TensorView seq_lens,          // GPU
+    const int32_t compress_ratio) {
+  auto B = SymbolicSize{"batch_size"};
+  auto device_ = SymbolicDevice{};
+  device_.set_options<kDLGPU>();
+
+  TensorMatcher({B})  //
+      .with_dtype<RID_T>()
+      .with_device(device_)
+      .verify(req_pool_indices);
+  TensorMatcher({B})  //
+      .with_dtype<IDX_T>()
+      .with_device(device_)
+      .verify(seq_lens);
+  RuntimeCheck(compress_ratio == 4 || compress_ratio == 128);
+
+  const auto batch_size = static_cast<uint32_t>(B.unwrap());
+  const auto device = device_.unwrap();
+  auto D = ffi::empty({batch_size, sizeof(PlanD)}, kDLUInt8, device);
+  const auto params = DecodeParamsLegacy{
+      .plan_d = static_cast<PlanD*>(D.data_ptr()),
+      .rid_ptr = static_cast<const RID_T*>(req_pool_indices.data_ptr()),
+      .seq_ptr = static_cast<const IDX_T*>(seq_lens.data_ptr()),
+      .batch_size = batch_size,
+      .compress_ratio = compress_ratio,
+  };
+  const auto block_size = 256;
+  const auto num_blocks = div_ceil(batch_size, block_size);
+  LaunchKernel(num_blocks, block_size, device)(plan_compress_decode_legacy_kernel, params);
+  return D;
+}
+
+}  // namespace host::compress
+
+using namespace host::compress;  // expose binding

python/sglang/jit_kernel/csrc/deepseek_v4/common.cuh

@@ -0,0 +1,208 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/deepseek_v4/compress.cuh>
+
+#include <dlpack/dlpack.h>
+
+namespace host::compress {
+
+using PlanResult = tvm::ffi::Tuple<uint32_t, uint32_t>;
+
+struct CompressParams {
+  PrefillPlan* __restrict__ compress_plan;
+  PrefillPlan* __restrict__ write_plan;
+  const int64_t* __restrict__ seq_lens;
+  const int64_t* __restrict__ extend_lens;
+  uint32_t batch_size;
+  uint32_t num_tokens;
+  uint32_t compress_ratio;
+  bool is_overlap;
+};
+
+inline constexpr uint32_t kBlockSize = 1024;
+
+#define PLAN_KERNEL __global__ __launch_bounds__(kBlockSize, 1) inline
+
+PLAN_KERNEL void plan_prefill_cuda(const __grid_constant__ CompressParams params) {
+  const auto &[
+    compress_plan, write_plan, seq_lens, extend_lens, // pointers
+    batch_size, num_tokens, compress_ratio, is_overlap // values
+  ] = params;
+
+  __shared__ uint32_t compress_counter;
+  __shared__ uint32_t write_counter;
+
+  uint32_t batch_id = 0;
+  uint32_t counter = 0;
+  uint32_t extend_len = extend_lens[0];
+
+  const auto tid = threadIdx.x;
+  if (tid == 0) {
+    compress_counter = 0;
+    write_counter = 0;
+  }
+  __syncthreads();
+
+  for (uint32_t i = tid; i < num_tokens; i += blockDim.x) {
+    const uint32_t ragged_id = i;
+    uint32_t j = ragged_id - counter;
+    while (j >= extend_len) {
+      j -= extend_len;
+      batch_id += 1;
+      if (batch_id >= batch_size) [[unlikely]]
+        break;
+      counter += extend_len;
+      extend_len = extend_lens[batch_id];
+    }
+    if (batch_id >= batch_size) [[unlikely]]
+      break;
+    const uint32_t seq_len = seq_lens[batch_id];
+    const uint32_t extend_len = extend_lens[batch_id];
+    const uint32_t prefix_len = seq_len - extend_len;
+    const uint32_t ratio = compress_ratio * (1 + is_overlap);
+    const uint32_t window_len = j + 1 < ratio ? ratio - (j + 1) : 0;
+    const uint32_t position = prefix_len + j;
+    const auto plan = PrefillPlan{
+        .ragged_id = ragged_id,
+        .batch_id = batch_id,
+        .position = position,
+        .window_len = window_len,
+    };
+    const uint32_t start_write_pos = [seq_len, compress_ratio, is_overlap] {
+      const uint32_t pos = seq_len / compress_ratio * compress_ratio;
+      if (!is_overlap) return pos;
+      return pos >= compress_ratio ? pos - compress_ratio : 0;
+    }();
+    if ((position + 1) % compress_ratio == 0) {
+      const auto write_pos = atomicAdd(&compress_counter, 1);
+      compress_plan[write_pos] = plan;
+    }
+    if (position >= start_write_pos) {
+      const auto write_pos = atomicAdd(&write_counter, 1);
+      write_plan[write_pos] = plan;
+    }
+  }
+  __syncthreads();
+  constexpr auto kInvalid = static_cast<uint32_t>(-1);
+  const auto kInvalidPlan = PrefillPlan{kInvalid, kInvalid, kInvalid, kInvalid};
+  const auto compress_count = compress_counter;
+  const auto write_count = write_counter;
+  for (uint32_t i = compress_count + tid; i < num_tokens; i += blockDim.x) {
+    compress_plan[i] = kInvalidPlan;
+  }
+  for (uint32_t i = write_count + tid; i < num_tokens; i += blockDim.x) {
+    write_plan[i] = kInvalidPlan;
+  }
+}
+
+inline PlanResult plan_prefill_host(const CompressParams& params, const bool use_cuda_graph) {
+  const auto &[
+    compress_ptr, write_ptr, seq_lens_ptr, extend_lens_ptr, // pointers
+    batch_size, num_tokens, compress_ratio, is_overlap // values
+  ] = params;
+
+  uint32_t counter = 0;
+  uint32_t compress_counter = 0;
+  uint32_t write_counter = 0;
+  const auto ratio = compress_ratio * (1 + is_overlap);
+  for (const auto i : irange(batch_size)) {
+    const uint32_t seq_len = seq_lens_ptr[i];
+    const uint32_t extend_len = extend_lens_ptr[i];
+    const uint32_t prefix_len = seq_len - extend_len;
+    RuntimeCheck(0 < extend_len && extend_len <= seq_len);
+    /// NOTE: `start_write_pos` must be a multiple of `compress_ratio`
+    const uint32_t start_write_pos = [seq_len, compress_ratio, is_overlap] {
+      const uint32_t pos = seq_len / compress_ratio * compress_ratio;
+      if (!is_overlap) return pos;
+      /// NOTE: to avoid unsigned integer underflow, don't use `pos - compress_ratio`
+      return pos >= compress_ratio ? pos - compress_ratio : 0;
+    }();
+    /// NOTE: `position` is within [prefix_len, seq_len)
+    for (const auto j : irange(extend_len)) {
+      const uint32_t position = prefix_len + j;
+      const auto plan = PrefillPlan{
+          .ragged_id = counter + j,
+          .batch_id = i,
+          .position = position,
+          .window_len = ratio - std::min(j + 1, ratio),
+      };
+      RuntimeCheck(plan.is_valid(compress_ratio, is_overlap), "Internal error!");
+      if ((position + 1) % compress_ratio == 0) {
+        compress_ptr[compress_counter++] = plan;
+      }
+      if (position >= start_write_pos) {
+        write_ptr[write_counter++] = plan;
+      }
+    }
+    counter += extend_len;
+  }
+  RuntimeCheck(counter == num_tokens, "input size ", counter, " != num_q_tokens ", num_tokens);
+  if (!use_cuda_graph) return PlanResult{compress_counter, write_counter};
+  constexpr auto kInvalid = static_cast<uint32_t>(-1);
+  constexpr auto kInvalidPlan = PrefillPlan{kInvalid, kInvalid, kInvalid, kInvalid};
+  for (const auto i : irange(compress_counter, num_tokens)) {
+    compress_ptr[i] = kInvalidPlan;
+  }
+  for (const auto i : irange(write_counter, num_tokens)) {
+    write_ptr[i] = kInvalidPlan;
+  }
+  return PlanResult{num_tokens, num_tokens};
+}
+
+inline PlanResult plan_prefill(
+    const tvm::ffi::TensorView extend_lens,
+    const tvm::ffi::TensorView seq_lens,
+    const tvm::ffi::TensorView compress_plan,
+    const tvm::ffi::TensorView write_plan,
+    const uint32_t compress_ratio,
+    const bool is_overlap,  // for overlap transform, we have to keep 1 more extra window
+    const bool use_cuda_graph) {
+  auto N = SymbolicSize{"batch_size"};
+  auto M = SymbolicSize{"num_tokens"};
+  auto device = SymbolicDevice{};
+  const bool is_cuda = [&] {
+    if (extend_lens.device().device_type == kDLCUDA) {
+      device.set_options<kDLCUDA>();
+      return true;
+    } else {
+      device.set_options<kDLCPU, kDLCUDAHost>();
+      return false;
+    }
+  }();
+  TensorMatcher({N})  // extend_lens and seq_lens
+      .with_dtype<int64_t>()
+      .with_device(device)
+      .verify(extend_lens)
+      .verify(seq_lens);
+  TensorMatcher({M, kPrefillPlanDim})  // compress_plan and write_plan
+      .with_dtype<PrefillPlanTensorDtype>()
+      .with_device(device)
+      .verify(compress_plan)
+      .verify(write_plan);
+
+  const auto params = CompressParams{
+      .compress_plan = static_cast<PrefillPlan*>(compress_plan.data_ptr()),
+      .write_plan = static_cast<PrefillPlan*>(write_plan.data_ptr()),
+      .seq_lens = static_cast<const int64_t*>(seq_lens.data_ptr()),
+      .extend_lens = static_cast<const int64_t*>(extend_lens.data_ptr()),
+      .batch_size = static_cast<uint32_t>(N.unwrap()),
+      .num_tokens = static_cast<uint32_t>(M.unwrap()),
+      .compress_ratio = compress_ratio,
+      .is_overlap = is_overlap,
+  };
+
+  if (!is_cuda) return plan_prefill_host(params, use_cuda_graph);
+  /// NOTE: cuda kernel plan is naturally compatible with cuda graph
+  LaunchKernel(1, kBlockSize, device.unwrap())(plan_prefill_cuda, params);
+  return PlanResult{params.num_tokens, params.num_tokens};
+}
+
+}  // namespace host::compress
+
+namespace {
+
+[[maybe_unused]]
+constexpr auto& plan_compress_prefill = host::compress::plan_prefill;
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/fp8_wo_a_group_major_quant.cuh

@@ -0,0 +1,169 @@
+// DeepSeek-V4 wo_a activation quantization for DeepGEMM fp8_einsum.
+//
+// This is intentionally narrower than the generic per_token_group_quant_8bit_v2
+// kernel: input is a [T, G, D] view with contiguous hidden groups, output_q is
+// contiguous [T, G, D], group_size is fixed to 128, scales are fp32 UE8M0
+// power-of-two values, and output_s is a logical [T, G, D/128] view backed by
+// group-major [G, T, D/128] storage.
+//
+// The generic kernel cannot read the strided DSV4 view while producing
+// contiguous [T, G, D] codes and group-major scales without an extra full-tensor
+// copy.
+#include <sgl_kernel/tensor.h>  // TensorMatcher, SymbolicSize/Device
+#include <sgl_kernel/utils.h>   // RuntimeCheck
+
+#include <sgl_kernel/utils.cuh>  // fp8 aliases, PDL helpers
+#include <sgl_kernel/warp.cuh>   // warp::reduce_max
+
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>  // UE8M0 and FP8 helpers
+
+#include <tvm/ffi/container/tensor.h>  // tvm::ffi::TensorView
+
+#include <cstdint>
+#include <cuda_fp8.h>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::inv_scale_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+constexpr float LOCAL_ABSMAX_ABS = 1e-10f;
+constexpr uint32_t GROUP_SIZE = 128;
+constexpr uint32_t THREADS_PER_GROUP = 8;
+constexpr uint32_t SUBWARPS_PER_BLOCK = 16;
+constexpr uint32_t INPUT_VEC_NUM_BYTES = 32;
+constexpr uint32_t INPUT_INT4_SIZE = INPUT_VEC_NUM_BYTES / sizeof(int4);
+
+template <int THREADS_PER_SUBWARP>
+SGL_DEVICE float GroupReduceMax(float val) {
+  static_assert(
+      (THREADS_PER_SUBWARP & (THREADS_PER_SUBWARP - 1)) == 0 && THREADS_PER_SUBWARP <= 16 && THREADS_PER_SUBWARP >= 1,
+      "THREADS_PER_SUBWARP must be 1, 2, 4, 8, or 16");
+  // Tail subwarps can be inactive at the bounds check, so reduce with only the
+  // current subgroup's lanes rather than a full-warp mask.
+  constexpr device::warp::mask_t kSub = (device::warp::mask_t{1} << THREADS_PER_SUBWARP) - 1;
+  const device::warp::mask_t mask = kSub << (THREADS_PER_SUBWARP * ((threadIdx.x % 32) / THREADS_PER_SUBWARP));
+  return device::warp::reduce_max<THREADS_PER_SUBWARP>(val, mask);
+}
+
+template <typename T, bool kUsePDL>
+__global__ void fp8_wo_a_group_major_quant_ue8m0_kernel(
+    const T* __restrict__ input,
+    fp8_e4m3_t* __restrict__ output_q,
+    float* __restrict__ output_s,
+    int64_t total_scale_groups,
+    int64_t num_tokens,
+    int hidden_dim_groups,
+    int num_outer_groups,
+    int64_t input_stride_t) {
+  device::PDLWaitPrimary<kUsePDL>();
+
+  const int64_t subwarp_id = threadIdx.x / THREADS_PER_GROUP;
+  const int lane_id = threadIdx.x % THREADS_PER_GROUP;
+  const int64_t group_id = static_cast<int64_t>(blockIdx.x) * SUBWARPS_PER_BLOCK + subwarp_id;
+  if (group_id < total_scale_groups) {
+    const int hidden_group = group_id % hidden_dim_groups;
+    const int64_t token_outer = group_id / hidden_dim_groups;
+    const int outer_idx = token_outer % num_outer_groups;
+    const int64_t token_idx = token_outer / num_outer_groups;
+
+    constexpr uint32_t INPUT_VEC_SIZE = INPUT_VEC_NUM_BYTES / sizeof(T);
+    static_assert(INPUT_VEC_SIZE * THREADS_PER_GROUP == GROUP_SIZE);
+
+    const int64_t input_group_start_offset =
+        token_idx * input_stride_t + outer_idx * GROUP_SIZE * hidden_dim_groups + hidden_group * GROUP_SIZE;
+    const int64_t output_group_start_offset = group_id * GROUP_SIZE;
+
+    int4 input_int4[INPUT_INT4_SIZE];
+    T* input_vec = reinterpret_cast<T*>(input_int4);
+
+#pragma unroll
+    for (uint32_t j = 0; j < INPUT_INT4_SIZE; ++j) {
+      input_int4[j] = reinterpret_cast<const int4*>(input + input_group_start_offset + lane_id * INPUT_VEC_SIZE)[j];
+    }
+
+    float local_absmax = LOCAL_ABSMAX_ABS;
+#pragma unroll
+    for (uint32_t j = 0; j < INPUT_VEC_SIZE; ++j) {
+      const float val = static_cast<float>(input_vec[j]);
+      local_absmax = fmaxf(local_absmax, fabsf(val));
+    }
+
+    local_absmax = GroupReduceMax<THREADS_PER_GROUP>(local_absmax);
+
+    constexpr float kFp8MaxInv = 1.0f / kFP8E4M3Max;
+    const int32_t scale_ue8m0 = cast_to_ue8m0(local_absmax * kFp8MaxInv);
+    const float y_scale = inv_scale_ue8m0(scale_ue8m0);
+    const float y_scale_inv = __uint_as_float(static_cast<uint32_t>(scale_ue8m0) << 23);
+
+    int4 output_buf;
+    auto* output_buf_ptr = reinterpret_cast<fp8x2_e4m3_t*>(&output_buf);
+#pragma unroll
+    for (uint32_t j = 0; j < INPUT_VEC_SIZE; j += 2) {
+      output_buf_ptr[j / 2] =
+          pack_fp8(static_cast<float>(input_vec[j]) * y_scale, static_cast<float>(input_vec[j + 1]) * y_scale);
+    }
+
+    *reinterpret_cast<int4*>(output_q + output_group_start_offset + lane_id * INPUT_VEC_SIZE) = output_buf;
+
+    if (lane_id == 0) {
+      output_s[(outer_idx * num_tokens + token_idx) * hidden_dim_groups + hidden_group] = y_scale_inv;
+    }
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename T, bool kUsePDL>
+struct FP8WoAGroupMajorQuantUE8M0Kernel {
+  static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    auto TSize = SymbolicSize{"num_tokens"};
+    auto GSize = SymbolicSize{"num_outer_groups"};
+    auto DSize = SymbolicSize{"hidden_dim"};
+    auto SSize = SymbolicSize{"hidden_dim_groups"};
+
+    TensorMatcher({TSize, GSize, DSize}).with_strides({-1, DSize, 1}).with_dtype<T>().with_device(device).verify(input);
+    TensorMatcher({TSize, GSize, DSize}).with_dtype<fp8_e4m3_t>().with_device(device).verify(output_q);
+    TensorMatcher({GSize, TSize, SSize}).with_dtype<float>().with_device(device).verify(output_s);
+
+    const auto num_tokens = TSize.unwrap();
+    const auto num_outer_groups = GSize.unwrap();
+    const auto hidden_dim = DSize.unwrap();
+    const auto hidden_dim_groups = SSize.unwrap();
+    const auto input_stride_t = input.stride(0);
+    constexpr int64_t kInputAlignElements = sizeof(int4) / sizeof(T);
+
+    RuntimeCheck(hidden_dim % GROUP_SIZE == 0, "hidden_dim must be divisible by 128");
+    RuntimeCheck(hidden_dim_groups == hidden_dim / GROUP_SIZE, "output_s hidden dim mismatch");
+    RuntimeCheck(
+        reinterpret_cast<uintptr_t>(input.data_ptr()) % sizeof(int4) == 0,
+        "input base pointer must be 16-byte aligned");
+    RuntimeCheck(
+        num_tokens <= 1 || input_stride_t % kInputAlignElements == 0,
+        "input token stride must preserve 16-byte vector-load alignment");
+
+    const int64_t total_scale_groups = num_tokens * num_outer_groups * hidden_dim_groups;
+    if (total_scale_groups == 0) return;
+
+    const auto grid = dim3((total_scale_groups + SUBWARPS_PER_BLOCK - 1) / SUBWARPS_PER_BLOCK);
+    const auto block = dim3(SUBWARPS_PER_BLOCK * THREADS_PER_GROUP);
+    host::LaunchKernel(grid, block, device.unwrap())
+        .enable_pdl(kUsePDL)(
+            fp8_wo_a_group_major_quant_ue8m0_kernel<T, kUsePDL>,
+            static_cast<const T*>(input.data_ptr()),
+            static_cast<fp8_e4m3_t*>(output_q.data_ptr()),
+            static_cast<float*>(output_s.data_ptr()),
+            total_scale_groups,
+            static_cast<int64_t>(num_tokens),
+            static_cast<int>(hidden_dim_groups),
+            static_cast<int>(num_outer_groups),
+            static_cast<int64_t>(input_stride_t));
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope.cuh

@@ -0,0 +1,254 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.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/compress.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <type_traits>
+
+namespace {
+
+using Plan = device::compress::PrefillPlan;
+
+/// \brief common block size for memory-bound kernel
+constexpr uint32_t kBlockSize = 128;
+constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
+
+struct FusedNormRopeParams {
+  void* __restrict__ input;
+  const void* __restrict__ weight;
+  float eps;
+  uint32_t num_works;
+  const void* __restrict__ handle;
+  const float* __restrict__ freqs_cis;
+  uint32_t compress_ratio;
+};
+
+enum class ForwardMode {
+  CompressExtend = 0,
+  CompressDecode = 1,
+  DefaultForward = 2,
+};
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, ForwardMode kMode, bool kUsePDL>
+__global__ void fused_norm_rope(const __grid_constant__ FusedNormRopeParams params) {
+  using namespace device;
+  using enum ForwardMode;
+
+  constexpr int64_t kMaxVecSize = 16 / sizeof(DType);
+  constexpr int64_t kVecSize = std::min(kMaxVecSize, kHeadDim / kWarpThreads);
+  constexpr int64_t kLocalSize = kHeadDim / (kWarpThreads * kVecSize);
+  constexpr int64_t kRopeVecSize = kRopeDim / (kWarpThreads * 2);
+  constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
+  static_assert(kHeadDim % (kWarpThreads * kVecSize) == 0);
+  static_assert(kLocalSize * kVecSize * kWarpThreads == kHeadDim);
+  static_assert(kRopeDim % (kWarpThreads * 2) == 0);
+  static_assert(kRopeDim % (kVecSize * kLocalSize) == 0);
+  static_assert(kRopeSize <= kWarpThreads);
+  static_assert(kRopeVecSize == 1, "only support rope dim = 64");
+
+  const auto& [
+    _input, _weight, eps, num_works, // norm
+    handle, freqs_cis, compress_ratio // rope
+  ] = params;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto work_id = blockIdx.x * kNumWarps + warp_id;
+
+  if (work_id >= num_works) return;
+
+  DType* input;
+  int32_t position;
+  if constexpr (kMode == CompressExtend) {
+    const auto plan = static_cast<const Plan*>(handle)[work_id];
+    input = static_cast<DType*>(_input) + plan.ragged_id * kHeadDim;
+    position = plan.position + 1 - compress_ratio;
+    if (plan.ragged_id == 0xFFFFFFFF) [[unlikely]]
+      return;
+  } else if constexpr (kMode == CompressDecode) {
+    input = static_cast<DType*>(_input) + work_id * kHeadDim;
+    const auto seq_len = static_cast<const int32_t*>(handle)[work_id];
+    if (seq_len % compress_ratio != 0) return;
+    position = seq_len - compress_ratio;
+  } else if constexpr (kMode == DefaultForward) {
+    input = static_cast<DType*>(_input) + work_id * kHeadDim;
+    position = static_cast<const int64_t*>(handle)[work_id];
+  } else {
+    static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
+  }
+
+  using Storage = AlignedVector<DType, kVecSize>;
+  __shared__ Storage s_rope_input[kNumWarps][kRopeSize];
+
+  // prefetch freq
+  const auto mem_freq = tile::Memory<fp32x2_t>::warp();
+  const auto freq = mem_freq.load(freqs_cis + position * kRopeDim);
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // part 1: norm
+  {
+    const auto gmem = tile::Memory<Storage>::warp();
+    Storage input_vec[kLocalSize];
+    Storage weight_vec[kLocalSize];
+#pragma unroll
+    for (int i = 0; i < kLocalSize; ++i) {
+      input_vec[i] = gmem.load(input, i);
+    }
+
+#pragma unroll
+    for (int i = 0; i < kLocalSize; ++i) {
+      weight_vec[i] = gmem.load(_weight, i);
+    }
+
+    float sum_of_squares = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kLocalSize; ++i) {
+#pragma unroll
+      for (int j = 0; j < kVecSize; ++j) {
+        const auto fp32_input = cast<float>(input_vec[i][j]);
+        sum_of_squares += fp32_input * fp32_input;
+      }
+    }
+
+    sum_of_squares = warp::reduce_sum(sum_of_squares);
+    const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + eps);
+
+#pragma unroll
+    for (int i = 0; i < kLocalSize; ++i) {
+#pragma unroll
+      for (int j = 0; j < kVecSize; ++j) {
+        const auto fp32_input = cast<float>(input_vec[i][j]);
+        const auto fp32_weight = cast<float>(weight_vec[i][j]);
+        input_vec[i][j] = cast<DType>(fp32_input * norm_factor * fp32_weight);
+      }
+    }
+
+    const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
+
+#pragma unroll
+    for (int i = 0; i < kLocalSize; ++i) {
+      if (i == kLocalSize - 1 && is_rope_lane) {
+        const auto rope_id = lane_id - (kWarpThreads - kRopeSize);
+        s_rope_input[warp_id][rope_id] = input_vec[i];
+      } else {
+        gmem.store(input, input_vec[i], i);
+      }
+    }
+
+    __syncwarp();
+  }
+
+  // part 2: rope
+  {
+    // mem elem = DType x 2
+    using DTypex2_t = packed_t<DType>;
+    const auto mem_elem = tile::Memory<DTypex2_t>::warp();
+    const auto elem = mem_elem.load(s_rope_input[warp_id]);
+    const auto [x_real, x_imag] = cast<fp32x2_t>(elem);
+    const auto [freq_real, freq_imag] = freq;
+    const fp32x2_t output = {
+        x_real * freq_real - x_imag * freq_imag,
+        x_real * freq_imag + x_imag * freq_real,
+    };
+    mem_elem.store(input + (kHeadDim - kRopeDim), cast<DTypex2_t>(output));
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, bool kUsePDL>
+struct FusedNormRopeKernel {
+  template <ForwardMode kMode>
+  static constexpr auto fused_kernel = fused_norm_rope<DType, kHeadDim, kRopeDim, kMode, kUsePDL>;
+
+  static void forward(
+      const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView handle,
+      const tvm::ffi::TensorView freqs_cis,
+      int32_t _mode,
+      float eps,
+      uint32_t compress_ratio) {
+    using namespace host;
+    using enum ForwardMode;
+
+    const auto mode = static_cast<ForwardMode>(_mode);
+
+    auto B = SymbolicSize{"num_q_tokens"};
+    auto N = SymbolicSize{"num_compress_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, kHeadDim})  // input
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(input);
+    TensorMatcher({kHeadDim})  // weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(weight);
+    TensorMatcher({-1, kRopeDim})  // freqs_cis
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(freqs_cis);
+    switch (mode) {
+      case CompressExtend:
+        TensorMatcher({N, compress::kPrefillPlanDim})  // plan
+            .with_dtype<compress::PrefillPlanTensorDtype>()
+            .with_device(device_)
+            .verify(handle);
+        RuntimeCheck(compress_ratio > 0);
+        break;
+      case CompressDecode:
+        TensorMatcher({N})  // seq_len
+            .with_dtype<int32_t>()
+            .with_device(device_)
+            .verify(handle);
+        RuntimeCheck(compress_ratio > 0);
+        break;
+      case DefaultForward:
+        TensorMatcher({N})  // position
+            .with_dtype<int64_t>()
+            .with_device(device_)
+            .verify(handle);
+        RuntimeCheck(compress_ratio == 0);
+        break;
+      default:
+        Panic("unsupported forward mode: ", static_cast<int>(mode));
+    }
+
+    // launch kernel
+    const auto num_compress_tokens = static_cast<uint32_t>(N.unwrap());
+    if (num_compress_tokens == 0) return;
+    const auto params = FusedNormRopeParams{
+        .input = input.data_ptr(),
+        .weight = weight.data_ptr(),
+        .eps = eps,
+        .num_works = num_compress_tokens,
+        .handle = handle.data_ptr(),
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .compress_ratio = compress_ratio,
+    };
+    const auto num_blocks = div_ceil(num_compress_tokens, kNumWarps);
+    using KernelType = std::decay_t<decltype(fused_norm_rope<DType, kHeadDim, kRopeDim, CompressExtend, kUsePDL>)>;
+    static constexpr KernelType kernel_table[3] = {
+        [static_cast<int>(CompressExtend)] = fused_kernel<CompressExtend>,
+        [static_cast<int>(CompressDecode)] = fused_kernel<CompressDecode>,
+        [static_cast<int>(DefaultForward)] = fused_kernel<DefaultForward>,
+    };
+    const auto kernel = kernel_table[static_cast<int>(mode)];
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/fused_norm_rope_v2.cuh

@@ -0,0 +1,663 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/tile.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/compress_v2.cuh>
+#include <sgl_kernel/deepseek_v4/fp8_utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+using PlanC = device::compress::CompressPlan;
+using PlanD = device::compress::DecodePlan;
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::inv_scale_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+SGL_DEVICE uint8_t quant_fp4_e2m1(float x) {
+  const float ax = fminf(fabsf(x), 6.0f);
+  uint8_t idx = 0;
+  idx += ax > 0.25f;
+  idx += ax > 0.75f;
+  idx += ax > 1.25f;
+  idx += ax > 1.75f;
+  idx += ax > 2.5f;
+  idx += ax > 3.5f;
+  idx += ax > 5.0f;
+  if (x < 0.0f && idx != 0) idx |= 0x8;
+  return idx;
+}
+
+constexpr uint32_t kBlockSize = 256;
+constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
+
+struct FusedNormRopeStoreParams {
+  void* __restrict__ input;
+  const void* __restrict__ handle;  // plan decode / compress
+  const void* __restrict__ weight;
+  const float* __restrict__ freqs_cis;
+  const int64_t* __restrict__ out_loc;
+  uint8_t* __restrict__ kvcache;
+  float eps;
+  uint32_t compress_ratio;
+  uint32_t num_tokens;
+};
+
+enum class ForwardMode : bool {
+  CompressExtend = 0,
+  CompressDecode = 1,
+};
+
+#define INDEXER_KERNEL __global__ __launch_bounds__(kBlockSize, 8)
+#define FLASHMLA_KERNEL __global__ __launch_bounds__(kBlockSize, 8)
+
+// ----------------------------------------------------------------------------
+// Indexer variant: kHeadDim = 128, 1 token per *warp* (8 tokens per block).
+// Each warp's 32 lanes cover the full 128-elem head_dim (kVecSize = 4 each).
+// Cache layout: 132 bytes/token (128 fp8 nope + 4 fp32 scale).
+// ----------------------------------------------------------------------------
+template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
+INDEXER_KERNEL void fused_norm_rope_indexer(const __grid_constant__ FusedNormRopeStoreParams params) {
+  using namespace device;
+  using enum ForwardMode;
+
+  constexpr int64_t kHeadDim = 128;
+  constexpr int64_t kRopeDim = 64;
+  constexpr int64_t kVecSize = 4;
+  constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
+  constexpr int64_t kPageBytes = 132ll << kPageBits;
+  static_assert(kHeadDim == kWarpThreads * kVecSize);
+  static_assert(kRopeDim == kWarpThreads * 2);
+  static_assert(kRopeSize <= kWarpThreads);
+  using Storage = AlignedVector<DType, kVecSize>;
+  using Float4 = AlignedVector<float, kVecSize>;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto work_id = blockIdx.x * kNumWarps + warp_id;
+  // Lanes whose 4-elem pack lies in the rope tail (= last `kRopeSize` packs).
+  const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
+
+  if (work_id >= params.num_tokens) return;
+
+  const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
+  int32_t position;
+  int64_t out_loc;
+  if constexpr (kMode == CompressExtend) {
+    const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
+    if (plan.is_invalid()) return;
+    position = plan.seq_len - params.compress_ratio;
+    out_loc = params.out_loc[plan.ragged_id];
+  } else if constexpr (kMode == CompressDecode) {
+    const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
+    if (plan.seq_len % params.compress_ratio != 0) return;
+    position = plan.seq_len - params.compress_ratio;
+    out_loc = params.out_loc[work_id];
+  } else {
+    static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
+  }
+  const auto freqs_cis = params.freqs_cis + position * kRopeDim;
+
+  PDLWaitPrimary<kUsePDL>();
+  Float4 data, freq;
+
+  // part 1: norm
+  {
+    Storage input_vec, weight_vec;
+    input_vec.load(input, lane_id);
+    weight_vec.load(params.weight, lane_id);
+    if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
+
+    float sum_of_squares = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto fp32_input = cast<float>(input_vec[i]);
+      sum_of_squares += fp32_input * fp32_input;
+    }
+
+    sum_of_squares = warp::reduce_sum(sum_of_squares);
+    const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
+
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto fp32_input = cast<float>(input_vec[i]);
+      const auto fp32_weight = cast<float>(weight_vec[i]);
+      data[i] = fp32_input * norm_factor * fp32_weight;
+    }
+  }
+
+  // part 2: rope (rope-lane only, 4 elems per lane = 2 (real, imag) pairs)
+  if (is_rope_lane) {
+    const auto x_real = data[0];
+    const auto x_imag = data[1];
+    const auto y_real = data[2];
+    const auto y_imag = data[3];
+    const auto freq_x_real = freq[0];
+    const auto freq_x_imag = freq[1];
+    const auto freq_y_real = freq[2];
+    const auto freq_y_imag = freq[3];
+    data[0] = x_real * freq_x_real - x_imag * freq_x_imag;
+    data[1] = x_real * freq_x_imag + x_imag * freq_x_real;
+    data[2] = y_real * freq_y_real - y_imag * freq_y_imag;
+    data[3] = y_real * freq_y_imag + y_imag * freq_y_real;
+  }
+
+  // part 3: hadamard transform
+  {
+    // Stage 1: butterfly (data[0], data[1]) and (data[2], data[3]).
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a1;
+      data[1] = a0 - a1;
+      data[2] = a2 + a3;
+      data[3] = a2 - a3;
+    }
+    // Stage 2: butterfly (data[0], data[2]) and (data[1], data[3]).
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a2;
+      data[1] = a1 + a3;
+      data[2] = a0 - a2;
+      data[3] = a1 - a3;
+    }
+    // Stages 3..7: cross-lane butterflies. Lower-lane (mask bit clear) keeps
+    // the sum, upper-lane (mask bit set) keeps the difference. shfl_xor is
+    // unsynchronized across early-returned lanes, but invalid-plan returns
+    // happen above for *all* lanes of a warp (work_id is warp-uniform), so
+    // the warp is intact here.
+#pragma unroll
+    for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
+#pragma unroll
+      for (int i = 0; i < kVecSize; ++i) {
+#ifndef USE_ROCM
+        const float other = __shfl_xor_sync(kFullMask, data[i], mask, kWarpThreads);
+#else
+        const float other = __shfl_xor(data[i], mask, kWarpThreads);
+#endif
+        data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
+      }
+    }
+    const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i)
+      data[i] *= kHadamardScale;
+  }
+
+  // part 4: per-warp UE8M0 quant + store. The whole warp emits one fp8 group
+  // (= 128 elements) plus a single fp32 scale, matching the indexer cache
+  // layout (`fused_store_indexer_cache`).
+  {
+    using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
+    float local_max = math::abs(data[0]);
+#pragma unroll
+    for (int i = 1; i < kVecSize; ++i) {
+      local_max = math::max(local_max, math::abs(data[i]));
+    }
+    const auto abs_max = warp::reduce_max(local_max);
+    const auto scale = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
+    const auto inv_scale = 1.0f / scale;
+    const int64_t page = out_loc >> kPageBits;
+    const int64_t offset = out_loc & ((1 << kPageBits) - 1);
+    const auto page_ptr = params.kvcache + page * kPageBytes;
+    const auto value_ptr = page_ptr + offset * 128;
+    const auto scale_ptr = page_ptr + (128 << kPageBits) + offset * 4;
+    OutStorage result;
+    result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
+    result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
+    PDLTriggerSecondary<kUsePDL>();
+    result.store(value_ptr, lane_id);
+    // The single fp32 scale is identical across all lanes -- write from any lane.
+    if (lane_id == 0) reinterpret_cast<float*>(scale_ptr)[0] = scale;
+  }
+}
+
+template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL>
+INDEXER_KERNEL void fused_norm_rope_indexer_fp4(const __grid_constant__ FusedNormRopeStoreParams params) {
+  using namespace device;
+  using enum ForwardMode;
+
+  constexpr int64_t kHeadDim = 128;
+  constexpr int64_t kRopeDim = 64;
+  constexpr int64_t kVecSize = 4;
+  constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
+  constexpr int64_t kPageBytes = 68ll << kPageBits;
+  static_assert(kHeadDim == kWarpThreads * kVecSize);
+  static_assert(kRopeDim == kWarpThreads * 2);
+  static_assert(kRopeSize <= kWarpThreads);
+  using Storage = AlignedVector<DType, kVecSize>;
+  using Float4 = AlignedVector<float, kVecSize>;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto work_id = blockIdx.x * kNumWarps + warp_id;
+  const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
+
+  if (work_id >= params.num_tokens) return;
+
+  const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
+  int32_t position;
+  int64_t out_loc;
+  if constexpr (kMode == CompressExtend) {
+    const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
+    if (plan.is_invalid()) return;
+    position = plan.seq_len - params.compress_ratio;
+    out_loc = params.out_loc[plan.ragged_id];
+  } else if constexpr (kMode == CompressDecode) {
+    const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
+    if (plan.seq_len % params.compress_ratio != 0) return;
+    position = plan.seq_len - params.compress_ratio;
+    out_loc = params.out_loc[work_id];
+  } else {
+    static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
+  }
+  const auto freqs_cis = params.freqs_cis + position * kRopeDim;
+
+  PDLWaitPrimary<kUsePDL>();
+  Float4 data, freq;
+
+  {
+    Storage input_vec, weight_vec;
+    input_vec.load(input, lane_id);
+    weight_vec.load(params.weight, lane_id);
+    if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
+
+    float sum_of_squares = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto fp32_input = cast<float>(input_vec[i]);
+      sum_of_squares += fp32_input * fp32_input;
+    }
+
+    sum_of_squares = warp::reduce_sum(sum_of_squares);
+    const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
+
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto fp32_input = cast<float>(input_vec[i]);
+      const auto fp32_weight = cast<float>(weight_vec[i]);
+      data[i] = fp32_input * norm_factor * fp32_weight;
+    }
+  }
+
+  if (is_rope_lane) {
+    const auto x_real = data[0];
+    const auto x_imag = data[1];
+    const auto y_real = data[2];
+    const auto y_imag = data[3];
+    const auto freq_x_real = freq[0];
+    const auto freq_x_imag = freq[1];
+    const auto freq_y_real = freq[2];
+    const auto freq_y_imag = freq[3];
+    data[0] = x_real * freq_x_real - x_imag * freq_x_imag;
+    data[1] = x_real * freq_x_imag + x_imag * freq_x_real;
+    data[2] = y_real * freq_y_real - y_imag * freq_y_imag;
+    data[3] = y_real * freq_y_imag + y_imag * freq_y_real;
+  }
+
+  {
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a1;
+      data[1] = a0 - a1;
+      data[2] = a2 + a3;
+      data[3] = a2 - a3;
+    }
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a2;
+      data[1] = a1 + a3;
+      data[2] = a0 - a2;
+      data[3] = a1 - a3;
+    }
+#pragma unroll
+    for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
+#pragma unroll
+      for (int i = 0; i < kVecSize; ++i) {
+#ifndef USE_ROCM
+        const float other = __shfl_xor_sync(kFullMask, data[i], mask, kWarpThreads);
+#else
+        const float other = __shfl_xor(data[i], mask, kWarpThreads);
+#endif
+        data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
+      }
+    }
+    const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i)
+      data[i] *= kHadamardScale;
+  }
+
+  {
+    float local_max = math::abs(data[0]);
+#pragma unroll
+    for (int i = 1; i < kVecSize; ++i) {
+      local_max = math::max(local_max, math::abs(data[i]));
+    }
+    local_max = warp::reduce_max<8>(local_max);
+
+    const auto scale_raw = fmaxf(1e-4f, local_max) / 6.0f;
+    const auto scale_ue8m0 = static_cast<uint8_t>(cast_to_ue8m0(scale_raw));
+    const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
+
+    const uint8_t packed0 = quant_fp4_e2m1(data[0] * inv_scale) | (quant_fp4_e2m1(data[1] * inv_scale) << 4);
+    const uint8_t packed1 = quant_fp4_e2m1(data[2] * inv_scale) | (quant_fp4_e2m1(data[3] * inv_scale) << 4);
+    const uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8);
+
+    const int64_t page = out_loc >> kPageBits;
+    const int64_t offset = out_loc & ((1 << kPageBits) - 1);
+    const auto page_ptr = params.kvcache + page * kPageBytes;
+    const auto value_ptr = page_ptr + offset * 64;
+    const auto scale_ptr = page_ptr + (64 << kPageBits) + offset * 4;
+
+    PDLTriggerSecondary<kUsePDL>();
+    reinterpret_cast<uint16_t*>(value_ptr)[lane_id] = packed;
+    if ((lane_id & 7) == 0) static_cast<uint8_t*>(scale_ptr)[lane_id >> 3] = scale_ue8m0;
+  }
+}
+
+// ----------------------------------------------------------------------------
+// FlashMLA variant: kHeadDim = 512, 1 token per *block* (256 threads).
+// Each thread loads kVecSize=2 BF16, so 256 threads cover the full 512 elems.
+// Cache layout: 584 bytes/token = 448 fp8 nope + 64 (=32 bf16x2) rope + 8 scale.
+// ----------------------------------------------------------------------------
+template <typename DType, ForwardMode kMode, int32_t kPageBits, bool kUsePDL, bool kBf16Store = false>
+FLASHMLA_KERNEL void fused_norm_rope_flashmla(const __grid_constant__ FusedNormRopeStoreParams params) {
+  using namespace device;
+  using enum ForwardMode;
+
+  constexpr int64_t kHeadDim = 512;
+  constexpr int64_t kRopeDim = 64;
+  constexpr int64_t kVecSize = 2;
+  // Last warp owns the rope tail. The remaining 7 warps each emit one
+  // 64-element fp8 group (own UE8M0 scale).
+  constexpr uint32_t kRopeWarp = kNumWarps - 1;
+  // kBf16Store: write the whole head_dim as plain BF16 (no fp8 / no scale) into a
+  // [num_slots, head_dim] bf16 cache (page_size==1) at row out_loc
+  constexpr int64_t kPageBytes =
+      kBf16Store ? ((kHeadDim * 2ll) << kPageBits) : host::div_ceil(584ll << kPageBits, 576) * 576;
+  static_assert(kHeadDim == kBlockSize * kVecSize);
+  static_assert(kRopeDim == kWarpThreads * kVecSize);
+  static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
+  using Storage = AlignedVector<DType, kVecSize>;
+  using Float2 = AlignedVector<float, kVecSize>;
+
+  const auto tx = threadIdx.x;
+  const auto warp_id = tx / kWarpThreads;
+  const auto lane_id = tx % kWarpThreads;
+  const auto work_id = blockIdx.x;
+
+  if (work_id >= params.num_tokens) return;
+
+  const auto input = static_cast<DType*>(params.input) + work_id * kHeadDim;
+  int32_t position;
+  int64_t out_loc;
+  if constexpr (kMode == CompressExtend) {
+    const auto plan = static_cast<const PlanC*>(params.handle)[work_id];
+    if (plan.is_invalid()) return;
+    position = plan.seq_len - params.compress_ratio;
+    out_loc = params.out_loc[plan.ragged_id];
+  } else if constexpr (kMode == CompressDecode) {
+    const auto plan = static_cast<const PlanD*>(params.handle)[work_id];
+    if (plan.seq_len % params.compress_ratio != 0) return;
+    position = plan.seq_len - params.compress_ratio;
+    out_loc = params.out_loc[work_id];
+  } else {
+    static_assert(host::dependent_false_v<DType>, "Unsupported Mode");
+  }
+  const auto freqs_cis = params.freqs_cis + position * kRopeDim;
+
+  PDLWaitPrimary<kUsePDL>();
+  Float2 data, freq;
+
+  // part 1: norm. Each thread owns one 2-elem pack (`tx`-th pack of input).
+  // Sum of squares is reduced across the whole block via per-warp partials.
+  {
+    __shared__ float partial_sums[kNumWarps];
+
+    Storage input_vec, weight_vec;
+    input_vec.load(input, tx);
+    weight_vec.load(params.weight, tx);
+    if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
+
+    float sum_of_squares = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto fp32_input = cast<float>(input_vec[i]);
+      sum_of_squares += fp32_input * fp32_input;
+    }
+
+    const auto warp_sum = warp::reduce_sum(sum_of_squares);
+    if (lane_id == 0) partial_sums[warp_id] = warp_sum;
+    __syncthreads();
+    // Replicate the per-warp partial sums to a full warp and reduce. Every
+    // lane-group of `kNumWarps` lanes ends up with the global sum.
+    sum_of_squares = warp::reduce_sum<kNumWarps>(partial_sums[lane_id % kNumWarps]);
+    const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
+
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto fp32_input = cast<float>(input_vec[i]);
+      const auto fp32_weight = cast<float>(weight_vec[i]);
+      data[i] = fp32_input * norm_factor * fp32_weight;
+    }
+  }
+
+  const int64_t page = out_loc >> kPageBits;
+  const int64_t offset = out_loc & ((1 << kPageBits) - 1);
+  const auto page_ptr = params.kvcache + page * kPageBytes;
+  const auto value_ptr = page_ptr + offset * (kBf16Store ? (kHeadDim * 2) : 576);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // part 2: rope on the rope warp (BF16 store), or per-warp FP8 quant + store.
+  if constexpr (kBf16Store) {
+    Float2 d = data;
+    if (warp_id == kRopeWarp) {
+      const auto x_real = data[0];
+      const auto x_imag = data[1];
+      const auto freq_real = freq[0];
+      const auto freq_imag = freq[1];
+      d[0] = x_real * freq_real - x_imag * freq_imag;
+      d[1] = x_real * freq_imag + x_imag * freq_real;
+    }
+    reinterpret_cast<bf16x2_t*>(value_ptr)[tx] = cast<bf16x2_t>(fp32x2_t{d[0], d[1]});
+  } else if (warp_id == kRopeWarp) {
+    // Each rope-warp lane owns exactly one (real, imag) pair within the rope
+    // tail. Apply rotation, downcast to BF16, write to the slot's rope region.
+    const auto x_real = data[0];
+    const auto x_imag = data[1];
+    const auto freq_real = freq[0];
+    const auto freq_imag = freq[1];
+    data[0] = x_real * freq_real - x_imag * freq_imag;
+    data[1] = x_real * freq_imag + x_imag * freq_real;
+    const auto result = cast<bf16x2_t>(fp32x2_t{data[0], data[1]});
+    const auto rope_ptr = value_ptr + 448;
+    reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
+  } else {
+    // Non-rope warp: per-warp UE8M0 group (64 elems -> 64 fp8 + 1 scale byte).
+    // BF16 round-trip to match the precision of the non-fused path
+    // (which goes through quant_to_nope_fp8_rope_bf16_pack_triton with bf16 input).
+    const auto x = cast<float>(cast<bf16_t>(data[0]));
+    const auto y = cast<float>(cast<bf16_t>(data[1]));
+    const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
+    const auto scale_raw = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
+    const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
+    const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
+    const auto result = pack_fp8(x * inv_scale, y * inv_scale);
+    const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
+    reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
+    // All lanes in this warp produce the same scale byte; let lane 0 publish.
+    if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
+  }
+}
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL, bool kBf16Store = false>
+struct FusedNormRopeKernel {
+  static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
+  static constexpr bool kIsIndexer = (kHeadDim == 128);
+  static_assert(!(kIsIndexer && kBf16Store), "bf16 store only for flashmla head_dim=512");
+  static constexpr int64_t kIndexerBytes = 132 * kPageSize;
+  static constexpr int64_t kFlashMLABytes = host::div_ceil(584 * kPageSize, 576) * 576;
+  static constexpr int64_t kBf16Bytes = kHeadDim * 2 * kPageSize;  // plain bf16 cache
+  static constexpr int64_t kPageBytes = kBf16Store ? kBf16Bytes : (kIsIndexer ? kIndexerBytes : kFlashMLABytes);
+
+  /// TODO: Let's fix the config for now.
+  static_assert(kRopeDim == 64 && (kHeadDim == 128 || kHeadDim == 512));
+  static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
+
+  template <ForwardMode kMode>
+  static constexpr auto select_kernel() {
+    if constexpr (kIsIndexer) {
+      return fused_norm_rope_indexer<DType, kMode, kLogPageSize, kUsePDL>;
+    } else {
+      return fused_norm_rope_flashmla<DType, kMode, kLogPageSize, kUsePDL, kBf16Store>;
+    }
+  }
+
+  template <ForwardMode kMode>
+  static constexpr auto select_fp4_kernel() {
+    static_assert(kIsIndexer, "FP4 fused store is only defined for the indexer");
+    return fused_norm_rope_indexer_fp4<DType, kMode, kLogPageSize, kUsePDL>;
+  }
+
+  static void forward(
+      const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView plan,
+      const tvm::ffi::TensorView weight,
+      const float eps,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView out_loc,
+      const tvm::ffi::TensorView kvcache,
+      const bool is_decode,
+      const uint32_t compress_ratio) {
+    using namespace host;
+    using enum ForwardMode;
+
+    const auto mode = static_cast<ForwardMode>(is_decode);
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLGPU>();
+
+    TensorMatcher({N, kHeadDim})  // input
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(input);
+    TensorMatcher({kHeadDim})  // weight
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(weight);
+    TensorMatcher({-1, kRopeDim})  // freqs_cis
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(freqs_cis);
+    TensorMatcher({-1})  // out_loc
+        .with_dtype<int64_t>()
+        .with_device(device_)
+        .verify(out_loc);
+    TensorMatcher({-1, -1})  // cache
+        .with_strides({kPageBytes, 1})
+        .with_dtype<uint8_t>()
+        .with_device(device_)
+        .verify(kvcache);
+
+    switch (mode) {
+      case CompressExtend:
+        compress::verify_plan_c(plan, N, device_);
+        RuntimeCheck(out_loc.size(0) >= N.unwrap());
+        break;
+      case CompressDecode:
+        compress::verify_plan_d(plan, N, device_);
+        RuntimeCheck(out_loc.size(0) == N.unwrap());
+        break;
+    }
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    if (num_tokens == 0) return;
+    const auto params = FusedNormRopeStoreParams{
+        .input = input.data_ptr(),
+        .handle = plan.data_ptr(),
+        .weight = weight.data_ptr(),
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
+        .kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
+        .eps = eps,
+        .compress_ratio = compress_ratio,
+        .num_tokens = num_tokens,
+    };
+    // Indexer packs `kNumWarps` tokens per block (warp-major); FlashMLA uses
+    // a whole block per token (cta-major sum-reduce over head_dim=512).
+    const uint32_t num_blocks = kIsIndexer ? div_ceil(num_tokens, kNumWarps) : num_tokens;
+    const auto device = device_.unwrap();
+    const auto kernel = mode == CompressExtend ? select_kernel<CompressExtend>() : select_kernel<CompressDecode>();
+    LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+  }
+
+  static void forward_fp4(
+      const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView plan,
+      const tvm::ffi::TensorView weight,
+      const float eps,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView out_loc,
+      const tvm::ffi::TensorView kvcache,
+      const bool is_decode,
+      const uint32_t compress_ratio) {
+    using namespace host;
+    using enum ForwardMode;
+
+    static_assert(kIsIndexer, "FP4 fused store is only defined for the indexer");
+    constexpr int64_t kFp4PageBytes = 68 * kPageSize;
+    const auto mode = static_cast<ForwardMode>(is_decode);
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({N, kHeadDim}).with_dtype<DType>().with_device(device_).verify(input);
+    TensorMatcher({kHeadDim}).with_dtype<DType>().with_device(device_).verify(weight);
+    TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis);
+    TensorMatcher({-1}).with_dtype<int64_t>().with_device(device_).verify(out_loc);
+    TensorMatcher({-1, -1}).with_strides({kFp4PageBytes, 1}).with_dtype<uint8_t>().with_device(device_).verify(kvcache);
+
+    switch (mode) {
+      case CompressExtend:
+        compress::verify_plan_c(plan, N, device_);
+        RuntimeCheck(out_loc.size(0) >= N.unwrap());
+        break;
+      case CompressDecode:
+        compress::verify_plan_d(plan, N, device_);
+        RuntimeCheck(out_loc.size(0) == N.unwrap());
+        break;
+    }
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    if (num_tokens == 0) return;
+    const auto params = FusedNormRopeStoreParams{
+        .input = input.data_ptr(),
+        .handle = plan.data_ptr(),
+        .weight = weight.data_ptr(),
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .out_loc = static_cast<const int64_t*>(out_loc.data_ptr()),
+        .kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
+        .eps = eps,
+        .compress_ratio = compress_ratio,
+        .num_tokens = num_tokens,
+    };
+    const uint32_t num_blocks = div_ceil(num_tokens, kNumWarps);
+    const auto device = device_.unwrap();
+    const auto kernel =
+        mode == CompressExtend ? select_fp4_kernel<CompressExtend>() : select_fp4_kernel<CompressDecode>();
+    LaunchKernel(num_blocks, kBlockSize, device).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/hash_topk.cuh

@@ -0,0 +1,214 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cmath>
+#include <cstdint>
+
+namespace {
+
+[[maybe_unused]]
+SGL_DEVICE float act_sqrt_softplus(float x) {
+  const float softplus = fmaxf(x, 0.0f) + log1pf(expf(-fabsf(x)));
+  return sqrtf(softplus);
+}
+
+struct MoEHashTopKParams {
+  const float* __restrict__ router_logits;
+  const int64_t* __restrict__ input_id;
+  const int32_t* __restrict__ tid2eid;
+  int32_t* __restrict__ topk_ids;
+  float* __restrict__ topk_weights;
+  uint32_t num_tokens;
+  uint32_t topk;
+  uint32_t num_routed_experts;
+  uint32_t num_shared_experts;
+  float routed_scaling_factor;
+};
+
+template <auto Fn, bool kUsePDL>
+__global__ void moe_hash_topk_fused(const MoEHashTopKParams __grid_constant__ params) {
+  using namespace device;
+  const auto& [
+    router_logits, input_id, tid2eid, topk_ids, topk_weights, // pointers
+    num_tokens, topk, num_routed_experts, num_shared_experts, routed_scaling_factor] =
+      params;
+
+  const uint32_t topk_fused = topk + num_shared_experts;
+  const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t warp_id = tid / kWarpThreads;
+  const uint32_t lane_id = tid % kWarpThreads;
+  if (warp_id >= num_tokens) return;
+  // we can safely prefetch the token id
+  const auto token_id = input_id[warp_id];
+
+  PDLWaitPrimary<kUsePDL>();
+
+  float routed_weight = 0.0f;
+  int32_t expert_id = 0;
+  if (lane_id < topk) {
+    expert_id = tid2eid[token_id * topk + lane_id];
+    routed_weight = Fn(router_logits[warp_id * num_routed_experts + expert_id]);
+  }
+
+  const auto routed_sum = device::warp::reduce_sum(routed_weight);
+  if (lane_id < topk_fused) {
+    const bool is_shared = lane_id >= topk;
+    const auto output_offset = warp_id * topk_fused + lane_id;
+    topk_ids[output_offset] = is_shared ? num_routed_experts + lane_id - topk : expert_id;
+    topk_weights[output_offset] = is_shared ? 1.0f / routed_scaling_factor : routed_weight / routed_sum;
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+struct TopKParams {
+  int32_t* __restrict__ topk_ids;
+  // Exactly one is active: ntn_ptr == nullptr means use ntn_value.
+  const int32_t* __restrict__ ntn_ptr;
+  int32_t ntn_value;
+  int64_t stride;
+  uint32_t topk;
+  uint32_t num_tokens;
+};
+
+__global__ void mask_topk_ids_padded_region(const TopKParams __grid_constant__ params) {
+  const uint32_t tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const uint32_t warp_id = tid / device::kWarpThreads;
+  const uint32_t lane_id = tid % device::kWarpThreads;
+  if (warp_id >= params.num_tokens || lane_id >= params.topk) return;
+  device::PDLWaitPrimary<true>();
+  const uint32_t num = (params.ntn_ptr != nullptr)  //
+                           ? static_cast<uint32_t>(params.ntn_ptr[0])
+                           : static_cast<uint32_t>(params.ntn_value);
+  if (warp_id >= num) params.topk_ids[warp_id * params.stride + lane_id] = -1;
+  device::PDLTriggerSecondary<true>();
+}
+
+template <auto Fn, bool kUsePDL>
+struct HashTopKKernel {
+  static constexpr auto kernel = moe_hash_topk_fused<Fn, kUsePDL>;
+
+  static void
+  run(const tvm::ffi::TensorView router_logits,
+      const tvm::ffi::TensorView input_id,
+      const tvm::ffi::TensorView tid2eid,
+      const tvm::ffi::TensorView topk_weights,
+      const tvm::ffi::TensorView topk_ids,
+      float routed_scaling_factor) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto E = SymbolicSize{"num_routed_experts"};
+    auto K = SymbolicSize{"topk_fused"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, E})  //
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(router_logits);
+    TensorMatcher({N})  //
+        .with_dtype<int64_t>()
+        .with_device(device)
+        .verify(input_id);
+    TensorMatcher({-1, -1})  //
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(tid2eid);
+    TensorMatcher({N, K})  //
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(topk_weights);
+    TensorMatcher({N, K})  //
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(topk_ids);
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto topk_fused = static_cast<uint32_t>(K.unwrap());
+    const auto topk = static_cast<uint32_t>(tid2eid.size(1));
+    const auto shared_experts = topk_fused - topk;
+    RuntimeCheck(topk <= topk_fused, "HashTopKKernel requires topk <= topk_fused");
+    RuntimeCheck(topk_fused <= device::kWarpThreads, "HashTopKKernel requires topk_fused <= warp size");
+
+    const auto params = MoEHashTopKParams{
+        .router_logits = static_cast<const float*>(router_logits.data_ptr()),
+        .input_id = static_cast<const int64_t*>(input_id.data_ptr()),
+        .tid2eid = static_cast<const int32_t*>(tid2eid.data_ptr()),
+        .topk_ids = static_cast<int32_t*>(topk_ids.data_ptr()),
+        .topk_weights = static_cast<float*>(topk_weights.data_ptr()),
+        .num_tokens = num_tokens,
+        .topk = topk,
+        .num_routed_experts = static_cast<uint32_t>(E.unwrap()),
+        .num_shared_experts = shared_experts,
+        .routed_scaling_factor = routed_scaling_factor,
+    };
+    const auto kBlockSize = 128u;
+    const auto kNumWarps = kBlockSize / device::kWarpThreads;
+    const auto num_blocks = div_ceil(num_tokens, kNumWarps);
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+// TODO this may not be related to *hash* topk, thus may move
+struct MaskKernel {
+  static constexpr auto kernel = mask_topk_ids_padded_region;
+
+  static void run(tvm::ffi::TensorView topk_ids, tvm::ffi::TensorView num_token_non_padded) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto K = SymbolicSize{"topk"};
+    auto D = SymbolicSize{"stride"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N, K})  //
+        .with_strides({D, 1})
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(topk_ids);
+    RuntimeCheck(num_token_non_padded.numel() == 1, "num_token_non_padded should be a scalar");
+    RuntimeCheck(K.unwrap() <= device::kWarpThreads, "MaskKernel requires topk <= warp size");
+    const int32_t* ntn_ptr = nullptr;
+    int32_t ntn_value = 0;
+    const auto ntn_dev = num_token_non_padded.device().device_type;
+    if (ntn_dev == kDLCUDA) {
+      RuntimeCheck(is_type<int32_t>(num_token_non_padded.dtype()), "num_token_non_padded on CUDA must be int32");
+      ntn_ptr = static_cast<const int32_t*>(num_token_non_padded.data_ptr());
+    } else if (ntn_dev == kDLCPU) {
+      if (is_type<int32_t>(num_token_non_padded.dtype())) {
+        ntn_value = *static_cast<const int32_t*>(num_token_non_padded.data_ptr());
+      } else if (is_type<int64_t>(num_token_non_padded.dtype())) {
+        ntn_value = static_cast<int32_t>(*static_cast<const int64_t*>(num_token_non_padded.data_ptr()));
+      } else {
+        RuntimeCheck(false, "num_token_non_padded on CPU must be int32 or int64");
+      }
+    } else {
+      RuntimeCheck(false, "num_token_non_padded must be on CPU or CUDA");
+    }
+
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = TopKParams{
+        .topk_ids = static_cast<int32_t*>(topk_ids.data_ptr()),
+        .ntn_ptr = ntn_ptr,
+        .ntn_value = ntn_value,
+        .stride = static_cast<int64_t>(D.unwrap()),
+        .topk = static_cast<uint32_t>(K.unwrap()),
+        .num_tokens = num_tokens,
+    };
+    const auto kBlockSize = 128u;
+    const auto kNumWarps = kBlockSize / device::kWarpThreads;
+    const auto num_blocks = div_ceil(num_tokens, kNumWarps);
+    LaunchKernel(num_blocks, kBlockSize, device.unwrap())  //
+        .enable_pdl(true)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/main_norm_rope.cuh

@@ -0,0 +1,882 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.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 <tvm/ffi/container/tensor.h>
+
+#include <bit>
+#include <cstdint>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::inv_scale_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+SGL_DEVICE uint8_t quant_fp4_e2m1(float x) {
+  const float ax = fminf(fabsf(x), 6.0f);
+  uint8_t idx = 0;
+  idx += ax > 0.25f;
+  idx += ax > 0.75f;
+  idx += ax > 1.25f;
+  idx += ax > 1.75f;
+  idx += ax > 2.5f;
+  idx += ax > 3.5f;
+  idx += ax > 5.0f;
+  if (x < 0.0f && idx != 0) idx |= 0x8;
+  return idx;
+}
+
+// 4 warps per block: warp-per-(token, head) work-item dispatch (Q kernel).
+constexpr uint32_t kFusedQBlockSize = 128;
+constexpr uint32_t kFusedQNumWarps = kFusedQBlockSize / device::kWarpThreads;
+
+// 8 warps per block: block-per-token work-item dispatch (K kernel).
+constexpr uint32_t kFusedKBlockSize = 256;
+constexpr uint32_t kFusedKNumWarps = kFusedKBlockSize / device::kWarpThreads;
+
+#define Q_KERNEL __global__ __launch_bounds__(kFusedQBlockSize, 16)
+#define K_KERNEL __global__ __launch_bounds__(kFusedKBlockSize, 8)
+
+template <int64_t kRopeDim>
+SGL_DEVICE device::AlignedVector<float, 4>
+load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id) {
+  constexpr int64_t kHalfRopeDim = kRopeDim / 2;
+  const int32_t pair0 = lane_id * 2;
+  const int32_t pair1 = pair0 + 1;
+  device::AlignedVector<float, 4> freq;
+  freq[0] = cos_sin_cache[pair0];
+  freq[1] = cos_sin_cache[kHalfRopeDim + pair0];
+  freq[2] = cos_sin_cache[pair1];
+  freq[3] = cos_sin_cache[kHalfRopeDim + pair1];
+  return freq;
+}
+
+// ============================================================================
+// Q kernel: warp-per-(token, head) rmsnorm-self + RoPE + write to q_out.
+// ============================================================================
+
+struct FusedQNormRopeParams {
+  const void* __restrict__ q_input;     // (B, num_q_heads, kHeadDim) DType
+  void* __restrict__ q_output;          // (B, num_q_heads, kHeadDim) DType
+  const float* __restrict__ freqs_cis;  // (max_pos, kRopeDim) fp32 (re/im interleaved)
+  const void* __restrict__ positions;   // (B,) PosT
+  int64_t q_input_stride_batch;
+  int64_t q_output_stride_batch;
+  uint32_t batch_size;
+  uint32_t num_q_heads;
+  float eps;
+};
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, bool kUsePDL>
+Q_KERNEL void fused_q_norm_rope(const __grid_constant__ FusedQNormRopeParams params) {
+  using namespace device;
+
+  constexpr int64_t kMaxVecSize = 16 / sizeof(DType);
+  constexpr int64_t kVecSize = std::min(kMaxVecSize, kHeadDim / kWarpThreads);
+  constexpr int64_t kLocalSize = kHeadDim / (kWarpThreads * kVecSize);
+  constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
+  static_assert(kHeadDim % (kWarpThreads * kVecSize) == 0);
+  static_assert(kLocalSize * kVecSize * kWarpThreads == kHeadDim);
+  static_assert(kRopeDim % kVecSize == 0);
+  static_assert(kRopeSize <= kWarpThreads);
+  static_assert(kRopeDim == kWarpThreads * 2, "1 (real, imag) pair per lane");
+
+  using Storage = AlignedVector<DType, kVecSize>;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
+
+  const uint32_t total_works = params.batch_size * params.num_q_heads;
+  if (work_id >= total_works) return;
+
+  const uint32_t batch_id = work_id / params.num_q_heads;
+  const uint32_t head_id = work_id % params.num_q_heads;
+  const auto input_ptr =
+      static_cast<const DType*>(params.q_input) + batch_id * params.q_input_stride_batch + head_id * kHeadDim;
+  const auto output_ptr =
+      static_cast<DType*>(params.q_output) + batch_id * params.q_output_stride_batch + head_id * kHeadDim;
+  const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
+
+  __shared__ Storage s_rope[kFusedQNumWarps][kRopeSize];
+
+  // Prefetch this lane's freq pair before the PDL gate so the wait happens
+  // outside the dependency chain on `position`.
+  const auto mem_freq = tile::Memory<fp32x2_t>{lane_id, kWarpThreads};
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // part 1: rmsnorm-self (no weight).
+  const auto gmem = tile::Memory<Storage>{lane_id, kWarpThreads};
+  Storage input_vec[kLocalSize];
+#pragma unroll
+  for (int i = 0; i < kLocalSize; ++i) {
+    input_vec[i] = gmem.load(input_ptr, i);
+  }
+
+  const auto freq = mem_freq.load(params.freqs_cis + position * kRopeDim);
+
+  float sum_of_squares = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kLocalSize; ++i) {
+#pragma unroll
+    for (int j = 0; j < kVecSize; ++j) {
+      const auto x = cast<float>(input_vec[i][j]);
+      sum_of_squares += x * x;
+    }
+  }
+  sum_of_squares = warp::reduce_sum(sum_of_squares);
+  const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
+
+#pragma unroll
+  for (int i = 0; i < kLocalSize; ++i) {
+#pragma unroll
+    for (int j = 0; j < kVecSize; ++j) {
+      const auto x = cast<float>(input_vec[i][j]);
+      input_vec[i][j] = cast<DType>(x * norm_factor);
+    }
+  }
+
+  // Stash the rope tail (last kRopeSize lanes' last tile) into shared memory;
+  // write nope tiles to gmem directly.
+  const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
+#pragma unroll
+  for (int i = 0; i < kLocalSize; ++i) {
+    if (i == kLocalSize - 1 && is_rope_lane) {
+      const auto rope_id = lane_id - (kWarpThreads - kRopeSize);
+      s_rope[warp_id][rope_id] = input_vec[i];
+    } else {
+      gmem.store(output_ptr, input_vec[i], i);
+    }
+  }
+  __syncwarp();
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // part 2: RoPE on all 32 lanes -- one (real, imag) bf16x2 pair per lane.
+  using DType2 = packed_t<DType>;
+  const auto mem_elem = tile::Memory<DType2>{lane_id, kWarpThreads};
+  const auto elem = mem_elem.load(s_rope[warp_id]);
+  const auto [x_real, x_imag] = cast<fp32x2_t>(elem);
+  const auto [freq_real, freq_imag] = freq;
+  const fp32x2_t rotated = {
+      x_real * freq_real - x_imag * freq_imag,
+      x_real * freq_imag + x_imag * freq_real,
+  };
+  mem_elem.store(output_ptr + (kHeadDim - kRopeDim), cast<DType2>(rotated));
+}
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, bool kUsePDL>
+struct FusedQNormRopeKernel {
+  template <typename PosT>
+  static constexpr auto kernel = fused_q_norm_rope<DType, kHeadDim, kRopeDim, PosT, kUsePDL>;
+
+  static void forward(
+      const tvm::ffi::TensorView q_input,
+      const tvm::ffi::TensorView q_output,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView positions,
+      float eps) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto H = SymbolicSize{"num_q_heads"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, H, kHeadDim})  //
+        .with_strides({-1, kHeadDim, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q_input);
+    TensorMatcher({B, H, kHeadDim})  //
+        .with_strides({-1, kHeadDim, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q_output);
+    TensorMatcher({-1, kRopeDim})  //
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(freqs_cis);
+    auto pos_dtype = SymbolicDType{};
+    TensorMatcher({B})  //
+        .with_dtype<int32_t, int64_t>(pos_dtype)
+        .with_device(device_)
+        .verify(positions);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto num_q_heads = static_cast<uint32_t>(H.unwrap());
+    if (batch_size == 0) return;
+
+    const auto params = FusedQNormRopeParams{
+        .q_input = q_input.data_ptr(),
+        .q_output = q_output.data_ptr(),
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .positions = positions.data_ptr(),
+        .q_input_stride_batch = q_input.stride(0),
+        .q_output_stride_batch = q_output.stride(0),
+        .batch_size = batch_size,
+        .num_q_heads = num_q_heads,
+        .eps = eps,
+    };
+    const auto total_works = batch_size * num_q_heads;
+    const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
+    const auto k_int32 = kernel<int32_t>;
+    const auto k_int64 = kernel<int64_t>;
+    const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
+    LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(k, params);
+  }
+};
+
+// ============================================================================
+// K kernel: block-per-token rmsnorm (with kv_weight) + RoPE + FlashMLA store.
+// ============================================================================
+
+struct FusedKNormRopeFlashMLAParams {
+  const void* __restrict__ kv;          // (B, kHeadDim) DType
+  const void* __restrict__ kv_weight;   // (kHeadDim,) DType
+  const float* __restrict__ freqs_cis;  // (max_pos, kRopeDim) fp32
+  const void* __restrict__ positions;   // (B,) PosT
+  const int32_t* __restrict__ out_loc;  // (B,) int32 -> cache slot id
+  uint8_t* __restrict__ kvcache;        // (npages, kPageBytes) uint8
+  // Row stride for `kv` in elements. Required because the upstream caller often
+  // passes `qkv_a[..., q_lora_rank:]`, a non-contiguous slice whose stride[0]
+  // equals `q_lora_rank + kHeadDim` rather than `kHeadDim`.
+  int64_t kv_stride_batch;
+  uint32_t batch_size;
+  float eps;
+};
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, typename PosT, int32_t kPageBits, bool kUsePDL>
+K_KERNEL void fused_k_norm_rope_flashmla(const __grid_constant__ FusedKNormRopeFlashMLAParams params) {
+  using namespace device;
+
+  constexpr int64_t kVecSize = 2;
+  constexpr uint32_t kRopeWarp = kFusedKNumWarps - 1;
+  constexpr int64_t kPageBytes = host::div_ceil(584ll << kPageBits, 576) * 576;
+  static_assert(kHeadDim == kFusedKBlockSize * kVecSize);
+  static_assert(kRopeDim == kWarpThreads * kVecSize);
+  static_assert(kHeadDim - kRopeDim == kRopeWarp * kWarpThreads * kVecSize);
+  using Storage = AlignedVector<DType, kVecSize>;
+  using Float2 = AlignedVector<float, kVecSize>;
+
+  const auto tx = threadIdx.x;
+  const auto warp_id = tx / kWarpThreads;
+  const auto lane_id = tx % kWarpThreads;
+  const auto work_id = blockIdx.x;
+  if (work_id >= params.batch_size) return;
+
+  const auto input_ptr = static_cast<const DType*>(params.kv) + work_id * params.kv_stride_batch;
+  const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[work_id]);
+  const auto out_loc = params.out_loc[work_id];
+  const auto freqs_cis = params.freqs_cis + position * kRopeDim;
+
+  PDLWaitPrimary<kUsePDL>();
+  Float2 data, freq;
+
+  // part 1: norm. Each thread owns one 2-elem pack (the `tx`-th).
+  // Sum-of-squares is reduced block-wide via per-warp partials.
+  {
+    __shared__ float partial_sums[kFusedKNumWarps];
+
+    Storage input_vec, weight_vec;
+    input_vec.load(input_ptr, tx);
+    weight_vec.load(params.kv_weight, tx);
+    if (warp_id == kRopeWarp) freq.load(freqs_cis, lane_id);
+
+    float sum_of_squares = 0.0f;
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto x = cast<float>(input_vec[i]);
+      sum_of_squares += x * x;
+    }
+    const auto warp_sum = warp::reduce_sum(sum_of_squares);
+    if (lane_id == 0) partial_sums[warp_id] = warp_sum;
+    __syncthreads();
+    // Replicate the per-warp partial sums onto all lanes of one warp and
+    // reduce. Every group of `kBlockItemNumWarps` lanes ends up with the
+    // global sum.
+    sum_of_squares = warp::reduce_sum<kFusedKNumWarps>(partial_sums[lane_id % kFusedKNumWarps]);
+    const auto norm_factor = math::rsqrt(sum_of_squares / kHeadDim + params.eps);
+
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      const auto x = cast<float>(input_vec[i]);
+      const auto w = cast<float>(weight_vec[i]);
+      data[i] = x * norm_factor * w;
+    }
+  }
+
+  // A negative out_loc marks a slot with no KV write target (e.g. the -1
+  // sentinel from the full->SWA translation for out-of-window tokens or
+  // padded rows); skip the row instead of writing out of bounds. Checked
+  // here, not at the load, so the out_loc prefetch overlaps the norm above.
+  if (out_loc < 0) return;
+
+  const int32_t page = out_loc >> kPageBits;
+  const int32_t offset = out_loc & ((1 << kPageBits) - 1);
+  const auto page_ptr = params.kvcache + page * kPageBytes;
+  const auto value_ptr = page_ptr + offset * 576;
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // part 2: rope on warp 7 (BF16 store), per-warp UE8M0 quant + store on warps 0..6.
+  if (warp_id == kRopeWarp) {
+    const auto x_real = data[0];
+    const auto x_imag = data[1];
+    const auto freq_real = freq[0];
+    const auto freq_imag = freq[1];
+    data[0] = x_real * freq_real - x_imag * freq_imag;
+    data[1] = x_real * freq_imag + x_imag * freq_real;
+    const auto result = cast<bf16x2_t>(fp32x2_t{data[0], data[1]});
+    const auto rope_ptr = value_ptr + 448;
+    reinterpret_cast<bf16x2_t*>(rope_ptr)[lane_id] = result;
+  } else {
+    const auto x = data[0];
+    const auto y = data[1];
+    const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
+    const auto scale_raw = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
+    const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
+    const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
+    const auto result = pack_fp8(x * inv_scale, y * inv_scale);
+    const auto scale_ptr = page_ptr + (576 << kPageBits) + offset * 8;
+    reinterpret_cast<fp8x2_e4m3_t*>(value_ptr)[tx] = result;
+    if (lane_id == 0) static_cast<uint8_t*>(scale_ptr)[warp_id] = scale_ue8m0;
+  }
+}
+
+template <typename DType, int64_t kHeadDim, int64_t kRopeDim, uint32_t kPageSize, bool kUsePDL>
+struct FusedKNormRopeFlashMLAKernel {
+  static constexpr int32_t kLogPageSize = std::countr_zero(kPageSize);
+  static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
+  static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
+  static_assert(1 << kLogPageSize == kPageSize);
+  static_assert(kHeadDim == 512 && kRopeDim == 64, "FlashMLA layout requires (512, 64)");
+
+  template <typename PosT>
+  static constexpr auto kernel = fused_k_norm_rope_flashmla<DType, kHeadDim, kRopeDim, PosT, kLogPageSize, kUsePDL>;
+
+  static void forward(
+      const tvm::ffi::TensorView kv,
+      const tvm::ffi::TensorView kv_weight,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView positions,
+      const tvm::ffi::TensorView out_loc,
+      const tvm::ffi::TensorView kvcache,
+      float eps) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, kHeadDim})  //
+        .with_strides({-1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(kv);
+    TensorMatcher({kHeadDim})  //
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(kv_weight);
+    TensorMatcher({-1, kRopeDim})  //
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(freqs_cis);
+    auto pos_dtype = SymbolicDType{};
+    TensorMatcher({B})  //
+        .with_dtype<int32_t, int64_t>(pos_dtype)
+        .with_device(device_)
+        .verify(positions);
+    TensorMatcher({B})  //
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(out_loc);
+    TensorMatcher({-1, -1})  //
+        .with_strides({kPageBytes, 1})
+        .with_dtype<uint8_t>()
+        .with_device(device_)
+        .verify(kvcache);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    if (batch_size == 0) return;
+
+    const auto params = FusedKNormRopeFlashMLAParams{
+        .kv = kv.data_ptr(),
+        .kv_weight = kv_weight.data_ptr(),
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .positions = positions.data_ptr(),
+        .out_loc = static_cast<const int32_t*>(out_loc.data_ptr()),
+        .kvcache = static_cast<uint8_t*>(kvcache.data_ptr()),
+        .kv_stride_batch = kv.stride(0),
+        .batch_size = batch_size,
+        .eps = eps,
+    };
+    const auto k_int32 = kernel<int32_t>;
+    const auto k_int64 = kernel<int64_t>;
+    const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
+    LaunchKernel(batch_size, kFusedKBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(k, params);
+  }
+};
+
+// ============================================================================
+// Indexer Q kernel: warp-per-(token, head) RoPE + Hadamard + fp8 act-quant.
+// ============================================================================
+
+struct FusedQIndexerRopeHadamardQuantParams {
+  const void* __restrict__ q_input;  // (B, num_heads, 128) DType
+  void* __restrict__ q_fp8;          // (B, num_heads, 128) fp8_e4m3
+  // weights_out[b, h] = weight[b, h] * weight_scale * q_scale[b, h].
+  // q_scale is computed internally and not exposed -- the only consumer of
+  // it is `weights_out`.
+  const void* __restrict__ weight;  // (B, num_heads) DType
+  float* __restrict__ weights_out;  // (B, num_heads) fp32 (== (B, H, 1) flat)
+  float weight_scale;               // scalar c4_indexer.weight_scale
+  // Template-dependent layout:
+  //   kRopeFirst=false: (max_pos, 64) fp32 interleaved [cos0, sin0, ...]
+  //   kRopeFirst=true : (max_pos, 64) fp32 halves [cos..., sin...]
+  const float* __restrict__ rope_cache;
+  const void* __restrict__ positions;  // (B,) PosT
+  // Row stride for `weight` (caller passes the non-contiguous wk slice directly).
+  int64_t weight_stride_batch;
+  uint32_t batch_size;
+  uint32_t num_heads;
+};
+
+template <typename DType, typename PosT, bool kUsePDL, bool kRopeFirst = false, bool kHadamard = true>
+Q_KERNEL void fused_q_indexer_rope_hadamard_quant(const __grid_constant__ FusedQIndexerRopeHadamardQuantParams params) {
+  using namespace device;
+
+  constexpr int64_t kHeadDim = 128;
+  constexpr int64_t kRopeDim = 64;
+  constexpr int64_t kVecSize = 4;
+  constexpr uint32_t kRopeSize = kRopeDim / kVecSize;  // = 16
+  static_assert(kHeadDim == kWarpThreads * kVecSize);
+  static_assert(kRopeDim == kWarpThreads * 2);
+  static_assert(kRopeSize <= kWarpThreads);
+
+  using Storage = AlignedVector<DType, kVecSize>;
+  using Float4 = AlignedVector<float, kVecSize>;
+  using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;  // 4 fp8 / lane
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
+  // V4 ropes the trailing kRopeDim dims (kRopeFirst=false); V3.2 ropes the
+  // leading kRopeDim dims (kRopeFirst=true). Select the owning lanes per layout.
+  const bool is_rope_lane = kRopeFirst ? (lane_id < kRopeSize) : (lane_id >= kWarpThreads - kRopeSize);
+
+  const uint32_t total_works = params.batch_size * params.num_heads;
+  if (work_id >= total_works) return;
+
+  const uint32_t batch_id = work_id / params.num_heads;
+  const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
+  const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
+  const auto rope_cache = params.rope_cache + position * kRopeDim;
+
+  // Lane 0 prefetches the weight scalar for this (token, head) work item.
+  // Weight is (B, num_heads) DType; we need one scalar per warp -- offload
+  // the load to lane 0 only. The multiply + store happens once the q_scale
+  // is known (part 4).
+
+  PDLWaitPrimary<kUsePDL>();
+  Float4 data, freq;
+  const uint32_t head_id = work_id - batch_id * params.num_heads;
+  const auto weight_val =
+      cast<float>(static_cast<const DType*>(params.weight)[batch_id * params.weight_stride_batch + head_id]);
+
+  // part 1: load (no norm). Each lane owns a 4-elem pack.
+  {
+    Storage input_vec;
+    input_vec.load(input_ptr, lane_id);
+    if (is_rope_lane) {
+      if constexpr (kRopeFirst) {
+        freq = load_rope_first_cos_sin<kRopeDim>(rope_cache, lane_id);
+      } else {
+        freq.load(rope_cache, lane_id - (kWarpThreads - kRopeSize));
+      }
+    }
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      data[i] = cast<float>(input_vec[i]);
+    }
+  }
+
+  // part 2: rope on rope lanes only (4 elems / lane = 2 (real, imag) pairs).
+  if (is_rope_lane) {
+    const auto x_real = data[0];
+    const auto x_imag = data[1];
+    const auto y_real = data[2];
+    const auto y_imag = data[3];
+    const auto fxr = freq[0];
+    const auto fxi = freq[1];
+    const auto fyr = freq[2];
+    const auto fyi = freq[3];
+    data[0] = x_real * fxr - x_imag * fxi;
+    data[1] = x_real * fxi + x_imag * fxr;
+    data[2] = y_real * fyr - y_imag * fyi;
+    data[3] = y_real * fyi + y_imag * fyr;
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // part 3: 128-point Hadamard (2 local stages + 5 cross-lane shfl_xor stages).
+  // Same recipe as `fused_norm_rope_indexer`; see comments there for the
+  // butterfly invariants and the early-return safety argument. V3.2 omits the
+  // rotation (kHadamard=false): it is logit-preserving (H orthonormal, applied
+  // to both q and k), so dropping it only trades fp8 quant accuracy.
+  if constexpr (kHadamard) {
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a1;
+      data[1] = a0 - a1;
+      data[2] = a2 + a3;
+      data[3] = a2 - a3;
+    }
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a2;
+      data[1] = a1 + a3;
+      data[2] = a0 - a2;
+      data[3] = a1 - a3;
+    }
+#pragma unroll
+    for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
+#pragma unroll
+      for (int i = 0; i < kVecSize; ++i) {
+        const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads);
+        data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
+      }
+    }
+    const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i)
+      data[i] *= kHadamardScale;
+  }
+
+  {
+    float local_max = math::abs(data[0]);
+#pragma unroll
+    for (int i = 1; i < kVecSize; ++i) {
+      local_max = math::max(local_max, math::abs(data[i]));
+    }
+    const auto abs_max = warp::reduce_max(local_max);
+    const auto scale = fmaxf(1e-4f, abs_max) / math::FP8_E4M3_MAX;
+    const auto inv_scale = 1.0f / scale;
+    OutStorage result;
+    result[0] = pack_fp8(data[0] * inv_scale, data[1] * inv_scale);
+    result[1] = pack_fp8(data[2] * inv_scale, data[3] * inv_scale);
+
+    // q_fp8 row pointer: 128 fp8 / row = 32 OutStorage / row, one per lane.
+    auto out_row = static_cast<uint8_t*>(params.q_fp8) + work_id * kHeadDim;
+    result.store(out_row, lane_id);
+    params.weights_out[work_id] = weight_val * params.weight_scale * scale;
+  }
+}
+
+template <typename DType, bool kUsePDL, bool kRopeFirst = false, bool kHadamard = true>
+struct FusedQIndexerRopeHadamardQuantKernel {
+  template <typename PosT>
+  static constexpr auto kernel = fused_q_indexer_rope_hadamard_quant<DType, PosT, kUsePDL, kRopeFirst, kHadamard>;
+
+  static void forward(
+      const tvm::ffi::TensorView q_input,
+      const tvm::ffi::TensorView q_fp8,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView weights_out,
+      double weight_scale,
+      const tvm::ffi::TensorView rope_cache,
+      const tvm::ffi::TensorView positions) {
+    using namespace host;
+    constexpr int64_t kHeadDim = 128;
+    constexpr int64_t kRopeDim = 64;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto H = SymbolicSize{"num_heads"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    // Caller path is `wq_b(q_lora).view(-1, H, D)` -> contiguous; the kernel
+    // assumes a flat `(B*H, kHeadDim)` layout for both q_input and q_fp8.
+    // Pin the head/innermost strides; assert the batch stride below.
+    TensorMatcher({B, H, kHeadDim})  //
+        .with_strides({-1, kHeadDim, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q_input);
+    TensorMatcher({B, H, kHeadDim})  //
+        .with_strides({-1, kHeadDim, 1})
+        .with_dtype<fp8_e4m3_t>()
+        .with_device(device_)
+        .verify(q_fp8);
+    TensorMatcher({B, H})  //
+        .with_strides({-1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(weight);
+    TensorMatcher({B, H, 1})  //
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(weights_out);
+    TensorMatcher({-1, kRopeDim})  //
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(rope_cache);
+    auto pos_dtype = SymbolicDType{};
+    TensorMatcher({B})  //
+        .with_dtype<int32_t, int64_t>(pos_dtype)
+        .with_device(device_)
+        .verify(positions);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto num_heads = static_cast<uint32_t>(H.unwrap());
+    if (batch_size == 0) return;
+
+    // The kernel computes row pointers as `base + work_id * kHeadDim`, so
+    // both inputs must be contiguous in (batch, head, elem) order.
+    const int64_t expected_batch_stride = static_cast<int64_t>(num_heads) * kHeadDim;
+    RuntimeCheck(
+        q_input.stride(0) == expected_batch_stride,
+        "q_input must be contiguous (B, H, kHeadDim); got stride[0]=",
+        q_input.stride(0));
+    RuntimeCheck(
+        q_fp8.stride(0) == expected_batch_stride,
+        "q_fp8 must be contiguous (B, H, kHeadDim); got stride[0]=",
+        q_fp8.stride(0));
+
+    const auto params = FusedQIndexerRopeHadamardQuantParams{
+        .q_input = q_input.data_ptr(),
+        .q_fp8 = q_fp8.data_ptr(),
+        .weight = weight.data_ptr(),
+        .weights_out = static_cast<float*>(weights_out.data_ptr()),
+        .weight_scale = static_cast<float>(weight_scale),
+        .rope_cache = static_cast<const float*>(rope_cache.data_ptr()),
+        .positions = positions.data_ptr(),
+        .weight_stride_batch = weight.stride(0),
+        .batch_size = batch_size,
+        .num_heads = num_heads,
+    };
+    const auto total_works = batch_size * num_heads;
+    const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
+    const auto k_int32 = kernel<int32_t>;
+    const auto k_int64 = kernel<int64_t>;
+    const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
+    LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(k, params);
+  }
+};
+
+struct FusedQIndexerRopeHadamardFp4QuantParams {
+  const void* __restrict__ q_input;
+  void* __restrict__ q_fp4;
+  int32_t* __restrict__ q_sf;
+  const void* __restrict__ weight;
+  float* __restrict__ weights_out;
+  float weight_scale;
+  const float* __restrict__ freqs_cis;
+  const void* __restrict__ positions;
+  uint32_t batch_size;
+  uint32_t num_heads;
+};
+
+template <typename DType, typename PosT, bool kUsePDL>
+Q_KERNEL void
+fused_q_indexer_rope_hadamard_fp4_quant(const __grid_constant__ FusedQIndexerRopeHadamardFp4QuantParams params) {
+  using namespace device;
+
+  constexpr int64_t kHeadDim = 128;
+  constexpr int64_t kRopeDim = 64;
+  constexpr int64_t kVecSize = 4;
+  constexpr uint32_t kRopeSize = kRopeDim / kVecSize;
+  static_assert(kHeadDim == kWarpThreads * kVecSize);
+  static_assert(kRopeDim == kWarpThreads * 2);
+  static_assert(kRopeSize <= kWarpThreads);
+
+  using Storage = AlignedVector<DType, kVecSize>;
+  using Float4 = AlignedVector<float, kVecSize>;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto work_id = blockIdx.x * kFusedQNumWarps + warp_id;
+  const bool is_rope_lane = lane_id >= kWarpThreads - kRopeSize;
+
+  const uint32_t total_works = params.batch_size * params.num_heads;
+  if (work_id >= total_works) return;
+
+  const uint32_t batch_id = work_id / params.num_heads;
+  const auto input_ptr = static_cast<const DType*>(params.q_input) + work_id * kHeadDim;
+  const auto position = static_cast<int32_t>(static_cast<const PosT*>(params.positions)[batch_id]);
+  const auto freqs_cis = params.freqs_cis + position * kRopeDim;
+
+  PDLWaitPrimary<kUsePDL>();
+  Float4 data, freq;
+  const auto weight_val = cast<float>(static_cast<const DType*>(params.weight)[work_id]);
+
+  {
+    Storage input_vec;
+    input_vec.load(input_ptr, lane_id);
+    if (is_rope_lane) freq.load(freqs_cis, lane_id - (kWarpThreads - kRopeSize));
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i) {
+      data[i] = cast<float>(input_vec[i]);
+    }
+  }
+
+  if (is_rope_lane) {
+    const auto x_real = data[0];
+    const auto x_imag = data[1];
+    const auto y_real = data[2];
+    const auto y_imag = data[3];
+    const auto fxr = freq[0];
+    const auto fxi = freq[1];
+    const auto fyr = freq[2];
+    const auto fyi = freq[3];
+    data[0] = x_real * fxr - x_imag * fxi;
+    data[1] = x_real * fxi + x_imag * fxr;
+    data[2] = y_real * fyr - y_imag * fyi;
+    data[3] = y_real * fyi + y_imag * fyr;
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i)
+      data[i] = cast<float>(cast<DType>(data[i]));
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  {
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a1;
+      data[1] = a0 - a1;
+      data[2] = a2 + a3;
+      data[3] = a2 - a3;
+    }
+    {
+      const float a0 = data[0], a1 = data[1], a2 = data[2], a3 = data[3];
+      data[0] = a0 + a2;
+      data[1] = a1 + a3;
+      data[2] = a0 - a2;
+      data[3] = a1 - a3;
+    }
+#pragma unroll
+    for (uint32_t mask = 1; mask < kWarpThreads; mask <<= 1) {
+#pragma unroll
+      for (int i = 0; i < kVecSize; ++i) {
+        const float other = __shfl_xor_sync(0xFFFFFFFFu, data[i], mask, kWarpThreads);
+        data[i] = (lane_id & mask) ? (other - data[i]) : (data[i] + other);
+      }
+    }
+    const float kHadamardScale = math::rsqrt(static_cast<float>(kHeadDim));
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i)
+      data[i] *= kHadamardScale;
+#pragma unroll
+    for (int i = 0; i < kVecSize; ++i)
+      data[i] = cast<float>(cast<DType>(data[i]));
+  }
+
+  {
+    float local_max = math::abs(data[0]);
+#pragma unroll
+    for (int i = 1; i < kVecSize; ++i) {
+      local_max = math::max(local_max, math::abs(data[i]));
+    }
+    local_max = warp::reduce_max<8>(local_max);
+    const auto scale_raw = fmaxf(1e-4f, local_max) / 6.0f;
+    const auto scale_ue8m0 = static_cast<uint8_t>(cast_to_ue8m0(scale_raw));
+    const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
+    const uint8_t packed0 = quant_fp4_e2m1(data[0] * inv_scale) | (quant_fp4_e2m1(data[1] * inv_scale) << 4);
+    const uint8_t packed1 = quant_fp4_e2m1(data[2] * inv_scale) | (quant_fp4_e2m1(data[3] * inv_scale) << 4);
+    const uint16_t packed = static_cast<uint16_t>(packed0) | (static_cast<uint16_t>(packed1) << 8);
+    auto out_row = static_cast<uint8_t*>(params.q_fp4) + work_id * (kHeadDim / 2);
+    reinterpret_cast<uint16_t*>(out_row)[lane_id] = packed;
+    if ((lane_id & 7) == 0) {
+      reinterpret_cast<uint8_t*>(params.q_sf + work_id)[lane_id >> 3] = scale_ue8m0;
+    }
+    params.weights_out[work_id] = weight_val * params.weight_scale;
+  }
+}
+
+template <typename DType, bool kUsePDL>
+struct FusedQIndexerRopeHadamardFp4QuantKernel {
+  template <typename PosT>
+  static constexpr auto kernel = fused_q_indexer_rope_hadamard_fp4_quant<DType, PosT, kUsePDL>;
+
+  static void forward(
+      const tvm::ffi::TensorView q_input,
+      const tvm::ffi::TensorView q_fp4,
+      const tvm::ffi::TensorView q_sf,
+      const tvm::ffi::TensorView weight,
+      const tvm::ffi::TensorView weights_out,
+      double weight_scale,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView positions) {
+    using namespace host;
+    constexpr int64_t kHeadDim = 128;
+    constexpr int64_t kRopeDim = 64;
+    constexpr int64_t kFp4Dim = kHeadDim / 2;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto H = SymbolicSize{"num_heads"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, H, kHeadDim})
+        .with_strides({-1, kHeadDim, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q_input);
+    TensorMatcher({B, H, kFp4Dim})
+        .with_strides({-1, kFp4Dim, 1})
+        .with_dtype<int8_t>()
+        .with_device(device_)
+        .verify(q_fp4);
+    TensorMatcher({B, H}).with_dtype<int32_t>().with_device(device_).verify(q_sf);
+    TensorMatcher({B, H}).with_dtype<DType>().with_device(device_).verify(weight);
+    TensorMatcher({B, H, 1}).with_dtype<float>().with_device(device_).verify(weights_out);
+    TensorMatcher({-1, kRopeDim}).with_dtype<float>().with_device(device_).verify(freqs_cis);
+    auto pos_dtype = SymbolicDType{};
+    TensorMatcher({B}).with_dtype<int32_t, int64_t>(pos_dtype).with_device(device_).verify(positions);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto num_heads = static_cast<uint32_t>(H.unwrap());
+    if (batch_size == 0) return;
+
+    const int64_t expected_q_stride = static_cast<int64_t>(num_heads) * kHeadDim;
+    const int64_t expected_fp4_stride = static_cast<int64_t>(num_heads) * kFp4Dim;
+    RuntimeCheck(q_input.stride(0) == expected_q_stride, "q_input must be contiguous");
+    RuntimeCheck(q_fp4.stride(0) == expected_fp4_stride, "q_fp4 must be contiguous");
+    RuntimeCheck(q_sf.stride(0) == static_cast<int64_t>(num_heads) && q_sf.stride(1) == 1, "q_sf must be contiguous");
+
+    const auto params = FusedQIndexerRopeHadamardFp4QuantParams{
+        .q_input = q_input.data_ptr(),
+        .q_fp4 = q_fp4.data_ptr(),
+        .q_sf = static_cast<int32_t*>(q_sf.data_ptr()),
+        .weight = weight.data_ptr(),
+        .weights_out = static_cast<float*>(weights_out.data_ptr()),
+        .weight_scale = static_cast<float>(weight_scale),
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .positions = positions.data_ptr(),
+        .batch_size = batch_size,
+        .num_heads = num_heads,
+    };
+    const auto total_works = batch_size * num_heads;
+    const auto num_blocks = div_ceil(total_works, kFusedQNumWarps);
+    const auto k_int32 = kernel<int32_t>;
+    const auto k_int64 = kernel<int64_t>;
+    const auto k = pos_dtype.is_type<int32_t>() ? k_int32 : k_int64;
+    LaunchKernel(num_blocks, kFusedQBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(k, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh

@@ -0,0 +1,221 @@
+#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 <cstdint>
+#include <cuda_fp8.h>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+struct MegaMoEPreDispatchParams {
+  const bf16_t* __restrict__ x;            // [num_tokens, hidden]
+  const int32_t* __restrict__ topk_idx;    // [num_tokens, top_k]
+  const float* __restrict__ topk_weights;  // [num_tokens, top_k]
+
+  fp8_e4m3_t* __restrict__ buf_x;        // [padded_max, hidden]
+  int32_t* __restrict__ buf_x_sf;        // contiguous int32 [P, G/4]; see layout comment
+  int64_t* __restrict__ buf_topk_idx;    // [padded_max, top_k]
+  float* __restrict__ buf_topk_weights;  // [padded_max, top_k]
+
+  uint32_t num_tokens;
+  uint32_t padded_max;
+  uint32_t hidden;
+  uint32_t num_groups;  // hidden / group_size
+  uint32_t top_k;
+};
+
+// kGroupSize must match sglang_per_token_group_quant_fp8_ue8m0(group_size=).
+template <uint32_t kGroupSize, bool kUsePDL>
+__global__ __launch_bounds__(1024, 2) void  //
+    mega_moe_pre_dispatch_kernel(const MegaMoEPreDispatchParams __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 bid = blockIdx.x;
+  const uint32_t tid = threadIdx.x;
+
+  PDLWaitPrimary<kUsePDL>();
+  if (bid < params.num_tokens) {
+    // ---- Quantize path: one CTA per valid token ----
+
+    const uint32_t token_id = bid;
+    const auto token_in = params.x + static_cast<uint64_t>(token_id) * params.hidden;
+    const auto token_out = params.buf_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)));
+    }
+
+    // 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);
+    // 2^-ue8m0_exp as fp32 (equivalent to 1 / __uint_as_float(ue8m0 << 23)).
+    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);
+
+    // One thread per group writes its UE8M0 byte into the contiguous
+    // row-major int32-packed layout: byte address = t*num_groups + g
+    // (see layout comment at the top of the file).
+    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.buf_x_sf)[byte_off] = static_cast<uint8_t>(ue8m0_exp);
+    }
+
+    // Copy this token's topk row (no alignment assumptions; top_k is small).
+    if (tid < params.top_k) {
+      const uint32_t off = token_id * params.top_k + tid;
+      params.buf_topk_idx[off] = params.topk_idx[off];
+      params.buf_topk_weights[off] = params.topk_weights[off];
+    }
+  } else {
+    // ---- Pad path: trailing blocks fill [num_tokens, padded_max) with (-1, 0) ----
+    const uint32_t copy_bid = bid - params.num_tokens;
+    const uint32_t pad_base = params.num_tokens * params.top_k;
+    const uint32_t slot = pad_base + copy_bid * blockDim.x + tid;
+    const uint32_t total_slots = params.padded_max * params.top_k;
+
+    if (slot < total_slots) {
+      params.buf_topk_idx[slot] = -1;
+      params.buf_topk_weights[slot] = 0.0f;
+    }
+  }
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// ---- Host wrapper
+// ------------------------------------------------------------------------------------------------------------------------
+
+template <int64_t kGroupSize, bool kUsePDL>
+struct MegaMoEPreDispatchKernel {
+  static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, "unsupported group_size");
+  static constexpr auto kernel = mega_moe_pre_dispatch_kernel<static_cast<uint32_t>(kGroupSize), kUsePDL>;
+
+  static void
+  run(const tvm::ffi::TensorView x,
+      const tvm::ffi::TensorView topk_idx,
+      const tvm::ffi::TensorView topk_weights,
+      const tvm::ffi::TensorView buf_x,
+      const tvm::ffi::TensorView buf_x_sf,
+      const tvm::ffi::TensorView buf_topk_idx,
+      const tvm::ffi::TensorView buf_topk_weights) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    auto M = SymbolicSize{"num_tokens"};
+    auto P = SymbolicSize{"padded_max"};
+    auto H = SymbolicSize{"hidden"};
+    auto K = SymbolicSize{"top_k"};
+    auto G4 = SymbolicSize{"num_groups_div_4"};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, H})  // input x
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(x);
+    TensorMatcher({M, K})  // topk_idx
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(topk_idx);
+    TensorMatcher({M, K})  // topk_weights
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(topk_weights);
+    // DeepGEMM versions expose this fp8 dispatch buffer either as raw int8
+    // storage or as torch.float8_e4m3fn; the kernel writes fp8 bytes in both.
+    TensorMatcher({P, H})  // buf.x
+        .with_dtype<int8_t, fp8_e4m3_t>()
+        .with_device(device)
+        .verify(buf_x);
+    // buf.x_sf is the contiguous row-major int32 view from DeepGEMM's mega
+    // symm buffer (DeepGEMM/csrc/apis/mega.hpp): shape (P, G/4), strides
+    // (G/4, 1). No explicit strides required -> TensorMatcher enforces
+    // is_contiguous().
+    TensorMatcher({P, G4})  // buf_x_sf
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(buf_x_sf);
+    TensorMatcher({P, K})  // buf.topk_idx
+        .with_dtype<int64_t>()
+        .with_device(device)
+        .verify(buf_topk_idx);
+    TensorMatcher({P, K})  // buf.topk_weights
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(buf_topk_weights);
+
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    const auto padded_max = static_cast<uint32_t>(P.unwrap());
+    const auto hidden = static_cast<uint32_t>(H.unwrap());
+    const auto top_k = static_cast<uint32_t>(K.unwrap());
+    const auto num_groups_div_4 = static_cast<uint32_t>(G4.unwrap());
+
+    RuntimeCheck(num_tokens <= padded_max, "num_tokens must not exceed padded_max");
+    RuntimeCheck(hidden % kGroupSize == 0, "hidden must be a multiple of group_size");
+    const auto num_groups = hidden / static_cast<uint32_t>(kGroupSize);
+    RuntimeCheck(num_groups == num_groups_div_4 * 4u, "num_groups must be a multiple of 4");
+    RuntimeCheck(hidden % 8u == 0, "hidden must be a multiple of 8 (16B bf16 loads)");
+    const auto num_threads = hidden / 8u;
+    RuntimeCheck(num_threads <= 1024, "hidden too large for single-block-per-row quant");
+    RuntimeCheck(num_threads >= top_k, "top_k must fit into one quant CTA");
+
+    const auto pad_slots = (padded_max - num_tokens) * top_k;
+    const uint32_t num_pad_blocks = pad_slots == 0 ? 0u : ((pad_slots + num_threads - 1u) / num_threads);
+    const auto num_total_blocks = num_tokens + num_pad_blocks;
+
+    const auto params = MegaMoEPreDispatchParams{
+        .x = static_cast<const bf16_t*>(x.data_ptr()),
+        .topk_idx = static_cast<const int32_t*>(topk_idx.data_ptr()),
+        .topk_weights = static_cast<const float*>(topk_weights.data_ptr()),
+        .buf_x = static_cast<fp8_e4m3_t*>(buf_x.data_ptr()),
+        .buf_x_sf = static_cast<int32_t*>(buf_x_sf.data_ptr()),
+        .buf_topk_idx = static_cast<int64_t*>(buf_topk_idx.data_ptr()),
+        .buf_topk_weights = static_cast<float*>(buf_topk_weights.data_ptr()),
+        .num_tokens = num_tokens,
+        .padded_max = padded_max,
+        .hidden = hidden,
+        .num_groups = num_groups,
+        .top_k = top_k,
+    };
+
+    if (num_total_blocks == 0) return;
+    LaunchKernel(num_total_blocks, num_threads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/online_c128_mtp.cuh

@@ -0,0 +1,387 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+SGL_DEVICE int64_t clamp_accept_len(int64_t delta, int64_t max_accept) {
+  if (delta < 0) return 0;
+  return delta < max_accept ? delta : max_accept;
+}
+
+template <typename TSeq, typename TReq>
+struct OnlineC128MTPWritePrefixParams {
+  const float* __restrict__ kv_score_input;
+  const TSeq* __restrict__ seq_lens;
+  const TReq* __restrict__ req_pool_indices;
+  const int32_t* __restrict__ req_to_token;
+  const float* __restrict__ ape;
+  float* __restrict__ state;
+  int64_t kv_score_stride_b;
+  int64_t req_to_token_stride_b;
+  int64_t ape_stride_r;
+  int64_t state_stride_b;
+  int64_t layer_bs;
+  int64_t num_verify_tokens;
+  int64_t state_slot_stride;
+};
+
+template <typename TSeq, typename TReq>
+struct OnlineC128MTPMarkPendingParams {
+  const TSeq* __restrict__ seq_lens;
+  const TReq* __restrict__ req_pool_indices;
+  int64_t* __restrict__ pending_seq_lens;
+  int64_t bs;
+  int64_t max_num_reqs;
+};
+
+template <typename TSeq, typename TReq>
+struct OnlineC128MTPCommitPendingParams {
+  const TSeq* __restrict__ cur_seq_lens;
+  const TReq* __restrict__ cur_req_pool_indices;
+  const int32_t* __restrict__ req_to_token;
+  const int64_t* __restrict__ pending_seq_lens;
+  float* __restrict__ state;
+  int64_t cur_bs;
+  int64_t req_to_token_stride_b;
+  int64_t state_stride_b;
+  int64_t num_verify_tokens;
+  int64_t state_slot_stride;
+  int64_t max_num_reqs;
+};
+
+__global__ void online_c128_mtp_clear_all_pending_kernel(int64_t* pending_seq_lens, int64_t max_num_reqs) {
+  const int64_t idx = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (idx < max_num_reqs) pending_seq_lens[idx] = -1;
+}
+
+template <typename TSeq, typename TReq>
+__global__ void online_c128_mtp_mark_pending_kernel(const OnlineC128MTPMarkPendingParams<TSeq, TReq> params) {
+  const int64_t bid = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (bid >= params.bs) return;
+  const int64_t req = static_cast<int64_t>(params.req_pool_indices[bid]);
+  if (req >= 0 && req < params.max_num_reqs) {
+    params.pending_seq_lens[req] = static_cast<int64_t>(params.seq_lens[bid]);
+  }
+}
+
+template <int64_t kHeadDim, typename TSeq, typename TReq>
+__global__ void online_c128_mtp_commit_pending_kernel(const OnlineC128MTPCommitPendingParams<TSeq, TReq> params) {
+  const int64_t bid = static_cast<int64_t>(blockIdx.x);
+  if (bid >= params.cur_bs) return;
+
+  const int64_t req = static_cast<int64_t>(params.cur_req_pool_indices[bid]);
+  if (req < 0 || req >= params.max_num_reqs) return;
+  const int64_t old_seq = params.pending_seq_lens[req];
+  if (old_seq < 0) return;
+
+  const int64_t cur_seq = static_cast<int64_t>(params.cur_seq_lens[bid]);
+  const int64_t accept = clamp_accept_len(cur_seq - old_seq, params.num_verify_tokens);
+  if (accept <= 0) return;
+
+  const int64_t final_seq = old_seq + accept;
+  if ((final_seq & 127) == 0) return;
+
+  const int64_t slot = req;
+  const float* const src = params.state + (slot + accept * params.state_slot_stride) * params.state_stride_b;
+  float* const dst = params.state + slot * params.state_stride_b;
+
+  for (int64_t d = static_cast<int64_t>(threadIdx.x); d < kHeadDim * 3; d += blockDim.x) {
+    dst[d] = src[d];
+  }
+}
+
+template <int64_t kHeadDim, typename TSeq, typename TReq>
+__global__ void online_c128_mtp_write_prefix_kernel(const OnlineC128MTPWritePrefixParams<TSeq, TReq> params) {
+  const int64_t bid = static_cast<int64_t>(blockIdx.x);
+  if (bid >= params.layer_bs) return;
+
+  const int64_t seq_before = static_cast<int64_t>(params.seq_lens[bid]);
+  const int64_t req_idx = static_cast<int64_t>(params.req_pool_indices[bid]);
+  const int64_t start_pos = seq_before & 127;
+  const bool has_partial = seq_before > 0 && start_pos != 0;
+
+  int64_t init_slot = 0;
+  if (has_partial) {
+    init_slot = req_idx;
+  }
+
+  const int64_t d = static_cast<int64_t>(threadIdx.x);
+  float run_max = 0.0f;
+  float run_sum = 0.0f;
+  float run_kv = 0.0f;
+  if (has_partial) {
+    const float* const init = params.state + init_slot * params.state_stride_b;
+    run_max = init[d];
+    run_sum = init[kHeadDim + d];
+    run_kv = init[kHeadDim * 2 + d];
+  }
+
+  constexpr int kMaxVerifyTokens = 8;
+  float kv_steps[kMaxVerifyTokens];
+  float score_steps[kMaxVerifyTokens];
+
+#pragma unroll
+  for (int step = 0; step < kMaxVerifyTokens; ++step) {
+    if (step >= params.num_verify_tokens) break;
+
+    const int64_t pos = (start_pos + step) & 127;
+    const float* const kv = params.kv_score_input + (bid * params.num_verify_tokens + step) * params.kv_score_stride_b;
+    kv_steps[step] = kv[d];
+    score_steps[step] = kv[kHeadDim + d] + params.ape[pos * params.ape_stride_r + d];
+  }
+
+#pragma unroll
+  for (int step = 0; step < kMaxVerifyTokens; ++step) {
+    if (step >= params.num_verify_tokens) break;
+
+    const int64_t pos = (start_pos + step) & 127;
+    const float kv_step = kv_steps[step];
+    const float score_step = score_steps[step];
+    if (pos == 0) {
+      run_kv = kv_step;
+      run_max = score_step;
+      run_sum = 1.0f;
+    } else {
+      const float new_max = fmaxf(run_max, score_step);
+      const float old_sum_scaled = run_sum * __expf(run_max - new_max);
+      const float new_exp = __expf(score_step - new_max);
+      const float new_sum = old_sum_scaled + new_exp;
+      run_kv = (run_kv * old_sum_scaled + kv_step * new_exp) / new_sum;
+      run_max = new_max;
+      run_sum = new_sum;
+    }
+
+    const int64_t final_seq = seq_before + step + 1;
+    if ((final_seq & 127) != 0) {
+      const int64_t slot = req_idx + (step + 1) * params.state_slot_stride;
+      float* const out = params.state + slot * params.state_stride_b;
+      out[d] = run_max;
+      out[kHeadDim + d] = run_sum;
+      out[kHeadDim * 2 + d] = run_kv;
+    }
+
+    if (pos == 127) {
+      run_kv = 0.0f;
+      run_max = 0.0f;
+      run_sum = 0.0f;
+    }
+  }
+}
+
+template <int64_t kHeadDim, typename TSeq, typename TReq>
+struct OnlineC128MTPWritePrefixKernel {
+  static void launch(
+      tvm::ffi::TensorView kv_score_input,
+      tvm::ffi::TensorView seq_lens,
+      tvm::ffi::TensorView req_pool_indices,
+      tvm::ffi::TensorView req_to_token,
+      tvm::ffi::TensorView ape,
+      tvm::ffi::TensorView state,
+      int64_t layer_bs,
+      int64_t num_verify_tokens,
+      int64_t state_slot_stride,
+      DLDevice device) {
+    using namespace host;
+
+    const auto params = OnlineC128MTPWritePrefixParams<TSeq, TReq>{
+        .kv_score_input = static_cast<const float*>(kv_score_input.data_ptr()),
+        .seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
+        .req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
+        .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
+        .ape = static_cast<const float*>(ape.data_ptr()),
+        .state = static_cast<float*>(state.data_ptr()),
+        .kv_score_stride_b = kv_score_input.stride(0),
+        .req_to_token_stride_b = req_to_token.stride(0),
+        .ape_stride_r = ape.stride(0),
+        .state_stride_b = state.stride(0),
+        .layer_bs = layer_bs,
+        .num_verify_tokens = num_verify_tokens,
+        .state_slot_stride = state_slot_stride,
+    };
+
+    static_assert(kHeadDim == 512, "online c128 MTP write-prefix only supports head_dim=512");
+    constexpr uint32_t kThreads = static_cast<uint32_t>(kHeadDim);
+    LaunchKernel(static_cast<uint32_t>(layer_bs), kThreads, device)(
+        online_c128_mtp_write_prefix_kernel<kHeadDim, TSeq, TReq>, params);
+  }
+
+  static void
+  run(tvm::ffi::TensorView kv_score_input,
+      tvm::ffi::TensorView seq_lens,
+      tvm::ffi::TensorView req_pool_indices,
+      tvm::ffi::TensorView req_to_token,
+      tvm::ffi::TensorView ape,
+      tvm::ffi::TensorView state,
+      int64_t layer_bs,
+      int64_t num_verify_tokens,
+      int64_t state_slot_stride) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({-1, kHeadDim * 2}).with_dtype<float>().with_device(device).verify(kv_score_input);
+    TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(seq_lens);
+    TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
+    TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
+    TensorMatcher({128, kHeadDim}).with_dtype<float>().with_device(device).verify(ape);
+    TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
+
+    if (layer_bs <= 0) return;
+    RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens);
+    RuntimeCheck(state_slot_stride > 0, "state_slot_stride must be positive");
+    RuntimeCheck(layer_bs <= seq_lens.shape()[0], "layer_bs exceeds seq_lens rows");
+    RuntimeCheck(layer_bs <= req_pool_indices.shape()[0], "layer_bs exceeds req_pool_indices rows");
+    RuntimeCheck(layer_bs * num_verify_tokens <= kv_score_input.shape()[0], "kv_score_input is too small");
+
+    launch(
+        kv_score_input,
+        seq_lens,
+        req_pool_indices,
+        req_to_token,
+        ape,
+        state,
+        layer_bs,
+        num_verify_tokens,
+        state_slot_stride,
+        device.unwrap());
+  }
+};
+
+template <int64_t kHeadDim, typename TSeq, typename TReq>
+struct OnlineC128MTPMarkPendingKernel {
+  static void launch(
+      tvm::ffi::TensorView seq_lens,
+      tvm::ffi::TensorView req_pool_indices,
+      tvm::ffi::TensorView pending_seq_lens,
+      int64_t bs,
+      int64_t max_num_reqs,
+      DLDevice device) {
+    using namespace host;
+
+    const auto params = OnlineC128MTPMarkPendingParams<TSeq, TReq>{
+        .seq_lens = static_cast<const TSeq*>(seq_lens.data_ptr()),
+        .req_pool_indices = static_cast<const TReq*>(req_pool_indices.data_ptr()),
+        .pending_seq_lens = static_cast<int64_t*>(pending_seq_lens.data_ptr()),
+        .bs = bs,
+        .max_num_reqs = max_num_reqs,
+    };
+
+    constexpr uint32_t kThreads = 256;
+    const uint32_t clear_blocks = host::div_ceil(static_cast<uint32_t>(max_num_reqs), kThreads);
+    LaunchKernel(clear_blocks, kThreads, device)(
+        online_c128_mtp_clear_all_pending_kernel, params.pending_seq_lens, max_num_reqs);
+    const uint32_t mark_blocks = host::div_ceil(static_cast<uint32_t>(bs), kThreads);
+    LaunchKernel(mark_blocks, kThreads, device)(online_c128_mtp_mark_pending_kernel<TSeq, TReq>, params);
+  }
+
+  static void
+  run(tvm::ffi::TensorView seq_lens,
+      tvm::ffi::TensorView req_pool_indices,
+      tvm::ffi::TensorView pending_seq_lens,
+      int64_t bs,
+      int64_t max_num_reqs) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(seq_lens);
+    TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(req_pool_indices);
+    TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
+
+    if (bs <= 0) return;
+    RuntimeCheck(bs <= seq_lens.shape()[0], "bs exceeds seq_lens rows");
+    RuntimeCheck(bs <= req_pool_indices.shape()[0], "bs exceeds req_pool_indices rows");
+    RuntimeCheck(max_num_reqs <= pending_seq_lens.shape()[0], "max_num_reqs exceeds pending rows");
+
+    launch(seq_lens, req_pool_indices, pending_seq_lens, bs, max_num_reqs, device.unwrap());
+  }
+};
+
+template <int64_t kHeadDim, typename TSeq, typename TReq>
+struct OnlineC128MTPCommitPendingKernel {
+  static void launch(
+      tvm::ffi::TensorView cur_seq_lens,
+      tvm::ffi::TensorView cur_req_pool_indices,
+      tvm::ffi::TensorView req_to_token,
+      tvm::ffi::TensorView pending_seq_lens,
+      tvm::ffi::TensorView state,
+      int64_t cur_bs,
+      int64_t num_verify_tokens,
+      int64_t state_slot_stride,
+      int64_t max_num_reqs,
+      DLDevice device) {
+    using namespace host;
+
+    const auto params = OnlineC128MTPCommitPendingParams<TSeq, TReq>{
+        .cur_seq_lens = static_cast<const TSeq*>(cur_seq_lens.data_ptr()),
+        .cur_req_pool_indices = static_cast<const TReq*>(cur_req_pool_indices.data_ptr()),
+        .req_to_token = static_cast<const int32_t*>(req_to_token.data_ptr()),
+        .pending_seq_lens = static_cast<const int64_t*>(pending_seq_lens.data_ptr()),
+        .state = static_cast<float*>(state.data_ptr()),
+        .cur_bs = cur_bs,
+        .req_to_token_stride_b = req_to_token.stride(0),
+        .state_stride_b = state.stride(0),
+        .num_verify_tokens = num_verify_tokens,
+        .state_slot_stride = state_slot_stride,
+        .max_num_reqs = max_num_reqs,
+    };
+
+    constexpr uint32_t kThreads = 256;
+    LaunchKernel(static_cast<uint32_t>(cur_bs), kThreads, device)(
+        online_c128_mtp_commit_pending_kernel<kHeadDim, TSeq, TReq>, params);
+  }
+
+  static void
+  run(tvm::ffi::TensorView cur_seq_lens,
+      tvm::ffi::TensorView cur_req_pool_indices,
+      tvm::ffi::TensorView req_to_token,
+      tvm::ffi::TensorView pending_seq_lens,
+      tvm::ffi::TensorView state,
+      int64_t cur_bs,
+      int64_t num_verify_tokens,
+      int64_t state_slot_stride,
+      int64_t max_num_reqs) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({-1}).with_dtype<TSeq>().with_device(device).verify(cur_seq_lens);
+    TensorMatcher({-1}).with_dtype<TReq>().with_device(device).verify(cur_req_pool_indices);
+    TensorMatcher({-1, -1}).with_dtype<int32_t>().with_device(device).verify(req_to_token);
+    TensorMatcher({-1}).with_dtype<int64_t>().with_device(device).verify(pending_seq_lens);
+    TensorMatcher({-1, kHeadDim * 3}).with_dtype<float>().with_device(device).verify(state);
+
+    if (cur_bs <= 0) return;
+    RuntimeCheck(num_verify_tokens > 0 && num_verify_tokens <= 8, "unsupported num_verify_tokens=", num_verify_tokens);
+    RuntimeCheck(state_slot_stride > 0, "state_slot_stride must be positive");
+    RuntimeCheck(cur_bs <= cur_seq_lens.shape()[0], "cur_bs exceeds seq_lens rows");
+    RuntimeCheck(cur_bs <= cur_req_pool_indices.shape()[0], "cur_bs exceeds req rows");
+    RuntimeCheck(max_num_reqs <= pending_seq_lens.shape()[0], "max_num_reqs exceeds pending rows");
+
+    launch(
+        cur_seq_lens,
+        cur_req_pool_indices,
+        req_to_token,
+        pending_seq_lens,
+        state,
+        cur_bs,
+        num_verify_tokens,
+        state_slot_stride,
+        max_num_reqs,
+        device.unwrap());
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/paged_mqa_metadata.cuh

@@ -0,0 +1,119 @@
+#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>
+
+namespace {
+
+constexpr uint32_t kBlockSize = 1024;
+constexpr uint32_t kSplitKV = 256;  // const for both SM90 and SM100
+
+struct MetadataParams {
+  /// NOTE: batch_size > 0
+  uint32_t batch_size;
+  uint32_t num_sm;
+  const uint32_t* __restrict__ context_lens;
+  uint32_t* __restrict__ schedule_metadata;
+  bool use_smem = true;
+};
+
+__global__ __launch_bounds__(kBlockSize, 1)  //
+    void smxx_paged_mqa_logits_metadata(const MetadataParams params) {
+  using namespace device;
+  extern __shared__ uint32_t s_length[];
+  static constexpr auto kNumWarps = kBlockSize / kWarpThreads;
+  static_assert(kNumWarps == kWarpThreads);
+
+  const auto tx = threadIdx.x;
+  const auto lane_id = tx % kWarpThreads;
+  const auto warp_id = tx / kWarpThreads;
+
+  __shared__ uint32_t s_warp_sum[kNumWarps];
+
+  uint32_t local_sum = 0;
+  for (uint32_t i = tx; i < params.batch_size; i += kBlockSize) {
+    const auto length = params.context_lens[i];
+    local_sum += (length + kSplitKV - 1) / kSplitKV;
+    if (params.use_smem) s_length[i] = length;
+  }
+
+  s_warp_sum[warp_id] = warp::reduce_sum(local_sum);
+  __syncthreads();
+
+  const auto global_sum = warp::reduce_sum(s_warp_sum[lane_id]);
+  if (lane_id != 0) return;
+
+  const auto length_ptr = params.use_smem ? s_length : params.context_lens;
+
+  const auto avg = global_sum / params.num_sm;
+  const auto ret = global_sum % params.num_sm;
+  uint32_t q = 0;
+  uint32_t num_work = (length_ptr[0] + kSplitKV - 1) / kSplitKV;
+  uint32_t sum_work = num_work;
+  for (auto i = warp_id; i <= params.num_sm; i += kNumWarps) {
+    const auto target = i * avg + min(i, ret);
+    while (sum_work <= target) {
+      if (++q >= params.batch_size) break;
+      num_work = (length_ptr[q] + kSplitKV - 1) / kSplitKV;
+      sum_work += num_work;
+    }
+    if (q >= params.batch_size) {
+      params.schedule_metadata[2 * i + 0] = params.batch_size;
+      params.schedule_metadata[2 * i + 1] = 0;
+    } else {
+      // sum > target && (sum - length) <= target
+      params.schedule_metadata[2 * i + 0] = q;
+      params.schedule_metadata[2 * i + 1] = target - (sum_work - num_work);
+    }
+  }
+}
+
+template <auto* f, size_t kMaxDynamicSMEM>
+void setup_kernel_smem_once(host::DebugInfo where = {}) {
+  [[maybe_unused]]
+  static const auto result = [] {
+    const auto fptr = std::bit_cast<const void*>(f);
+    return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
+  }();
+  host::RuntimeDeviceCheck(result, where);
+}
+
+struct IndexerMetadataKernel {
+  static constexpr auto kMaxBatchSizeInSmem = 16384 * 2;  // 128 KB smeme
+  static void run(tvm::ffi::TensorView seq_lens, tvm::ffi::TensorView metadata) {
+    using namespace host;
+    auto N = SymbolicSize{"batch_size"};
+    auto M = SymbolicSize{"num_sm"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N})  //
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(seq_lens);
+    TensorMatcher({M, 2})  //
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(metadata);
+    const auto batch_size = static_cast<uint32_t>(N.unwrap());
+    const auto num_sm = static_cast<uint32_t>(M.unwrap()) - 1;
+    RuntimeCheck(num_sm <= 1024);
+    const auto use_smem = batch_size <= kMaxBatchSizeInSmem;
+    const auto params = MetadataParams{
+        .batch_size = batch_size,
+        .num_sm = num_sm,
+        .context_lens = static_cast<uint32_t*>(seq_lens.data_ptr()),
+        .schedule_metadata = static_cast<uint32_t*>(metadata.data_ptr()),
+        .use_smem = use_smem,
+    };
+    constexpr auto kernel = smxx_paged_mqa_logits_metadata;
+    setup_kernel_smem_once<kernel, (kMaxBatchSizeInSmem + 1) * sizeof(uint32_t)>();
+    const auto smem = use_smem ? (batch_size + 1) * sizeof(uint32_t) : 0;
+    LaunchKernel(1, kBlockSize, device.unwrap(), smem)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/rope.cuh

@@ -0,0 +1,169 @@
+#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 <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+
+namespace {
+
+using DType = bf16_t;
+constexpr int64_t kRopeDim = 64;
+constexpr uint32_t kBlockSize = 128;
+constexpr uint32_t kNumWarps = kBlockSize / device::kWarpThreads;
+
+struct FusedQKRopeParams {
+  void* __restrict__ q;
+  void* __restrict__ k;
+  const float* __restrict__ freqs_cis;
+  const void* __restrict__ positions;
+  int64_t q_stride_batch;
+  int64_t k_stride_batch;
+  int64_t q_stride_head;
+  int64_t k_stride_head;
+  uint32_t num_q_heads;
+  uint32_t num_k_heads;
+  uint32_t batch_size;
+};
+
+template <bool kUsePDL, bool kInverse, typename IndexType>
+__global__ __launch_bounds__(kBlockSize, 16)  //
+    void deepseek_rope_kernel(const __grid_constant__ FusedQKRopeParams param) {
+  using namespace device;
+  using DType2 = packed_t<DType>;
+
+  const auto warp_id = threadIdx.x / kWarpThreads;
+  const auto lane_id = threadIdx.x % kWarpThreads;
+  const auto global_warp_id = blockIdx.x * kNumWarps + warp_id;
+
+  const auto& [
+    q, k, freqs_cis, positions, //
+    q_stride_batch, k_stride_batch, q_stride_head, k_stride_head, //
+    num_q_heads, num_k_heads, batch_size
+  ] = param;
+
+  const auto num_total_heads = num_q_heads + num_k_heads;
+  const auto head_id = global_warp_id % num_total_heads;
+  const auto batch_id = global_warp_id / num_total_heads;
+  if (batch_id >= batch_size) return;
+
+  const auto position = static_cast<const IndexType*>(positions)[batch_id];
+  const auto is_q = head_id < num_q_heads;
+  const auto local_head = is_q ? head_id : (head_id - num_q_heads);
+  const auto stride_batch = is_q ? q_stride_batch : k_stride_batch;
+  const auto stride_head = is_q ? q_stride_head : k_stride_head;
+  const auto base_ptr = is_q ? q : k;
+  const auto input = static_cast<DType2*>(pointer::offset(base_ptr, batch_id * stride_batch, local_head * stride_head));
+
+  const auto freq_ptr = reinterpret_cast<const fp32x2_t*>(freqs_cis + position * kRopeDim);
+  const auto [f_real, f_imag] = freq_ptr[lane_id];
+  PDLWaitPrimary<kUsePDL>();
+
+  const auto data = input[lane_id];
+  const auto [x_real, x_imag] = cast<fp32x2_t>(data);
+  fp32x2_t output;
+  if constexpr (kInverse) {
+    // (a + bi) * (c - di) = (ac + bd) + (bc - ad)i
+    output = {
+        x_real * f_real + x_imag * f_imag,
+        x_imag * f_real - x_real * f_imag,
+    };
+  } else {
+    // (a + bi) * (c + di) = (ac - bd) + (ad + bc)i
+    output = {
+        x_real * f_real - x_imag * f_imag,
+        x_real * f_imag + x_imag * f_real,
+    };
+  }
+  input[lane_id] = cast<DType2>(output);
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <bool kUsePDL>
+struct FusedQKRopeKernel {
+  // 4 kernel variants: {forward, inverse} x {int32, int64}
+  static constexpr auto kernel_fwd_i32 = deepseek_rope_kernel<kUsePDL, false, int32_t>;
+  static constexpr auto kernel_fwd_i64 = deepseek_rope_kernel<kUsePDL, false, int64_t>;
+  static constexpr auto kernel_inv_i32 = deepseek_rope_kernel<kUsePDL, true, int32_t>;
+  static constexpr auto kernel_inv_i64 = deepseek_rope_kernel<kUsePDL, true, int64_t>;
+
+  static void forward(
+      const tvm::ffi::TensorView q,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> k,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView positions,
+      bool inverse) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto Q = SymbolicSize{"num_q_heads"};
+    auto K = SymbolicSize{"num_k_heads"};
+    constexpr auto D = kRopeDim;
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, Q, D})  //
+        .with_strides({-1, -1, 1})
+        .with_dtype<DType>()
+        .with_device(device_)
+        .verify(q);
+    if (k.has_value()) {
+      TensorMatcher({B, K, D})  //
+          .with_strides({-1, -1, 1})
+          .with_dtype<DType>()
+          .with_device(device_)
+          .verify(k.value());
+    } else {
+      K.set_value(0);
+    }
+    TensorMatcher({-1, D})  //
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(freqs_cis);
+
+    auto pos_dtype = SymbolicDType{};
+    TensorMatcher({B})  //
+        .with_dtype<int32_t, int64_t>(pos_dtype)
+        .with_device(device_)
+        .verify(positions);
+    const bool pos_i32 = pos_dtype.is_type<int32_t>();
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    if (batch_size == 0) return;
+
+    const auto num_q_heads = static_cast<uint32_t>(Q.unwrap());
+    const auto num_k_heads = static_cast<uint32_t>(K.unwrap());
+    const auto num_total_heads = num_q_heads + num_k_heads;
+    const auto total_warps = batch_size * num_total_heads;
+    const auto num_blocks = div_ceil(total_warps, kNumWarps);
+
+    const auto elem_size = static_cast<int64_t>(sizeof(DType));
+    const auto params = FusedQKRopeParams{
+        .q = q.data_ptr(),
+        .k = k ? k.value().data_ptr() : nullptr,
+        .freqs_cis = static_cast<const float*>(freqs_cis.data_ptr()),
+        .positions = positions.data_ptr(),
+        .q_stride_batch = q.stride(0) * elem_size,
+        .k_stride_batch = k ? k.value().stride(0) * elem_size : 0,
+        .q_stride_head = q.stride(1) * elem_size,
+        .k_stride_head = k ? k.value().stride(1) * elem_size : 0,
+        .num_q_heads = num_q_heads,
+        .num_k_heads = num_k_heads,
+        .batch_size = batch_size,
+    };
+
+    // dispatch: {inverse} x {pos_i32}
+    using KernelType = decltype(kernel_fwd_i32);
+    const KernelType kernel =
+        inverse ? (pos_i32 ? kernel_inv_i32 : kernel_inv_i64) : (pos_i32 ? kernel_fwd_i32 : kernel_fwd_i64);
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/silu_and_mul_masked_post_quant.cuh

@@ -0,0 +1,540 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/math.cuh>
+#include <sgl_kernel/tile.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 <cstdint>
+#include <cuda_fp8.h>
+#include <type_traits>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+struct SiluMulQuantVarlenParams {
+  const bf16_t* __restrict__ input;
+  fp8_e4m3_t* __restrict__ output;
+  float* __restrict__ output_scale;
+  const int32_t* __restrict__ masked_m;
+  float swiglu_limit;  // only read when kApplySwigluLimit=true
+  int64_t hidden_dim;
+  uint32_t num_tokens;
+  uint32_t num_experts;
+};
+
+constexpr uint32_t kMaxExperts = 256;
+
+struct alignas(16) CTAWork {
+  uint32_t expert_id;
+  uint32_t expert_token_id;
+  bool valid;
+};
+
+SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) {
+  static_assert(device::kWarpThreads == 32);
+#pragma unroll
+  for (uint32_t offset = 1; offset < 32; offset *= 2) {
+    uint32_t n = __shfl_up_sync(0xFFFFFFFF, val, offset);
+    if (lane_id >= offset) val += n;
+  }
+  return val;
+}
+
+template <bool kApplySwigluLimit, bool kPrecise = true, typename DType2>
+SGL_DEVICE fp32x2_t silu_and_mul(DType2 gate, DType2 up, float limit) {
+  using namespace device;
+  // refer to as implementation. TL;DR: must clamp in bf16
+  // https://github.com/deepseek-ai/DeepGEMM/blob/7f2a703ed51ac1f7af07f5e1453b2d3267d37d50/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh#L984-L997
+  if constexpr (kApplySwigluLimit) {
+    static_assert(std::is_same_v<DType2, bf16x2_t>);
+    gate = __hmin2(gate, {limit, limit});
+    up = __hmax2(up, {-limit, -limit});
+    up = __hmin2(up, {limit, limit});
+  }
+  const auto [g0, g1] = cast<fp32x2_t>(gate);
+  const auto [u0, u1] = cast<fp32x2_t>(up);
+  const auto silu0 = g0 / (1.0f + __expf(-g0));
+  const auto silu1 = g1 / (1.0f + __expf(-g1));
+  const float val0 = silu0 * u0;
+  const float val1 = silu1 * u1;
+  if constexpr (kPrecise) {  // I don't know if we should enable this?
+    return {val0, val1};
+  } else {
+    return cast<fp32x2_t>(cast<bf16x2_t>(fp32x2_t{val0, val1}));
+  }
+}
+
+[[maybe_unused]]
+SGL_DEVICE CTAWork get_work(const SiluMulQuantVarlenParams& params) {
+  // Preconditions:
+  // 1. blockDim.x >= params.num_experts
+  // 2. params.num_experts <= kMaxExperts
+  using namespace device;
+  static_assert(kWarpThreads == 32);
+
+  static __shared__ uint32_t s_warp_sum[32];
+  static __shared__ CTAWork result;
+
+  result.valid = false;
+
+  const uint32_t tx = threadIdx.x;
+  const uint32_t lane_id = tx % kWarpThreads;
+  const uint32_t warp_id = tx / kWarpThreads;
+
+  const uint32_t val = tx < params.num_experts ? params.masked_m[tx] : 0u;
+
+  // Per-warp inclusive scan of masked_m.
+  const uint32_t warp_inclusive = warp_inclusive_sum(lane_id, val);
+  const uint32_t warp_exclusive = warp_inclusive - val;
+
+  // Write each warp total.
+  if (lane_id == kWarpThreads - 1) s_warp_sum[warp_id] = warp_inclusive;
+  __syncthreads();
+  const auto tmp_val = lane_id < warp_id ? s_warp_sum[lane_id] : 0u;
+  const auto prefix_exclusive = warp::reduce_sum(tmp_val) + warp_exclusive;
+  const auto bx = blockIdx.x;
+  if (prefix_exclusive <= bx && bx < prefix_exclusive + val) {
+    result = {tx, bx - prefix_exclusive, true};
+  }
+  __syncthreads();
+  return result;
+}
+
+template <bool kScaleUE8M0, bool kTransposed, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
+__global__ __launch_bounds__(1024, 2) void  // maximize occupancy
+    silu_mul_quant_varlen_kernel(const SiluMulQuantVarlenParams __grid_constant__ params) {
+  using namespace device;
+
+  constexpr uint32_t kGroupSize = 128u;
+  constexpr uint32_t kWorkThreads = 16u;
+  // each thread will handle 8 elements
+  using InputVec = AlignedVector<bf16x2_t, 4>;
+  using OutputVec = AlignedVector<fp8x2_e4m3_t, 4>;
+  static_assert(8 * kWorkThreads == 128, "Invalid tiling");
+  static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0");
+
+  const auto [expert_id, token_id, valid] = get_work(params);
+
+  if (!valid) return;
+
+  const auto work_id = threadIdx.x / kWorkThreads;
+
+  const auto offset = expert_id * params.num_tokens + token_id;
+  const auto input = params.input + offset * params.hidden_dim * 2;
+  const auto output = params.output + offset * params.hidden_dim;
+  [[maybe_unused]]
+  const auto output_scale = [&] {
+    const auto num_groups = params.hidden_dim / kGroupSize;
+    if constexpr (kTransposed) {
+      const auto base = reinterpret_cast<uint8_t*>(params.output_scale);
+      // Physical layout is [E, G//4, N] int32.  Each int32 packs 4 consecutive
+      // group scales for the same token, so the byte address is:
+      //   expert_offset + (group/4)*N*4 + token*4 + group%4
+      return base + expert_id * num_groups * params.num_tokens + (work_id / 4u) * (params.num_tokens * 4u) +
+             token_id * 4u + (work_id % 4u);
+    } else {
+      return params.output_scale + offset * num_groups + work_id;
+    }
+  }();
+
+  PDLWaitPrimary<kUsePDL>();
+
+  InputVec gate_vec, up_vec;
+  if constexpr (kSwizzle) {
+    // gran=8 interleaved: every 16-element chunk on the N axis is
+    // [gate[0..7], up[0..7]]. Each thread handles 8 consecutive output
+    // elements, so its gate chunk lives at vec index 2*threadIdx.x and its
+    // up chunk at 2*threadIdx.x+1.
+    gate_vec.load(input, threadIdx.x * 2);
+    up_vec.load(input, threadIdx.x * 2 + 1);
+  } else {
+    gate_vec.load(input, threadIdx.x);
+    up_vec.load(input, threadIdx.x + blockDim.x);
+  }
+
+  float local_max = 0.0f;
+  float results[8];
+
+#pragma unroll
+  for (uint32_t i = 0; i < 4; ++i) {
+    const auto [x, y] = silu_and_mul<kApplySwigluLimit>(gate_vec[i], up_vec[i], params.swiglu_limit);
+    results[2 * i + 0] = x;
+    results[2 * i + 1] = y;
+    local_max = fmaxf(local_max, fmaxf(fabsf(x), fabsf(y)));
+  }
+
+  local_max = warp::reduce_max<kWorkThreads>(local_max);
+
+  const float absmax = fmaxf(local_max, 1e-10f);
+  float scale;
+  uint32_t ue8m0_exp;
+
+  if constexpr (kScaleUE8M0) {
+    const float raw_scale = absmax / math::FP8_E4M3_MAX;
+    ue8m0_exp = cast_to_ue8m0(raw_scale);
+    scale = __uint_as_float(ue8m0_exp << 23);
+  } else {
+    scale = absmax / math::FP8_E4M3_MAX;
+  }
+  const auto inv_scale = 1.0f / scale;
+
+  OutputVec out_vec;
+#pragma unroll
+  for (uint32_t i = 0; i < 4; ++i) {
+    const float scaled_val0 = results[2 * i + 0] * inv_scale;
+    const float scaled_val1 = results[2 * i + 1] * inv_scale;
+    out_vec[i] = pack_fp8(scaled_val0, scaled_val1);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  out_vec.store(output, threadIdx.x);
+  if constexpr (kTransposed) {
+    *output_scale = ue8m0_exp;
+  } else {
+    *output_scale = scale;
+  }
+}
+
+struct SiluAndMulClampParams {
+  const void* __restrict__ input;
+  void* __restrict__ output;
+  float swiglu_limit;
+};
+
+template <typename DType, bool kUsePDL>
+__global__ __launch_bounds__(1024, 2) void  // maximize occupancy
+    silu_mul_clamp_kernel(const SiluAndMulClampParams __grid_constant__ params) {
+  using namespace device;
+  static_assert(sizeof(DType) == 2, "only fp16/bf16 supported");
+  using DType2 = packed_t<DType>;
+  constexpr auto kVecSize = 16 / sizeof(DType);
+  static_assert(kVecSize % 2 == 0 && kVecSize > 0);
+  using Vec = AlignedVector<DType2, kVecSize / 2>;
+  const auto bid = blockIdx.x;
+  const auto tile = tile::Memory<Vec>::cta();
+  const float limit = params.swiglu_limit;
+
+  PDLWaitPrimary<kUsePDL>();
+  const auto gate = tile.load(params.input, bid * 2 + 0);
+  const auto up = tile.load(params.input, bid * 2 + 1);
+  Vec out;
+
+#pragma unroll
+  for (uint32_t i = 0; i < kVecSize / 2; ++i) {
+    out[i] = cast<DType2>(silu_and_mul<true>(cast<bf16x2_t>(gate[i]), cast<bf16x2_t>(up[i]), limit));
+  }
+
+  tile.store(params.output, out, bid);
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+// ---- Host wrapper
+// ------------------------------------------------------------------------------------------------------------------------
+
+template <int64_t kGroupSize, bool kScaleUE8M0, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
+struct SiluAndMulMaskedPostQuantKernel {
+  static_assert(kGroupSize == 128);
+  static constexpr auto kernel_normal =
+      silu_mul_quant_varlen_kernel<kScaleUE8M0, false, kSwizzle, kUsePDL, kApplySwigluLimit>;
+  static constexpr auto kernel_transposed =
+      silu_mul_quant_varlen_kernel<true, true, kSwizzle, kUsePDL, kApplySwigluLimit>;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView output,
+      const tvm::ffi::TensorView output_scale,
+      const tvm::ffi::TensorView masked_m,
+      const uint32_t topk,
+      const bool transposed,
+      const double swiglu_limit) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    auto E = SymbolicSize{"num_experts"};
+    auto T = SymbolicSize{"num_tokens_padded"};
+    auto D = SymbolicSize{"hidden_dim x 2"};
+    auto N = SymbolicSize{"hidden_dim"};
+    auto G = SymbolicSize{"num_groups"};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({E, T, D})  // input
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({E, T, N})  // output
+        .with_dtype<fp8_e4m3_t>()
+        .with_device(device)
+        .verify(output);
+    if (!transposed) {
+      TensorMatcher({E, T, G})  //
+          .with_dtype<fp32_t>()
+          .with_device(device)
+          .verify(output_scale);
+    } else {
+      RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true");
+      auto G_ = SymbolicSize{"G // 4"};
+      TensorMatcher({E, G_, T})  //
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(output_scale);
+      G.set_value(G_.unwrap() * 4);
+    }
+    TensorMatcher({E})  //
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(masked_m);
+
+    const auto num_experts = static_cast<uint32_t>(E.unwrap());
+    const auto num_tokens = static_cast<uint32_t>(T.unwrap());
+    const auto num_groups = static_cast<uint32_t>(G.unwrap());
+    const auto hidden_dim = N.unwrap();
+
+    RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension");
+    RuntimeCheck(hidden_dim % kGroupSize == 0);
+    RuntimeCheck(num_experts <= kMaxExperts, "num_experts exceeds maximum (256)");
+    RuntimeCheck(num_groups * kGroupSize == hidden_dim, "invalid num_groups");
+
+    const auto params = SiluMulQuantVarlenParams{
+        .input = static_cast<const bf16_t*>(input.data_ptr()),
+        .output = static_cast<fp8_e4m3_t*>(output.data_ptr()),
+        .output_scale = static_cast<float*>(output_scale.data_ptr()),
+        .masked_m = static_cast<const int32_t*>(masked_m.data_ptr()),
+        .swiglu_limit = static_cast<float>(swiglu_limit),
+        .hidden_dim = hidden_dim,
+        .num_tokens = num_tokens,
+        .num_experts = num_experts,
+    };
+
+    const auto num_threads = hidden_dim / 8;
+    RuntimeCheck(num_threads % device::kWarpThreads == 0);
+    RuntimeCheck(num_threads >= num_experts);
+    const auto kernel = transposed ? kernel_transposed : kernel_normal;
+    LaunchKernel(num_tokens * topk, num_threads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+template <typename DType, bool kUsePDL>
+struct SiluAndMulClampKernel {
+  static constexpr auto kernel = silu_mul_clamp_kernel<DType, kUsePDL>;
+
+  static void run(const tvm::ffi::TensorView input, const tvm::ffi::TensorView output, const double swiglu_limit) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    auto M = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"gate_up_dim"};  // 2 * out_dim
+    auto H = SymbolicSize{"out_dim"};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, D})  // input  (gate || up)
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({M, H})  // output
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(output);
+    RuntimeCheck(D.unwrap() == 2 * H.unwrap(), "input last dim must be 2 * output last dim");
+
+    constexpr uint32_t kVecSize = 16 / sizeof(DType);
+    const auto out_dim = static_cast<uint32_t>(H.unwrap());
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+    RuntimeCheck(out_dim % kVecSize == 0, "out_dim must be divisible by vector size");
+    const auto num_threads = out_dim / kVecSize;
+    RuntimeCheck(num_threads <= 1024, "out_dim too large for single-block-per-row launch");
+
+    const auto params = SiluAndMulClampParams{
+        .input = input.data_ptr(),
+        .output = output.data_ptr(),
+        .swiglu_limit = static_cast<float>(swiglu_limit),
+    };
+    LaunchKernel(num_tokens, num_threads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+struct SiluMulQuantContigParams {
+  const bf16_t* __restrict__ input;
+  fp8_e4m3_t* __restrict__ output;
+  float* __restrict__ output_scale;
+  float swiglu_limit;  // only read when kApplySwigluLimit=true
+  int64_t hidden_dim;
+  uint32_t num_tokens;
+  uint32_t scale_row_stride_int32;  // only used when kTransposed=true
+};
+
+template <bool kScaleUE8M0, bool kTransposed, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
+__global__ __launch_bounds__(1024, 2) void  // maximize occupancy
+    silu_mul_quant_contig_kernel(const SiluMulQuantContigParams __grid_constant__ params) {
+  using namespace device;
+
+  constexpr uint32_t kGroupSize = 128u;
+  constexpr uint32_t kWorkThreads = 16u;
+  using InputVec = AlignedVector<bf16x2_t, 4>;
+  using OutputVec = AlignedVector<fp8x2_e4m3_t, 4>;
+  static_assert(8 * kWorkThreads == 128, "Invalid tiling");
+  static_assert(!(kTransposed && !kScaleUE8M0), "transposed layout only supports ue8m0");
+
+  const auto token_id = blockIdx.x;
+  const auto work_id = threadIdx.x / kWorkThreads;
+
+  const auto input = params.input + token_id * params.hidden_dim * 2;
+  const auto output = params.output + token_id * params.hidden_dim;
+  [[maybe_unused]]
+  const auto output_scale = [&] {
+    const auto num_groups = params.hidden_dim / kGroupSize;
+    if constexpr (kTransposed) {
+      // Physical layout is (G//4_pad, M_pad) int32; each int32 packs 4
+      // consecutive UE8M0 exponents for the same token. Byte address:
+      //   (work_id / 4) * M_pad * 4  +  token * 4  +  (work_id % 4).
+      const auto base = reinterpret_cast<uint8_t*>(params.output_scale);
+      return base + (work_id / 4u) * (params.scale_row_stride_int32 * 4u) + token_id * 4u + (work_id % 4u);
+    } else {
+      return params.output_scale + token_id * num_groups + work_id;
+    }
+  }();
+
+  PDLWaitPrimary<kUsePDL>();
+
+  InputVec gate_vec, up_vec;
+  if constexpr (kSwizzle) {
+    gate_vec.load(input, threadIdx.x * 2);
+    up_vec.load(input, threadIdx.x * 2 + 1);
+  } else {
+    gate_vec.load(input, threadIdx.x);
+    up_vec.load(input, threadIdx.x + blockDim.x);
+  }
+
+  float local_max = 0.0f;
+  float results[8];
+
+#pragma unroll
+  for (uint32_t i = 0; i < 4; ++i) {
+    const auto [x, y] = silu_and_mul<kApplySwigluLimit>(gate_vec[i], up_vec[i], params.swiglu_limit);
+    results[2 * i + 0] = x;
+    results[2 * i + 1] = y;
+    local_max = fmaxf(local_max, fmaxf(fabsf(x), fabsf(y)));
+  }
+
+  local_max = warp::reduce_max<kWorkThreads>(local_max);
+
+  const float absmax = fmaxf(local_max, 1e-10f);
+  float scale;
+  uint32_t ue8m0_exp;
+
+  if constexpr (kScaleUE8M0) {
+    const float raw_scale = absmax / math::FP8_E4M3_MAX;
+    ue8m0_exp = cast_to_ue8m0(raw_scale);
+    scale = __uint_as_float(ue8m0_exp << 23);
+  } else {
+    scale = absmax / math::FP8_E4M3_MAX;
+  }
+  const auto inv_scale = 1.0f / scale;
+
+  OutputVec out_vec;
+#pragma unroll
+  for (uint32_t i = 0; i < 4; ++i) {
+    const float scaled_val0 = results[2 * i + 0] * inv_scale;
+    const float scaled_val1 = results[2 * i + 1] * inv_scale;
+    out_vec[i] = pack_fp8(scaled_val0, scaled_val1);
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  out_vec.store(output, threadIdx.x);
+  if constexpr (kTransposed) {
+    *output_scale = ue8m0_exp;
+  } else {
+    *output_scale = scale;
+  }
+}
+
+template <int64_t kGroupSize, bool kScaleUE8M0, bool kSwizzle, bool kUsePDL, bool kApplySwigluLimit>
+struct SiluAndMulContigPostQuantKernel {
+  static_assert(kGroupSize == 128);
+  static constexpr auto kernel_normal =
+      silu_mul_quant_contig_kernel<kScaleUE8M0, false, kSwizzle, kUsePDL, kApplySwigluLimit>;
+  static constexpr auto kernel_transposed =
+      silu_mul_quant_contig_kernel<true, true, kSwizzle, kUsePDL, kApplySwigluLimit>;
+
+  static void
+  run(const tvm::ffi::TensorView input,
+      const tvm::ffi::TensorView output,
+      const tvm::ffi::TensorView output_scale,
+      const bool transposed,
+      const double swiglu_limit) {
+    using namespace host;
+
+    auto device = SymbolicDevice{};
+    auto M = SymbolicSize{"num_tokens"};
+    auto D = SymbolicSize{"hidden_dim x 2"};
+    auto N = SymbolicSize{"hidden_dim"};
+    auto G = SymbolicSize{"num_groups"};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({M, D})  // input (gate/up, natural or gran=8 interleaved on last dim)
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(input);
+    TensorMatcher({M, N})  // fp8 output
+        .with_dtype<fp8_e4m3_t>()
+        .with_device(device)
+        .verify(output);
+
+    const auto hidden_dim = N.unwrap();
+    RuntimeCheck(D.unwrap() == 2 * hidden_dim, "invalid dimension");
+    RuntimeCheck(hidden_dim % kGroupSize == 0);
+    const auto num_groups = static_cast<uint32_t>(hidden_dim / kGroupSize);
+
+    uint32_t scale_row_stride_int32 = 0;
+    if (!transposed) {
+      G.set_value(num_groups);
+      TensorMatcher({M, G})  // (M, G) fp32 natural row-major
+          .with_dtype<fp32_t>()
+          .with_device(device)
+          .verify(output_scale);
+    } else {
+      RuntimeCheck(kScaleUE8M0, "transposed layout only supports scale_ue8m0=true");
+      RuntimeCheck(num_groups % 4 == 0, "transposed layout requires num_groups % 4 == 0");
+      auto G_ = SymbolicSize{"G // 4"};
+      G_.set_value(num_groups / 4);
+      auto M_pad = SymbolicSize{"M padded"};
+      TensorMatcher({M, G_})                  // `.transpose(-1,-2)[:M,:]` view of (G//4_pad, M_pad) int32
+          .with_strides({int64_t{1}, M_pad})  // col-major transposed
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(output_scale);
+      scale_row_stride_int32 = static_cast<uint32_t>(M_pad.unwrap());
+    }
+
+    const auto num_tokens = static_cast<uint32_t>(M.unwrap());
+
+    const auto params = SiluMulQuantContigParams{
+        .input = static_cast<const bf16_t*>(input.data_ptr()),
+        .output = static_cast<fp8_e4m3_t*>(output.data_ptr()),
+        .output_scale = static_cast<float*>(output_scale.data_ptr()),
+        .swiglu_limit = static_cast<float>(swiglu_limit),
+        .hidden_dim = hidden_dim,
+        .num_tokens = num_tokens,
+        .scale_row_stride_int32 = scale_row_stride_int32,
+    };
+
+    const auto num_threads = hidden_dim / 8;
+    RuntimeCheck(num_threads % device::kWarpThreads == 0);
+    const auto kernel = transposed ? kernel_transposed : kernel_normal;
+    LaunchKernel(num_tokens, num_threads, device.unwrap())  //
+        .enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/store.cuh

@@ -0,0 +1,205 @@
+#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 <bit>
+#include <cstdint>
+#include <cuda_fp8.h>
+
+namespace {
+
+using deepseek_v4::fp8::cast_to_ue8m0;
+using deepseek_v4::fp8::inv_scale_ue8m0;
+using deepseek_v4::fp8::pack_fp8;
+
+struct FusedStoreCacheParam {
+  const void* __restrict__ input;
+  void* __restrict__ cache;
+  const void* __restrict__ indices;
+  uint32_t num_tokens;
+};
+
+template <typename Float, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
+__global__ void fused_store_flashmla_cache(const __grid_constant__ FusedStoreCacheParam param) {
+  using namespace device;
+
+  /// NOTE: 584 = 576 + 8
+  constexpr int64_t kPageBytes = host::div_ceil(584 << kPageBits, 576) * 576;
+
+  // each warp handles 64 elements, 8 warps, each block handles 1 row
+  const auto& [input, cache, indices, num_tokens] = param;
+  const uint32_t bid = blockIdx.x;
+  const uint32_t tid = threadIdx.x;
+  const uint32_t wid = tid / 32;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // prefetch the index
+  const auto index = static_cast<const IndicesT*>(indices)[bid];
+  // always load the value from input (don't store if invalid)
+  using Float2 = packed_t<Float>;
+  const auto elems = static_cast<const Float2*>(input)[tid + bid * 256];
+  if (wid != 7) {
+    const auto [x, y] = cast<fp32x2_t>(elems);
+    const auto abs_max = warp::reduce_max(fmaxf(fabs(x), fabs(y)));
+    const auto scale_raw = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
+    const auto scale_ue8m0 = cast_to_ue8m0(scale_raw);
+    const auto inv_scale = inv_scale_ue8m0(scale_ue8m0);
+    const auto result = pack_fp8(x * inv_scale, y * inv_scale);
+    const int32_t page = index >> kPageBits;
+    const int32_t offset = index & ((1 << kPageBits) - 1);
+    const auto page_ptr = pointer::offset(cache, page * kPageBytes);
+    const auto value_ptr = pointer::offset(page_ptr, offset * 576);
+    const auto scale_ptr = pointer::offset(page_ptr, 576 << kPageBits, offset * 8);
+    static_cast<fp8x2_e4m3_t*>(value_ptr)[tid] = result;
+    static_cast<uint8_t*>(scale_ptr)[wid] = scale_ue8m0;
+  } else {
+    const auto result = cast<bf16x2_t>(elems);
+    const int32_t page = index >> kPageBits;
+    const int32_t offset = index & ((1 << kPageBits) - 1);
+    const auto page_ptr = pointer::offset(cache, page * kPageBytes);
+    const auto value_ptr = pointer::offset(page_ptr, offset * 576, 448);
+    static_cast<bf16x2_t*>(value_ptr)[tid - 7 * 32] = result;
+  }
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename Float, typename IndicesT, uint32_t kPageBits, bool kUsePDL>
+__global__ void fused_store_indexer_cache(const __grid_constant__ FusedStoreCacheParam param) {
+  using namespace device;
+
+  /// NOTE: 132 = 128 + 4
+  constexpr int64_t kPageBytes = 132 << kPageBits;
+
+  // each warp handles 128 elements, 1 warp, each block handles multiple rows
+  const auto& [input, cache, indices, num_tokens] = param;
+  const auto global_tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const auto global_wid = global_tid / 32;
+  const auto lane_id = threadIdx.x % 32;
+
+  if (global_wid >= num_tokens) return;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // prefetch the index
+  const auto index = static_cast<const IndicesT*>(indices)[global_wid];
+  // always load the value from input (don't store if invalid)
+  using Float2 = packed_t<Float>;
+  using InStorage = AlignedVector<Float2, 2>;
+  using OutStorage = AlignedVector<fp8x2_e4m3_t, 2>;
+  const auto elems = static_cast<const InStorage*>(input)[global_tid];
+  const auto [x0, x1] = cast<fp32x2_t>(elems[0]);
+  const auto [y0, y1] = cast<fp32x2_t>(elems[1]);
+  const auto local_max = fmaxf(fmaxf(fabs(x0), fabs(x1)), fmaxf(fabs(y0), fabs(y1)));
+  const auto abs_max = warp::reduce_max(local_max);
+  // use normal fp32 scale
+  const auto scale = fmaxf(1e-4f, abs_max) / kFP8E4M3Max;
+  const auto inv_scale = 1.0f / scale;
+  const int32_t page = index >> kPageBits;
+  const int32_t offset = index & ((1 << kPageBits) - 1);
+  const auto page_ptr = pointer::offset(cache, page * kPageBytes);
+  const auto value_ptr = pointer::offset(page_ptr, offset * 128);
+  const auto scale_ptr = pointer::offset(page_ptr, 128 << kPageBits, offset * 4);
+  OutStorage result;
+  result[0] = pack_fp8(x0 * inv_scale, x1 * inv_scale);
+  result[1] = pack_fp8(y0 * inv_scale, y1 * inv_scale);
+  static_cast<OutStorage*>(value_ptr)[lane_id] = result;
+  static_cast<float*>(scale_ptr)[0] = scale;
+
+  PDLTriggerSecondary<kUsePDL>();
+}
+
+template <typename Float, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
+struct FusedStoreCacheFlashMLAKernel {
+  static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
+  static constexpr int64_t kPageBytes = host::div_ceil(584 * kPageSize, 576) * 576;
+  static constexpr auto kernel = fused_store_flashmla_cache<Float, IndicesT, kLogSize, kUsePDL>;
+
+  static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
+  static_assert(1 << kLogSize == kPageSize);
+
+  static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    TensorMatcher({N, 512})  // input
+        .with_dtype<Float>()
+        .with_device(device_)
+        .verify(input);
+    TensorMatcher({-1, -1})  // cache
+        .with_strides({kPageBytes, 1})
+        .with_dtype<uint8_t>()
+        .with_device(device_)
+        .verify(cache);
+    TensorMatcher({N})  // indices
+        .with_dtype<IndicesT>()
+        .with_device(device_)
+        .verify(indices);
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = FusedStoreCacheParam{
+        .input = input.data_ptr(),
+        .cache = cache.data_ptr(),
+        .indices = indices.data_ptr(),
+        .num_tokens = num_tokens,
+    };
+    const auto kBlockSize = 256;
+    const auto num_blocks = num_tokens;
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+template <typename Float, typename IndicesT, uint32_t kPageSize, bool kUsePDL>
+struct FusedStoreCacheIndexerKernel {
+  static constexpr int32_t kLogSize = std::countr_zero(kPageSize);
+  static constexpr int64_t kPageBytes = 132 * kPageSize;
+  static constexpr auto kernel = fused_store_indexer_cache<Float, IndicesT, kLogSize, kUsePDL>;
+
+  static_assert(std::has_single_bit(kPageSize), "kPageSize must be a power of 2");
+  static_assert(1 << kLogSize == kPageSize);
+
+  static void run(tvm::ffi::TensorView input, tvm::ffi::TensorView cache, tvm::ffi::TensorView indices) {
+    using namespace host;
+
+    auto N = SymbolicSize{"num_tokens"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+    TensorMatcher({N, 128})  // input
+        .with_dtype<Float>()
+        .with_device(device_)
+        .verify(input);
+    TensorMatcher({-1, -1})  // cache
+        .with_strides({kPageBytes, 1})
+        .with_dtype<uint8_t>()
+        .with_device(device_)
+        .verify(cache);
+    TensorMatcher({N})  // indices
+        .with_dtype<IndicesT>()
+        .with_device(device_)
+        .verify(indices);
+    const auto num_tokens = static_cast<uint32_t>(N.unwrap());
+    const auto params = FusedStoreCacheParam{
+        .input = input.data_ptr(),
+        .cache = cache.data_ptr(),
+        .indices = indices.data_ptr(),
+        .num_tokens = num_tokens,
+    };
+    const auto kBlockSize = 128;
+    const auto num_blocks = div_ceil(num_tokens * 32, kBlockSize);
+    LaunchKernel(num_blocks, kBlockSize, device_.unwrap()).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/topk_v1.cuh

@@ -0,0 +1,340 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <bit>
+#include <cstdint>
+
+namespace {
+
+#ifndef SGL_TOPK
+#define SGL_TOPK 512
+#endif
+
+constexpr uint32_t kTopK = SGL_TOPK;
+constexpr uint32_t kTopKBlockSize = SGL_TOPK;
+constexpr uint32_t kSMEM = 16 * 1024 * sizeof(uint32_t);  // 64KB (bytes)
+
+struct TopKParams {
+  const float* __restrict__ scores;
+  const int32_t* __restrict__ seq_lens;
+  const int32_t* __restrict__ page_table;
+  int32_t* __restrict__ page_indices;
+  int32_t* __restrict__ raw_indices;  // optional: output raw abs position indices before page transform
+  const int64_t score_stride;
+  const int64_t page_table_stride;
+  uint32_t page_bits;
+};
+
+SGL_DEVICE uint8_t convert_to_uint8(float x) {
+  __half h = __float2half_rn(x);
+  uint16_t bits = __half_as_ushort(h);
+  uint16_t key = (bits & 0x8000) ? static_cast<uint16_t>(~bits) : static_cast<uint16_t>(bits | 0x8000);
+  return static_cast<uint8_t>(key >> 8);
+}
+
+SGL_DEVICE uint32_t convert_to_uint32(float x) {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+SGL_DEVICE int32_t page_to_indices(const int32_t* __restrict__ page_table, uint32_t i, uint32_t page_bits) {
+  const uint32_t mask = (1u << page_bits) - 1u;
+  return (page_table[i >> page_bits] << page_bits) | (i & mask);
+}
+
+[[maybe_unused]]
+SGL_DEVICE void naive_transform(
+    const float* __restrict__,  // unused
+    const int32_t* __restrict__ page_table,
+    int32_t* __restrict__ indices,
+    int32_t* __restrict__ raw_indices,  // optional: output raw abs position indices
+    const uint32_t length,
+    const uint32_t page_bits) {
+  static_assert(kTopK <= kTopKBlockSize);
+  if (const auto tx = threadIdx.x; tx < length) {
+    indices[tx] = page_to_indices(page_table, tx, page_bits);
+    if (raw_indices != nullptr) {
+      raw_indices[tx] = tx;
+    }
+  } else if (kTopK == kTopKBlockSize || tx < kTopK) {
+    indices[tx] = -1;  // fill invalid indices to -1
+    if (raw_indices != nullptr) {
+      raw_indices[tx] = -1;
+    }
+  }
+}
+
+[[maybe_unused]]
+SGL_DEVICE void radix_topk(const float* __restrict__ input, int32_t* __restrict__ output, const uint32_t length) {
+  constexpr uint32_t RADIX = 256;
+  constexpr uint32_t BLOCK_SIZE = kTopKBlockSize;
+  constexpr uint32_t SMEM_INPUT_SIZE = kSMEM / (2 * sizeof(int32_t));
+
+  alignas(128) __shared__ uint32_t _s_histogram_buf[2][RADIX + 32];
+  alignas(128) __shared__ uint32_t s_counter;
+  alignas(128) __shared__ uint32_t s_threshold_bin_id;
+  alignas(128) __shared__ uint32_t s_num_input[2];
+  alignas(128) __shared__ int32_t s_last_remain;
+
+  extern __shared__ uint32_t s_input_idx[][kSMEM / (2 * sizeof(int32_t))];
+
+  const uint32_t tx = threadIdx.x;
+  uint32_t remain_topk = kTopK;
+  auto& s_histogram = _s_histogram_buf[0];
+
+  const auto run_cumsum = [&] {
+#pragma unroll 8
+    for (int32_t i = 0; i < 8; ++i) {
+      static_assert(1 << 8 == RADIX);
+      if (tx < RADIX) {
+        const auto j = 1 << i;
+        const auto k = i & 1;
+        auto value = _s_histogram_buf[k][tx];
+        if (tx + j < RADIX) {
+          value += _s_histogram_buf[k][tx + j];
+        }
+        _s_histogram_buf[k ^ 1][tx] = value;
+      }
+      __syncthreads();
+    }
+  };
+
+  // stage 1: 8bit coarse histogram
+  if (tx < RADIX + 1) s_histogram[tx] = 0;
+  __syncthreads();
+  for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
+    const auto bin = convert_to_uint8(input[idx]);
+    ::atomicAdd(&s_histogram[bin], 1);
+  }
+  __syncthreads();
+  run_cumsum();
+  if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
+    s_threshold_bin_id = tx;
+    s_num_input[0] = 0;
+    s_counter = 0;
+  }
+  __syncthreads();
+
+  const auto threshold_bin = s_threshold_bin_id;
+  remain_topk -= s_histogram[threshold_bin + 1];
+  if (remain_topk == 0) {
+    for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const uint32_t bin = convert_to_uint8(input[idx]);
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        output[pos] = idx;
+      }
+    }
+    __syncthreads();
+    return;
+  } else {
+    __syncthreads();
+    if (tx < RADIX + 1) {
+      s_histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (uint32_t idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const float raw_input = input[idx];
+      const uint32_t bin = convert_to_uint8(raw_input);
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        output[pos] = idx;
+      } else if (bin == threshold_bin) {
+        const auto pos = ::atomicAdd(&s_num_input[0], 1);
+        if (pos < SMEM_INPUT_SIZE) {
+          [[likely]] s_input_idx[0][pos] = idx;
+          const auto bin = convert_to_uint32(raw_input);
+          const auto sub_bin = (bin >> 24) & 0xFF;
+          ::atomicAdd(&s_histogram[sub_bin], 1);
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // stage 2: refine with 8bit radix passes
+#pragma unroll 4
+  for (int round = 0; round < 4; ++round) {
+    const auto r_idx = round % 2;
+
+    // clip here to prevent overflow
+    const auto raw_num_input = s_num_input[r_idx];
+    const auto num_input = raw_num_input < SMEM_INPUT_SIZE ? raw_num_input : SMEM_INPUT_SIZE;
+
+    run_cumsum();
+    if (tx < RADIX && s_histogram[tx] > remain_topk && s_histogram[tx + 1] <= remain_topk) {
+      s_threshold_bin_id = tx;
+      s_num_input[r_idx ^ 1] = 0;
+      s_last_remain = remain_topk - s_histogram[tx + 1];
+    }
+    __syncthreads();
+
+    const auto threshold_bin = s_threshold_bin_id;
+    remain_topk -= s_histogram[threshold_bin + 1];
+
+    if (remain_topk == 0) {
+      for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto offset = 24 - round * 8;
+        const auto bin = (convert_to_uint32(input[idx]) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          output[pos] = idx;
+        }
+      }
+      __syncthreads();
+      break;
+    } else {
+      __syncthreads();
+      if (tx < RADIX + 1) {
+        s_histogram[tx] = 0;
+      }
+      __syncthreads();
+      for (uint32_t i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto raw_input = input[idx];
+        const auto offset = 24 - round * 8;
+        const auto bin = (convert_to_uint32(raw_input) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          output[pos] = idx;
+        } else if (bin == threshold_bin) {
+          if (round == 3) {
+            const auto pos = ::atomicAdd(&s_last_remain, -1);
+            if (pos > 0) {
+              output[kTopK - pos] = idx;
+            }
+          } else {
+            const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
+            if (pos < SMEM_INPUT_SIZE) {
+              /// NOTE: (dark) fuse the histogram computation here
+              [[likely]] s_input_idx[r_idx ^ 1][pos] = idx;
+              const auto bin = convert_to_uint32(raw_input);
+              const auto sub_bin = (bin >> (offset - 8)) & 0xFF;
+              ::atomicAdd(&s_histogram[sub_bin], 1);
+            }
+          }
+        }
+      }
+      __syncthreads();
+    }
+  }
+}
+
+template <bool kUsePDL>
+__global__ void topk_transform_kernel(const __grid_constant__ TopKParams params) {
+  const auto &[
+    scores, seq_lens, page_table, page_indices, raw_indices, // pointers
+    score_stride, page_table_stride, page_bits // sizes
+  ] = params;
+  const uint32_t work_id = blockIdx.x;
+
+  /// NOTE: dangerous prefetch seq_len before PDL wait
+  const uint32_t seq_len = seq_lens[work_id];
+  const auto score_ptr = scores + work_id * score_stride;
+  const auto page_ptr = page_table + work_id * page_table_stride;
+  const auto indices_ptr = page_indices + work_id * kTopK;
+  const auto raw_indices_ptr = raw_indices != nullptr ? raw_indices + work_id * kTopK : nullptr;
+
+  device::PDLWaitPrimary<kUsePDL>();
+
+  if (seq_len <= kTopK) {
+    naive_transform(score_ptr, page_ptr, indices_ptr, raw_indices_ptr, seq_len, page_bits);
+  } else {
+    __shared__ int32_t s_topk_indices[kTopK];
+    radix_topk(score_ptr, s_topk_indices, seq_len);
+    static_assert(kTopK <= kTopKBlockSize);
+    const auto tx = threadIdx.x;
+    if (kTopK == kTopKBlockSize || tx < kTopK) {
+      indices_ptr[tx] = page_to_indices(page_ptr, s_topk_indices[tx], page_bits);
+      if (raw_indices_ptr != nullptr) {
+        raw_indices_ptr[tx] = s_topk_indices[tx];
+      }
+    }
+  }
+
+  device::PDLTriggerSecondary<kUsePDL>();
+}
+
+template <auto* f, size_t kMaxDynamicSMEM>
+void setup_kernel_smem_once(host::DebugInfo where = {}) {
+  [[maybe_unused]]
+  static const auto result = [] {
+    const auto fptr = std::bit_cast<const void*>(f);
+    return ::cudaFuncSetAttribute(fptr, ::cudaFuncAttributeMaxDynamicSharedMemorySize, kMaxDynamicSMEM);
+  }();
+  host::RuntimeDeviceCheck(result, where);
+}
+
+template <bool kUsePDL>
+struct TopKKernel {
+  static constexpr auto kernel = topk_transform_kernel<kUsePDL>;
+
+  static void transform(
+      const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::TensorView page_table,
+      const tvm::ffi::TensorView page_indices,
+      const uint32_t page_size,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
+    using namespace host;
+    auto B = SymbolicSize{"batch_size"};
+    auto S = SymbolicSize{"score_stride"};
+    auto P = SymbolicSize{"page_table_stride"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B, -1})  // strided scores
+        .with_strides({S, 1})
+        .with_dtype<float>()
+        .with_device(device)
+        .verify(scores);
+    TensorMatcher({B})  // seq_lens, must be contiguous
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(seq_lens);
+    TensorMatcher({B, -1})  // strided page table
+        .with_strides({P, 1})
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(page_table);
+    TensorMatcher({B, kTopK})  // output, must be contiguous
+        .with_dtype<int32_t>()
+        .with_device(device)
+        .verify(page_indices);
+
+    int32_t* raw_indices_ptr = nullptr;
+    if (raw_indices.has_value()) {
+      TensorMatcher({B, kTopK})  // optional raw indices output, must be contiguous
+          .with_dtype<int32_t>()
+          .with_device(device)
+          .verify(raw_indices.value());
+      raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
+    }
+
+    RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
+    const auto page_bits = static_cast<uint32_t>(std::countr_zero(page_size));
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto params = TopKParams{
+        .scores = static_cast<float*>(scores.data_ptr()),
+        .seq_lens = static_cast<int32_t*>(seq_lens.data_ptr()),
+        .page_table = static_cast<int32_t*>(page_table.data_ptr()),
+        .page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
+        .raw_indices = raw_indices_ptr,
+        .score_stride = S.unwrap(),
+        .page_table_stride = P.unwrap(),
+        .page_bits = page_bits,
+    };
+    constexpr auto kSMEM_ = kSMEM + sizeof(int32_t);  // align up a little
+    setup_kernel_smem_once<kernel, kSMEM_>();
+    LaunchKernel(batch_size, kTopKBlockSize, device.unwrap(), kSMEM_).enable_pdl(kUsePDL)(kernel, params);
+  }
+};
+
+}  // namespace

python/sglang/jit_kernel/csrc/deepseek_v4/topk_v2.cuh

@@ -0,0 +1,458 @@
+/**
+ * \file topk_v2.cuh
+ * \brief TopK kernel for DeepSeek v4.
+ * Adapted from
+ * 1:
+ *   https://github.com/vllm-project/vllm/blob/a8c6ee9b787d273916206a29b77feebadb80c368/csrc/persistent_topk.cuh
+ * 2:
+ *   https://github.com/flashinfer-ai/flashinfer/blob/c2b4db2b1a84448d802f0e6ac445243312bd6a4c/include/flashinfer/topk.cuh
+ * DarkSharpness never took a detailed look at these 2 implementation, but his claude code did.
+ * So we add credit to the reference implementations.
+ */
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <sgl_kernel/deepseek_v4/topk_impl.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <bit>
+#include <cstdint>
+#include <iterator>
+
+namespace {
+
+namespace impl = device::topk;
+using impl::TopKProblem;
+
+using Register2 = impl::TopKRegister<2>;  // <= 8192, register-resident, 1 read
+using Register4 = impl::TopKRegister<4>;  // <= 16384, register-resident, 1 read
+using Streaming = impl::TopKStreaming;
+using Cluster = impl::TopKCluster<8>;
+
+constexpr uint32_t kBlockSize = impl::TopKConfig::kBlockSize;
+constexpr uint32_t kOccupancy = impl::TopKConfig::kOccupancy;
+constexpr uint32_t kMaxTopK = impl::TopKConfig::kMaxTopK;
+constexpr uint32_t kClusterSize = Cluster::kClusterSize;
+constexpr uint32_t kReg2MaxSeqLen = Register2::kMaxSeqLen;  // 8192
+constexpr uint32_t kReg4MaxSeqLen = Register4::kMaxSeqLen;  // 16384
+
+#define TOPK_KERNEL __global__ __launch_bounds__(kBlockSize, kOccupancy)
+#define CLUSTER_TOPK_KERNEL TOPK_KERNEL __cluster_dims__(1, kClusterSize, 1)
+
+constexpr uint32_t kClusterFloor = 65536;
+constexpr uint32_t kClusterMaxBatch = 512;
+constexpr uint32_t kNumPersistentClusters = 15 * kOccupancy;
+
+/// Metadata tensor rows (each 8 B / 2 int32). Row 0 is the global plan result;
+/// rows 1..N are the (batch_id, seq_len) of items routed to the cluster pool.
+struct alignas(8) GlobalMetadata {
+  uint32_t cluster_threshold;
+  uint32_t num_cluster_items;  // N = number of items routed to the cluster pool
+};
+struct alignas(8) PlanItem {
+  uint32_t batch_id;
+  uint32_t seq_len;
+};
+static_assert(sizeof(GlobalMetadata) == 2 * sizeof(int32_t) && sizeof(PlanItem) == sizeof(GlobalMetadata));
+
+struct TopKLaunchParams {
+  const float* __restrict__ scores;
+  const int32_t* __restrict__ seq_lens;
+  const int32_t* __restrict__ page_table;
+  int32_t* __restrict__ page_indices;
+  int32_t* __restrict__ raw_indices;      // optional raw (pre-transform) indices output; nullptr if unused
+  const PlanItem* __restrict__ metadata;  // [0]=GlobalMetadata, [1+i]=PlanItem
+  int64_t score_stride;
+  int64_t page_table_stride;
+  uint32_t topk;
+  uint32_t page_bits;
+  uint32_t cluster_floor;  // seq_len > this routes to the cluster path (batch-aware, host-set)
+
+  SGL_DEVICE const GlobalMetadata& global() const {
+    return *reinterpret_cast<const GlobalMetadata*>(metadata);
+  }
+  SGL_DEVICE uint32_t cluster_threshold() const {
+    return global().cluster_threshold;
+  }
+  SGL_DEVICE const PlanItem& item(uint32_t i) const {
+    return metadata[1 + i];
+  }
+  SGL_DEVICE int32_t* get_output_ptr(uint32_t batch_id) const {
+    return page_indices + batch_id * static_cast<int64_t>(topk);
+  }
+  SGL_DEVICE TopKProblem problem(uint32_t batch_id, uint32_t seq_len) const {
+    const auto k = static_cast<int64_t>(topk);
+    return TopKProblem{
+        .in = scores + batch_id * score_stride,
+        .out = page_indices + batch_id * k,
+        .raw_out = raw_indices != nullptr ? raw_indices + batch_id * k : nullptr,
+        .page_table = page_table + batch_id * page_table_stride,
+        .topk = topk,
+        .seq_len = seq_len,
+        .page_bits = page_bits,
+    };
+  }
+  SGL_DEVICE TopKProblem problem(uint32_t batch_id) const {
+    return this->problem(batch_id, static_cast<uint32_t>(seq_lens[batch_id]));
+  }
+};
+
+/**
+ * \brief Persistent cluster kernel for the long items. It will handle long inputs.
+ * The short items are handled by the separate topk_kernel.
+ */
+template <bool kPDL>
+CLUSTER_TOPK_KERNEL void topk_persistent_cluster_kernel(const __grid_constant__ TopKLaunchParams params) {
+  device::enable_smem_spilling();
+  __shared__ impl::MaxSmem<Cluster::Smem> smem;
+  const uint32_t num_cluster_items = params.global().num_cluster_items;
+  device::PDLWaitPrimary<kPDL>();
+  device::PDLTriggerSecondary<kPDL>();
+#pragma unroll 1
+  for (uint32_t w = blockIdx.x; w < num_cluster_items; w += kNumPersistentClusters) {
+    const auto it = params.item(w);
+    const auto problem = params.problem(it.batch_id, it.seq_len);
+    Cluster::forward<false>(problem, &smem);
+    __syncthreads();
+  }
+}
+
+template <typename F>
+SGL_DEVICE void for_each_item(uint32_t topk, const F& f) {
+  constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
+#pragma unroll
+  for (uint32_t i = 0; i < kNumElems; ++i) {
+    if (const auto tx = i * kBlockSize + threadIdx.x; tx < topk) {
+      __builtin_assume(tx < kMaxTopK);
+      f(tx, i);
+    }
+  }
+}
+
+template <bool kPDL>
+SGL_DEVICE void trivial_transform(const TopKProblem& problem) {
+  device::PDLWaitPrimary<kPDL>();
+  device::PDLTriggerSecondary<kPDL>();
+  for_each_item(problem.topk, [&](uint32_t tx, uint32_t) {
+    problem.transform_output(tx, tx < problem.seq_len ? static_cast<int32_t>(tx) : -1);
+  });
+}
+
+SGL_DEVICE void problem_transform(TopKProblem& problem, int32_t* output_ptr) {
+  static_assert(kMaxTopK % kBlockSize == 0);
+  constexpr uint32_t kNumElems = kMaxTopK / kBlockSize;
+  int32_t source_index[kNumElems];
+  for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { source_index[i] = problem.out[tx]; });
+  problem.out = output_ptr;
+  for_each_item(problem.topk, [&](uint32_t tx, uint32_t i) { problem.transform_output(tx, source_index[i]); });
+}
+
+/**
+ * \brief Main kernel for the short items and epilogue of long items.
+ * \tparam kPDL whether to use PDL to synchronize with the cluster kernel (if any)
+ * \tparam kLevel:
+ * - Level 0: max_seq_len <= 8192           -> trivial + register<2>
+ * - Level 1: max_seq_len <= 16384          -> trivial + register<4>
+ * - Level 2: max_seq_len <= cluster_floor  -> trivial + register<4> + streaming
+ * - Level 3: max_seq_len > cluster_floor   -> + epilogue process of cluster path
+ */
+template <bool kPDL, int kLevel>
+TOPK_KERNEL void topk_main_kernel(const __grid_constant__ TopKLaunchParams params) {
+  device::enable_smem_spilling();
+  auto problem = params.problem(blockIdx.x);
+  constexpr uint32_t kU32Max = std::numeric_limits<uint32_t>::max();
+  __shared__ impl::MaxSmem<Register2::Smem, Register4::Smem, Streaming::Smem> smem;
+  if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
+  __shared__ int32_t topk_indices[kMaxTopK];
+  problem.out = topk_indices;
+
+  constexpr bool kHandleCluster = (kLevel == 3);
+  // non-trivial path: dispatch based on level and seq_len
+  const auto cluster_threshold = kHandleCluster ? params.cluster_threshold() : kU32Max;
+  if constexpr (kLevel == 0) {
+    __builtin_assume(problem.seq_len <= kReg2MaxSeqLen);
+    Register2::forward<kPDL>(problem, &smem);
+  } else if constexpr (kLevel == 1) {
+    __builtin_assume(problem.seq_len <= kReg4MaxSeqLen);
+    Register4::forward<kPDL>(problem, &smem);  // max_seq_len <= 16384 guarantees seq <= 16384
+  } else {
+    static_assert(kLevel == 2 || kLevel == 3, "we only support level = 0,1,2,3 now");
+    // if using cluster, we can delay the PDL wait
+    constexpr bool kPDLEarly = kPDL && !kHandleCluster;
+    constexpr bool kPDLFinal = kPDL && kHandleCluster;
+    if (problem.seq_len <= kReg4MaxSeqLen) {
+      Register4::forward<kPDLEarly>(problem, &smem);
+    } else if (problem.seq_len <= cluster_threshold) {
+      Streaming::forward<kPDLEarly>(problem, &smem);
+    } else {  // cluster path do nothing here
+      problem.out = params.get_output_ptr(blockIdx.x);
+    }
+    device::PDLWaitPrimary<kPDLFinal>();
+  }
+
+  // page-table transform pass (gathers kept out of the hot scatter loop),
+  // then trigger the dependent kernel only after the full output is written.
+  device::PDLTriggerSecondary<kPDL>();
+  __syncthreads();
+  problem_transform(problem, params.get_output_ptr(blockIdx.x));
+}
+
+template <bool kPDL>
+CLUSTER_TOPK_KERNEL void topk_small_batch_kernel(const __grid_constant__ TopKLaunchParams params) {
+  device::enable_smem_spilling();
+  auto problem = params.problem(blockIdx.x);
+  __shared__ impl::MaxSmem<Streaming::Smem, Cluster::Smem> smem;
+  if (problem.seq_len <= problem.topk) return trivial_transform<kPDL>(problem);
+  __shared__ int32_t topk_indices[kMaxTopK];
+  problem.out = topk_indices;
+
+  // randomly elect one worker rank to avoid workload imbalance
+  const auto worker_rank = blockIdx.x % kClusterSize;
+
+  // for small batch, we will fuse in the cluster case
+  if (problem.seq_len <= kReg4MaxSeqLen) {
+    if (blockIdx.y == worker_rank) Register4::forward<kPDL>(problem, &smem);
+  } else if (problem.seq_len <= params.cluster_floor) {
+    if (blockIdx.y == worker_rank) Streaming::forward<kPDL>(problem, &smem);
+  } else {
+    auto cluster = cooperative_groups::this_cluster();
+    problem.out = cluster.map_shared_rank(topk_indices, worker_rank);
+    Cluster::forward<kPDL>(problem, &smem);  // write to peer's output shared memory
+    cluster.sync();
+  }
+
+  device::PDLWaitPrimary<kPDL>();
+  __syncthreads();
+  if (blockIdx.y == worker_rank) problem_transform(problem, params.get_output_ptr(blockIdx.x));
+}
+
+// --- Plan: choose cluster_threshold from the seq_len distribution -----------
+__global__ __launch_bounds__(kBlockSize, 1) void topk_plan(
+    const uint32_t* __restrict__ seq_lens,
+    PlanItem* __restrict__ metadata,  // [0]=GlobalMetadata, [1+i]=PlanItem
+    const uint32_t batch_size,
+    const uint32_t static_cluster_threshold) {
+  // Candidate (threshold T_j, cap_j) pairs, T strictly increasing. The plan lowers
+  // cluster_threshold to T_j while #(items with seq_len > T_j) <= cap_j, so cap_j
+  // bounds how many long items go to the persistent pool. The pool runs N items in
+  // ceil(N / kNumPersistentClusters) waves; the longer the seq the more waves pay
+  // off (streaming a single block over a long item is very slow), so cap_j is the
+  // measured cluster-vs-streaming crossover (B200, occ2) and GROWS with T -- a flat
+  // cap = pool size only fits the shortest (~98K, one-wave) bucket. (Plan is tunable.)
+  struct Pair {
+    uint32_t threshold;
+    uint32_t max_batch_size;
+  };
+  constexpr Pair kCandidates[] = {
+      {65536, 30},    // (65536,98304]:    ~1 pool wave, streams beyond 30
+      {98304, 48},    // (98304,131072]
+      {131072, 60},   // (131072,196608]
+      {196608, 80},   // (196608,262144]
+      {262144, 112},  // (262144,393216]
+      {393216, 128},  // (393216,inf):     longest -- worth many pool waves; a top
+                      // threshold here lets overloaded ~280-393K batches still stream
+  };
+  constexpr uint32_t kNumCandidates = std::size(kCandidates);
+  static_assert(kCandidates[0].threshold == kClusterFloor);
+
+  __shared__ uint32_t s_counts[kNumCandidates];
+  __shared__ uint32_t s_threshold;
+  __shared__ uint32_t s_count;
+
+  const auto tx = threadIdx.x;
+  if (tx < kNumCandidates) s_counts[tx] = 0;
+  if (tx == 0) s_count = 0;
+  __syncthreads();
+
+  if (static_cluster_threshold > 0) {
+    if (tx == 0) s_threshold = static_cluster_threshold;
+  } else {
+    for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
+      const uint32_t sl = seq_lens[i];
+      uint32_t count = 0;
+#pragma unroll
+      for (uint32_t j = 0; j < kNumCandidates; ++j) {
+        count += (sl > kCandidates[j].threshold ? 1 : 0);
+      }
+      if (count > 0) atomicAdd(&s_counts[count - 1], 1);
+    }
+    __syncthreads();
+    if (tx == 0) {
+      uint32_t accum = 0;
+      uint32_t chosen = kCandidates[kNumCandidates - 1].threshold;
+#pragma unroll
+      for (uint32_t i = 0; i < kNumCandidates; ++i) {
+        const auto j = kNumCandidates - 1 - i;
+        accum += s_counts[j];  // # items with seq_len > kCandidates[j].threshold
+        if (accum > kCandidates[j].max_batch_size) break;
+        chosen = kCandidates[j].threshold;
+      }
+      s_threshold = chosen;
+    }
+  }
+  __syncthreads();
+  const auto cluster_threshold = max(s_threshold, kClusterFloor);
+
+  // Compact items with seq_len > threshold into metadata[1..N]: their batch ids
+  // are the work list the persistent cluster pool fetches.
+  for (uint32_t i = tx; i < batch_size; i += kBlockSize) {
+    const uint32_t sl = seq_lens[i];
+    if (sl > cluster_threshold) {
+      const auto pos = atomicAdd(&s_count, 1);
+      metadata[1 + pos] = {i, sl};
+    }
+  }
+  __syncthreads();
+  if (tx == 0) {
+    auto* g = reinterpret_cast<GlobalMetadata*>(metadata);
+    *g = {.cluster_threshold = cluster_threshold, .num_cluster_items = s_count};
+  }
+}
+
+struct TopKKernel {
+  static void plan(  //
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::TensorView metadata,
+      const uint32_t static_cluster_threshold) {
+    using namespace host;
+    auto B = SymbolicSize{"batch_size"};
+    auto Bp1 = SymbolicSize{"batch_size_plus_1"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B})  // seq_lens
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(seq_lens);
+    TensorMatcher({Bp1, 2})  // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len)
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(metadata);
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
+    const auto device = device_.unwrap();
+    LaunchKernel(1, kBlockSize, device)(  //
+        topk_plan,
+        static_cast<const uint32_t*>(seq_lens.data_ptr()),
+        static_cast<PlanItem*>(metadata.data_ptr()),
+        batch_size,
+        static_cluster_threshold);
+  }
+
+  static void transform(
+      const tvm::ffi::TensorView scores,
+      const tvm::ffi::TensorView seq_lens,
+      const tvm::ffi::TensorView page_table,
+      const tvm::ffi::TensorView page_indices,
+      const uint32_t page_size,
+      const tvm::ffi::TensorView metadata,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> raw_indices) {
+    using namespace host;
+    auto B = SymbolicSize{"batch_size"};
+    auto Bp1 = SymbolicSize{"batch_size_plus_1"};
+    auto L = SymbolicSize{"max_seq_len"};
+    auto S = SymbolicSize{"score_stride"};
+    auto P = SymbolicSize{"page_table_stride"};
+    auto K = SymbolicSize{"topk"};
+    auto device_ = SymbolicDevice{};
+    device_.set_options<kDLCUDA>();
+
+    TensorMatcher({B, L})  // score
+        .with_strides({S, 1})
+        .with_dtype<float>()
+        .with_device(device_)
+        .verify(scores);
+    TensorMatcher({B})  // seq_lens
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(seq_lens);
+    TensorMatcher({B, -1})  // page_table
+        .with_strides({P, 1})
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(page_table);
+    TensorMatcher({B, K})  // page_indices
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(page_indices);
+    TensorMatcher({Bp1, 2})  // metadata: [0]=GlobalMetadata, [1..N]=PlanItem(batch_id, seq_len)
+        .with_dtype<int32_t>()
+        .with_device(device_)
+        .verify(metadata);
+
+    int32_t* raw_indices_ptr = nullptr;
+    if (raw_indices.has_value()) {
+      TensorMatcher({B, K}).with_dtype<int32_t>().with_device(device_).verify(raw_indices.value());
+      raw_indices_ptr = static_cast<int32_t*>(raw_indices.value().data_ptr());
+    }
+
+    RuntimeCheck(std::has_single_bit(page_size), "page_size must be power of 2");
+    RuntimeCheck(S.unwrap() % 4 == 0, "score_stride must be a multiple of 4 (16-byte vectorized load)");
+    RuntimeCheck(Bp1.unwrap() == B.unwrap() + 1, "invalid metadata shape");
+    const auto topk = static_cast<uint32_t>(K.unwrap());
+    RuntimeCheck(topk > 0 && topk <= kMaxTopK, "topk must be in (0, 2048]");
+
+    const auto page_bits = static_cast<uint32_t>(std::countr_zero(page_size));
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+    const auto max_seq_len = static_cast<uint32_t>(L.unwrap());
+    const auto device = device_.unwrap();
+
+    // The fused kernel runs one 8-block cluster per batch element, and B200 fits one
+    // wave of exactly 15 such clusters (occ2). For batch <= 15 it stays latency-bound,
+    // so the 8-way split beats streaming from a much lower seq (measured crossover
+    // ~36-40K); batch 16 spills into a 2nd wave (+25%) and keeps the 64K floor.
+    // The floor is chosen on the host per launch.
+    constexpr uint32_t kClusterFloorSmall = 32768;
+    constexpr uint32_t kSmallBatchLowFloor = 15;
+    const auto params = TopKLaunchParams{
+        .scores = static_cast<const float*>(scores.data_ptr()),
+        .seq_lens = static_cast<const int32_t*>(seq_lens.data_ptr()),
+        .page_table = static_cast<const int32_t*>(page_table.data_ptr()),
+        .page_indices = static_cast<int32_t*>(page_indices.data_ptr()),
+        .raw_indices = raw_indices_ptr,
+        .metadata = static_cast<const PlanItem*>(metadata.data_ptr()),
+        .score_stride = S.unwrap(),
+        .page_table_stride = P.unwrap(),
+        .topk = topk,
+        .page_bits = page_bits,
+        .cluster_floor = (batch_size <= kSmallBatchLowFloor) ? kClusterFloorSmall : kClusterFloor,
+    };
+
+    const bool use_cluster = (max_seq_len > params.cluster_floor) && (batch_size <= kClusterMaxBatch);
+    constexpr bool kUsePDL = true;
+    if (use_cluster) {
+      if (batch_size <= kNumPersistentClusters) {
+        LaunchKernel({batch_size, kClusterSize}, kBlockSize, device)
+            .config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
+            .launch(topk_small_batch_kernel<kUsePDL>, params);
+      } else {
+        const uint32_t num_clusters = std::min(batch_size, kNumPersistentClusters);
+        LaunchKernel({num_clusters, kClusterSize}, kBlockSize, device)
+            .config({.use_pdl = kUsePDL, .cluster_dim = dim3{1, kClusterSize}})
+            .launch(topk_persistent_cluster_kernel<kUsePDL>, params);
+        LaunchKernel(batch_size, kBlockSize, device)
+            .config({.use_pdl = kUsePDL})
+            .launch(topk_main_kernel<kUsePDL, /*kLevel=*/3>, params);
+      }
+    } else if (max_seq_len <= kReg2MaxSeqLen) {
+      LaunchKernel(batch_size, kBlockSize, device)
+          .config({.use_pdl = kUsePDL})
+          .launch(topk_main_kernel<kUsePDL, /*kLevel=*/0>, params);
+    } else if (max_seq_len <= kReg4MaxSeqLen) {
+      LaunchKernel(batch_size, kBlockSize, device)
+          .config({.use_pdl = kUsePDL})
+          .launch(topk_main_kernel<kUsePDL, /*kLevel=*/1>, params);
+    } else {
+      LaunchKernel(batch_size, kBlockSize, device)
+          .config({.use_pdl = kUsePDL})
+          .launch(topk_main_kernel<kUsePDL, /*kLevel=*/2>, params);
+    }
+  }
+};
+
+}  // namespace