[Spec] Budget the DFLASH draft KV pool from its own attention geometry (#34234)

Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>

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,

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

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(

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,