diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index 351636574..3675a6e39 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/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; diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 986fd8dca..54eabcb7e 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/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, diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 68bc1eddd..c227cb0c9 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/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 diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 1e9550d77..16e8714d0 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/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, diff --git a/test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_draft_pad.py b/test/registered/kernels/ops/attention/test_deepseek_v4_compress_plan_draft_pad.py new file mode 100644 index 000000000..8765cb7c8 --- /dev/null +++ b/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() diff --git a/test/registered/unit/mem_cache/test_dsv4_compress_write_pad.py b/test/registered/unit/mem_cache/test_dsv4_compress_write_pad.py new file mode 100644 index 000000000..2f5d1007c --- /dev/null +++ b/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()