diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 63c16baf4..9ba6ecc10 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -81,14 +81,16 @@ def _dflash_draft_cell_size(kvc: KVCacheConfigurator) -> int: 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. + a flat additive term. Under DCP, the target pool is sharded while the draft + pool spans the allocator's widened virtual location space, so the draft + term is replicated across DCP ranks. """ 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) + return int(cell_size) * get_parallel().attn_dcp_size def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]: diff --git a/test/registered/unit/model_executor/test_pool_configurator.py b/test/registered/unit/model_executor/test_pool_configurator.py index ebea13970..9d7df2df0 100644 --- a/test/registered/unit/model_executor/test_pool_configurator.py +++ b/test/registered/unit/model_executor/test_pool_configurator.py @@ -678,6 +678,37 @@ class TestDflashDraftKvBudget(unittest.TestCase): 0, ) + def test_dcp_replication_scales_draft_budget(self): + """The replicated draft pool spans every DCP virtual location.""" + draft_kv_per_token = 10_240 + mr = _make_model_runner() + 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, + ) + + target_kv_per_token = 4 * (64 + 64) * 32 * KV_SIZE + for dcp_size in (1, 8): + with self.subTest(dcp_size=dcp_size): + # TP=8 makes both topologies valid. The mock deliberately keeps + # target geometry fixed so this assertion isolates the draft term. + with ( + mock_cpu_env(tp_size=8), + get_parallel().override(attn_dcp_size=dcp_size), + ): + from sglang.srt.model_executor.pool_configurator import ( + create_memory_pool_configurator, + ) + + cfg = create_memory_pool_configurator(mr) + + self.assertEqual( + cfg._cell_size, + target_kv_per_token + draft_kv_per_token * dcp_size, + ) + 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