[omirosh]

Motivation

This ROCm bugfix, with no speed change, turns off the fused K/V allocation from #32272 on ROCm, because with it the MI35x PD disaggregation test test_disaggregation_basic.py returns garbage on every AMD PR.

#32272 changed MHATokenToKVPool to store each layer's K and V as block-aligned halves of one allocation, so v_buffer[i] now starts inside the allocation of k_buffer[i]. Only the trtllm_mha fmha_v2 prefill on CUDA reads that layout, but the allocation change has no platform gate. The test serves Qwen3-8B with Mooncake, prefill and decode on one MI35x node. The decode server now returns nonsense with no transfer error:

FAIL: test_gsm8k (__main__.TestDisaggregationAccuracy.test_gsm8k)
AssertionError: 0.005 not greater than 0.7

The MI35x 8-GPU disaggregation job ran on many PRs on 2026-10-27. Grouped by the main commit each run merged onto:

main base	jobs	test_disaggregation_basic.py
db90f584d2, 2b36c718cb, e4d4a9a7de (before #32272)	9	passed in all 9 (example)
1a120754e0 (#32272) and every later base seen (d3a157a2f2, 61a676125a, 449617bc94, 3927a73a34, 8a455f1b90, a6f5bd397b)	9	failed in all 9, GSM8K accuracy 0.0 to 0.01 (first)

Seven of the nine earlier jobs failed in test_disaggregation_pp.py, which also fails on main's scheduled runs. test_disaggregation_basic.py passed in each of them.

Three other commits sit between e4d4a9a7de and 1a120754e0. They are #31286 (removes an unused HiCache field), #46366 (gfx1250 diffusion GEMM) and #33743 (flashinfer TRT-LLM MoE weight reload). None of them touches the KV pool or PD. Every job ran on a different runner, so one bad machine does not explain it.

Mooncake's HIP transport carries the same-node transfer. It registers each KV buffer with its own hipIpcGetMemHandle and maps peer addresses as dest - entry.addr + shm_addr. With the fused layout, both handles of a layer point into one allocation. I did not find which step then puts V at the wrong address, so the evidence here is the bisect, not a mechanism.

Modifications

• In memory_pool.py, the fused allocation now also needs not _is_hip. ROCm gets back the two torch.zeros per layer it had before [Perf] fmha_v2: use non-interleaved paged KV input for trtllm sm90/120 GQA prefill #32272.
• In test_mha_pool_fused_kv_layout.py, a new CPU test patches _is_hip and checks that K and V do not share storage. The existing fused-layout test still runs on the CPU runner, where _is_hip is false. Both tests build the pool with a shared _make_pool helper.

Off ROCm this change does nothing. On ROCm nothing reads the fused layout, because fmha_v2 runs only on CUDA, so the only effect there is the allocation from before #32272.

Accuracy Tests

I have not run it on MI35x. The AMD disaggregation job of this PR is the end-to-end check, and I will post its result here.

The CPU unit test on this branch and on main:

	test_v_is_block_aligned_offset_from_k	test_rocm_keeps_k_and_v_in_separate_allocations
this PR	pass	pass
main	pass	fail (AssertionError: 4573659136 == 4573659136 : layer 0)

Untested: same-node PD on CUDA through an IPC transport, which could meet the same layout.

Speed Tests and Profiling

None. ROCm goes back to the allocation it used before #32272. CUDA is unchanged, so it keeps the fmha_v2 speedup from #32272.

Checklist

• Format your code according to the Format code with pre-commit. (isort 7.0.0, ruff 0.15.1 check and format)
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. Not needed, no user-facing option changes.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. Pending the MI35x PD job of this PR.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): ⏳ Run #40656882185
Latest PR Test (Extra): ❌ Run #40656882229
Latest PR Test (AMD ROCm 10): ⏳ Run #40656882098

[1am9trash]

LGTM.
Add a guard to resolve PD acc issue.
Only impact on amd code path.