From b51bf9ec9e5ccce8ab05e7d4ce555b1fb51df3db Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 10 Aug 2026 01:28:49 -0700 Subject: [PATCH] [Spec] Budget the DFLASH draft KV pool from its own attention geometry (#34234) Co-authored-by: Brayden Zhong --- python/sglang/srt/mem_cache/kv_cache_dtype.py | 2 +- .../spec_aux_hidden_state.py | 48 ++++++++++++++ .../srt/model_executor/pool_configurator.py | 20 ++++++ python/sglang/srt/speculative/dflash_utils.py | 16 +++++ .../model_executor/test_pool_configurator.py | 62 +++++++++++++++++++ 5 files changed, 147 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py index 789dbe713..2bbdb499a 100644 --- a/python/sglang/srt/mem_cache/kv_cache_dtype.py +++ b/python/sglang/srt/mem_cache/kv_cache_dtype.py @@ -22,7 +22,7 @@ TORCH_DTYPE_TO_KV_CACHE_STR = { def configure_kv_cache_dtype( *, server_args_kv_cache_dtype: str, - model: nn.Module, + model: nn.Module | None, model_dtype: torch.dtype, is_draft_worker: bool, is_dflash: bool, diff --git a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py index 95c887636..08dd87afd 100644 --- a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py +++ b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py @@ -21,6 +21,8 @@ class SpecAuxHiddenStateConfig(msgspec.Struct, kw_only=True): dflash_use_aux_hidden_state: bool = False dflash_draft_num_layers: Optional[int] = None dflash_target_layer_ids: Any = None + # DFLASH draft KV bytes/token; None when unresolved. + dflash_draft_cell_size_per_token: int | None = None def resolve_spec_aux_hidden_state_config( @@ -163,3 +165,49 @@ def _resolve_dflash_aux_hidden_state( config.dflash_use_aux_hidden_state = True config.dflash_draft_num_layers = int(draft_num_layers) config.dflash_target_layer_ids = target_layer_ids + config.dflash_draft_cell_size_per_token = _resolve_dflash_draft_cell_size( + server_args=server_args, + draft_model_config=draft_model_config, + draft_num_layers=int(draft_num_layers), + ) + + +def _resolve_dflash_draft_cell_size( + *, + server_args: ServerArgs, + draft_model_config: ModelConfig, + draft_num_layers: int, +) -> int | None: + """Bytes/token the DFLASH draft KV pool will cost the target's pool budget. + + Resolved from the draft's own attention geometry and the KV dtype the draft + worker will actually resolve. Returns None if anything is unresolvable, + leaving callers on layer-count scaling. + """ + from sglang.srt.mem_cache.kv_cache_dtype import configure_kv_cache_dtype + from sglang.srt.speculative.dflash_utils import dflash_draft_cell_size_per_token + + try: + _, draft_kv_cache_dtype = configure_kv_cache_dtype( + server_args_kv_cache_dtype=server_args.kv_cache_dtype, + model=None, + model_dtype=draft_model_config.dtype, + is_draft_worker=True, + is_dflash=True, + speculative_draft_attention_backend=( + server_args.speculative_draft_attention_backend + ), + ) + return dflash_draft_cell_size_per_token( + draft_model_config=draft_model_config, + draft_num_layers=draft_num_layers, + draft_kv_cache_dtype=draft_kv_cache_dtype, + tp_size=server_args.tp_size, + ) + except Exception as e: # noqa: BLE001 + logger.warning( + "Could not resolve DFLASH draft KV bytes/token (%s); falling back to " + "layer-count scaling for the KV pool budget.", + e, + ) + return None diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index c227cb0c9..63c16baf4 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -76,6 +76,21 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +def _dflash_draft_cell_size(kvc: KVCacheConfigurator) -> int: + """Bytes/token the DFLASH draft KV pool adds to the target's budget, 0 if none. + + Unlike an EAGLE draft, which reuses the target's attention config and is + therefore priced by layer count, a DFLASH draft has its own geometry and is + a flat additive term. + """ + if kvc.is_draft_worker or not kvc.spec_algorithm.is_dflash_family(): + return 0 + cell_size = kvc.spec_aux_config.dflash_draft_cell_size_per_token + if cell_size is None or int(cell_size) <= 0: + return 0 + return int(cell_size) + + def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]: dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower() if dtype_name in ("float32", "fp32"): @@ -182,6 +197,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator): target_num_layers=int(num_layers), draft_num_layers=int(draft_num_layers) * get_parallel().attn_dcp_size, + draft_cell_size_per_token=_dflash_draft_cell_size(kvc) or None, ) def _compute_cell_size(self, kvc: KVCacheConfigurator, num_layers: int) -> int: @@ -412,6 +428,8 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): self._draft_swa_full_layers_num = banded_depths self._draft_full_layers_num = draft_layers - banded_depths + self._draft_cell_size = _dflash_draft_cell_size(kvc) + # Bytes per token of max_total_num_tokens. # # Hybrid (full_layers > 0): max_total = full_tokens, so cell_size accounts @@ -426,6 +444,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): self._swa_per_token * self._swa_layers_num + self._full_per_token * self._draft_full_layers_num + self._swa_per_token * self._draft_swa_full_layers_num + + self._draft_cell_size ) else: self._cell_size = ( @@ -435,6 +454,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator): + self._swa_full_tokens_ratio * self._swa_per_token * self._swa_layers_num + + self._draft_cell_size ) def _solve_pool_sizes( diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index e30eb759e..e4c13a6bc 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -65,6 +65,22 @@ def is_dflash_sampling_verify_available() -> bool: return _DFLASH_SAMPLING_VERIFY_AVAILABLE +def dflash_draft_cell_size_per_token( + *, + draft_model_config: Any, + draft_num_layers: int, + draft_kv_cache_dtype: torch.dtype, + tp_size: int, +) -> int: + """Exact bytes/token of the DFLASH draft KV pool.""" + if draft_num_layers <= 0: + return 0 + num_kv_heads = draft_model_config.get_num_kv_heads(tp_size) + kv_dim_per_head = draft_model_config.head_dim + draft_model_config.v_head_dim + dtype_size = torch._utils._element_size(draft_kv_cache_dtype) + return int(num_kv_heads * kv_dim_per_head * int(draft_num_layers) * dtype_size) + + def scale_kv_cell_size_per_token_for_dflash( *, target_cell_size_per_token: int, diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index b0f58df64..ebea13970 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -645,5 +645,67 @@ class TestFactory(unittest.TestCase): self.assertNotIsInstance(_cfg(None), SWAChunkCapPoolConfigurator) +class TestDflashDraftKvBudget(unittest.TestCase): + """DFLASH draft KV pool as a flat bytes/token term on the target's budget.""" + + def test_bytes_per_token_from_draft_geometry(self): + import torch + + from sglang.srt.speculative.dflash_utils import ( + dflash_draft_cell_size_per_token, + ) + + draft = SimpleNamespace( + get_num_kv_heads=lambda tp: 4, head_dim=128, v_head_dim=128 + ) + # 4 kv heads * (128 + 128) dims * 5 layers * 2 bytes + self.assertEqual( + dflash_draft_cell_size_per_token( + draft_model_config=draft, + draft_num_layers=5, + draft_kv_cache_dtype=torch.bfloat16, + tp_size=1, + ), + 10240, + ) + self.assertEqual( + dflash_draft_cell_size_per_token( + draft_model_config=draft, + draft_num_layers=0, + draft_kv_cache_dtype=torch.bfloat16, + tp_size=1, + ), + 0, + ) + + def test_hybrid_swa_budget_shrinks_by_draft_pool(self): + """HybridSWA carried no draft term, so the draft pool fell outside the budget.""" + available = 10_000_000 + + def _tokens(draft_kv_per_token): + mr = _make_model_runner( + is_hybrid_swa=True, + full_attention_layer_ids=list(range(16)), + swa_attention_layer_ids=list(range(16, 32)), + swa_num_kv_heads=4, + ) + mr.spec_algorithm.is_dflash_family.return_value = True + mr.spec_aux_config = SimpleNamespace( + eagle_draft_num_layers=None, + dflash_draft_num_layers=5, + dflash_draft_cell_size_per_token=draft_kv_per_token, + ) + with mock_cpu_env(): + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + cfg = create_memory_pool_configurator(mr) + config = cfg.calculate_pool_sizes(available, mr.server_args.page_size) + return config.full_max_total_num_tokens + + self.assertLess(_tokens(10240), _tokens(None)) + + if __name__ == "__main__": unittest.main()