[Spec] Budget the DFLASH/DSPARK draft KV pool by attn_tp_size, not tp_size

[Baichuan7]

Fixes #38202

Problem

With DP attention on, the target KV pool eats essentially the whole
--mem-fraction-static budget and the draft worker OOMs allocating its own pool when serving Kimi-K3.

A DFLASH/DSPARK draft never sizes itself: KVCacheConfigurator.configure
short-circuits for draft workers and reuses the target's resolved
MemoryPoolConfig, then allocates its own token_to_kv_pool at the target's
max_total_num_tokens. Room for it is reserved by inflating the target's
per-token cell size in pool_configurator.py, which bottoms out in a single
head-count input:

# model_runner_components/spec_aux_hidden_state.py
tp_size=get_parallel().tp_size,   # -> draft_model_config.get_num_kv_heads(tp_size)

Every KV pool builder, however, shards by attn_tp_size:

• kv_cache_configurator._build_mha_kv_pool
• kv_cache_configurator._build_hybrid_linear_kv_pool
• the target's own pool_configurator._compute_cell_size

With --tp-size 16 --dp-size 16 --enable-dp-attention, attn_tp_size is
16 // 16 // 1 = 1 while the reservation still divides by 16, so the draft term
is under-counted by min(attn_dp_size, total_kv_heads). Without DP attention the
two are equal, which is why this went unnoticed. This is the only place in the KV
sizing path (model_executor/, mem_cache/) that reads the raw tp_size.

Fix

Derive the attention width from the configured leaves via
derive_attention_widths, and thread server_args in to reach them.

Reading get_parallel().attn_tp_size directly does not work here, and fails
in a way worth calling out: this runs in ModelRunner.init_spec_aux_hidden_state(),
which is invoked before init_torch_distributed(), so the attention TP group does
not exist yet. The property raises
attention tensor model parallel group is not initialized, the surrounding
except Exception swallows it, and the budget silently degrades to the
layer-count fallback — turning a wrong reservation into no reservation at all.
derive_attention_widths is a pure function over the config leaves and is
already used this way by dp_attention.compute_dp_attention_world_info.

Test plan

Manually verified on a PD decode node (Kimi-K3 target + DSPARK draft,
--tp-size 16 --dp-size 16 --enable-dp-attention, --mem-fraction-static 0.90):
before the fix the draft pool OOMs during alloc_memory_pool; after it, the node
starts cleanly.

Not addressed here

The except Exception in _resolve_dflash_draft_cell_size degrades to a
zero-byte draft reservation on the hybrid-SWA path: HybridSWAPoolConfigurator's
layer-count fallback only covers EAGLE/STANDALONE, so a hybrid-SWA DFLASH/DSPARK
target reserves nothing when resolution fails. That is what turned the bug above
into an OOM rather than an error. Happy to fold a fallback (or a loud failure)
into this PR if reviewers prefer.

---

CI States

Latest PR Test (Base): ❌ Run #34024742578
Latest PR Test (Extra): ❌ Run #34024742435
Latest PR Test (AMD ROCm 10): ❌ Run #34024742593

[he-yufeng]

