fix(dflash): account for DCP in draft KV pool sizing (#33912)

Signed-off-by: Alexandre Milesi <milesial@users.noreply.github.com>

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]:

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