[Kernel] GLM 5.3 Flash related kernels (ported from #36507) (#37477)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/kernels/jit/csrc/dsa/kpool_topk_transform.cuh

@@ -0,0 +1,412 @@
+// Radix top-k core adapted from https://github.com/tile-ai/tilelang/blob/main/examples/deepseek_v32/topk_selector.py.
+// This JIT variant adds pool expansion and page-table/ragged-offset transforms.
+#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 <cstddef>
+#include <cstdint>
+#include <cuda_fp16.h>
+
+namespace sglang {
+namespace {
+
+#ifndef C10_LIKELY
+#define C10_LIKELY(expr) (__builtin_expect(static_cast<bool>(expr), 1))
+#endif
+
+#ifndef SGL_GROUP_TOPK
+#define SGL_GROUP_TOPK 256
+#endif
+
+inline constexpr int kGroupTopK = SGL_GROUP_TOPK;
+inline constexpr int kThreadsPerBlock = 1024;
+
+inline constexpr std::size_t kSmem = 8 * 1024 * sizeof(uint32_t);  // 32KB (bytes)
+
+struct FastTopKParams {
+  const float* __restrict__ input;         // [B, input_stride]
+  const int32_t* __restrict__ row_starts;  // [B] or nullptr
+  int32_t* __restrict__ indices;           // unused here (kept for layout parity)
+  const int32_t* __restrict__ lengths;     // [B]
+  int64_t input_stride;
+};
+
+__device__ __forceinline__ auto convert_to_uint8(float x) -> uint8_t {
+  __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);
+}
+
+__device__ __forceinline__ auto convert_to_uint32(float x) -> uint32_t {
+  uint32_t bits = __float_as_uint(x);
+  return (bits & 0x80000000u) ? ~bits : (bits | 0x80000000u);
+}
+
+template <int K>
+__device__ void
+fast_topk_cuda_tl_impl(const float* __restrict__ input, int* __restrict__ index, int row_start, int length) {
+  // We assume length > K here, or it will crash
+  int topk = K;
+  constexpr auto BLOCK_SIZE = 1024;
+  constexpr auto RADIX = 256;
+  constexpr auto SMEM_INPUT_SIZE = kSmem / (2 * sizeof(int));
+
+  alignas(128) __shared__ int s_histogram_buf[2][RADIX + 128];
+  alignas(128) __shared__ int s_counter;
+  alignas(128) __shared__ int s_threshold_bin_id;
+  alignas(128) __shared__ int s_num_input[2];
+
+  auto& s_histogram = s_histogram_buf[0];
+  extern __shared__ int s_input_idx[][SMEM_INPUT_SIZE];
+
+  const int tx = threadIdx.x;
+
+  if (tx < RADIX + 1) s_histogram[tx] = 0;
+  __syncthreads();
+
+  for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+    const auto bin = convert_to_uint8(input[idx + row_start]);
+    ::atomicAdd(&s_histogram[bin], 1);
+  }
+  __syncthreads();
+
+  const auto run_cumsum = [&] {
+#pragma unroll 8
+    for (int i = 0; i < 8; ++i) {
+      static_assert(1 << 8 == RADIX);
+      if (C10_LIKELY(tx < RADIX)) {
+        const auto j = 1 << i;
+        const auto k = i & 1;
+        auto value = s_histogram_buf[k][tx];
+        if (tx < RADIX - j) {
+          value += s_histogram_buf[k][tx + j];
+        }
+        s_histogram_buf[k ^ 1][tx] = value;
+      }
+      __syncthreads();
+    }
+  };
+
+  run_cumsum();
+  if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
+    s_threshold_bin_id = tx;
+    s_num_input[0] = 0;
+    s_counter = 0;
+  }
+  __syncthreads();
+
+  const auto threshold_bin = s_threshold_bin_id;
+  topk -= s_histogram[threshold_bin + 1];
+
+  if (topk == 0) {
+    for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const auto bin = static_cast<int>(convert_to_uint8(input[idx + row_start]));
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        index[pos] = idx;
+      }
+    }
+    __syncthreads();
+    return;
+  } else {
+    __syncthreads();
+    if (tx < RADIX + 1) {
+      s_histogram[tx] = 0;
+    }
+    __syncthreads();
+
+    for (int idx = tx; idx < length; idx += BLOCK_SIZE) {
+      const auto raw_input = input[idx + row_start];
+      const auto bin = static_cast<int>(convert_to_uint8(raw_input));
+      if (bin > threshold_bin) {
+        const auto pos = ::atomicAdd(&s_counter, 1);
+        index[pos] = idx;
+      } else if (bin == threshold_bin) {
+        const auto pos = ::atomicAdd(&s_num_input[0], 1);
+        if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
+          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();
+  }
+
+#pragma unroll 4
+  for (int round = 0; round < 4; ++round) {
+    __shared__ int s_last_remain;
+    const auto r_idx = round % 2;
+
+    const auto _raw_num_input = s_num_input[r_idx];
+    const auto num_input = (_raw_num_input < int(SMEM_INPUT_SIZE)) ? _raw_num_input : int(SMEM_INPUT_SIZE);
+
+    run_cumsum();
+    if (tx < RADIX && s_histogram[tx] > topk && s_histogram[tx + 1] <= topk) {
+      s_threshold_bin_id = tx;
+      s_num_input[r_idx ^ 1] = 0;
+      s_last_remain = topk - s_histogram[tx + 1];
+    }
+    __syncthreads();
+
+    const auto threshold_bin = s_threshold_bin_id;
+    topk -= s_histogram[threshold_bin + 1];
+
+    if (topk == 0) {
+      for (int 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 + row_start]) >> offset) & 0xFF;
+        if (bin > threshold_bin) {
+          const auto pos = ::atomicAdd(&s_counter, 1);
+          index[pos] = idx;
+        }
+      }
+      __syncthreads();
+      break;
+    } else {
+      __syncthreads();
+      if (tx < RADIX + 1) {
+        s_histogram[tx] = 0;
+      }
+      __syncthreads();
+      for (int i = tx; i < num_input; i += BLOCK_SIZE) {
+        const auto idx = s_input_idx[r_idx][i];
+        const auto raw_input = input[idx + row_start];
+        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);
+          index[pos] = idx;
+        } else if (bin == threshold_bin) {
+          if (round == 3) {
+            const auto pos = ::atomicAdd(&s_last_remain, -1);
+            if (pos > 0) {
+              index[K - pos] = idx;
+            }
+          } else {
+            const auto pos = ::atomicAdd(&s_num_input[r_idx ^ 1], 1);
+            if (C10_LIKELY(pos < SMEM_INPUT_SIZE)) {
+              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();
+    }
+  }
+}
+
+__device__ __forceinline__ int32_t transform_kpool_token(
+    int32_t raw_token,
+    const int32_t* __restrict__ page_table_entry,
+    const int32_t* __restrict__ topk_indices_offset,
+    int32_t offset) {
+  if (page_table_entry != nullptr) {
+    return page_table_entry[raw_token];
+  }
+  if (topk_indices_offset != nullptr) {
+    return raw_token + offset;
+  }
+  return raw_token;
+}
+
+template <int K>
+__global__ __launch_bounds__(kThreadsPerBlock) void kpool_topk_transform_kernel(
+    const __grid_constant__ FastTopKParams params,
+    int32_t* __restrict__ dst_token_indices,
+    const int64_t dst_stride,
+    const int32_t pool_size,
+    const int32_t token_topk,
+    const int32_t out_cols,
+    const int32_t* __restrict__ page_table,
+    const int64_t page_table_stride,
+    const int32_t* __restrict__ page_table_row_index,
+    const int32_t* __restrict__ topk_indices_offset,
+    const int32_t* __restrict__ seq_lens) {
+  const auto& [input, row_starts, _, lengths, input_stride] = params;
+  const auto bid = static_cast<uint64_t>(blockIdx.x);
+  const auto tid = threadIdx.x;
+  const auto row_start = row_starts == nullptr ? 0 : row_starts[bid];
+  const auto length = lengths[bid];
+  const auto score = input + bid * input_stride;
+  const auto dst = dst_token_indices + bid * dst_stride;
+  const auto page_table_row = page_table_row_index == nullptr ? bid : static_cast<uint64_t>(page_table_row_index[bid]);
+  const auto page_table_entry = page_table == nullptr ? nullptr : page_table + page_table_row * page_table_stride;
+  const auto offset = topk_indices_offset == nullptr ? 0 : topk_indices_offset[bid];
+  const bool append_tail = seq_lens != nullptr;
+  const auto full_pool_token_len = length * pool_size;
+  const auto history_len = full_pool_token_len < token_topk ? full_pool_token_len : token_topk;
+  const auto tail_count = append_tail ? seq_lens[bid] % pool_size : 0;
+
+  if (length <= K) {
+    for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
+      if (col < history_len) {
+        const auto group_rank = col / pool_size;
+        const auto slot = col % pool_size;
+        const auto raw_token = group_rank * pool_size + slot;
+        dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+      } else if (append_tail && col < history_len + tail_count) {
+        const auto raw_token = length * pool_size + (col - history_len);
+        dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+      } else {
+        dst[col] = -1;
+      }
+    }
+    return;
+  }
+
+  __shared__ int s_indices[K];
+  fast_topk_cuda_tl_impl<K>(score, s_indices, row_start, length);
+  for (int col = tid; col < out_cols; col += kThreadsPerBlock) {
+    if (col < history_len) {
+      const auto group_rank = col / pool_size;
+      const auto group_id = s_indices[group_rank];
+      const auto slot = col % pool_size;
+      const auto raw_token = group_id * pool_size + slot;
+      dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+    } else if (append_tail && col < history_len + tail_count) {
+      const auto raw_token = length * pool_size + (col - history_len);
+      dst[col] = transform_kpool_token(raw_token, page_table_entry, topk_indices_offset, offset);
+    } else {
+      dst[col] = -1;
+    }
+  }
+}
+
+template <auto* f, std::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 <typename T>
+const T* optional_data_ptr(const tvm::ffi::Optional<tvm::ffi::TensorView>& opt) {
+  if (!opt.has_value()) {
+    return nullptr;
+  }
+  return static_cast<const T*>(opt.value().data_ptr());
+}
+
+struct KpoolTopKTransformKernel {
+  static constexpr auto kernel = kpool_topk_transform_kernel<kGroupTopK>;
+
+  static void transform(
+      const tvm::ffi::TensorView score,
+      const tvm::ffi::TensorView lengths,
+      const tvm::ffi::TensorView dst_token_indices,
+      const int64_t pool_size,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> page_table_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> topk_indices_offset_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> row_starts_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> seq_lens_opt,
+      const tvm::ffi::Optional<tvm::ffi::TensorView> page_table_row_index_opt) {
+    using namespace host;
+
+    auto B = SymbolicSize{"batch_size"};
+    auto S = SymbolicSize{"score_stride"};
+    auto out_cols_sym = SymbolicSize{"out_cols"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({B, -1}).with_strides({S, 1}).with_dtype<float>().with_device(device).verify(score);
+    TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(lengths);
+    TensorMatcher({B, out_cols_sym}).with_dtype<int32_t>().with_device(device).verify(dst_token_indices);
+
+    RuntimeCheck(pool_size > 1, "pool_size must be > 1, got ", pool_size);
+    RuntimeCheck(
+        !(page_table_opt.has_value() && topk_indices_offset_opt.has_value()),
+        "page_table and topk_indices_offset are mutually exclusive");
+    RuntimeCheck(
+        !page_table_row_index_opt.has_value() || page_table_opt.has_value(),
+        "page_table_row_index requires page_table");
+
+    const auto out_cols = static_cast<int32_t>(out_cols_sym.unwrap());
+    const auto tail_cols = seq_lens_opt.has_value() ? static_cast<int32_t>(pool_size) - 1 : 0;
+    RuntimeCheck(out_cols > tail_cols, "dst_token_indices columns ", out_cols, " must exceed tail ", tail_cols);
+    const auto token_topk = out_cols - tail_cols;
+    RuntimeCheck(token_topk % static_cast<int32_t>(pool_size) == 0, "token_topk must be a multiple of pool_size");
+    RuntimeCheck(
+        token_topk / static_cast<int32_t>(pool_size) == kGroupTopK,
+        "this module is built for group_topk=",
+        kGroupTopK,
+        " but got ",
+        token_topk / static_cast<int32_t>(pool_size));
+
+    const auto batch_size = static_cast<uint32_t>(B.unwrap());
+
+    int64_t page_table_stride = 0;
+    const int32_t* page_table_ptr = nullptr;
+    if (page_table_opt.has_value()) {
+      auto P = SymbolicSize{"page_table_stride"};
+      if (page_table_row_index_opt.has_value()) {
+        auto page_table_rows = SymbolicSize{"page_table_rows"};
+        TensorMatcher({page_table_rows, -1})
+            .with_strides({P, 1})
+            .with_dtype<int32_t>()
+            .with_device(device)
+            .verify(page_table_opt.value());
+      } else {
+        TensorMatcher({B, -1}).with_strides({P, 1}).with_dtype<int32_t>().with_device(device).verify(
+            page_table_opt.value());
+      }
+      page_table_ptr = static_cast<const int32_t*>(page_table_opt.value().data_ptr());
+      page_table_stride = static_cast<int64_t>(P.unwrap());
+    }
+
+    if (topk_indices_offset_opt.has_value()) {
+      TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(topk_indices_offset_opt.value());
+    }
+    if (row_starts_opt.has_value()) {
+      TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(row_starts_opt.value());
+    }
+    if (seq_lens_opt.has_value()) {
+      TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(seq_lens_opt.value());
+    }
+    if (page_table_row_index_opt.has_value()) {
+      TensorMatcher({B}).with_dtype<int32_t>().with_device(device).verify(page_table_row_index_opt.value());
+    }
+
+    const auto params = FastTopKParams{
+        .input = static_cast<const float*>(score.data_ptr()),
+        .row_starts = optional_data_ptr<int32_t>(row_starts_opt),
+        .indices = nullptr,
+        .lengths = static_cast<const int32_t*>(lengths.data_ptr()),
+        .input_stride = static_cast<int64_t>(S.unwrap()),
+    };
+
+    setup_kernel_smem_once<kernel, kSmem>();
+    LaunchKernel(batch_size, kThreadsPerBlock, device.unwrap(), kSmem)(
+        kernel,
+        params,
+        static_cast<int32_t*>(dst_token_indices.data_ptr()),
+        static_cast<int64_t>(dst_token_indices.strides()[0]),
+        static_cast<int32_t>(pool_size),
+        token_topk,
+        out_cols,
+        page_table_ptr,
+        page_table_stride,
+        optional_data_ptr<int32_t>(page_table_row_index_opt),
+        optional_data_ptr<int32_t>(topk_indices_offset_opt),
+        optional_data_ptr<int32_t>(seq_lens_opt));
+  }
+};
+
+}  // namespace
+
+}  // namespace sglang

