[DSV4] Fix silent KV corruption when speculative draft tokens > 4 (#34189)

Co-authored-by: Deleter-D <867909454@qq.com>

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

@@ -47,6 +47,8 @@ struct Prefill0Params {
   uint32_t num_q_tokens;
   int32_t compress_ratio;
   int32_t swa_page_size;
+  /// \brief Trailing tokens the write plan keeps resident in the compress state ring.
+  /// Derived from the ring in `plan_compress_prefill`; see the bound there.
   int32_t mtp_pad;
 };
 
@@ -509,13 +511,19 @@ inline PrefillPlan plan_compress_prefill(
   RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens);
   // `swa_page_size` >= `ring_size` >= `compress_ratio`
   RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0);
+  // Write pad: trailing tokens kept resident so a verify batch's committed tail survives
+  // any accept length. Zero without speculation -- nothing rolls back, and the ring is
+  // then exactly one window wide. Otherwise the ring bounds it: a write at `w` aliases
+  // onto `w - ring_size`, and the earliest position a future compression still needs is
+  // `prefix_len - window_size + 2` (the next batch commits >= 1 token, and `run_prefill`
+  // launches the compress kernel before the write kernel, so a batch's own compressions
+  // read the pre-write ring). Padding past the extend range is harmless: the loops only
+  // span `[prefix_len, seq_len)`.
+  const auto mtp_pad = ring_size > window_size ? ring_size - window_size + 2 : 0;
 
   const auto device = device_.unwrap();
   const auto stream = LaunchKernel::resolve_device(device);
 
