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 d05dd0819..7b5f01280 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -78,6 +78,19 @@ def get_swa_ring_size(sliding_window: int, is_speculative: bool = False) -> int: return sliding_window + spec_extra +def resolve_unified_kv_fp8(unified_fp8: Optional[bool] = None) -> bool: + """Per-pool fp8 layout. None follows SGLANG_DSV4_UNIFIED_KV_FP8. + + A caller may pass False so this pool keeps the bf16 ring while the env + stays on (target fused-Q still keys off the global switch). + """ + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, + ) + + return is_unified_kv_fp8() if unified_fp8 is None else bool(unified_fp8) + + def _num_dsv4_physical_kv_pages( size: int, physical_page_size: int, logical_page_size: int ) -> int: @@ -867,6 +880,7 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): enable_hisparse: bool = False, online_mtp_max_draft_tokens: int = 0, num_req_slots: Optional[int] = None, + unified_fp8: Optional[bool] = None, kv_source_layers: Sequence[int] = (), full_size: Optional[int] = None, is_draft_worker: bool = False, @@ -893,13 +907,6 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): self.compressed_kv_layout_option = compressed_kv_layout c4_logical_size = c128_size * 32 - logger.info( - "Initialize DeepSeekV4TokenToKVPool with " - f"{max_num_reqs=} {swa_size=} {c4_size=} " - f"{c4_logical_size=} {c128_size=} " - f"{c4_state_pool_size=} {c128_state_pool_size=}" - ) - self.max_num_reqs = max_num_reqs # PD preallocation can exceed max_num_reqs; # the SWA ring must cover every addressable req_pool_idx. @@ -915,7 +922,19 @@ class DeepSeekV4TokenToKVPool(BaseSWAKVPool): # Resolve the unified-kv gate before any sizing so the two cannot drift. self._unified_kv = is_unified_kv_triton() - self._unified_kv_fp8 = is_unified_kv_fp8() + self._unified_kv_fp8 = resolve_unified_kv_fp8(unified_fp8) + logger.info( + "Initialize DeepSeekV4TokenToKVPool with " + f"{max_num_reqs=} {swa_size=} {c4_size=} " + f"{c4_logical_size=} {c128_size=} " + f"{c4_state_pool_size=} {c128_state_pool_size=} " + f"unified={self._unified_kv} unified_fp8={self._unified_kv_fp8}" + ) + if is_unified_kv_fp8() and not self._unified_kv_fp8: + logger.info( + "SGLANG_DSV4_UNIFIED_KV_FP8 is on; this pool stays bf16 " + "(unified_fp8=False)" + ) # Uniform 512-dim e4m3 layout for the trtllm attention backend self.uniform_fp8 = ( not self._unified_kv diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 6025de05c..2bd16e700 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -133,6 +133,16 @@ def _get_dsv4_compress_state_dtypes() -> tuple[torch.dtype, torch.dtype]: _is_npu = is_npu() +def unified_fp8_for_dsv4_pool(*, is_draft_worker: bool, spec_algorithm) -> bool: + """Per-pool fp8 layout. DSpark draft writers scatter bf16, so that pool + stays a bf16 ring; MTP/EAGLE NextN follows the env.""" + from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( + is_unified_kv_fp8, + ) + + return is_unified_kv_fp8() and not (is_draft_worker and spec_algorithm.is_dspark()) + + def _should_enable_lazy_compaction() -> bool: """Lazy compaction default — ON unless `SGLANG_DISABLE_LAZY_COMPACTION=1` (escape hatch for A/B / rollback). @@ -1365,6 +1375,11 @@ class KVCacheConfigurator: kv_layout=kv_layout, compressed_kv_layout=compressed_kv_layout ) + unified_fp8 = unified_fp8_for_dsv4_pool( + is_draft_worker=self.is_draft_worker, + spec_algorithm=self.spec_algorithm, + ) + token_to_kv_pool = pool_cls( max_num_reqs=max_running_requests, # SWA ring is indexed by req_pool_idx; PD decode inflates req_to_token @@ -1392,6 +1407,7 @@ class KVCacheConfigurator: end_layer=self.layer_info.end_layer, enable_hisparse=get_memory().enable_hisparse, online_mtp_max_draft_tokens=(max_speculative_num_draft_tokens() or 0), + unified_fp8=unified_fp8, kv_source_layers=kv_source_layers, full_size=full_max_total_num_tokens, **({"is_draft_worker": self.is_draft_worker} if not _is_npu else {}), diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 89dfb6972..d70e248cf 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -968,6 +968,9 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): is the request-scoped fixed pools that do not scale with full_token. """ + # object.__new__ stubs (SWA floor tests) skip __init__ + _dspark_draft_on_bf16 = False + def __init__(self, kvc: KVCacheConfigurator): self.kv_cache_dtype_str = kvc.kv_cache_dtype_str cfg = kvc.model_config @@ -986,6 +989,11 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): # Resolve the unified-kv gate before any sizing so the two cannot drift. self._unified = is_unified_kv_triton() self._unified_fp8 = is_unified_kv_fp8() + # DSpark draft still allocates a bf16 ring; target fp8 * (T+1)/T would + # under-count that ring (640 vs 1024). MTP keeps the old inflation. + self._dspark_draft_on_bf16 = bool( + self._unified_fp8 and kvc.spec_algorithm.is_dspark() + ) # Row width across both unified pools: 1024 B bf16, 640 B fp8. self._unified_row_bytes = dsv4_unified_row_bytes( self.qk_nope_head_dim, self.qk_rope_head_dim, self._unified_fp8 @@ -1395,20 +1403,32 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): def _fixed_swa_bytes(self, max_running_requests: int) -> int: """Unified_kv SWA is a fixed per-request ring, sized by concurrency - (num_req_slots) rather than by full_token. Return its byte footprint - across all full layers, inflated for the draft worker the same way as the - per-token coeff. Returns 0 on the non-unified path (where SWA is already - accounted per-token).""" + (num_req_slots) rather than by full_token. MTP inflates the target ring + by _spec_infl; DSpark+fp8 adds a bf16 draft ring instead (640 vs 1024). + Returns 0 on the non-unified path (SWA already counted per-token).""" if not self._unified: return 0 num_req_slots = self._get_num_req_slots(max_running_requests) - ring_bytes = ( + target_ring = ( num_req_slots * self._swa_ring_size * self._unified_row_bytes * self.num_layers_total ) - return int(ring_bytes * self._spec_infl) + if self._dspark_draft_on_bf16: + from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + dsv4_unified_row_bytes, + ) + + draft_row = dsv4_unified_row_bytes( + self.qk_nope_head_dim, self.qk_rope_head_dim, fp8=False + ) + # 1 layer is what the shipped DSpark drafts allocate. A multi-stage + # draft would under-count by ~9 MB/layer (128-wide window, ~65 req + # slots), which the (T+1)/T on bytes_per_full_token already covers. + draft_ring = num_req_slots * self._swa_ring_size * draft_row + return int(target_ring + draft_ring) + return int(target_ring * self._spec_infl) def _to_config(self, sizes: _DSV4PoolSizes) -> MemoryPoolConfig: full = sizes.full_max_total_num_tokens @@ -1486,6 +1506,7 @@ class DSV4PoolConfigurator(MemoryPoolConfigurator): logger.info( f"DSV4 memory calculation: unified={self._unified}, " f"unified_fp8={self._unified_fp8}, " + f"dspark_draft_bf16={self._dspark_draft_on_bf16}, " f"bytes_per_full_token={self.bytes_per_full_token:.2f}, " f"available_bytes={available_bytes / (1 << 30):.2f} GB, " f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, " diff --git a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py index 7e82d66d3..82255abe2 100644 --- a/test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py +++ b/test/registered/e2e/dsv4/test_dsv4_unified_fp8_scatter.py @@ -244,6 +244,23 @@ class TestUnifiedFp8SwaScatter(CustomTestCase): with self.assertRaises(AssertionError): _store(kv, pool_nope, self.state_slot, self.positions) + def test_scatter_bf16_still_writes_a_bf16_draft_pool(self): + """Cut A: DSpark draft stays on the bf16 ring; scatter_bf16 must not + dtype-assert against that pool even when the fp8 env is on.""" + kv = torch.randn( + self.n_rows, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16 + ).contiguous() + loc = ( + self.state_slot.long() * RING_STRIDE + self.positions.long() % RING_STRIDE + ).to(torch.int32) + loc[1] = -1 + pool = torch.zeros(N_PAGES, BF16_LATENT, device=DEVICE, dtype=torch.bfloat16) + expected = pool.clone() + keep = loc >= 0 + expected[loc[keep].long()] = kv[keep] + runtime.scatter_bf16_into_unified(kv=kv, loc=loc, unified_kv=pool) + self.assertTrue(torch.equal(pool, expected)) + def test_empty_batch_is_a_noop(self): empty_slot = torch.zeros(0, device=DEVICE, dtype=torch.int32) kv_nope, _ = _packed_nope(0) diff --git a/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py b/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py index 19d51cacb..7e72aa515 100644 --- a/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py +++ b/test/registered/unit/mem_cache/test_dsv4_unified_fp8_pool.py @@ -1,5 +1,7 @@ import contextlib import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch import torch @@ -9,7 +11,10 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( DeepSeekV4TokenToKVPool, DeepSeekV4UnifiedKVPool, dsv4_unified_row_bytes, + resolve_unified_kv_fp8, ) +from sglang.srt.mem_cache.kv_cache_configurator import unified_fp8_for_dsv4_pool +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -237,5 +242,213 @@ class TestDSV4UnifiedRegionBuffers(CustomTestCase): self.assertAlmostEqual((nope + rope) / whole, 0.625) +class TestResolveUnifiedKvFp8(CustomTestCase): + def test_override_false_wins_over_env(self): + env_mod = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate" + with patch(f"{env_mod}.is_unified_kv_fp8", return_value=True): + self.assertFalse(resolve_unified_kv_fp8(False)) + self.assertTrue(resolve_unified_kv_fp8(True)) + self.assertTrue(resolve_unified_kv_fp8(None)) + + def test_none_follows_env_off(self): + env_mod = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate" + with patch(f"{env_mod}.is_unified_kv_fp8", return_value=False): + self.assertFalse(resolve_unified_kv_fp8(None)) + self.assertFalse(resolve_unified_kv_fp8(False)) + self.assertTrue(resolve_unified_kv_fp8(True)) + + +class TestDsv4PoolFp8Gate(CustomTestCase): + _ENV = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate" + + def _layout(self, *, is_draft_worker, algo, env_on=True): + with patch(f"{self._ENV}.is_unified_kv_fp8", return_value=env_on): + return unified_fp8_for_dsv4_pool( + is_draft_worker=is_draft_worker, spec_algorithm=algo + ) + + def test_dspark_draft_stays_bf16_when_env_on(self): + self.assertFalse( + self._layout( + is_draft_worker=True, algo=SpeculativeAlgorithm.DSPARK, env_on=True + ) + ) + + def test_eagle_draft_stays_two_pool_when_env_on(self): + self.assertTrue( + self._layout( + is_draft_worker=True, algo=SpeculativeAlgorithm.EAGLE, env_on=True + ) + ) + + def test_target_stays_two_pool_under_dspark_and_eagle(self): + for algo in (SpeculativeAlgorithm.DSPARK, SpeculativeAlgorithm.EAGLE): + with self.subTest(algo=algo): + self.assertTrue( + self._layout(is_draft_worker=False, algo=algo, env_on=True) + ) + + def test_env_off_is_bf16_for_every_worker(self): + for draft, algo in ( + (True, SpeculativeAlgorithm.DSPARK), + (True, SpeculativeAlgorithm.EAGLE), + (False, SpeculativeAlgorithm.DSPARK), + ): + with self.subTest(draft=draft, algo=algo): + self.assertFalse( + self._layout(is_draft_worker=draft, algo=algo, env_on=False) + ) + + +class TestBuildDsv4KvPoolPassesGate(CustomTestCase): + class _RecPool: + last = None + + def __init__(self, **kwargs): + type(self).last = kwargs + self._unified_kv = False + self._unified_kv_fp8 = kwargs.get("unified_fp8") + + def _kvc(self, *, is_draft_worker, spec_algorithm): + from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator + + kvc = object.__new__(KVCacheConfigurator) + kvc.is_draft_worker = is_draft_worker + kvc.spec_algorithm = spec_algorithm + kvc.layer_info = SimpleNamespace( + num_effective_layers=1, start_layer=0, end_layer=1 + ) + kvc.model_config = SimpleNamespace( + compress_ratios=[0], + window_size=256, + qk_nope_head_dim=NOPE_DIM, + qk_rope_head_dim=ROPE_DIM, + index_head_dim=128, + hf_config=SimpleNamespace(kv_source_layer_ids=[]), + ) + kvc.kv_cache_dtype = torch.bfloat16 + kvc.device = "cpu" + return kvc + + def _build(self, *, is_draft_worker, spec_algorithm): + kvc = self._kvc(is_draft_worker=is_draft_worker, spec_algorithm=spec_algorithm) + sched = MagicMock() + sched.page_size = 256 + exec_cfg = MagicMock() + exec_cfg.features.enable_memory_saver = False + mem = MagicMock() + mem.enable_hisparse = False + par = MagicMock() + par.attn_dcp_size = 1 + req = SimpleNamespace(req_to_token=torch.zeros(4, 1)) + rec = self._RecPool + rec.last = None + env = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate" + with ( + patch( + "sglang.srt.mem_cache.kv_cache_configurator.DeepSeekV4TokenToKVPool", + rec, + ), + patch( + "sglang.srt.mem_cache.kv_cache_configurator.get_schedule", + return_value=sched, + ), + patch( + "sglang.srt.mem_cache.kv_cache_configurator.get_exec", + return_value=exec_cfg, + ), + patch( + "sglang.srt.mem_cache.kv_cache_configurator.get_memory", + return_value=mem, + ), + patch( + "sglang.srt.mem_cache.kv_cache_configurator.get_parallel", + return_value=par, + ), + patch( + "sglang.srt.mem_cache.kv_cache_configurator.max_speculative_num_draft_tokens", + return_value=0, + ), + patch(f"{env}.is_unified_kv_fp8", return_value=True), + ): + kvc._build_dsv4_kv_pool( + max_running_requests=2, + full_max_total_num_tokens=256, + swa_max_total_num_tokens=256, + c4_max_total_num_tokens=0, + c128_max_total_num_tokens=1, + c4_state_pool_size=0, + c128_state_pool_size=0, + c4_state_dtype=None, + c128_state_dtype=None, + req_to_token_pool=req, + ) + return rec.last + + def test_dspark_draft_ctor_gets_unified_fp8_false(self): + kw = self._build( + is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.DSPARK + ) + self.assertIsNotNone(kw) + self.assertFalse(kw["unified_fp8"]) + + def test_eagle_draft_ctor_gets_unified_fp8_true(self): + kw = self._build( + is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.EAGLE + ) + self.assertIsNotNone(kw) + self.assertTrue(kw["unified_fp8"]) + + def test_target_ctor_gets_unified_fp8_true(self): + kw = self._build( + is_draft_worker=False, spec_algorithm=SpeculativeAlgorithm.DSPARK + ) + self.assertIsNotNone(kw) + self.assertTrue(kw["unified_fp8"]) + + +class TestUnifiedKvPoolFollowsCtorFp8(CustomTestCase): + def _pool(self, fp8): + return DeepSeekV4UnifiedKVPool( + stage_ratios=[0], + num_slots=2, + num_blocks=1, + page_size=256, + qk_nope_head_dim=NOPE_DIM, + qk_rope_head_dim=ROPE_DIM, + device="cpu", + memory_saver_adapter=_StubMemorySaver(), + custom_mem_pool=None, + swa_ring_size=8, + fp8=fp8, + ) + + def test_dspark_draft_layout_has_no_rope_pool(self): + env = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate" + with patch(f"{env}.is_unified_kv_fp8", return_value=True): + fp8 = unified_fp8_for_dsv4_pool( + is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.DSPARK + ) + self.assertFalse(fp8) + pool = self._pool(fp8) + buf = pool.kv_buffer[0] + self.assertEqual(buf.dtype, torch.bfloat16) + self.assertEqual(buf.shape[1], NOPE_DIM + ROPE_DIM) + self.assertIsNone(pool.kv_buffer_rope[0]) + + def test_eagle_draft_layout_has_rope_pool(self): + env = "sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate" + with patch(f"{env}.is_unified_kv_fp8", return_value=True): + fp8 = unified_fp8_for_dsv4_pool( + is_draft_worker=True, spec_algorithm=SpeculativeAlgorithm.EAGLE + ) + self.assertTrue(fp8) + pool = self._pool(fp8) + buf, rope = pool.kv_buffer[0], pool.kv_buffer_rope[0] + self.assertEqual(buf.dtype, torch.float8_e4m3fn) + self.assertEqual(rope.dtype, torch.bfloat16) + self.assertEqual(rope.shape[1], ROPE_DIM) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index fe7f3bb6c..7fbf006f8 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1223,6 +1223,7 @@ class TestSWAPoolFloor(CustomTestCase): cfg.request_window_bytes = 0 cfg.bytes_per_swa_token = 0.0 cfg._unified_fp8 = False + cfg._dspark_draft_on_bf16 = False # object.__new__ skips __init__; bf16 unified row is 2B * latent cfg._unified_row_bytes = cfg.attn_head_dim * 2 return cfg @@ -1291,6 +1292,38 @@ class TestSWAPoolFloor(CustomTestCase): self.assertEqual(sizes.swa_max_total_num_tokens, 3072) self.assertEqual(sizes.c4_state_pool_size, 0) + def test_dsv4_fp8_dspark_swa_ring_includes_bf16_draft(self): + """Target fp8 ring + one-layer bf16 draft ring, not (T+1)/T * fp8.""" + from sglang.srt.mem_cache.deepseek_v4_memory_pool import dsv4_unified_row_bytes + + cfg = self._dsv4_configurator_for_budget() + cfg._unified_fp8 = True + cfg._unified_row_bytes = dsv4_unified_row_bytes(448, 64, fp8=True) + cfg.qk_nope_head_dim, cfg.qk_rope_head_dim = 448, 64 + cfg._dspark_draft_on_bf16 = True + cfg._spec_infl = (cfg.num_layers_total + 1) / cfg.num_layers_total + mrr = 32 + slots = cfg._get_num_req_slots(mrr) + got = cfg._fixed_swa_bytes(mrr) + target = ( + slots * cfg._swa_ring_size * cfg._unified_row_bytes * cfg.num_layers_total + ) + draft = slots * cfg._swa_ring_size * dsv4_unified_row_bytes(448, 64, False) + self.assertEqual(got, target + draft) + mtp_formula = int(target * cfg._spec_infl) + self.assertGreater(got, mtp_formula) + + def test_dsv4_fp8_mtp_swa_ring_keeps_spec_inflation(self): + cfg = self._dsv4_configurator_for_budget() + cfg._unified_fp8 = True + cfg._unified_row_bytes = 640 + cfg._dspark_draft_on_bf16 = False + cfg._spec_infl = (cfg.num_layers_total + 1) / cfg.num_layers_total + mrr = 32 + slots = cfg._get_num_req_slots(mrr) + target = slots * cfg._swa_ring_size * 640 * cfg.num_layers_total + self.assertEqual(cfg._fixed_swa_bytes(mrr), int(target * cfg._spec_infl)) + if __name__ == "__main__": unittest.main()