python/sglang/kernels/ops/attention/__init__.py

@@ -95,6 +95,7 @@ for _mod, _fn in [
     ("dsa.triton_sparse_mla", "triton_sparse_mla_fwd"),
     ("dsa.transform_index", "transform_index_page_table_prefill"),
     ("dsa.transform_index", "transform_index_page_table_decode"),
+    ("dsa.transform_index", "prepare_trtllm_nope_sparse_metadata"),
     ("dsa.cp_split", "dsa_cp_round_robin_split_q_seqs_kernel"),
     ("dsv4.fp4_indexer", "quantize_fp4_indexer_tensor"),
     ("dsv4.fp4_indexer", "store_fp4_index_k_cache"),

python/sglang/kernels/ops/attention/dsa/transform_index.py

@@ -14,6 +14,58 @@ def transform_index_page_table_decode(**kwargs):
     return transform_index_page_table_decode_fast(**kwargs)
 
 
+@triton.jit
+def prepare_trtllm_nope_sparse_metadata_kernel(
+    page_table_ptr: torch.Tensor,
+    topk_lens_ptr: torch.Tensor,
+    row_stride: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_TOPK: tl.constexpr,
+):
+    row = tl.program_id(0)
+    offsets = tl.arange(0, BLOCK_TOPK)
+    valid = offsets < TOPK
+    indices = tl.load(
+        page_table_ptr + row * row_stride + offsets,
+        mask=valid,
+        other=-1,
+    )
+    topk_len = tl.sum((valid & (indices >= 0)).to(tl.int32), axis=0)
+
+    # TRTLLM-GEN's native H512 dynamic sparse kernel produces NaNs for an
+    # empty row. CUDA-graph padding rows are never consumed, so point them at
+    # a valid dummy token and run a one-element attention instead.
+    is_empty = topk_len == 0
+    tl.store(page_table_ptr + row * row_stride, 0, mask=is_empty)
+    tl.store(topk_lens_ptr + row, tl.maximum(topk_len, 1))
+
+
+def prepare_trtllm_nope_sparse_metadata(
+    page_table: torch.Tensor,
+) -> torch.Tensor:
+    """Build per-query active top-k lengths for native H512 TRTLLM-GEN MLA.
+
+    ``page_table`` must contain packed valid token locations followed by ``-1``
+    padding. The tensor is only modified for fully empty CUDA-graph padding
+    rows, whose first entry is replaced with the valid dummy token location 0.
+    """
+    assert page_table.ndim == 2
+    assert page_table.dtype == torch.int32
+    assert page_table.is_contiguous()
+    num_rows, topk = page_table.shape
+    topk_lens = torch.empty(num_rows, dtype=torch.int32, device=page_table.device)
+    block_topk = triton.next_power_of_2(topk)
+    prepare_trtllm_nope_sparse_metadata_kernel[(num_rows,)](
+        page_table,
+        topk_lens,
+        page_table.stride(0),
+        TOPK=topk,
+        BLOCK_TOPK=block_topk,
+        num_warps=8,
+    )
+    return topk_lens
+
+
 def _allocate_prefill_result(
     topk_indices: torch.Tensor,
     real_num_tokens: int,

python/sglang/kernels/ops/attention/fla/kda.py

@@ -1155,6 +1155,7 @@ def chunk_kda_fwd(
         cu_seqlens=cu_seqlens,
         chunk_size=chunk_size,
         chunk_indices=chunk_indices,
+        safe_gate=lower_bound is not None,
         fuse_diagonal=_small_grid,
         fuse_recompute=_small_grid,
     )
@@ -1210,6 +1211,7 @@ def chunk_kda(
     dt_bias: Optional[torch.Tensor] = None,
     lower_bound: Optional[float] = None,
     output_intermediate_states: bool = False,
+    beta_is_raw: bool = False,
     **kwargs,
 ):
     if scale is None:
@@ -1219,6 +1221,9 @@ def chunk_kda(
         q = l2norm_fwd(q.contiguous())
         k = l2norm_fwd(k.contiguous())
 
+    if beta_is_raw:
+        beta = beta.float().sigmoid()
+
     # Returns o [B, T, H, V] when output_intermediate_states=False, or (o, h [B, NT, H, V, K]) when output_intermediate_states=True.
     return chunk_kda_fwd(
         q=q,

python/sglang/kernels/ops/attention/helion/kda_prefill.py

@@ -1314,6 +1314,7 @@ def chunk_kda(
     dt_bias: torch.Tensor | None = None,
     lower_bound: float | None = None,
     output_intermediate_states: bool = False,
+    beta_is_raw: bool = False,
     **kwargs: object,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
     """Match the public forward contract of SGLang's Triton ``chunk_kda``."""
@@ -1327,6 +1328,8 @@ def chunk_kda(
         raise ValueError("g and beta must cover every q token")
     g = g[:, :num_tokens]
     beta = beta[:, :num_tokens]
+    if beta_is_raw:
+        beta = beta.float().sigmoid()
     if num_tokens == 1:
         # Tracing constant-folds size-one dimensions, but the resulting kernel
         # can share a cache entry with longer inputs. Keep T=1 on Triton so a

python/sglang/kernels/ops/attention/utils.py

@@ -180,6 +180,27 @@ def mla_quantize_and_rope_for_fp8(
     return q_out, k_nope_out, k_rope_out
 
 
+def mla_quantize_for_fp8_no_rope(
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    k_nope: torch.Tensor,
+    k_rope: torch.Tensor,
+    kv_lora_rank: int,
+    qk_rope_head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    attn_dtype = torch.float8_e4m3fn
+    q_len, num_heads = q_rope.shape[:2]
+    q_out = q_rope.new_empty(
+        q_len,
+        num_heads,
+        kv_lora_rank + qk_rope_head_dim,
+        dtype=attn_dtype,
+    )
+    q_out[..., :kv_lora_rank] = q_nope.to(attn_dtype)
+    q_out[..., kv_lora_rank:] = q_rope.to(attn_dtype)
+    return q_out, k_nope.to(attn_dtype), k_rope.to(attn_dtype)
+
+
 def mla_quantize_without_rope_for_fp8(
     q_nope: torch.Tensor,
     q_rope: torch.Tensor,

python/sglang/kernels/ops/kvcache/mla_buffer.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import triton
 import triton.language as tl
@@ -81,6 +83,40 @@ def set_mla_kv_buffer_kernel(
         tl.extra.cuda.gdc_launch_dependents()
 
 
+@triton.jit
+def set_mla_kv_buffer_kernel_norope(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+    BLOCK: tl.constexpr,
+    USE_GDC: tl.constexpr = False,
+):
+    pid_loc = tl.program_id(0)
+    pid_blk = tl.program_id(1)
+
+    base = pid_blk * BLOCK
+    offs = base + tl.arange(0, BLOCK)
+    mask = offs < nope_dim
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_wait()
+
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs
+
+    src = tl.load(
+        cache_k_nope_ptr + pid_loc * nope_stride + offs,
+        mask=mask,
+    )
+    tl.store(dst_ptr, src, mask=mask)
+
+    if USE_GDC:
+        tl.extra.cuda.gdc_launch_dependents()
+
+
 # Above this loc count the TMA bulk-store path overtakes the single-CTA-per-loc
 # Triton kernel. Below it, Triton with BLOCK = next_pow2(total_dim) (one CTA
 # does the whole row in one tile, no boundary fan-out) is the winning fallback.
@@ -92,7 +128,7 @@ def _set_mla_kv_buffer_impl(
     kv_buffer: torch.Tensor,
     loc: torch.Tensor,
     cache_k_nope: torch.Tensor,
-    cache_k_rope: torch.Tensor,
+    cache_k_rope: Optional[torch.Tensor] = None,
     *,
     reserved_skip_index: int,
     dcp_world_size: int,
@@ -127,6 +163,28 @@ def _set_mla_kv_buffer_impl(
     Shared body of the two entry points below; the owner rule reaches it as
     ``1, 0`` (nothing to select) or as the live topology.
     """
+    has_rope = cache_k_rope is not None and cache_k_rope.numel() > 0
+    n_loc = loc.numel()
+    nope_dim = cache_k_nope.shape[-1]
+
+    if not has_rope:
+        BLOCK = triton.next_power_of_2(nope_dim)
+        grid = (n_loc, 1)
+        pdl_kwargs = (
+            {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
+        )
+        set_mla_kv_buffer_kernel_norope[grid](
+            kv_buffer,
+            cache_k_nope,
+            loc,
+            kv_buffer.stride(0),
+            cache_k_nope.stride(0),
+            nope_dim,
+            BLOCK=BLOCK,
+            **pdl_kwargs,
+        )
+        return
+
     from sglang.kernels.ops.kvcache.set_mla_kv_buffer import (
         can_use_set_mla_kv_buffer,
     )
@@ -134,7 +192,6 @@ def _set_mla_kv_buffer_impl(
         set_mla_kv_buffer as jit_set_mla_kv_buffer,
     )
 
-    n_loc = loc.numel()
     nope_bytes = cache_k_nope.shape[-1] * cache_k_nope.element_size()
     rope_bytes = cache_k_rope.shape[-1] * cache_k_rope.element_size()
     if (
@@ -157,7 +214,6 @@ def _set_mla_kv_buffer_impl(
     # ``set_mla_kv_buffer_kernel`` handles the over-allocation past total_dim
     # via the offs<total_dim mask). Beats BLOCK=128 by 60-2700 ns across the
     # 2 <= bs <= 512 range on GB300.
-    nope_dim = cache_k_nope.shape[-1]
     rope_dim = cache_k_rope.shape[-1]
     total_dim = nope_dim + rope_dim
     BLOCK = triton.next_power_of_2(total_dim)
@@ -437,18 +493,51 @@ def get_mla_kv_buffer_kernel(
     )
 
 
+@triton.jit
+def get_mla_kv_buffer_kernel_norope(
+    kv_buffer_ptr,
+    cache_k_nope_ptr,
+    loc_ptr,
+    buffer_stride: tl.constexpr,
+    nope_stride: tl.constexpr,
+    nope_dim: tl.constexpr,
+):
+    pid_loc = tl.program_id(0)
+    loc = tl.load(loc_ptr + pid_loc).to(tl.int64)
+    loc_src_ptr = kv_buffer_ptr + loc * buffer_stride
+
+    nope_offs = tl.arange(0, nope_dim)
+    nope_src = tl.load(loc_src_ptr + nope_offs)
+    tl.store(
+        cache_k_nope_ptr + pid_loc * nope_stride + nope_offs,
+        nope_src,
+    )
+
+
 def get_mla_kv_buffer_triton(
     kv_buffer: torch.Tensor,
     loc: torch.Tensor,
     cache_k_nope: torch.Tensor,
-    cache_k_rope: torch.Tensor,
+    cache_k_rope: Optional[torch.Tensor] = None,
 ):
     # The source data type will be implicitly converted to the target data type.
     nope_dim = cache_k_nope.shape[-1]  # 512
-    rope_dim = cache_k_rope.shape[-1]  # 64
     n_loc = loc.numel()
     grid = (n_loc,)
 
+    has_rope = cache_k_rope is not None and cache_k_rope.numel() > 0
+    if not has_rope:
+        get_mla_kv_buffer_kernel_norope[grid](
+            kv_buffer,
+            cache_k_nope,
+            loc,
+            kv_buffer.stride(0),
+            cache_k_nope.stride(0),
+            nope_dim,
+        )
+        return
+
+    rope_dim = cache_k_rope.shape[-1]  # 64
     get_mla_kv_buffer_kernel[grid](
         kv_buffer,
         cache_k_nope,

python/sglang/kernels/ops/layernorm/mhc.py

@@ -1734,6 +1734,190 @@ def mhc_fused_post_pre(
     )
 
 
+def hc_expand(x: torch.Tensor, n: int) -> torch.Tensor:
+    return x.repeat(1, n)
+
+
+def hc_contract(x: torch.Tensor, n: int) -> torch.Tensor:
+    return x.unflatten(-1, (n, -1)).mean(dim=-2)
+
+
+def _mhc_pre_torch(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    import torch.nn.functional as F
+
+    s, n, h = residual.shape
+    dtype = residual.dtype
+
+    x_flat = residual.view(s, n * h).float()
+    rsqrt = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + rms_eps)
+    mixes = F.linear(x_flat, fn) * rsqrt
+
+    pre_raw = mixes[:, :n]
+    post_raw = mixes[:, n : 2 * n]
+    comb_raw = mixes[:, 2 * n :].view(s, n, n)
+    pre_base = hc_base[:n]
+    post_base = hc_base[n : 2 * n]
+    comb_base = hc_base[2 * n :].view(n, n)
+
+    pre = torch.sigmoid(pre_raw * hc_scale[0] + pre_base) + hc_pre_eps
+    post = hc_post_mult_value * torch.sigmoid(post_raw * hc_scale[1] + post_base)
+    comb = comb_raw * hc_scale[2] + comb_base
+
+    comb = comb.softmax(-1) + hc_sinkhorn_eps
+    comb = comb / (comb.sum(-2, keepdim=True) + hc_sinkhorn_eps)
+    for _ in range(sinkhorn_repeat - 1):
+        comb = comb / (comb.sum(-1, keepdim=True) + hc_sinkhorn_eps)
+        comb = comb / (comb.sum(-2, keepdim=True) + hc_sinkhorn_eps)
+
+    layer_input = (pre.unsqueeze(-1) * residual.float()).sum(dim=1).to(dtype)
+    return post.unsqueeze(-1), comb, layer_input
+
+
+def _mhc_post_torch(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    out = post_layer_mix * x.unsqueeze(1) + (
+        comb_res_mix.unsqueeze(-1) * residual.unsqueeze(2)
+    ).sum(dim=1)
+    return out.type_as(x)
+
+
+@torch._dynamo.disable
+def _mhc_pre_dispatch(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float | None = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
+    assert residual.dim() == 3, f"residual must be (s, n, h); got {residual.shape}"
+    if not envs.SGLANG_OPT_USE_TILELANG_MHC_PRE.get():
+        post_mix, comb_mix, layer_input = _mhc_pre_torch(
+            residual=residual,
+            fn=fn,
+            hc_scale=hc_scale,
+            hc_base=hc_base,
+            rms_eps=rms_eps,
+            hc_pre_eps=hc_pre_eps,
+            hc_sinkhorn_eps=hc_sinkhorn_eps,
+            hc_post_mult_value=hc_post_mult_value,
+            sinkhorn_repeat=sinkhorn_repeat,
+        )
+        return post_mix, comb_mix, layer_input, False
+
+    post_mix, comb_mix, layer_input = mhc_pre(
+        residual=residual,
+        fn=fn,
+        hc_scale=hc_scale,
+        hc_base=hc_base,
+        rms_eps=rms_eps,
+        hc_pre_eps=hc_pre_eps,
+        hc_sinkhorn_eps=hc_sinkhorn_eps,
+        hc_post_mult_value=hc_post_mult_value,
+        sinkhorn_repeat=sinkhorn_repeat,
+        norm_weight=norm_weight,
+        norm_eps=norm_eps,
+    )
+    return post_mix, comb_mix, layer_input, norm_weight is not None
+
+
+@torch._dynamo.disable
+def _mhc_post_dispatch(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    assert x.dim() == 2 and residual.dim() == 3
+    assert post_layer_mix.dim() == 3 and comb_res_mix.dim() == 3
+    if not envs.SGLANG_OPT_USE_TILELANG_MHC_POST.get():
+        return _mhc_post_torch(x, residual, post_layer_mix, comb_res_mix)
+    return mhc_post(x, residual, post_layer_mix, comb_res_mix)
+
+
+def hc_pre(
+    x: torch.Tensor,
+    hc_fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    hc_mult: int,
+    rms_eps: float,
+    hc_eps: float,
+    sinkhorn_iters: int,
+    post_mult_value: float = 2.0,
+    hc_norm_weight: torch.Tensor | None = None,
+    out_norm_weight: torch.Tensor | None = None,
+    out_norm_eps: float | None = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, bool]:
+    s, total = x.shape
+    hidden_size = total // hc_mult
+    if x.numel() == 0:
+        empty_layer_input = x.new_zeros((s, hidden_size))
+        empty_h_res = torch.zeros(
+            (s, hc_mult * hc_mult), device=x.device, dtype=torch.float32
+        )
+        empty_h_post = torch.zeros((s, hc_mult), device=x.device, dtype=torch.float32)
+        return empty_layer_input, empty_h_res, empty_h_post, False
+
+    fn = hc_fn if hc_norm_weight is None else hc_fn * hc_norm_weight
+    residual_3d = x.view(s, hc_mult, hidden_size)
+    post_mix, comb_mix, layer_input, norm_fused = _mhc_pre_dispatch(
+        residual=residual_3d,
+        fn=fn,
+        hc_scale=hc_scale,
+        hc_base=hc_base,
+        rms_eps=rms_eps,
+        hc_pre_eps=hc_eps,
+        hc_sinkhorn_eps=hc_eps,
+        hc_post_mult_value=post_mult_value,
+        sinkhorn_repeat=sinkhorn_iters,
+        norm_weight=out_norm_weight,
+        norm_eps=out_norm_eps,
+    )
+    return (
+        layer_input,
+        comb_mix.reshape(s, hc_mult * hc_mult),
+        post_mix.reshape(s, hc_mult),
+        norm_fused,
+    )
+
+
+def hc_post(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    h_post: torch.Tensor,
+    h_res: torch.Tensor,
+    hc_mult: int,
+) -> torch.Tensor:
+    s, hidden_size = x.shape
+    if s == 0:
+        return x.new_zeros((s, hc_mult * hidden_size))
+    residual = residual.view(s, hc_mult, hidden_size)
+    h_post = h_post.view(s, hc_mult, 1)
+    h_res = h_res.view(s, hc_mult, hc_mult)
+    out = _mhc_post_dispatch(x, residual, h_post, h_res)
+    return out.view(s, -1)
+
+
 def npu_hc_pre(
     x: torch.Tensor,
     hc_fn: torch.Tensor,

python/sglang/kernels/ops/moe/kpool_topk_transform.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+SUPPORTED_GROUP_TOPK = (128, 160, 192, 224, 256, 512)
+
+
+@cache_once
+def _jit_kpool_topk_transform_module(group_topk: int) -> Module:
+    assert group_topk in SUPPORTED_GROUP_TOPK, (
+        "fast_kpool_topk_transform supports pool-level topk "
+        f"{SUPPORTED_GROUP_TOPK}, got {group_topk}"
+    )
+    return load_jit(
+        f"kpool_topk_transform_{group_topk}",
+        cuda_files=["dsa/kpool_topk_transform.cuh"],
+        cuda_wrappers=[("kpool_topk_transform", "KpoolTopKTransformKernel::transform")],
+        extra_cuda_cflags=[f"-DSGL_GROUP_TOPK={group_topk}"],
+    )
+
+
+def fast_kpool_topk_transform_fused(
+    score: torch.Tensor,
+    lengths: torch.Tensor,
+    pool_size: int,
+    topk: int,
+    page_table: Optional[torch.Tensor] = None,
+    topk_indices_offset: Optional[torch.Tensor] = None,
+    row_starts: Optional[torch.Tensor] = None,
+    seq_lens: Optional[torch.Tensor] = None,
+    page_table_row_index: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert topk % pool_size == 0
+    group_topk = topk // pool_size
+    assert group_topk in SUPPORTED_GROUP_TOPK, (
+        "fast_kpool_topk_transform supports pool-level topk "
+        f"{SUPPORTED_GROUP_TOPK}, got {group_topk}"
+    )
+    assert score.dim() == 2
+    assert page_table is None or topk_indices_offset is None
+    assert page_table_row_index is None or page_table is not None
+    if seq_lens is not None:
+        assert seq_lens.dim() == 1
+        assert seq_lens.shape[0] == score.shape[0]
+    if page_table_row_index is not None:
+        assert page_table_row_index.dim() == 1
+        assert page_table_row_index.shape[0] == score.shape[0]
+
+    out_cols = topk + (pool_size - 1 if seq_lens is not None else 0)
+    dst_token_indices = score.new_empty((score.shape[0], out_cols), dtype=torch.int32)
+
+    module = _jit_kpool_topk_transform_module(group_topk)
+    module.kpool_topk_transform(
+        score,
+        lengths,
+        dst_token_indices,
+        pool_size,
+        page_table,
+        topk_indices_offset,
+        row_starts,
+        seq_lens,
+        page_table_row_index,
+    )
+    return dst_token_indices

python/sglang/srt/layers/attention/dsa/kpool_plan.py

@@ -0,0 +1,852 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, List, NamedTuple, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.dsa.kpool_fp8_index import (
+    INDEX_HEAD_DIM,
+    build_pooled_page_table_64,
+    kpool_build_ragged_layout,
+    kpool_max_closed_pools,
+    update_kpool_write_plan_cuda_graph,
+)
+from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp
+from sglang.srt.model_executor.forward_context import get_req_to_token_pool
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import is_cuda
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod
+    from sglang.srt.layers.attention.dsa_backend import DSAMetadata
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
+
+
+_RAGGED_SCRATCH_K_U8: Optional[torch.Tensor] = None
+_RAGGED_SCRATCH_K_SCALE: Optional[torch.Tensor] = None
+
+
+def _get_ragged_scratch(
+    total_k_rows: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor]:
+    global _RAGGED_SCRATCH_K_U8, _RAGGED_SCRATCH_K_SCALE
+    cur = _RAGGED_SCRATCH_K_U8
+    grow = (
+        cur is None
+        or cur.device.type != device.type
+        or (device.index is not None and cur.device.index != device.index)
+        or cur.shape[0] < total_k_rows
+    )
+    if grow:
+        _RAGGED_SCRATCH_K_U8 = torch.empty(
+            (total_k_rows, INDEX_HEAD_DIM), dtype=torch.uint8, device=device
+        )
+        _RAGGED_SCRATCH_K_SCALE = torch.empty(
+            (total_k_rows,), dtype=torch.float32, device=device
+        )
+    return _RAGGED_SCRATCH_K_U8[:total_k_rows], _RAGGED_SCRATCH_K_SCALE[:total_k_rows]
+
+
+@dataclass(frozen=True)
+class PoolWriteRows:
+    req: torch.Tensor
+    pool_id: torch.Tensor
+    n_from_tail: torch.Tensor
+    chunk_src: torch.Tensor
+    tail_logical_base: torch.Tensor
+    write_loc: torch.Tensor
+
+    @property
+    def is_empty(self) -> bool:
+        return self.pool_id.shape[0] == 0
+
+
+@dataclass(frozen=True)
+class TailWriteRows:
+    req: torch.Tensor
+    dst_logical_start: torch.Tensor
+    chunk_src: torch.Tensor
+    n_write: torch.Tensor
+
+    @property
+    def is_empty(self) -> bool:
+        return self.req.shape[0] == 0
+
+
+@dataclass(frozen=True)
+class KPoolCpInfo:
+    size: int
+    rank: int
+    owner_rank: torch.Tensor
+    local_write_mask: torch.Tensor
+
+
+@dataclass(frozen=True)
+class KPoolExtendPlan:
+    writes: PoolWriteRows
+    tails: TailWriteRows
+    pooled_seq_lens_expanded: torch.Tensor
+    seq_lens_expanded: torch.Tensor
+    ragged_concat_page_table: torch.Tensor
+    ragged_q_ks: torch.Tensor
+    ragged_q_ke: torch.Tensor
+    ragged_total_k_rows: int
+    ragged_k_u8: Optional[torch.Tensor]
+    ragged_k_scale: Optional[torch.Tensor]
+    ragged_paged_page_table: Optional[torch.Tensor]
+    ragged_paged_page_table_row_index: Optional[torch.Tensor]
+    cp: Optional[KPoolCpInfo] = None
+
+
+@dataclass(frozen=True)
+class KPoolWritePlan:
+    """``write_loc[b, p]`` is the compression destination for candidate closed
+    pool ``base_pool[b] + p``; the kernel decides which candidates closed."""
+
+    req: torch.Tensor
+    write_start: torch.Tensor
+    tail_logical_start: torch.Tensor
+    write_loc: torch.Tensor  # int64 [B, max_closed_pools]
+    num_draft_tokens: int
+    pool_seqlens_per_q: Optional[torch.Tensor] = None
+    seqlens_per_q: Optional[torch.Tensor] = None
+    pool_schedule_metadata: Optional[torch.Tensor] = None
+    effective_n_per_batch: Optional[torch.Tensor] = None
+
+
+def _is_kpool_layout_enabled(pool_size: int, real_page_size: int) -> bool:
+    return pool_size > 1 and real_page_size == 64 and real_page_size % pool_size == 0
+
+
+@dataclass
+class _KPoolCpuPlan:
+    pool_batch_idx: List[int] = field(default_factory=list)
+    pool_req: List[int] = field(default_factory=list)
+    pool_pool_id: List[int] = field(default_factory=list)
+    pool_n_from_tail: List[int] = field(default_factory=list)
+    pool_chunk_src: List[int] = field(default_factory=list)
+    pool_tail_logical_base: List[int] = field(default_factory=list)
+
+    tail_req: List[int] = field(default_factory=list)
+    tail_dst_logical_start: List[int] = field(default_factory=list)
+    tail_chunk_src: List[int] = field(default_factory=list)
+    tail_n_write: List[int] = field(default_factory=list)
+
+    ragged_q_len: List[int] = field(default_factory=list)
+    ragged_pool_pages: List[int] = field(default_factory=list)
+    cu_pages_excl: List[int] = field(default_factory=list)
+    cu_q_len_excl: List[int] = field(default_factory=list)
+    total_pool_pages: int = 0
+
+
+class _KPoolDecompose(NamedTuple):
+    first_slot: int
+    base_pool: int
+    n_pool: int
+    tail_n_write: int
+
+
+def _decompose_compress(start: int, length: int, pool_size: int) -> _KPoolDecompose:
+    first_slot = start % pool_size
+    base_pool = start // pool_size
+    n_pool = (start + length) // pool_size - base_pool
+    consumed = max(0, n_pool * pool_size - first_slot)
+    tail_n = length - consumed
+    return _KPoolDecompose(
+        first_slot=first_slot,
+        base_pool=base_pool,
+        n_pool=n_pool,
+        tail_n_write=tail_n,
+    )
+
+
+def _append_compress_rows(
+    plan: _KPoolCpuPlan,
+    pool_size: int,
+    batch_size: int,
+    extend_seq_lens_cpu: List[int],
+    seq_lens_cpu: List[int],
+    req_pool_indices_cpu: List[int],
+) -> None:
+    q_offset = 0
+    for i in range(batch_size):
+        q_len = extend_seq_lens_cpu[i]
+        assert q_len > 0, f"extend_seq_lens_cpu[{i}] = {q_len}; expected > 0"
+
+        seq_len = seq_lens_cpu[i]
+        req = req_pool_indices_cpu[i]
+        d = _decompose_compress(seq_len - q_len, q_len, pool_size)
+
+        if d.n_pool > 0:
+            plan.pool_batch_idx.extend([i] * d.n_pool)
+            plan.pool_req.extend([req] * d.n_pool)
+            plan.pool_pool_id.extend(range(d.base_pool, d.base_pool + d.n_pool))
+            plan.pool_n_from_tail.append(d.first_slot)
+            plan.pool_n_from_tail.extend([0] * (d.n_pool - 1))
+            bulk_start = q_offset + pool_size - d.first_slot
+            plan.pool_chunk_src.append(q_offset)
+            plan.pool_chunk_src.extend(
+                range(bulk_start, bulk_start + (d.n_pool - 1) * pool_size, pool_size)
+            )
+            plan.pool_tail_logical_base.extend(
+                range(
+                    d.base_pool * pool_size,
+                    (d.base_pool + d.n_pool) * pool_size,
+                    pool_size,
+                )
+            )
+
+        if d.tail_n_write > 0:
+            consumed = q_len - d.tail_n_write
+            plan.tail_req.append(req)
+            plan.tail_dst_logical_start.append(seq_len - q_len + consumed)
+            plan.tail_chunk_src.append(q_offset + consumed)
+            plan.tail_n_write.append(d.tail_n_write)
+
+        q_offset += q_len
+
+
+def _append_local_rows(
+    plan: _KPoolCpuPlan,
+    pool_size: int,
+    slots_per_page: int,
+    local_extend_seq_lens_cpu: List[int],
+    local_seq_lens_cpu: List[int],
+) -> None:
+    q_offset = 0
+    for q_len, seq_len in zip(
+        local_extend_seq_lens_cpu, local_seq_lens_cpu, strict=True
+    ):
+        assert q_len > 0, f"local_extend_seq_lens_cpu has non-positive {q_len = }"
+        plan.ragged_q_len.append(q_len)
+        pool_seq_len = seq_len // pool_size
+        pool_pages_i = (pool_seq_len + slots_per_page - 1) // slots_per_page
+        plan.ragged_pool_pages.append(pool_pages_i)
+        plan.cu_pages_excl.append(plan.total_pool_pages)
+        plan.cu_q_len_excl.append(q_offset)
+        plan.total_pool_pages += pool_pages_i
+        q_offset += q_len
+
+
+def _kpool_cpu_plan(
+    forward_batch: ForwardBatch,
+    pool_size: int,
+    slots_per_page: int,
+    *,
+    local_extend_seq_lens_cpu: Optional[List[int]] = None,
+    local_seq_lens_cpu: Optional[List[int]] = None,
+) -> _KPoolCpuPlan:
+    plan = _KPoolCpuPlan()
+
+    extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
+    if isinstance(extend_seq_lens_cpu, torch.Tensor):
+        extend_seq_lens_cpu = extend_seq_lens_cpu.tolist()
+    seq_lens_cpu = forward_batch.seq_lens_cpu.tolist()
+    req_pool_indices_cpu = forward_batch.req_pool_indices.tolist()
+
+    _append_compress_rows(
+        plan,
+        pool_size,
+        forward_batch.batch_size,
+        extend_seq_lens_cpu,
+        seq_lens_cpu,
+        req_pool_indices_cpu,
+    )
+
+    if local_extend_seq_lens_cpu is None:
+        local_extend_seq_lens_cpu = extend_seq_lens_cpu
+        local_seq_lens_cpu = seq_lens_cpu
+
+    _append_local_rows(
+        plan,
+        pool_size,
+        slots_per_page,
+        local_extend_seq_lens_cpu,
+        local_seq_lens_cpu,
+    )
+    return plan
+
+
+def _kpool_plan_to_gpu(
+    cpu: _KPoolCpuPlan,
+    forward_batch: ForwardBatch,
+    full_real_page_table: torch.Tensor,
+    local_real_page_table: torch.Tensor,
+    local_seqlens_expanded: torch.Tensor,
+    local_req_pool_indices: torch.Tensor,
+    pool_size: int,
+    slots_per_page: int,
+    topk_transform_method: TopkTransformMethod,
+) -> KPoolExtendPlan:
+    from sglang.srt.layers.attention.dsa.dsa_topk_backend import TopkTransformMethod
+
+    device = forward_batch.seq_lens.device
+    n_pool = len(cpu.pool_pool_id)
+    n_tail = len(cpu.tail_req)
+    n_rag = len(cpu.ragged_q_len)
+
+    total_pool_pages = cpu.total_pool_pages
+    ragged_total_k_rows = total_pool_pages * slots_per_page
+
+    need_paged = (
+        topk_transform_method == TopkTransformMethod.PAGED
+        and envs.SGLANG_DSA_FUSE_TOPK.get()
+        and n_rag > 0
+    )
+
+    i64_total = 4 * n_pool + 2 * n_tail
+    if i64_total > 0:
+        i64_cpu = torch.tensor(
+            cpu.pool_req
+            + cpu.pool_pool_id
+            + cpu.pool_chunk_src
+            + cpu.pool_batch_idx
+            + cpu.tail_req
+            + cpu.tail_chunk_src,
+            dtype=torch.int64,
+            pin_memory=True,
+        )
+        i64_gpu = i64_cpu.to(device, non_blocking=True)
+        c = 0
+        pool_req_t = i64_gpu[c : c + n_pool]
+        c += n_pool
+        pool_pool_id_t = i64_gpu[c : c + n_pool]
+        c += n_pool
+        pool_chunk_src_t = i64_gpu[c : c + n_pool]
+        c += n_pool
+        pool_batch_idx_t = i64_gpu[c : c + n_pool]
+        c += n_pool
+        tail_req_t = i64_gpu[c : c + n_tail]
+        c += n_tail
+        tail_chunk_src_t = i64_gpu[c : c + n_tail]
+    else:
+        empty_i64 = torch.empty((0,), dtype=torch.int64, device=device)
+        pool_req_t = pool_pool_id_t = pool_chunk_src_t = pool_batch_idx_t = empty_i64
+        tail_req_t = tail_chunk_src_t = empty_i64
+
+    i32_total = 2 * n_pool + 2 * n_tail + 4 * n_rag
+    if i32_total > 0:
+        i32_cpu = torch.tensor(
+            cpu.pool_n_from_tail
+            + cpu.pool_tail_logical_base
+            + cpu.tail_dst_logical_start
+            + cpu.tail_n_write
+            + cpu.ragged_pool_pages
+            + cpu.ragged_q_len
+            + cpu.cu_pages_excl
+            + cpu.cu_q_len_excl,
+            dtype=torch.int32,
+            pin_memory=True,
+        )
+        i32_gpu = i32_cpu.to(device, non_blocking=True)
+        c = 0
+        pool_n_from_tail_t = i32_gpu[c : c + n_pool]
+        c += n_pool
+        pool_tail_logical_base_t = i32_gpu[c : c + n_pool]
+        c += n_pool
+        tail_dst_logical_start_t = i32_gpu[c : c + n_tail]
+        c += n_tail
+        tail_n_write_t = i32_gpu[c : c + n_tail]
+        c += n_tail
+        ragged_pool_pages_t = i32_gpu[c : c + n_rag]
+        c += n_rag
+        ragged_q_len_t = i32_gpu[c : c + n_rag]
+        c += n_rag
+        cu_pages_excl_t = i32_gpu[c : c + n_rag]
+        c += n_rag
+        cu_q_len_excl_t = i32_gpu[c : c + n_rag]
+    else:
+        empty_i32 = torch.empty((0,), dtype=torch.int32, device=device)
+        pool_n_from_tail_t = pool_tail_logical_base_t = empty_i32
+        tail_dst_logical_start_t = tail_n_write_t = empty_i32
+        ragged_pool_pages_t = ragged_q_len_t = empty_i32
+        cu_pages_excl_t = cu_q_len_excl_t = empty_i32
+
+    if n_pool > 0:
+        pool_page_group = torch.div(
+            pool_pool_id_t, slots_per_page, rounding_mode="floor"
+        )
+        token_page_row = pool_page_group * pool_size
+        packed_page = full_real_page_table[pool_batch_idx_t, token_page_row].to(
+            torch.int64
+        )
+        pool_write_locs = packed_page * slots_per_page + torch.remainder(
+            pool_pool_id_t, slots_per_page
+        )
+    else:
+        pool_write_locs = torch.empty((0,), dtype=torch.int64, device=device)
+
+    pooled_seq_lens_expanded = torch.div(
+        local_seqlens_expanded, pool_size, rounding_mode="floor"
+    ).to(torch.int32)
+
+    if n_rag > 0:
+        (
+            ragged_concat_page_table,
+            ragged_q_ks,
+            ragged_q_ke,
+        ) = kpool_build_ragged_layout(
+            full_page_table=local_real_page_table,
+            cu_pages_excl=cu_pages_excl_t,
+            ragged_pool_pages=ragged_pool_pages_t,
+            cu_q_len_excl=cu_q_len_excl_t,
+            ragged_q_len=ragged_q_len_t,
+            pooled_seq_lens_expanded=pooled_seq_lens_expanded,
+            slots_per_page=slots_per_page,
+            total_pool_pages=total_pool_pages,
+            total_q=pooled_seq_lens_expanded.shape[0],
+            pool_size=pool_size,
+        )
+    else:
+        empty_i32_dev = torch.empty((0,), dtype=torch.int32, device=device)
+        ragged_concat_page_table = empty_i32_dev
+        ragged_q_ks = empty_i32_dev
+        ragged_q_ke = empty_i32_dev
+
+    ragged_paged_page_table = None
+    ragged_paged_page_table_row_index = None
+    if need_paged:
+        req_to_token = get_req_to_token_pool().req_to_token
+        ragged_paged_page_table_row_index = torch.repeat_interleave(
+            local_req_pool_indices.to(torch.int32), ragged_q_len_t
+        )
+        ragged_paged_page_table = req_to_token
+
+    if ragged_total_k_rows > 0:
+        ragged_k_u8, ragged_k_scale = _get_ragged_scratch(ragged_total_k_rows, device)
+    else:
+        ragged_k_u8 = None
+        ragged_k_scale = None
+
+    return KPoolExtendPlan(
+        writes=PoolWriteRows(
+            req=pool_req_t,
+            pool_id=pool_pool_id_t,
+            n_from_tail=pool_n_from_tail_t,
+            chunk_src=pool_chunk_src_t,
+            tail_logical_base=pool_tail_logical_base_t,
+            write_loc=pool_write_locs,
+        ),
+        tails=TailWriteRows(
+            req=tail_req_t,
+            dst_logical_start=tail_dst_logical_start_t,
+            chunk_src=tail_chunk_src_t,
+            n_write=tail_n_write_t,
+        ),
+        pooled_seq_lens_expanded=pooled_seq_lens_expanded,
+        seq_lens_expanded=local_seqlens_expanded,
+        ragged_concat_page_table=ragged_concat_page_table,
+        ragged_q_ks=ragged_q_ks,
+        ragged_q_ke=ragged_q_ke,
+        ragged_total_k_rows=ragged_total_k_rows,
+        ragged_k_u8=ragged_k_u8,
+        ragged_k_scale=ragged_k_scale,
+        ragged_paged_page_table=ragged_paged_page_table,
+        ragged_paged_page_table_row_index=ragged_paged_page_table_row_index,
+        cp=_kpool_cp_owner_rank(forward_batch, n_pool, device),
+    )
+
+
+def _kpool_cp_owner_rank(
+    forward_batch: ForwardBatch,
+    n_pool: int,
+    device: torch.device,
+) -> Optional[KPoolCpInfo]:
+    if not dsa_use_prefill_cp(forward_batch):
+        return None
+
+    cp_size = get_parallel().attn_cp_size
+    if cp_size <= 1:
+        return None
+
+    cp_rank = get_parallel().attn_cp_rank
+    if n_pool > 0:
+        owner = torch.arange(n_pool, dtype=torch.int32, device=device) % cp_size
+        local_write_mask = owner == cp_rank
+    else:
+        owner = torch.empty((0,), dtype=torch.int32, device=device)
+        local_write_mask = torch.empty((0,), dtype=torch.bool, device=device)
+    return KPoolCpInfo(
+        size=cp_size,
+        rank=cp_rank,
+        owner_rank=owner,
+        local_write_mask=local_write_mask,
+    )
+
+
+def init_kpool_extend_metadata(
+    metadata: DSAMetadata,
+    forward_batch: ForwardBatch,
+    *,
+    pool_size: int,
+    real_page_size: int,
+    slots_per_page: int,
+    topk_transform_method: TopkTransformMethod,
+    full_real_page_table: torch.Tensor,
+    full_seqlens_expanded: torch.Tensor,
+    local_real_page_table: Optional[torch.Tensor] = None,
+    local_seqlens_expanded: Optional[torch.Tensor] = None,
+    local_extend_seq_lens_cpu: Optional[List[int]] = None,
+    local_seq_lens_cpu: Optional[List[int]] = None,
+    local_req_pool_indices: Optional[torch.Tensor] = None,
+) -> DSAMetadata:
+    mode = forward_batch.forward_mode
+    is_extend_like = mode.is_extend_without_speculative() or mode.is_draft_extend_v2()
+    if (
+        not _is_kpool_layout_enabled(pool_size, real_page_size)
+        or not is_extend_like
+        or forward_batch.extend_seq_lens_cpu is None
+        or forward_batch.seq_lens_cpu is None
+    ):
+        return metadata
+
+    if local_real_page_table is None:
+        local_real_page_table = full_real_page_table
+    if local_seqlens_expanded is None:
+        local_seqlens_expanded = full_seqlens_expanded
+    if local_req_pool_indices is None:
+        local_req_pool_indices = forward_batch.req_pool_indices
+
+    cpu = _kpool_cpu_plan(
+        forward_batch,
+        pool_size,
+        slots_per_page,
+        local_extend_seq_lens_cpu=local_extend_seq_lens_cpu,
+        local_seq_lens_cpu=local_seq_lens_cpu,
+    )
+    plan = _kpool_plan_to_gpu(
+        cpu,
+        forward_batch,
+        full_real_page_table,
+        local_real_page_table,
+        local_seqlens_expanded,
+        local_req_pool_indices,
+        pool_size,
+        slots_per_page,
+        topk_transform_method,
+    )
+    return dataclasses.replace(metadata, kpool_extend_plan=plan)
+
+
+_DEEP_GEMM_MODULE = None
+_DEEP_GEMM_IMPORT_FAILED = False
+
+
+def _get_deep_gemm():
+    # Cache import failure too, because this helper runs on every graph-replay
+    # metadata refresh.
+    global _DEEP_GEMM_MODULE, _DEEP_GEMM_IMPORT_FAILED
+    if _DEEP_GEMM_MODULE is None and not _DEEP_GEMM_IMPORT_FAILED:
+        try:
+            import deep_gemm
+        except (ImportError, ModuleNotFoundError):
+            _DEEP_GEMM_IMPORT_FAILED = True
+        else:
+            _DEEP_GEMM_MODULE = deep_gemm
+    return _DEEP_GEMM_MODULE
+
+
+def _compute_pool_schedule_metadata(
+    pool_seqlens: torch.Tensor,
+    *,
+    slots_per_page: int,
+) -> Optional[torch.Tensor]:
+    if not is_cuda():
+        return None
+    deep_gemm = _get_deep_gemm()
+    if deep_gemm is None:
+        return None
+    return deep_gemm.get_paged_mqa_logits_metadata(
+        pool_seqlens.contiguous().view(-1, 1).clamp(min=1),
+        slots_per_page,
+        deep_gemm.get_num_sms(),
+    )
+
+
+def init_pooled_paged_mqa_metadata(
+    metadata: DSAMetadata,
+    seqlens_32: torch.Tensor,
+    forward_mode: ForwardMode,
+    *,
+    pool_size: int,
+    real_page_size: int,
+    slots_per_page: int,
+    build_schedule_metadata: bool = True,
+) -> DSAMetadata:
+    if (
+        not _is_kpool_layout_enabled(pool_size, real_page_size)
+        or not is_cuda()
+        or not forward_mode.is_decode_or_idle()
+    ):
+        return metadata
+
+    pool_seqlens = torch.div(seqlens_32, pool_size, rounding_mode="floor").to(
+        torch.int32
+    )
+    pooled_page_table = build_pooled_page_table_64(
+        metadata.real_page_table, pool_size
+    ).contiguous()
+    schedule = (
+        _compute_pool_schedule_metadata(
+            pool_seqlens,
+            slots_per_page=slots_per_page,
+        )
+        if build_schedule_metadata
+        else None
+    )
+    return dataclasses.replace(
+        metadata,
+        pooled_index_kpool=pool_size,
+        pooled_cache_seqlens_int32=pool_seqlens,
+        pooled_real_page_table=pooled_page_table,
+        pooled_paged_mqa_schedule_metadata=schedule,
+    )
+
+
+def update_pooled_paged_mqa_metadata(
+    metadata: DSAMetadata,
+    seqlens_32: torch.Tensor,
+    forward_mode: ForwardMode,
+    *,
+    pool_size: int,
+    real_page_size: int,
+    slots_per_page: int,
+    build_schedule_metadata: bool = True,
+) -> None:
+    if (
+        not _is_kpool_layout_enabled(pool_size, real_page_size)
+        or not is_cuda()
+        or not forward_mode.is_decode_or_idle()
+    ):
+        return
+
+    if (
+        metadata.pooled_index_kpool != pool_size
+        or metadata.pooled_cache_seqlens_int32 is None
+        or metadata.pooled_real_page_table is None
+    ):
+        return
+
+    pool_seqlens = torch.div(seqlens_32, pool_size, rounding_mode="floor").to(
+        torch.int32
+    )
+    metadata.pooled_cache_seqlens_int32[: pool_seqlens.shape[0]].copy_(pool_seqlens)
+    pooled_page_table = build_pooled_page_table_64(
+        metadata.real_page_table, pool_size
+    ).contiguous()
+    metadata.pooled_real_page_table[
+        : pooled_page_table.shape[0], : pooled_page_table.shape[1]
+    ].copy_(pooled_page_table)
+
+    if (
+        build_schedule_metadata
+        and metadata.pooled_paged_mqa_schedule_metadata is not None
+    ):
+        new_schedule = _compute_pool_schedule_metadata(
+            metadata.pooled_cache_seqlens_int32,
+            slots_per_page=slots_per_page,
+        )
+        if new_schedule is not None:
+            metadata.pooled_paged_mqa_schedule_metadata.copy_(new_schedule)
+
+
+def _alloc_kpool_write_plan_buffers(
+    *,
+    max_bs: int,
+    num_draft_tokens: int,
+    pool_size: int,
+    device: torch.device,
+    is_verify: bool,
+    is_v2: bool = False,
+) -> KPoolWritePlan:
+    max_closed_pools = kpool_max_closed_pools(num_draft_tokens, pool_size)
+    verify_extras = {}
+    if is_verify:
+        n_rows = max_bs * num_draft_tokens
+        verify_extras = dict(
+            pool_seqlens_per_q=torch.zeros(n_rows, dtype=torch.int32, device=device),
+            seqlens_per_q=torch.zeros(n_rows, dtype=torch.int32, device=device),
+        )
+    if is_v2:
+        verify_extras["effective_n_per_batch"] = torch.zeros(
+            max_bs, dtype=torch.int32, device=device
+        )
+    return KPoolWritePlan(
+        req=torch.zeros(max_bs, dtype=torch.int64, device=device),
+        write_start=torch.zeros(max_bs, dtype=torch.int32, device=device),
+        tail_logical_start=torch.zeros(max_bs, dtype=torch.int32, device=device),
+        write_loc=torch.zeros(
+            max_bs, max_closed_pools, dtype=torch.int64, device=device
+        ),
+        num_draft_tokens=num_draft_tokens,
+        **verify_extras,
+    )
+
+
+def init_kpool_write_plan_capture(
+    metadata: DSAMetadata,
+    *,
+    max_bs: int,
+    pool_size: int,
+    real_page_size: int,
+    num_draft_tokens: int,
+    device: torch.device,
+    is_verify: bool,
+    slots_per_page: int,
+    is_v2: bool = False,
+    build_schedule_metadata: bool = True,
+) -> DSAMetadata:
+    if not _is_kpool_layout_enabled(pool_size, real_page_size) or num_draft_tokens == 0:
+        return metadata
+
+    plan = _alloc_kpool_write_plan_buffers(
+        max_bs=max_bs,
+        num_draft_tokens=num_draft_tokens,
+        pool_size=pool_size,
+        device=device,
+        is_verify=is_verify,
+        is_v2=is_v2,
+    )
+    if is_verify and build_schedule_metadata:
+        schedule = _compute_pool_schedule_metadata(
+            plan.pool_seqlens_per_q,
+            slots_per_page=slots_per_page,
+        )
+        plan = dataclasses.replace(plan, pool_schedule_metadata=schedule)
+    return dataclasses.replace(metadata, kpool_write_plan=plan)
+
+
+def update_kpool_write_plan(
+    metadata: DSAMetadata,
+    *,
+    write_start: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    real_page_table: torch.Tensor,
+    pool_size: int,
+    real_page_size: int,
+    num_draft_tokens: int,
+    forward_mode: ForwardMode,
+    slots_per_page: int,
+    effective_n_per_batch: Optional[torch.Tensor] = None,
+    include_deep_gemm_schedule: bool = True,
+) -> None:
+    if not _is_kpool_layout_enabled(pool_size, real_page_size) or not is_cuda():
+        return
+    is_verify = forward_mode.is_target_verify()
+    is_decode = forward_mode.is_decode_or_idle()
+    is_v2 = forward_mode.is_draft_extend_v2()
+    if not (is_verify or is_decode or is_v2):
+        return
+
+    plan = metadata.kpool_write_plan
+    assert plan is not None, "kpool_write_plan must be allocated before update"
+    update_kpool_write_plan_cuda_graph(
+        write_start=write_start,
+        req_pool_indices=req_pool_indices,
+        real_page_table=real_page_table,
+        req_out=plan.req,
+        write_start_out=plan.write_start,
+        tail_logical_start_out=plan.tail_logical_start,
+        write_loc_out=plan.write_loc,
+        pool_seqlens_per_q_out=plan.pool_seqlens_per_q,
+        seqlens_per_q_out=plan.seqlens_per_q,
+        pool_size=pool_size,
+        num_draft_tokens=num_draft_tokens,
+        slots_per_page=slots_per_page,
+    )
+
+    if (
+        is_v2
+        and effective_n_per_batch is not None
+        and plan.effective_n_per_batch is not None
+    ):
+        plan.effective_n_per_batch[: effective_n_per_batch.shape[0]].copy_(
+            effective_n_per_batch.to(torch.int32)
+        )
+
+    # In-graph replay updates plan lengths too late for host schedule construction;
+    # the caller rebuilds the schedule from raw seq_lens out of graph.
+    if include_deep_gemm_schedule and plan.pool_schedule_metadata is not None:
+        new_schedule = _compute_pool_schedule_metadata(
+            plan.pool_seqlens_per_q,
+            slots_per_page=slots_per_page,
+        )
+        if new_schedule is not None:
+            plan.pool_schedule_metadata.copy_(new_schedule)
+
+
+def refresh_kpool_pool_schedule_from(
+    metadata: DSAMetadata,
+    pool_seqlens_per_q: torch.Tensor,
+    *,
+    slots_per_page: int,
+) -> None:
+    """Use an explicit source because the captured plan buffer remains stale
+    until replay."""
+    plan = metadata.kpool_write_plan
+    if plan is None or plan.pool_schedule_metadata is None:
+        return
+    new_schedule = _compute_pool_schedule_metadata(
+        pool_seqlens_per_q,
+        slots_per_page=slots_per_page,
+    )
+    if new_schedule is not None:
+        plan.pool_schedule_metadata.copy_(new_schedule)
+
+
+def init_kpool_write_plan(
+    metadata: DSAMetadata,
+    forward_batch: ForwardBatch,
+    *,
+    pool_size: int,
+    real_page_size: int,
+    real_page_table: torch.Tensor,
+    num_draft_tokens: int,
+    write_start: torch.Tensor,
+    slots_per_page: int,
+    effective_n_per_batch: Optional[torch.Tensor] = None,
+    build_schedule_metadata: bool = True,
+) -> DSAMetadata:
+    forward_mode = forward_batch.forward_mode
+    is_verify = forward_mode.is_target_verify()
+    is_decode = forward_mode.is_decode_or_idle()
+    is_v2 = forward_mode.is_draft_extend_v2()
+    is_ring_write = is_verify or is_decode or is_v2
+    if not _is_kpool_layout_enabled(pool_size, real_page_size) or not is_ring_write:
+        return metadata
+
+    pool = getattr(forward_batch, "token_to_kv_pool", None)
+    if (is_verify or is_v2) and pool is not None:
+        assert pool.tail_extra_slots == num_draft_tokens, (
+            f"tail_extra_slots mismatch: pool={pool.tail_extra_slots}, "
+            f"forward={num_draft_tokens}"
+        )
+
+    metadata = init_kpool_write_plan_capture(
+        metadata,
+        max_bs=forward_batch.seq_lens.shape[0],
+        pool_size=pool_size,
+        real_page_size=real_page_size,
+        num_draft_tokens=num_draft_tokens,
+        device=forward_batch.seq_lens.device,
+        is_verify=is_verify or is_v2,
+        slots_per_page=slots_per_page,
+        is_v2=is_v2,
+        build_schedule_metadata=build_schedule_metadata,
+    )
+    update_kpool_write_plan(
+        metadata,
+        write_start=write_start,
+        req_pool_indices=forward_batch.req_pool_indices,
+        real_page_table=real_page_table,
+        pool_size=pool_size,
+        real_page_size=real_page_size,
+        num_draft_tokens=num_draft_tokens,
+        forward_mode=forward_mode,
+        slots_per_page=slots_per_page,
+        effective_n_per_batch=effective_n_per_batch,
+    )
+    return metadata

python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py

@@ -138,6 +138,8 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
         num_tokens = q_n.shape[0]
         g_in = g[0][:num_tokens]  # raw forget gate; activated inside chunk_kda_cutedsl
         beta_in = beta[0][:num_tokens].to(torch.float32)
+        if kwargs.get("beta_is_raw"):
+            beta_in = beta_in.sigmoid()
         cu_seqlens = query_start_loc.to(torch.int32)
 
         # Pool state I/O is fused into the h kernel's TMA load/store: pass the

python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py

@@ -39,6 +39,7 @@ def _triton_fallback(
     A_log=None,
     dt_bias=None,
     lower_bound=None,
+    beta_is_raw=False,
     return_intermediate_states=False,
 ):
     """Fall back to the Triton chunk_kda kernel (handles all preprocessing).
@@ -64,6 +65,7 @@ def _triton_fallback(
         A_log=A_log,
         dt_bias=dt_bias,
         lower_bound=lower_bound,
+        beta_is_raw=beta_is_raw,
         output_intermediate_states=return_intermediate_states,
     )
 
@@ -114,6 +116,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
         lower_bound: Optional[float] = None,
         extend_seq_lens_cpu: Optional[list] = None,
         is_spec_decode: bool = False,
+        beta_is_raw: bool = False,
         return_intermediate_states: bool = False,
         **kwargs,
     ) -> torch.Tensor:
@@ -136,6 +139,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
                 A_log=A_log,
                 dt_bias=dt_bias,
                 lower_bound=lower_bound,
+                beta_is_raw=beta_is_raw,
                 return_intermediate_states=return_intermediate_states,
             )
 
@@ -152,6 +156,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
                 A_log=A_log,
                 dt_bias=dt_bias,
                 lower_bound=lower_bound,
+                beta_is_raw=beta_is_raw,
             ),
             None,
         )
@@ -206,6 +211,7 @@ class FlashKDAKernel(LinearAttnKernelBase):
         A_log: Optional[torch.Tensor] = None,
         dt_bias: Optional[torch.Tensor] = None,
         lower_bound: Optional[float] = None,
+        beta_is_raw: bool = False,
     ) -> torch.Tensor:
         flash_kda = _load_flash_kda()
 
@@ -223,12 +229,11 @@ class FlashKDAKernel(LinearAttnKernelBase):
         v = v.contiguous()
         g = g.contiguous()
 
-        # KimiDeltaAttention.forward already applies sigmoid to beta on the
+        # FlashKDA applies sigmoid internally; invert only the already-activated
-        # prefill path, but flash_kda expects beta LOGITS (it sigmoids
+        # Kimi beta path.
-        # internally). Invert back so the kernel recovers the intended value:
+        if not beta_is_raw:
-        # sigmoid(logit(p)) == p. (triton/cuLA consume the post-sigmoid beta.)
+            beta = torch.logit(beta.float().clamp_(1e-7, 1.0 - 1e-7))
-        beta = torch.logit(beta.float().clamp_(1e-7, 1.0 - 1e-7)).to(torch.bfloat16)
+        beta = beta.to(torch.bfloat16).contiguous()
-        beta = beta.contiguous()
 
         # flash_kda wants A_log [H] fp32 and dt_bias [H, K] fp32. The model
         # stores A_log as [1, 1, H, 1] and dt_bias as 1D [H*K], so reshape both.

python/sglang/srt/layers/attention/linear/kernels/kda_triton.py

@@ -171,7 +171,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
         intermediate_states_buffer: torch.Tensor,
         intermediate_state_indices: torch.Tensor,
         cache_steps: int,
-        retrieve_parent_token: torch.Tensor,
+        retrieve_parent_token: Optional[torch.Tensor],
         lower_bound: Optional[float] = None,
         # fused ReplaySSM ring-write (dense verify only; off elsewhere).
         cache_ring: bool = False,
@@ -229,6 +229,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
         A_log: Optional[torch.Tensor] = None,
         dt_bias: Optional[torch.Tensor] = None,
         lower_bound: Optional[float] = None,
+        beta_is_raw: bool = False,
         return_intermediate_states: bool = False,
         **kwargs,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
@@ -245,5 +246,6 @@ class TritonKDAKernel(LinearAttnKernelBase):
             A_log=A_log,
             dt_bias=dt_bias,
             lower_bound=lower_bound,
+            beta_is_raw=beta_is_raw,
             output_intermediate_states=return_intermediate_states,
         )

test/registered/kernels/ops/attention/test_kda_helion.py

@@ -706,6 +706,7 @@ def _compare_prefill(
     A_log: torch.Tensor | None = None,
     dt_bias: torch.Tensor | None = None,
     lower_bound: float | None = None,
+    beta_is_raw: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     batch, tokens, heads, key_dim = q.shape
     value_dim = v.size(-1)
@@ -740,7 +741,8 @@ def _compare_prefill(
     k_rows = reference_k.view(batch * tokens, heads, key_dim)
     v_rows = v.view(batch * tokens, heads, value_dim).float()
     gate_rows = reference_gate.view(batch * tokens, heads, key_dim)
-    beta_rows = beta.view(batch * tokens, heads).float()
+    reference_beta = beta.float().sigmoid() if beta_is_raw else beta
+    beta_rows = reference_beta.view(batch * tokens, heads).float()
     out_rows = reference_out.view(batch * tokens, heads, value_dim)
 
     if cu_seqlens is None:
@@ -811,6 +813,7 @@ def _compare_prefill(
         A_log=A_log,
         dt_bias=dt_bias,
         lower_bound=lower_bound,
+        beta_is_raw=beta_is_raw,
     )
 
     assert helion_out.data_ptr() == helion_v.data_ptr()
@@ -852,6 +855,42 @@ def test_fixed_partial_prefill_and_state_pool_contract() -> None:
     assert torch.equal(helion_state[untouched], state[untouched])
 
 
+def test_raw_beta_prefill_contract() -> None:
+    torch.manual_seed(811)
+    batch, tokens, heads, key_dim, value_dim = 2, 17, 2, 32, 32
+    q = torch.randn(batch, tokens, heads, key_dim, device="cuda", dtype=torch.bfloat16)
+    # Keep the unnormalized recurrence numerically contractive while still
+    # exercising the no-QK-L2-normalization path. Unit-scale random keys make
+    # (I - beta * k k^T) expansive and obscure the raw-beta contract with
+    # exponentially amplified BF16 round-off.
+    k = torch.randn_like(q) * 0.05
+    v = torch.randn(
+        batch, tokens, heads, value_dim, device="cuda", dtype=torch.bfloat16
+    )
+    # Keep the recurrent decay contractive so the raw-beta check measures the
+    # sigmoid conversion instead of amplifying BF16 round-off exponentially.
+    gate = -torch.rand_like(q) * 0.2
+    raw_beta = torch.linspace(
+        -2,
+        2,
+        steps=batch * tokens * heads,
+        device="cuda",
+    ).reshape(batch, tokens, heads)
+    indices = torch.tensor([3, 1], device="cuda", dtype=torch.int32)
+    state = torch.randn(5, heads, value_dim, key_dim, device="cuda") * 0.01
+
+    _compare_prefill(
+        q,
+        k,
+        v,
+        gate,
+        raw_beta,
+        state,
+        indices,
+        beta_is_raw=True,
+    )
+
+
 @pytest.mark.parametrize("is_varlen", [False, True], ids=["fixed", "varlen"])
 def test_single_token_prefill_does_not_poison_later_shapes(
     is_varlen: bool,

test/registered/kernels/test_dsa_kpool_multi_pool.py

@@ -0,0 +1,213 @@
+"""CUDA regressions for DSA kpool speculative writes spanning multiple pools."""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.layers.attention.dsa.kpool_fp8_index import (
+    INDEX_HEAD_DIM,
+    kpool_assemble_softmax_rotate_write_cache,
+    kpool_max_closed_pools,
+    kpool_write_tail_and_maybe_compress,
+    update_kpool_write_plan_cuda_graph,
+)
+from sglang.srt.layers.attention.dsa.kpool_plan import (
+    _alloc_kpool_write_plan_buffers,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+
+@unittest.skipUnless(torch.cuda.is_available(), "Test requires CUDA")
+class TestDsaKpoolMultiPool(CustomTestCase):
+    POOL_SIZE = 4
+    PAGE_SIZE = 64
+    SLOTS_PER_PAGE = 64
+    NUM_DRAFT_TOKENS = 6
+
+    def _pool(self) -> SimpleNamespace:
+        return SimpleNamespace(
+            page_size=self.PAGE_SIZE,
+            index_head_dim=INDEX_HEAD_DIM,
+            slots_per_page=self.SLOTS_PER_PAGE,
+            index_kpool=self.POOL_SIZE,
+            tail_extra_slots=self.NUM_DRAFT_TOKENS,
+            quant_block_size=128,
+        )
+
+    def _empty_cache(self) -> torch.Tensor:
+        page_nbytes = self.SLOTS_PER_PAGE * INDEX_HEAD_DIM + self.SLOTS_PER_PAGE * 4
+        return torch.zeros((1, page_nbytes), dtype=torch.uint8, device="cuda")
+
+    def test_write_plan_records_every_candidate_pool(self):
+        batch_size = 2
+        num_draft_tokens = self.NUM_DRAFT_TOKENS
+        max_closed_pools = kpool_max_closed_pools(num_draft_tokens, self.POOL_SIZE)
+        self.assertEqual(max_closed_pools, 2)
+
+        plan = _alloc_kpool_write_plan_buffers(
+            max_bs=batch_size,
+            num_draft_tokens=num_draft_tokens,
+            pool_size=self.POOL_SIZE,
+            device=torch.device("cuda"),
+            is_verify=True,
+        )
+        self.assertEqual(plan.write_loc.shape, (batch_size, max_closed_pools))
+
+        write_start = torch.tensor([3, 255], dtype=torch.int32, device="cuda")
+        req_pool_indices = torch.tensor([7, 11], dtype=torch.int64, device="cuda")
+        real_page_table = torch.zeros(
+            (batch_size * num_draft_tokens, 8),
+            dtype=torch.int32,
+            device="cuda",
+        )
+        real_page_table[:num_draft_tokens, 0] = 2
+        real_page_table[:num_draft_tokens, 4] = 3
+        real_page_table[num_draft_tokens:, 0] = 5
+        real_page_table[num_draft_tokens:, 4] = 6
+
+        update_kpool_write_plan_cuda_graph(
+            write_start=write_start,
+            req_pool_indices=req_pool_indices,
+            real_page_table=real_page_table,
+            req_out=plan.req,
+            write_start_out=plan.write_start,
+            tail_logical_start_out=plan.tail_logical_start,
+            write_loc_out=plan.write_loc,
+            pool_seqlens_per_q_out=plan.pool_seqlens_per_q,
+            seqlens_per_q_out=plan.seqlens_per_q,
+            pool_size=self.POOL_SIZE,
+            num_draft_tokens=num_draft_tokens,
+            slots_per_page=self.SLOTS_PER_PAGE,
+        )
+
+        torch.testing.assert_close(plan.req, req_pool_indices)
+        torch.testing.assert_close(plan.write_start, write_start)
+        torch.testing.assert_close(
+            plan.tail_logical_start,
+            torch.tensor([0, 252], dtype=torch.int32, device="cuda"),
+        )
+        torch.testing.assert_close(
+            plan.write_loc,
+            torch.tensor(
+                [
+                    [2 * self.SLOTS_PER_PAGE, 2 * self.SLOTS_PER_PAGE + 1],
+                    [
+                        5 * self.SLOTS_PER_PAGE + 63,
+                        6 * self.SLOTS_PER_PAGE,
+                    ],
+                ],
+                dtype=torch.int64,
+                device="cuda",
+            ),
+        )
+
+    def _run_compress_case(self, effective_n: int, expected_closed_pools: int):
+        torch.manual_seed(42)
+        pool = self._pool()
+        num_draft_tokens = self.NUM_DRAFT_TOKENS
+        tail_size = self.POOL_SIZE + num_draft_tokens
+        write_start_value = 3
+
+        key = torch.randn(
+            num_draft_tokens, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+        )
+        score = torch.randn_like(key)
+        ape = torch.randn(
+            self.POOL_SIZE, INDEX_HEAD_DIM, dtype=torch.float32, device="cuda"
+        )
+        tail_k_initial = torch.randn(
+            1, tail_size, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+        )
+        tail_score_initial = torch.randn_like(tail_k_initial)
+
+        tail_k_expected = tail_k_initial.clone()
+        tail_score_expected = tail_score_initial.clone()
+        for i in range(num_draft_tokens):
+            physical_slot = (write_start_value + i) % tail_size
+            tail_k_expected[0, physical_slot] = key[i]
+            tail_score_expected[0, physical_slot] = score[i]
+
+        expected_cache = self._empty_cache()
+        dummy_chunk = torch.zeros(
+            1, INDEX_HEAD_DIM, dtype=torch.bfloat16, device="cuda"
+        )
+        kpool_assemble_softmax_rotate_write_cache(
+            pool=pool,
+            buf=expected_cache,
+            chunk_k=dummy_chunk,
+            chunk_score=dummy_chunk,
+            tail_k=tail_k_expected,
+            tail_score=tail_score_expected,
+            req_pool_idx=torch.zeros(
+                expected_closed_pools, dtype=torch.int64, device="cuda"
+            ),
+            n_from_tail=torch.full(
+                (expected_closed_pools,),
+                self.POOL_SIZE,
+                dtype=torch.int32,
+                device="cuda",
+            ),
+            chunk_src_start=torch.zeros(
+                expected_closed_pools, dtype=torch.int64, device="cuda"
+            ),
+            tail_logical_base=torch.arange(
+                0,
+                expected_closed_pools * self.POOL_SIZE,
+                self.POOL_SIZE,
+                dtype=torch.int32,
+                device="cuda",
+            ),
+            ape=ape,
+            loc=torch.arange(expected_closed_pools, dtype=torch.int64, device="cuda"),
+            round_scale=False,
+        )
+
+        actual_cache = self._empty_cache()
+        tail_k_actual = tail_k_initial.clone()
+        tail_score_actual = tail_score_initial.clone()
+        kpool_write_tail_and_maybe_compress(
+            pool=pool,
+            buf=actual_cache,
+            key=key,
+            score=score,
+            tail_k=tail_k_actual,
+            tail_score=tail_score_actual,
+            ape=ape,
+            req_pool_indices=torch.zeros(1, dtype=torch.int64, device="cuda"),
+            write_start=torch.tensor(
+                [write_start_value], dtype=torch.int32, device="cuda"
+            ),
+            tail_logical_start=torch.zeros(1, dtype=torch.int32, device="cuda"),
+            write_loc=torch.tensor([[0, 1]], dtype=torch.int64, device="cuda"),
+            out_cache_loc=torch.arange(
+                1, num_draft_tokens + 1, dtype=torch.int64, device="cuda"
+            ),
+            num_draft_tokens=num_draft_tokens,
+            round_scale=False,
+            effective_n_per_batch=torch.tensor(
+                [effective_n], dtype=torch.int32, device="cuda"
+            ),
+        )
+
+        torch.testing.assert_close(tail_k_actual, tail_k_expected, atol=0, rtol=0)
+        torch.testing.assert_close(
+            tail_score_actual, tail_score_expected, atol=0, rtol=0
+        )
+        torch.testing.assert_close(actual_cache, expected_cache, atol=0, rtol=0)
+
+    def test_compresses_two_pools_when_draft_window_closes_two(self):
+        self._run_compress_case(
+            effective_n=self.NUM_DRAFT_TOKENS,
+            expected_closed_pools=2,
+        )
+
+    def test_effective_n_only_compresses_accepted_pools(self):
+        self._run_compress_case(effective_n=2, expected_closed_pools=1)
+
+
+if __name__ == "__main__":
+    unittest.main()