[Spec] Read the DFLASH/DSPARK draft pool budget from the live attn_tp_size

[he-yufeng]

Motivation

Fixes #38202. Supersedes #38203.

Under DP attention the DFLASH/DSPARK draft KV pool is reserved at the raw tp_size width while the pool builders shard the draft KV heads by attn_tp_size, which is smaller by attn_dp_size. The reservation under-counts the draft pool by that factor and the draft worker OOMs out of the runtime slack, exactly as reported in the issue.

#38203 identified the same bug and proposed deriving the width from the configured leaves (tp_size, dp_size, enable_dp_attention) at the call site. That derivation no longer works on current main: resolution leaves the deprecated enable_dp_attention flag false, and the derive_attention_widths helper it imported is gone. More fundamentally, re-deriving topology at the call site duplicates logic that can drift from what the runtime actually stamps; initialize_dp_attention settles the effective attn_dp_size during init_parallel_runtime, before ModelRunner.__init__ ever reaches this resolver.

Modifications

• _resolve_dflash_draft_cell_size now reads get_parallel().attn_tp_size, the same width the draft KV pool builders shard by. The call site runs after init_parallel_runtime (init_distributed_environment -> initialize_model_parallel -> initialize_dp_attention), so the stamped value already carries any DP-attention correction.
• CPU regression test test_draft_budget_prices_the_attention_sharded_width: publishes a tp16/attn-dp16 topology and asserts the draft KV heads are priced at the sharded width.
• The test spells the topology as attn_dp_size=16 directly, since resolution leaves the deprecated enable_dp_attention flag false and override declarations do not re-run the deprecation mapping.

Co-authored with @Baichuan7, who reported the issue and wrote the first fix.

Accuracy Tests

No model-output path is touched; this only changes how much KV budget the draft pool reserves.

Speed Tests and Profiling

No hot path is touched; the resolver runs once at startup.

Checklist

• Format your code according to the Format code with pre-commit guide.
• Add unit tests according to the Run and add unit tests guide (test_pool_configurator.py, 52 passed locally on CPU).
• Update documentation according to the Write documentations guide. (No doc surface changes.)
• Provide accuracy and speed benchmark results. (Not applicable; see above.)
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #37583413583
Latest PR Test (Extra): ❌ Run #37583413529
Latest PR Test (AMD ROCm 10): ❌ Run #37583413643

[he-yufeng]

Rebased onto current main. The code half of this PR is already there: #42651 landed the same attn_tp_size fix in the reservation while backing the draft pool to post-capture sizing, so after the rebase the code line matches main and what remains here is:

• the longer comment at the reservation site, spelling out why the budget has to price attn_tp_size (under DP attention it is narrower than tp_size by attn_dp_size, and budgeting raw tp_size under-counts the draft pool by that factor), in place of the one-liner;
• the CPU regression test for the DP-attention budget arithmetic, which main does not cover (existing pool configurator tests all run attn_dp_size=1).

Re-verified locally after the rebase: the whole pool configurator suite passes, 52 tests plus 11 subtests, on a CPU-only run.

[he-yufeng]

Small status bump: with #42651 merged, this PR is now test-plus-comment only (the regression test pins the DP-attention budget arithmetic that the reservation fix relies on). The pr-gate failure above is just the missing run-ci label, which I can't add myself.

@hanming-lu, you merged #42651; could you tag run-ci here when you get a chance, or point me at whoever shepherds the DFlash test surface these days?