Repository navigation
fix: bound the DSA kpool ragged MQA logits allocation - #38469
SouthWest7 wants to merge 1 commit into
Conversation
|
Thanks for fixing the issue I encountered in #37712. I also found several related PRs, including #35759, #38086, and #38101. Since this issue appears across different accelerators and has been discussed in multiple PRs, it may be worth considering a follow-up redesign that combines budgeted workspaces with streaming top-k. Per-request Q chunking still requires a full K-length logits row, while other temporary allocations remain outside the workspace budget, so a more comprehensive memory-bounded design could help address the underlying problem. I would like to work on such re-design based on this PR if needed. |
|
Hi @SouthWest7 , can you please rebase on main, and reuse the util functions introduced in #39095, thanks~ |
|
@JackZeng0208 , we just merged #39095, which will be a starting point for fixing this OOM issue in an unified way. |
f1e32cf to
f804ff4
Compare
|
@Fridge003 Hi, Could you help review this? |
The kpool indexer's ragged extend path builds one dense fp32 logits tensor whose two dimensions are both summed over the whole batch, while each q row may only read its own request's column window. Peak memory is therefore quadratic in the request count and each row's useful area is ~1/B of the allocation, so long-context traffic with prefix-cache hits kills every TP rank inside deep_gemm.fp8_mqa_logits. The kpool path carried a dead, drifted copy of the budget predicate and never called it. Delete it, take the shared policy from attention/mqa_logits_utils.py (sgl-project#39095) instead, and split the kpool computation per request and then per q-row chunk under that budget. The single-call fast path is preserved unchanged when the dense tensor fits. Fixes sgl-project#37712
f804ff4 to
e3e631d
Compare
The kpool indexer's ragged extend path builds one dense fp32 logits tensor whose two dimensions are both summed over the whole batch, while each q row may only read its own request's column window. Peak memory is therefore quadratic in the request count and each row's useful area is ~1/B of the allocation, so long-context traffic with prefix-cache hits kills every TP rank inside deep_gemm.fp8_mqa_logits.
The non-kpool indexer already had a memory budget for the same tensor, but the kpool path carried a dead, drifted copy of the predicate and never called it. Move the policy to one module both flavors import, and split the kpool computation per request and then per q-row chunk under that budget. The single-call fast path is preserved unchanged when the dense tensor fits.
Fixes #37712
Motivation
Modifications
dsa/mqa_logits_budget.py— single home for the budget policy, since both indexer flavors need it and they are not in a subclass relationship. A second copy is how the kpool path drifted inthe first place.
dsa_indexer.py— drop the class constants and the three budget methods, import the module functions instead.dsa_indexer_kpool.py— delete the dead duplicate; consult the shared budget, and when the dense tensor does not fit, split per request (plan.ragged_groups) and then per q-row chunk. Everycolumn index rebases into the sliced K window. The single-call fast path is preserved unchanged when it fits.
kpool_plan.py— addRaggedGroupandKPoolExtendPlan.ragged_groups, the host-side geometry of the ragged concat layout, built from CPU lists already computed in the plan so there is nodevice→host sync.
Accuracy Tests
GLM-5.3-Flash on 8× B300-class GPUs, serving shape from the issue.
Chunking must not change which KV tokens the indexer selects. 4 passkey-retrieval prompts (~72 K tokens, needles at four depths, greedy), sent serially so batch composition is fixed, with
SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTIONlarge (dense path) vs0.001(~32 q-row chunks per step):Speed Tests and Profiling
12 × 132096-token contexts in one batch (18.14 GiB dense, i.e. above the budget so the fix does chunk), warm repeats with
/flush_cachebetween each:No measurable cost. Note the first long-context batch after server start is 3–20× slower from deep_gemm JIT; only warm repeats are comparable.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #35491432054
Latest PR Test (Extra): ❌ Run #35491431816
Latest PR Test (AMD ROCm 10): ❌ Run #35491431937