From 2e42bdb16119119c72cee1984333c204be86b596 Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 24 Sep 2026 17:27:56 +0800 Subject: [PATCH 1/4] fix: account for DSV4 FlashMLA KV page padding in token budgets --- .../srt/model_executor/pool_configurator.py | 28 ++++++++- .../model_executor/test_pool_configurator.py | 60 +++++++++++++++++++ 2 files changed, 85 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index fce87013641f..a72515e114b0 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -1280,6 +1280,28 @@ def _resolve_swa_cap_tokens(self) -> Optional[int]: headroom = self.swa_prefix_tails * (self.sliding_window_size + self.page_size) return ceil_align(cap + headroom, self.page_size) + def _get_paged_kv_bytes_per_token(self, compress_ratio: int = 0) -> float: + if ( + self._unified + or _is_hip + or _is_npu + or get_exec().kernel.dsv4_attn_backend == "trtllm" + ): + return self.kv_bytes + from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( + resolve_compressed_kv_layout, + select_dsv4_kv_layout, + ) + + layout, compressed_option = select_dsv4_kv_layout() + page_size = self.swa_page_size + if compress_ratio: + layout = resolve_compressed_kv_layout( + layout, compress_ratio, compressed_option + ) + page_size = self.page_size // compress_ratio + return layout.page_bytes(page_size) / page_size + def _get_bytes_per_swa_token(self) -> float: """Bytes one SWA slot costs across the stage. c4_state_pool_size = swa_tokens / swa_page_size * ring, so c4 compress state is priced per SWA slot too.""" @@ -1293,7 +1315,7 @@ def _get_bytes_per_swa_token(self) -> float: c4_state_ratio = self.c4_ring_size / self.swa_page_size return ( - self.kv_bytes * self.num_layers_total + self._get_paged_kv_bytes_per_token() * self.num_layers_total + c4_state_ratio * (c4_state_bytes + c4_indexer_state_bytes) * self.num_layers_ca4 @@ -1323,8 +1345,8 @@ def _get_bytes_per_full_token(self) -> float: return ( swa_ratio * self.bytes_per_swa_token + self.low_ratio_bytes_per_full_token - + c4_frac * self.kv_bytes * self.num_layers_ca4 - + 1 / 128 * self.kv_bytes * self.num_layers_ca128 + + c4_frac * self._get_paged_kv_bytes_per_token(4) * self.num_layers_ca4 + + 1 / 128 * self._get_paged_kv_bytes_per_token(128) * self.num_layers_ca128 + 1 / 4 * self.indexer_bytes_per_token * self.num_layers_ca4 + c128_state_ratio * c128_state_bytes * self.num_layers_ca128 ) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index a7c47c4609fc..9c19d1421b5f 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1277,6 +1277,66 @@ def test_dsv4_c2_state_fixed_bytes_mirror_pool(self): cfg.num_layers_c2_source = 0 self.assertEqual(cfg._get_c2_state_fixed_bytes(10), 0) + def test_dsv4_paged_kv_budget(self): + import torch + + from sglang.srt.environ import envs + from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4SingleKVPool + from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator + + _publish_config( + self, + dsv4_attn_backend="flashmla", + page_size=256, + max_running_requests=2, + chunked_prefill_size=256, + ) + kvc = SimpleNamespace( + kv_cache_dtype_str="fp8_e4m3", + model_config=SimpleNamespace( + qk_nope_head_dim=448, + qk_rope_head_dim=64, + index_head_dim=128, + context_len=131072, + compress_ratios=[0, 4, 128], + window_size=128, + hf_config=SimpleNamespace(kv_source_layer_ids=[]), + ), + layer_info=SimpleNamespace(start_layer=0, end_layer=3), + pp_size=1, + attn_dp_size=1, + sliding_window_size=128, + page_size=256, + spec_algorithm=SimpleNamespace( + is_dspark=lambda: False, is_none=lambda: True + ), + ) + with ( + envs.SGLANG_DSV4_KV_LAYOUT.override("v4"), + envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.override("float32"), + ): + planner = DSV4PoolConfigurator(kvc) + costs = [] + for ratio in (0, 4, 128): + page_size = 256 // (ratio or 1) + pool = DeepSeekV4SingleKVPool( + size=256, + page_size=page_size, + dtype=torch.float8_e4m3fn, + qk_nope_head_dim=448, + qk_rope_head_dim=64, + layer_num=1, + device="cpu", + enable_memory_saver=False, + ) + actual = pool.create_buffer(num_pages=1).nbytes / page_size + costs.append(actual) + self.assertEqual(planner.bytes_per_swa_token, 3 * costs[0] + 320) + self.assertEqual( + planner.bytes_per_full_token, + costs[1] / 4 + costs[2] / 128 + planner.indexer_bytes_per_token / 4, + ) + def test_dsv4_paged_dspark_budget_reserves_window_and_draft_layers(self): from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator From ee1d84930cf7afe972782288e2e87d79b9cc807d Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 24 Sep 2026 21:22:01 +0800 Subject: [PATCH 2/4] test: remove added DSV4 budget regression case --- .../model_executor/test_pool_configurator.py | 60 ------------------- 1 file changed, 60 deletions(-) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index 9c19d1421b5f..a7c47c4609fc 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1277,66 +1277,6 @@ def test_dsv4_c2_state_fixed_bytes_mirror_pool(self): cfg.num_layers_c2_source = 0 self.assertEqual(cfg._get_c2_state_fixed_bytes(10), 0) - def test_dsv4_paged_kv_budget(self): - import torch - - from sglang.srt.environ import envs - from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4SingleKVPool - from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator - - _publish_config( - self, - dsv4_attn_backend="flashmla", - page_size=256, - max_running_requests=2, - chunked_prefill_size=256, - ) - kvc = SimpleNamespace( - kv_cache_dtype_str="fp8_e4m3", - model_config=SimpleNamespace( - qk_nope_head_dim=448, - qk_rope_head_dim=64, - index_head_dim=128, - context_len=131072, - compress_ratios=[0, 4, 128], - window_size=128, - hf_config=SimpleNamespace(kv_source_layer_ids=[]), - ), - layer_info=SimpleNamespace(start_layer=0, end_layer=3), - pp_size=1, - attn_dp_size=1, - sliding_window_size=128, - page_size=256, - spec_algorithm=SimpleNamespace( - is_dspark=lambda: False, is_none=lambda: True - ), - ) - with ( - envs.SGLANG_DSV4_KV_LAYOUT.override("v4"), - envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.override("float32"), - ): - planner = DSV4PoolConfigurator(kvc) - costs = [] - for ratio in (0, 4, 128): - page_size = 256 // (ratio or 1) - pool = DeepSeekV4SingleKVPool( - size=256, - page_size=page_size, - dtype=torch.float8_e4m3fn, - qk_nope_head_dim=448, - qk_rope_head_dim=64, - layer_num=1, - device="cpu", - enable_memory_saver=False, - ) - actual = pool.create_buffer(num_pages=1).nbytes / page_size - costs.append(actual) - self.assertEqual(planner.bytes_per_swa_token, 3 * costs[0] + 320) - self.assertEqual( - planner.bytes_per_full_token, - costs[1] / 4 + costs[2] / 128 + planner.indexer_bytes_per_token / 4, - ) - def test_dsv4_paged_dspark_budget_reserves_window_and_draft_layers(self): from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator From da467a5eeb16a50c0852e8bf68b6e69b387cffbe Mon Sep 17 00:00:00 2001 From: alphabetc1 <2508695655@qq.com> Date: Thu, 24 Sep 2026 23:33:01 +0800 Subject: [PATCH 3/4] test: cover DSV4 paged KV budgeting with a compact fixture --- .../model_executor/test_pool_configurator.py | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index a7c47c4609fc..144e12209be3 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -1277,6 +1277,31 @@ def test_dsv4_c2_state_fixed_bytes_mirror_pool(self): cfg.num_layers_c2_source = 0 self.assertEqual(cfg._get_c2_state_fixed_bytes(10), 0) + def test_dsv4_paged_kv_budget(self): + from sglang.srt.environ import envs + + cfg = self._dsv4_configurator_for_budget() + cfg._unified = False + cfg.encoder_replay = False + cfg.page_size = cfg.swa_page_size = 256 + cfg.kv_bytes = 584 + cfg.attn_head_dim = 512 + cfg.num_layers_total = 3 + cfg.num_layers_ca4 = cfg.num_layers_ca128 = 1 + cfg.low_ratio_bytes_per_full_token = 0 + cfg.indexer_bytes_per_token = 132 + _publish_config(self, dsv4_attn_backend="flashmla") + with ( + envs.SGLANG_DSV4_KV_LAYOUT.override("v4"), + envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.override("float32"), + ): + cfg.bytes_per_swa_token = cfg._get_bytes_per_swa_token() + self.assertEqual(cfg.bytes_per_swa_token, 3 * 585 + 320) + self.assertEqual( + cfg._get_bytes_per_full_token(), + 0.1 * cfg.bytes_per_swa_token + 585 / 4 + 864 / 128 + 132 / 4, + ) + def test_dsv4_paged_dspark_budget_reserves_window_and_draft_layers(self): from sglang.srt.model_executor.pool_configurator import DSV4PoolConfigurator From 4e0982ea98882b6e73173b1f64a641d29f36b4a9 Mon Sep 17 00:00:00 2001 From: hnyls2002 Date: Thu, 24 Sep 2026 16:05:33 -0700 Subject: [PATCH 4/4] drop hip from paged kv pad --- python/sglang/srt/model_executor/pool_configurator.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index e24090e3e328..24ef6dfcf5ed 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -1213,12 +1213,9 @@ def _resolve_swa_cap_tokens(self) -> Optional[int]: return ceil_align(cap + headroom, self.page_size) def _get_paged_kv_bytes_per_token(self, compress_ratio: int = 0) -> float: - if ( - self._unified - or _is_hip - or _is_npu - or get_exec().kernel.dsv4_attn_backend == "trtllm" - ): + # Unified rings, the NPU pool and the trtllm uniform-FP8 pool do not go + # through DeepSeekV4SingleKVPool.create_buffer, so they carry no page pad. + if self._unified or _is_npu or get_exec().kernel.dsv4_attn_backend == "trtllm": return self.kv_bytes from sglang.srt.mem_cache.deepseek_v4_memory_pool import ( resolve_compressed_kv_layout,