-  constexpr int32_t kMaxMTPDraftTokens = 4;
-  const auto mtp_pad = std::min(ring_size - compress_ratio, kMaxMTPDraftTokens);
-
   if (cpu_or_gpu.unwrap().device_type == kDLGPU) {
     // GPU input path: kernel0 builds the (CPU-loop-equivalent) plan metadata directly
     // on device, padding to num_q_tokens with invalid; kernel_1 then finalizes the
@@ -575,7 +583,7 @@ inline PrefillPlan plan_compress_prefill(
     const int32_t extend_len = ext_ptr[i];
     const int32_t prefix_len = seq_len - extend_len;
     const int32_t last_c_pos = seq_len / compress_ratio * compress_ratio;
-    const int32_t first_w_pos = last_c_pos - (is_overlap ? compress_ratio : 0);
+    const int32_t first_w_pos = std::min(last_c_pos - (is_overlap ? compress_ratio : 0), seq_len - mtp_pad);
     RuntimeCheck(0 < extend_len && extend_len <= seq_len);
     const auto should_write = [=](int32_t position) {
       if (position >= first_w_pos) return true;

python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py

@@ -48,6 +48,14 @@ def get_compress_state_ring_size(
         return 8 if compress_ratio == 4 else 128
 
 
+def get_compress_state_write_pad(compress_ratio: int, ring_size: int) -> int:
+    """Largest draft-token count this ring can serve; mirrors `mtp_pad` in `c_plan.cuh`
+    (the bound is derived there). Zero for a non-speculative ring, which is exactly one
+    window wide."""
+    window_size = compress_ratio * (2 if compress_ratio == 4 else 1)
+    return ring_size - window_size + 2 if ring_size > window_size else 0
+
+
 class DeepSeekV4SingleKVPool(KVCache):
     def __init__(
         self,

python/sglang/srt/model_executor/pool_configurator.py

@@ -31,7 +31,10 @@ from sglang.srt.configs.model_config import (
 )
 from sglang.srt.environ import envs
 from sglang.srt.mem_cache.allocation_sizing import get_alloc_len_per_decode
-from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    get_compress_state_ring_size,
+    get_compress_state_write_pad,
+)
 from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils.common import (
@@ -667,6 +670,12 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
         self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4)
         self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128)
 
+        if self.is_speculative:
+            # Ring is sized once here, so it must serve the largest adaptive tier.
+            self._assert_ring_serves_draft_tokens(
+                kvc.server_args.max_speculative_num_draft_tokens or 0
+            )
+
         self.bytes_per_full_token = self._get_bytes_per_full_token()
         if self.is_speculative:
             # Reserve memory for the speculative draft worker by inflating
@@ -707,6 +716,26 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator):
                     "DSV4 compressed attention: online c128 enabled (ring_size=1)"
                 )
 
+    def _assert_ring_serves_draft_tokens(self, num_draft_tokens: int) -> None:
+        """A verify batch writes its whole optimistic tail into the ring, so ring
+        capacity bounds the draft count."""
+        for compress_ratio, ring_size, num_layers in (
+            (4, self.c4_ring_size, self.num_layers_ca4),
+            (128, self.c128_ring_size, self.num_layers_ca128),
+        ):
+            if num_layers == 0:
+                continue
+            if compress_ratio == 128 and envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+                # Online c128 keeps per-draft state instead of a ring; sized separately.
+                continue
+            max_draft_tokens = get_compress_state_write_pad(compress_ratio, ring_size)
+            assert num_draft_tokens <= max_draft_tokens, (
+                f"speculative_num_draft_tokens={num_draft_tokens} exceeds what the c{compress_ratio} "
+                f"compress state ring can keep resident (ring_size={ring_size} serves at most "
+                f"{max_draft_tokens} draft tokens). Lower the draft count, or grow the ring in "
+                f"get_compress_state_ring_size()."
+            )
+
     def _get_bytes_per_full_token(self) -> float:
         kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8
 

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -371,7 +371,7 @@ class MockDSV4ModelRunner:
             max_running_requests=None,
             pp_size=1,
             revision=None,
-            speculative_algorithm=None,
+            speculative_algorithm=("EAGLE" if speculative_num_draft_tokens else None),
             speculative_eagle_topk=speculative_eagle_topk,
             speculative_num_draft_tokens=speculative_num_draft_tokens,
             speculative_num_steps=max(0, speculative_num_draft_tokens - 1),
@@ -400,7 +400,7 @@ class MockDSV4ModelRunner:
             c4_state_pool_size=pool_batch_size,
             c128_state_pool_size=pool_batch_size,
             page_size=case.page_size,
-            swa_page_size=DSV4_SWA_WINDOW,
+            swa_page_size=case.page_size,
             dtype=torch.float8_e4m3fn,
             c4_state_dtype=dtype,
             c128_state_dtype=dtype,

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

@@ -0,0 +1,217 @@
+"""Kernel-level tests for the DSV4 compress write-plan (`plan_prefill`).
+
+`plan_w` decides which tokens' raw KV get persisted into the compress-state ring
+for a *future* compression window to read. A speculative verify batch plans from
+the optimistic `seq_len = prefix + num_draft_tokens` but rolls back to
+`prefix + accept_len`, so every committed token must stay resident whatever the
+accept length -- i.e. the plan must write all of `[prefix, seq_len)`.
+
+`c_plan.cuh` used to cap that pad at 4 (`kMaxMTPDraftTokens`), silently
+under-writing the ring for larger draft counts -- no IMA, no NaN, just wrong
+compressed state. The pad now comes from the ring itself
+(`ring_size - window_size + 2`), covering every draft count the ring can serve.
+Tests pin the invariant on both planner paths (CPU host loop and GPU
+`plan_compress_prefill_kernel0`) and both compress ratios.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kernels.deepseek_v4.common import make_paged_context, to_seq_extend
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+C4_RING_SIZE = 16  # get_compress_state_ring_size(4, is_speculative=True)
+C128_RING_SIZE = 256  # get_compress_state_ring_size(128, is_speculative=True)
+C4_RING_SIZE_NO_SPEC = 8  # get_compress_state_ring_size(4, is_speculative=False)
+C128_RING_SIZE_NO_SPEC = 128  # get_compress_state_ring_size(128, is_speculative=False)
+
+
+def _window_size(compress_ratio: int) -> int:
+    """Tokens read by one compression: c4 overlaps two chunks, c128 does not."""
+    return compress_ratio * (2 if compress_ratio == 4 else 1)
+
+
+def _max_draft_tokens(compress_ratio: int, ring_size: int) -> int:
+    """Largest draft count this ring serves; mirrors `mtp_pad` in c_plan.cuh."""
+    window = _window_size(compress_ratio)
+    return ring_size - window + 2 if ring_size > window else 0
+
+
+def _written_positions(plan_w: torch.Tensor, prefix_len: int) -> set[int]:
+    """Decode `plan_w` into the set of positions written, for a bs=1 plan.
+
+    `plan_w` is `[n, 8]` uint8 = (uint32 ragged_id, int32 write_loc). Stage 1
+    overwrites `write_loc` with the final state slot, but `ragged_id` survives and
+    equals the token's index within the ragged layout, so for a single request
+    `position = prefix_len + ragged_id`.
+    """
+    words = plan_w.cpu().view(torch.uint32).view(-1, 2)
+    ragged_ids = words[:, 0]
+    valid = ragged_ids != 0xFFFFFFFF
+    return {prefix_len + int(r) for r in ragged_ids[valid]}
+
+
+class TestCompressWritePlanDraftPad(CustomTestCase):
+    def _make_plan_positions(
+        self,
+        *,
+        compress_ratio: int,
+        ring_size: int,
+        prefix_len: int,
+        num_draft_tokens: int,
+        on_gpu: bool = False,
+    ) -> set[int]:
+        """Build a bs=1 verify plan and return the positions it writes.
+
+        `on_gpu=True` moves the planner inputs to device, which routes
+        `plan_prefill` to `plan_compress_prefill_kernel0` instead of the host loop.
+        """
+        ctx = make_paged_context(
+            bs=1, compress_ratio=compress_ratio, ring_size=ring_size
+        )
+        seq_lens, extend_lens, num_q = to_seq_extend(
+            [(prefix_len + num_draft_tokens, num_draft_tokens)]
+        )
+        if on_gpu:
+            seq_lens = seq_lens.to(ctx.req_to_token.device)
+            extend_lens = extend_lens.to(ctx.req_to_token.device)
+        plan = ctx.make_prefill_plan(seq_lens, extend_lens, num_q)
+        return _written_positions(plan.plan_w, prefix_len)
+
+    def _assert_ring_residency(self, compress_ratio: int, ring_size: int):
+        """Every committed token must be written, for each (D, sl mod cr) combo.
+
+        This is the sufficient condition, which is why there is no multi-step replay
+        test: if a step writes all of `[prefix, prefix + D)`, then whatever the accept
+        length, the tokens the next compression window needs are either from this step
+        (written here) or older (written by an earlier step, same invariant by
+        induction).
+        """
+        max_d = _max_draft_tokens(compress_ratio, ring_size)
+        # Vary `seq_len % compress_ratio`: that residue decides whether the unpadded rule
+        # alone would have sufficed. Four is enough -- with the pad in place it dominates
+        # `last_c_pos` for every residue, so the rest repeat one branch. Bases are
+        # page-aligned so the swa-page-boundary clause does not mask the pad.
+        bases = [512 + off for off in range(min(compress_ratio, 4))]
+        draft_counts = sorted(
+            {1, 2, 3, 4, 5, max_d - 1, max_d} & set(range(1, max_d + 1))
+        )
+        for num_draft_tokens in draft_counts:
+            for prefix_len in bases:
+                with self.subTest(
+                    cr=compress_ratio,
+                    D=num_draft_tokens,
+                    prefix=prefix_len,
+                ):
+                    written = self._make_plan_positions(
+                        compress_ratio=compress_ratio,
+                        ring_size=ring_size,
+                        prefix_len=prefix_len,
+                        num_draft_tokens=num_draft_tokens,
+                    )
+                    seq_len = prefix_len + num_draft_tokens
+                    missing = set(range(prefix_len, seq_len)) - written
+                    self.assertEqual(
+                        missing,
+                        set(),
+                        f"plan_w skipped committed positions {sorted(missing)}; "
+                        f"a later compression would read stale ring slots",
+                    )
+
+    def test_c4_ring_residency(self):
+        self._assert_ring_residency(4, C4_RING_SIZE)
+
+    def test_c128_ring_residency(self):
+        self._assert_ring_residency(128, C128_RING_SIZE)
+
+    def test_cpu_and_gpu_planner_agree(self):
+        """Both planner paths must emit the same write set.
+
+        The residency invariants above are checked on the host-loop plan; this
+        pins the GPU `plan_compress_prefill_kernel0` plan to it, so the pad fix
+        has to hold on both paths.
+        """
+        for compress_ratio, ring_size in ((4, C4_RING_SIZE), (128, C128_RING_SIZE)):
+            max_d = _max_draft_tokens(compress_ratio, ring_size)
+            for num_draft_tokens in (1, 4, max_d):
+                for prefix_len in (512, 513, 515):
+                    with self.subTest(
+                        cr=compress_ratio, D=num_draft_tokens, prefix=prefix_len
+                    ):
+                        kwargs = dict(
+                            compress_ratio=compress_ratio,
+                            ring_size=ring_size,
+                            prefix_len=prefix_len,
+                            num_draft_tokens=num_draft_tokens,
+                        )
+                        self.assertEqual(
+                            self._make_plan_positions(**kwargs, on_gpu=False),
+                            self._make_plan_positions(**kwargs, on_gpu=True),
+                        )
+
+    def test_plain_prefill_write_set(self):
+        """A non-speculative ring is exactly one window wide, so the pad is 0 and the
+        base write rule stands unchanged for both ratios."""
+        for compress_ratio, ring_size in (
+            (4, C4_RING_SIZE_NO_SPEC),
+            (128, C128_RING_SIZE_NO_SPEC),
+        ):
+            self.assertEqual(_max_draft_tokens(compress_ratio, ring_size), 0)
+            is_overlap = compress_ratio == 4
+            for seq_len in (512, 600, 777):
+                with self.subTest(cr=compress_ratio, sl=seq_len):
+                    ctx = make_paged_context(
+                        bs=1, compress_ratio=compress_ratio, ring_size=ring_size
+                    )
+                    seq_lens, extend_lens, num_q = to_seq_extend([(seq_len, seq_len)])
+                    plan = ctx.make_prefill_plan(seq_lens, extend_lens, num_q)
+                    written = _written_positions(plan.plan_w, 0)
+
+                    last_c_pos = seq_len // compress_ratio * compress_ratio
+                    first_w_pos = last_c_pos - (compress_ratio if is_overlap else 0)
+                    sps = ctx.swa_page_size
+                    expected = {
+                        p
+                        for p in range(seq_len)
+                        if p >= first_w_pos
+                        or (is_overlap and p % sps >= sps - compress_ratio)
+                    }
+                    self.assertEqual(written, expected)
+
+    def test_over_capacity_under_writes(self):
+        """Beyond the ring's capacity the plan silently under-writes.
+
+        The planner cannot tell an over-configured verify batch from an ordinary long
+        prefill, so it cannot fail loudly -- hence the startup check in
+        `DSV4PoolConfigurator._assert_ring_serves_draft_tokens`.
+        """
+        for compress_ratio, ring_size in ((4, C4_RING_SIZE), (128, C128_RING_SIZE)):
+            max_d = _max_draft_tokens(compress_ratio, ring_size)
+            # Far enough over that the `last_c_pos` term cannot cover the gap for any
+            # residue of `seq_len % compress_ratio`.
+            too_many = max_d + compress_ratio + 1
+            prefix_len = 512
+            with self.subTest(cr=compress_ratio, D=too_many):
+                written = self._make_plan_positions(
+                    compress_ratio=compress_ratio,
+                    ring_size=ring_size,
+                    prefix_len=prefix_len,
+                    num_draft_tokens=too_many,
+                )
+                missing = set(range(prefix_len, prefix_len + too_many)) - written
+                self.assertNotEqual(
+                    missing,
+                    set(),
+                    "expected the plan to under-write past the ring capacity; if this "
+                    "now covers everything, the startup bound can be relaxed",
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_dsv4_compress_write_pad.py

@@ -0,0 +1,45 @@
+import unittest
+
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
+    get_compress_state_ring_size,
+    get_compress_state_write_pad,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class TestCompressStateWritePad(CustomTestCase):
+    """The pad bounds how many speculative draft tokens a compress-state ring can serve.
+
+    Mirrors `mtp_pad` in `c_plan.cuh`; `DSV4PoolConfigurator` rejects a larger draft
+    count at startup.
+    """
+
+    def test_pad_is_zero_without_speculation(self):
+        """A non-speculative ring is exactly one window wide: nothing rolls back."""
+        for compress_ratio in (4, 128):
+            ring_size = get_compress_state_ring_size(compress_ratio, False)
+            with self.subTest(cr=compress_ratio, ring=ring_size):
+                self.assertEqual(
+                    get_compress_state_write_pad(compress_ratio, ring_size), 0
+                )
+
+    def test_pad_matches_speculative_ring_capacity(self):
+        """`ring_size - window_size + 2`, with window = 2*cr for the overlapping c4."""
+        for compress_ratio, expected in ((4, 10), (128, 130)):
+            ring_size = get_compress_state_ring_size(compress_ratio, True)
+            with self.subTest(cr=compress_ratio, ring=ring_size):
+                self.assertEqual(
+                    get_compress_state_write_pad(compress_ratio, ring_size), expected
+                )
+
+    def test_pad_is_zero_for_rings_below_one_window(self):
+        """Online c128 collapses the ring to 1; the pad must clamp instead of going
+        negative."""
+        self.assertEqual(get_compress_state_write_pad(128, 1), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()