[AMD] dsr1 mi355x disagg 8k1k mtp: fix gsm8k collapse at conc >= 1024 - #2466
[AMD] dsr1 mi355x disagg 8k1k mtp: fix gsm8k collapse at conc >= 1024#2466Duyi-Wang wants to merge 1 commit into
Conversation
The 2P1D entry of dsr1-fp4-mi355x-sglang-disagg-8k1k-mtp collapses on gsm8k at high concurrency. Measured on 2 prefill + 1 decode node, conc 2048, full gsm8k 5-shot: strict-match 0.0023. Cause: job.slurm defaults SGLANG_AITER_MLA_PERSIST to 0, selecting the non-persistent aiter MLA decode. That path leaves num_kv_splits to aiter's python-side auto-search in mla_decode_fwd; under CUDA graphs that python runs once at capture, so the split count -- and the partial buffers sized from it -- are frozen and reused for every replay. Above roughly 128 running requests per DP rank the frozen choice corrupts decode. Only this entry reaches that point (1024/8 = 128 and up); conc <= 640 stays below it and is unaffected. Single-node bisection, image 20260729, client concurrency 2048, varying only the server admission cap: cap peak batch/rank strict 256 32 0.9530 512 64 0.9500 1024 128 0.9530 2048 165 0.0356 Same failing point, one variable changed at a time: SGLANG_AITER_MLA_PERSIST=1 0.9530 --disable-cuda-graph 0.9492 speculative decoding off 0.0879 (still broken; not an MTP defect) bf16 KV instead of fp8 0.0265 (still broken; not a KV defect) Two prefill + one decode node, conc 2048, the configuration this entry runs: SGLANG_AITER_MLA_PERSIST=0 0.0023 SGLANG_AITER_MLA_PERSIST=1 0.9469 Scoped to this one entry rather than flipping the job.slurm default, because the persistent MLA reduce had its own defect (an LDS write-after-read race, wave 1 corrupting output channels [256, 512)) that aiter only fixed in cdd6628f9. Configs still pinned to an older image must keep the non-persistent path, so the image here is pinned to 20260729, which carries that fix.
|
Thanks for the contribution! Please reach out to respective companies' CODEOWNER to fill in the latest PR_REVIEW_CHECKLIST.md before pinging core maintainer on Slack for review. In order for the signoff PR check bot to trigger, you must follow the PR_REVIEW_CHECKLIST.md template correctly, including the phrase For PR verification, add the PR authors are responsible for ensuring that after merging, all GitHub Action jobs fully pass. A lot of the time, failures are just flakes and simply re-running the failed jobs will fix it. See GitHub's docs on re-running failed jobs 感谢你的贡献！请联系相应公司的 CODEOWNER 填写最新的 PR_REVIEW_CHECKLIST.md，然后再在 Slack 上联系核心维护者进行审阅。为了触发 signoff PR 检查机器人，你必须正确遵循 PR_REVIEW_CHECKLIST.md 模板，包括保留英文语句 如需进行 PR 验证，请为此 PR 添加 PR 作者有责任确保合并后所有 GitHub Action 任务完全通过。 很多时候失败只是偶发抖动（flake），重新运行失败的任务即可解决。参见 GitHub 关于重新运行失败任务的文档 |
1 participant
Opened as a draft — the change is validated on hardware, but I am not requesting merge until a maintainer confirms the scoping is what you want.
Problem
The 2P1D entry of
dsr1-fp4-mi355x-sglang-disagg-8k1k-mtpcollapses on gsm8k at high concurrency. On 2 prefill + 1 decode node, conc 2048, full gsm8k 5-shot: strict-match 0.0023. Generations are truncated after a handful of tokens rather than wrong, which is why it reads as a total failure rather than a quality regression.Cause
job.slurmdefaultsSGLANG_AITER_MLA_PERSISTto0, selecting the non-persistent aiter MLA decode. That path leavesnum_kv_splitsto aiter's python-side auto-search inmla_decode_fwd. Under CUDA graphs that python runs once, at capture — so the split count, and the partial buffers sized from it, are frozen and reused for every replay. Above roughly 128 running requests per DP rank the frozen choice corrupts decode.Only this entry reaches that point: with dp8, per-rank batch is conc/8, so conc 1024 → 128 and up. conc ≤ 640 stays below it and is unaffected.
Evidence
Single-node bisection, image
20260729, client concurrency 2048, varying only the server admission cap:Same failing point, one variable changed at a time:
SGLANG_AITER_MLA_PERSIST=1--disable-cuda-graphTwo prefill + one decode node, conc 2048 — the configuration this entry actually runs:
SGLANG_AITER_MLA_PERSIST=0SGLANG_AITER_MLA_PERSIST=1An accept-length of 1.997 (ceiling 2.0) in the failing runs initially looked like a speculative-decode verification defect; disabling speculative decoding entirely still fails, so that was a symptom, not the cause.
Why scoped rather than flipping the job.slurm default
The persistent MLA reduce had its own defect — an LDS write-after-read race in
mla_reduce_v1_impl_massive, where__builtin_amdgcn_s_barrier()emits nos_waitcnt lgkmcnt(0)and no memory fence, so wave 1 readsp_lds_lse_scalebefore wave 0's writes land and corrupts output channels[256, 512). aiter fixed that incdd6628f9. Configs still pinned to an older image must keep the non-persistent path, so flipping the global default would trade this bug for that one on every other AMD multi-node config.Hence: the image here is pinned to
20260729(which carriescdd6628f9) and the persistent path is enabled for this one entry only. The inline comment states that dependency so the two are not separated later.Note
dsr1-fp4-mi355x-sglang-disagg-mtphas a structurally identical 2P1D entry but a different image; I have deliberately not touched it.Verification
CI dispatched on this branch for
dsr1-fp4-mi355x-sglang-disagg-8k1k-mtp --conc 2048 --evals-only(2P1D, eval-only).🤖 Generated with Claude Code