[DSV4.1] Fuse DSpark verify compression, indexer and projections

python/sglang/kernels/jit/csrc/deepseek_v4/c2.cuh

@@ -50,12 +50,19 @@ constexpr uint32_t kC2VecSize = 2;
 /// An odd position completes a group with its even predecessor; an even one
 /// parks itself in the state.
 ///
+/// Target-verify runs that same schedule with `draft_len` consecutive positions
+/// per request instead of one, so a row's partner is usually the row before it
+/// in `kv_input` rather than the ring. That is the whole difference, and a 2D
+/// grid answers it without arithmetic: `blockIdx.x` is the position inside the
+/// block, `blockIdx.y` the request.
+///
 /// The three reductions below have three different widths and are not
 /// interchangeable: the RMSNorm statistic spans the row, an fp8 store scale
 /// spans 64 elements, an fp4 block spans 16. All asserted.
 template <
     bool kUsePDL,
     bool kStore,
+    bool kVerify,
     int64_t kHeadDim,
     int64_t kRopeDim,
     int32_t kPageBits,
@@ -85,7 +92,8 @@ __global__ __launch_bounds__(kHeadDim / kC2VecSize) void flash_c2_decode_kernel(
   using bf16_vec_t = AlignedVector<bf16x2_t, kVecSize / 2>;
 
   const auto tx = threadIdx.x;
-  const auto row = blockIdx.x;
+  // Verify gives each request a CTA column; decode a flat grid of one row each.
+  const auto row = kVerify ? blockIdx.y * gridDim.x + blockIdx.x : blockIdx.x;
   // Slots fit in int32 whatever width the scheduler hands them in.
   const auto raw_out_loc = static_cast<int32_t>(static_cast<const LocT*>(params.raw_out_loc)[row]);
   const auto pos = static_cast<const PosT*>(params.positions)[row];
@@ -101,8 +109,17 @@ __global__ __launch_bounds__(kHeadDim / kC2VecSize) void flash_c2_decode_kernel(
   score_new.load(params.kv_input + row * kStride, tx + kCTASize);
 
   fp32_vec_t kv_old, score_old;
-  kv_old.load(params.kv_state + read_row * kStride, tx);
-  score_old.load(params.kv_state + read_row * kStride, tx + kCTASize);
+  // Only a verify block's first row carries over from the ring; the rest pair
+  // with the row before them, which is already in `kv_input` under the same
+  // `| kv | score |` layout as the state, so this is a pointer swap. Taking the
+  // in-block partner from the input is also what keeps it race-free: the CTA
+  // that publishes that ring slot belongs to this very launch.
+  const float* partner = params.kv_state + read_row * kStride;
+  if constexpr (kVerify) {
+    if (blockIdx.x != 0) partner = params.kv_input + static_cast<int64_t>(row - 1) * kStride;
+  }
+  kv_old.load(partner, tx);
+  score_old.load(partner, tx + kCTASize);
 
   if ((pos & 1) == 0) {
     // padded case
@@ -256,14 +273,15 @@ struct FlashC2DecodeKernel {
   static constexpr uint32_t kBlockSize = kHeadDim / kC2VecSize;
   static constexpr int32_t kPageBits = std::bit_width(kPageSize) - 1;
   static constexpr int64_t kPageBytes = host::div_ceil(584ll * kPageSize, 576) * 576;
-  template <bool kStore, typename PosT, typename LocT>
-  static constexpr auto kernel = flash_c2_decode_kernel<kUsePDL, kStore, kHeadDim, kRopeDim, kPageBits, PosT, LocT>;
+  template <bool kStore, bool kVerify, typename PosT, typename LocT>
+  static constexpr auto kernel =
+      flash_c2_decode_kernel<kUsePDL, kStore, kVerify, kHeadDim, kRopeDim, kPageBits, PosT, LocT>;
 
   /// \brief The (`positions`, `raw_out_loc`) dtype pair, resolved at run time.
-  template <bool kStore>
+  template <bool kStore, bool kVerify>
   static auto select(const bool pos_i32, const bool loc_i32) {
-    if (pos_i32) return loc_i32 ? kernel<kStore, int32_t, int32_t> : kernel<kStore, int32_t, int64_t>;
-    return loc_i32 ? kernel<kStore, int64_t, int32_t> : kernel<kStore, int64_t, int64_t>;
+    if (pos_i32) return loc_i32 ? kernel<kStore, kVerify, int32_t, int32_t> : kernel<kStore, kVerify, int32_t, int64_t>;
+    return loc_i32 ? kernel<kStore, kVerify, int64_t, int32_t> : kernel<kStore, kVerify, int64_t, int64_t>;
   }
 
   // The sum of squares is reduced through a fixed-size shared array, so the CTA
@@ -312,6 +330,38 @@ struct FlashC2DecodeKernel {
     launch(kv_input, kv_state, kv_output, norm_weight, positions, req, raw_out_loc, eps, ring_size, freqs_cis, kvcache);
   }
 
+  /// \brief `run_decode_fusion` for a target-verify block.
+  ///
+  /// `draft_len` consecutive positions per request, request-major, which the
+  /// grid reproduces as `draft_len x batch`.
+  static void run_verify_fusion(
+      const tvm::ffi::TensorView kv_input,
+      const tvm::ffi::TensorView kv_state,
+      const tvm::ffi::TensorView kv_output,
+      const tvm::ffi::TensorView norm_weight,
+      const tvm::ffi::TensorView positions,
+      const tvm::ffi::TensorView req,
+      const tvm::ffi::TensorView raw_out_loc,
+      const float eps,
+      const tvm::ffi::TensorView freqs_cis,
+      const tvm::ffi::TensorView kvcache,
+      const int64_t ring_size,
+      const int64_t draft_len) {
+    launch(
+        kv_input,
+        kv_state,
+        kv_output,
+        norm_weight,
+        positions,
+        req,
+        raw_out_loc,
+        eps,
+        ring_size,
+        freqs_cis,
+        kvcache,
+        draft_len);
+  }
+
  private:
   using MaybeTensor = std::optional<tvm::ffi::TensorView>;
 
@@ -326,7 +376,8 @@ struct FlashC2DecodeKernel {
       const float eps,
       const int64_t ring_size,
       const MaybeTensor freqs_cis,
-      const MaybeTensor kvcache) {
+      const MaybeTensor kvcache,
+      const int64_t draft_len = 1) {
     using namespace host;
 
     auto N = SymbolicSize{"num_tokens"};
@@ -358,6 +409,19 @@ struct FlashC2DecodeKernel {
     const auto num_tokens = static_cast<uint32_t>(N.unwrap());
     if (num_tokens == 0) return;
     RuntimeCheck(ring_size > 0, "the pair-state ring must have at least one position");
+    RuntimeCheck(draft_len >= 1, "the draft length ", draft_len, " must be positive");
+    const auto is_verify = draft_len > 1;
+    RuntimeCheck(!is_verify || num_tokens % draft_len == 0, "verify rows must be a whole number of blocks");
+    // A block publishes its own even positions, so the slot its first row reads
+    // stays out of the launch's reach only while the ring is wider than the
+    // block. `get_compress_state_ring_size` satisfies this by construction.
+    RuntimeCheck(
+        !is_verify || ring_size > draft_len,
+        "the pair-state ring (",
+        ring_size,
+        ") must be wider than the draft length (",
+        draft_len,
+        ")");
 
     const auto params = C2Params{
         .kv_input = static_cast<const float*>(kv_input.data_ptr()),
@@ -375,12 +439,17 @@ struct FlashC2DecodeKernel {
     // `LaunchKernel` is move-only, so each arm builds its own.
     const auto pos_i32 = pos_dtype.is_type<int32_t>();
     const auto loc_i32 = loc_dtype.is_type<int32_t>();
-    if (store) {
-      const auto k = select<true>(pos_i32, loc_i32);
+    if (is_verify) {
+      const auto block = static_cast<uint32_t>(draft_len);
+      const auto k = select<true, true>(pos_i32, loc_i32);
+      LaunchKernel(dim3{block, num_tokens / block}, kBlockSize, device_.unwrap())  //
+          .enable_pdl(kUsePDL)(k, params);
+    } else if (store) {
+      const auto k = select<true, false>(pos_i32, loc_i32);
       LaunchKernel(num_tokens, kBlockSize, device_.unwrap())  //
           .enable_pdl(kUsePDL)(k, params);
     } else {
-      const auto k = select<false>(pos_i32, loc_i32);
+      const auto k = select<false, false>(pos_i32, loc_i32);
       LaunchKernel(num_tokens, kBlockSize, device_.unwrap())  //
           .enable_pdl(kUsePDL)(k, params);
     }

python/sglang/kernels/ops/attention/dsv4/c2.py

@@ -37,6 +37,7 @@ def _jit_c2_module(head_dim: int, rope_dim: int = 64, page_size: int = 128) -> M
         cuda_wrappers=[
             ("decode", f"FlashC2DecodeKernel<{args}>::run_decode"),
             ("decode_fusion", f"FlashC2DecodeKernel<{args}>::run_decode_fusion"),
+            ("verify_fusion", f"FlashC2DecodeKernel<{args}>::run_verify_fusion"),
         ],
     )
 
@@ -150,3 +151,59 @@ def c2_decode_norm_rope_store(
         int(ring_size),
     )
     return out
+
+
+def c2_verify_norm_rope_store(
+    kv_input: torch.Tensor,
+    kv_state: torch.Tensor,
+    norm_weight: torch.Tensor,
+    positions: torch.Tensor,
+    req: torch.Tensor,
+    raw_out_loc: torch.Tensor,
+    eps: float,
+    freqs_cis: torch.Tensor,
+    k_cache: torch.Tensor,
+    *,
+    page_size: int,
+    ring_size: int,
+    draft_len: int,
+    out: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """``c2_decode_norm_rope_store`` for a target-verify block.
+
+    A block is ``draft_len`` rows of one request at consecutive positions, laid
+    out request-major, so every row but the first pairs with the row before it
+    in ``kv_input`` rather than through the ring. Reading the in-block partner
+    from the input is what makes that safe: the row that publishes it into the
+    ring belongs to the same launch, with nothing ordering the two. The block's
+    first row does read the ring, at the slot before the block's own, which the
+    launch cannot reach while ``ring_size > draft_len`` -- a precondition the
+    kernel checks and ``get_compress_state_ring_size`` satisfies by
+    construction.
+
+    Nothing else changes, the pair state included: replaying a block one
+    position at a time through ``c2_decode_norm_rope_store`` gives the same
+    latents, the same ring and the same cache bytes.
+
+    :param draft_len: rows per request, ``speculative_num_draft_tokens``.
+    """
+    num_tokens, fused_dim = kv_input.shape
+    head_dim = fused_dim // 2
+    if out is None:
+        out = kv_input.new_empty((num_tokens, head_dim), dtype=torch.bfloat16)
+
+    _jit_c2_module(head_dim, freqs_cis.shape[-1], page_size).verify_fusion(
+        kv_input,
+        kv_state,
+        out,
+        norm_weight,
+        positions,
+        req,
+        raw_out_loc,
+        float(eps),
+        freqs_cis,
+        k_cache,
+        int(ring_size),
+        int(draft_len),
+    )
+    return out

python/sglang/kernels/ops/attention/dsv4/candidate_blocks.py

@@ -62,6 +62,29 @@ def _candidate_mask_kernel(
     tl.store(OUT + row * WIDTH + cols, values, cols < WIDTH)
 
 
+@triton.jit
+def _publish_candidate_mask_kernel(
+    INDICES,
+    VALUES,
+    KEEP,
+    WIDTH: tl.constexpr,
+    GROUP: tl.constexpr,
+    TOPK: tl.constexpr,
+    TILE: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    i = tl.program_id(1) * TILE + tl.arange(0, TILE)
+    selected = tl.load(INDICES + row * TOPK + i // GROUP, i < TOPK * GROUP, 0)
+    score = tl.load(VALUES + row * TOPK + i // GROUP, i < TOPK * GROUP, -float("inf"))
+    cols = selected * GROUP + i % GROUP
+    # torch.topk returns unique block indices: each output position has one writer.
+    tl.store(
+        KEEP + row * WIDTH + cols,
+        score > -float("inf"),
+        (i < TOPK * GROUP) & (cols < WIDTH),
+    )
+
+
 def candidate_block_logits(
     logits: torch.Tensor,
     seq_lens: torch.Tensor,
@@ -108,8 +131,19 @@ def candidate_block_logits(
         group_pad,
         tile,
     )
-    top = scores.topk(min(topk_blocks, blocks), dim=-1)
-    keep = torch.zeros_like(scores, dtype=torch.bool).scatter_(
-        -1, top.indices, top.values > -torch.inf
+    # Publication only needs membership; sorting the selected pairs is unused.
+    top = scores.topk(min(topk_blocks, blocks), dim=-1, sorted=False)
+    keep = torch.zeros((rows, width), dtype=torch.bool, device=logits.device)
+    _publish_candidate_mask_kernel[
+        (rows, triton.cdiv(top.indices.shape[1] * block_size, 256))
+    ](
+        top.indices,
+        top.values,
+        keep,
+        width,
+        block_size,
+        top.indices.shape[1],
+        256,
+        num_warps=4,
     )
-    return output, keep.repeat_interleave(block_size, dim=-1)[..., :width]
+    return output, keep

python/sglang/kernels/ops/attention/dsv4/indexer_postprocess.py

@@ -0,0 +1,74 @@
+"""Filter selected indexer scores and map logical positions to KV slots."""
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _filter_topk_pages(
+    SCORES,
+    INDICES,
+    PAGES,
+    OUT,
+    RAW,
+    WIDTH: tl.constexpr,
+    TOPK: tl.constexpr,
+    PAGE_SIZE: tl.constexpr,
+    SS: tl.constexpr,
+    SI: tl.constexpr,
+    SP: tl.constexpr,
+    SO: tl.constexpr,
+    SR: tl.constexpr,
+    WRITE_RAW: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    col = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
+    index = tl.load(INDICES + row * SI + col, col < TOPK, -1).to(tl.int64)
+    in_bounds = (index >= 0) & (index < WIDTH) & (col < TOPK)
+    score = tl.load(SCORES + row * SS + index, in_bounds, -float("inf"))
+    # This comparison also rejects NaN; +inf remains a valid score.
+    valid = in_bounds & (score > -float("inf"))
+    page = tl.load(PAGES + row * SP + index // PAGE_SIZE, valid, 0)
+    slot = (page * PAGE_SIZE).to(tl.int64) + index % PAGE_SIZE
+    tl.store(OUT + row * SO + col, tl.where(valid, slot, -1), col < TOPK)
+    if WRITE_RAW:
+        tl.store(RAW + row * SR + col, tl.where(valid, index, -1), col < TOPK)
+
+
+def filter_topk_pages(
+    scores: torch.Tensor,
+    indices: torch.Tensor,
+    page_table: torch.Tensor,
+    page_indices: torch.Tensor,
+    page_size: int,
+    raw_indices: torch.Tensor | None = None,
+) -> None:
+    """Preserve top-k order, write -1 for invalid scores, and map valid slots."""
+    rows, topk = indices.shape
+    assert scores.ndim == page_table.ndim == page_indices.ndim == 2
+    assert scores.shape[0] == page_table.shape[0] == page_indices.shape[0] == rows
+    assert page_indices.shape[1] == topk and scores.shape[1] > 0
+    assert page_table.shape[1] * page_size >= scores.shape[1]
+    assert all(t.stride(1) == 1 for t in (scores, indices, page_table, page_indices))
+    if raw_indices is not None:
+        assert raw_indices.shape == indices.shape and raw_indices.stride(1) == 1
+    _filter_topk_pages[(rows, triton.cdiv(topk, 256))](
+        scores,
+        indices,
+        page_table,
+        page_indices,
+        raw_indices,
+        scores.shape[1],
+        topk,
+        page_size,
+        scores.stride(0),
+        indices.stride(0),
+        page_table.stride(0),
+        page_indices.stride(0),
+        raw_indices.stride(0) if raw_indices is not None else 0,
+        raw_indices is not None,
+        256,
+        num_warps=4,
+    )