Repository navigation
[CI] [AMD] memory_pool: Keep separate K/V allocations on ROCm to fix Mooncake PD accuracy - #46875
Closed
michaelzhang-ai wants to merge 1 commit into
Closed
michaelzhang-ai wants to merge 1 commit into
michaelzhang-ai wants to merge 1 commit into
Conversation
sgl-project#32272 stores K and V as halves of one allocation per layer. On ROCm, Mooncake PD transfers then hand decode corrupted KV: Qwen3-8B GSM8K in test_disaggregation_basic.py drops from 0.955 to ~0.0 with no transfer errors, while NIXL and MORI are unaffected. The fused layout only serves the CUDA trtllm fmha_v2 kernels, so ROCm keeps separate buffers.
michaelzhang-ai
marked this pull request as ready for review
October 27, 2026 17:58
github-actions
Bot
added
memory-pool
and removed
run-ci
michaelzhang-ai
requested review from
Ying1123,
alphabetc1,
hanming-lu,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
xiezhq-hermann and
yizhang2077
as code owners
October 27, 2026 17:58
This was referenced Oct 27, 2026
Contributor
|
Hi @michaelzhang-ai , looks like a duplicate for #46861 |
Collaborator
|
Hi @michaelzhang-ai, The code change here is the same as #46861. Both add a |
Collaborator
Author
|
Closing as a duplicate of #46861, which made the same change and is merged. |
3 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Since #32272,
MHATokenToKVPoolstores each layer's K and V as two halves of one allocation. On ROCm this breaks Mooncake PD transfers: decode receives corrupted KV without any transfer error. Instage-b-test-large-8-gpu-mi35x-disaggregation-amd,test_disaggregation_basic.py(Qwen3-8B, Mooncake) dropped from 0.955 GSM8K accuracy to 0.0-0.01 on main (example job). NIXL and MORI tests in the same job are unaffected (NIXL GSM8K 0.765 before and after).The last passing main run was at 5884e4f; the failures start with the commit range that includes #32272. The other candidates in that range don't touch this path for a bf16 Qwen3 PD run.
Modifications
Skip the fused K/V allocation on ROCm (
and not _is_hip) and keep separate per-layer K and V buffers, as before #32272. The fused layout serves the CUDA trtllm fmha_v2 paged kernels, which don't run on ROCm. Other platforms are unchanged.Test plan
stage-b-test-large-8-gpu-mi35x-disaggregation-amdon rocm10, dispatched on this commit: all four files pass.test_disaggregation_basic.pyGSM8K 0.950 / 0.965,test_disaggregation_pp.py0.945-0.955, NIXL 0.770, MORI OK (job).CI States
Latest PR Test (Base): ❌ Run #40555506148
Latest PR Test (Extra): ❌ Run #40555506102
Latest PR Test (AMD ROCm 10): ✅ Run #40555506196