Verified the head of this PR against the issue chain. Your correction to my proposal is the right one: _resolve_dflash_draft_cell_size runs before init_torch_distributed, so get_parallel() is not live there and deriving attn_tp_size from the configured leaves via derive_attention_widths is the correct source. The arithmetic matches the allocation side (tp // attn_dp // attn_cp, with attn_dp = dp iff DP attention is on), the get_parallel import removal is clean (no other use left in the file), and the fallback-to-layer-count behavior on unresolvable inputs is unchanged.

One thing missing: nothing pins the derivation this fix stands on. A small unit test next to the other runtime_context tests would do it — feel free to take this verbatim:

def test_derive_attention_widths_matches_dp_attention_arithmetic():
    from sglang.srt.runtime_context import derive_attention_widths

    # DP attention on: attention shards by 1 under tp16/dp16 (#38202).
    assert derive_attention_widths(
        tp_size=16, attn_cp_size=1, dp_size=16, enable_dp_attention=True
    ) == (16, 1)
    # DP attention off: the full tp width is the attention width.
    assert derive_attention_widths(
        tp_size=16, attn_cp_size=1, dp_size=16, enable_dp_attention=False
    ) == (1, 16)
    # Context parallelism stacks on top of the DP split.
    assert derive_attention_widths(
        tp_size=16, attn_cp_size=2, dp_size=8, enable_dp_attention=True
    ) == (8, 1)

I can't push to your fork, so committing the test yourself or applying the patch above is equally good; happy to open a PR onto your branch instead if you prefer that route.

[AI-assisted]

[Baichuan7]

Thank you for your complement ! Just open a PR to my branch if you like.

[he-yufeng]

I couldn't push here directly (no write access to the fork), so the commits are on my fork at he-yufeng:fix/dflash-draft-budget-attn-tp (ee22440249): your two commits, a merge of current upstream main, and one commit from me. Fetching that branch and resetting yours to it brings the PR to the same state.

What my commit changes, and why:

The reservation now reads get_parallel().attn_tp_size instead of re-deriving the width from the configured leaves. The get_parallel() concern from the issue thread was right for the base this branch started from, but that base predates the runtime_context rewrite (#40343, #40707). On current main the resolve site is post-publish, _build_config_bags materializes the derived widths into the parallel bag (I verified attn_tp_size is in the bag's fields), and the reads two lines earlier (get_model(), `get_spec()") prove the bags are already projected. The "runs before init_torch_distributed" comment does not hold there.

The bigger reason is drift, not style. initialize_model_parallel stamps attn_tp_size from the widths the attention groups are actually built at, and initialize_dp_attention can settle an effective attn_dp_size that differs from the configured dp_size. The pool builders read the live value; a reservation re-derived from the configured leaves would miss that correction and the under-count would survive in exactly those DP configurations. Reading the same property the builders read makes the two sides agree by construction.

Regression test (test_draft_budget_prices_the_attention_sharded_width): publishes tp16/dp16 with DP attention and asserts the resolver hands 1, not 16, to get_num_kv_heads. That is the under-count behind the OOM.

Verification: the whole test_pool_configurator.py file passes on the merged result (47 tests), and the new test fails against the unfixed call site with seen == [16], so it actually guards the bug. CPU-only, no GPU needed.

[he-yufeng]

The file moved again since my last note (cell-size resolution now sits behind dflash_draft_cell_size_per_token on main), but the branch still applies clean: I just ran merge-tree of ee22440249 against current main (4a4e1ce014), zero conflicts.

You can pull it straight in:

git fetch https://github.com/he-yufeng/sglang.git fix/dflash-draft-budget-attn-tp
git merge FETCH_HEAD

or cherry-pick ee22440249 if you want only the budget fix and its CPU regression test. Happy to re-cut it against whatever head you push next.

[he-yufeng]

Heads-up on the CI state: the branch predates the MIN_BASE_SHA maintenance gate (8854857), so check-changes will fail until it is rebased onto current main. I do not have push access to the fork branch anymore, so that rebase needs the author's hand. The fix itself targets #38202: the draft KV pool reservation scales by attn_tp_size instead of tp_size, removing the attn_dp_size-fold over-allocation that OOMed DP attention with DFLASH/DSPARK enabled, verified against the reporter's repro at the time. Note #39225 attacks the same bug; whichever lands first is fine. Could a maintainer add the run-ci label once the branch is rebased so CI can actually exercise it?

[he-yufeng]

Thanks for the diagnosis and the first fix here. Since this branch has been quiet for a few weeks and main has moved on (the derive_attention_widths helper is gone and resolution now leaves enable_dp_attention false, so the leaf-based derivation no longer sees DP attention), I opened #42789 to carry this over the finish line: it reads the live attn_tp_size the pool builders shard by, keeps your authorship via the co-author trailer, and adds a CPU regression test for the DP-attention arithmetic. Closing this one in favor of #42789 would make the intent clear to reviewers.