From e9609232c7bf113a13c32a6a9d527f9f8f77fcad Mon Sep 17 00:00:00 2001 From: Olga Miroshnichenko Date: Tue, 27 Oct 2026 17:41:29 +0200 Subject: [PATCH] [Fix] [AMD] memory_pool: Keep K and V in separate allocations on ROCm Since #32272, MHATokenToKVPool puts K and V of each layer in one allocation, with V at a block offset. Only the CUDA trtllm fmha_v2 prefill reads that layout. On MI35x, the same-node Mooncake PD test test_disaggregation_basic.py now gets GSM8K accuracy 0.0 to 0.01. It passes on each main base before #32272 and fails on each base after it. Keep the fused layout off on ROCm. Add a CPU test that patches _is_hip and checks that K and V do not share storage. Co-authored-by: Cursor --- python/sglang/srt/mem_cache/memory_pool.py | 5 ++- .../test_mha_pool_fused_kv_layout.py | 42 +++++++++++++------ 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index f6d0b71ae3cb..141afa102296 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -2627,7 +2627,10 @@ def _zeros(shape): shape, dtype=self.store_dtype, device=self.device ) - if k_shape == v_shape: + # Only CUDA's trtllm fmha_v2 prefill reads the fused layout, + # and same-node Mooncake PD on ROCm decodes garbage when V + # sits inside K's allocation, so ROCm allocates them apart. + if k_shape == v_shape and not _is_hip: # Store K and V as block-aligned halves of one allocation # so paged kernels with a single pool base can reach V # via block offsets. Each half keeps its standalone diff --git a/test/registered/unit/mem_cache/test_mha_pool_fused_kv_layout.py b/test/registered/unit/mem_cache/test_mha_pool_fused_kv_layout.py index 447368463850..57a79cf4cace 100644 --- a/test/registered/unit/mem_cache/test_mha_pool_fused_kv_layout.py +++ b/test/registered/unit/mem_cache/test_mha_pool_fused_kv_layout.py @@ -2,15 +2,18 @@ trtllm_mha's fmha_v2 prefill reaches V from k_cache's base pointer via int32 block offsets and falls back to a full-pool copy (with only a warning) when -the layout breaks. This pins the invariant that path relies on. +the layout breaks. This pins the invariant that path relies on. ROCm keeps K +and V in separate allocations. python -m pytest test/registered/unit/mem_cache/test_mha_pool_fused_kv_layout.py -v """ import unittest +from unittest import mock import torch +from sglang.srt.mem_cache import memory_pool from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool from sglang.test.ci.ci_register import register_cpu_ci @@ -21,21 +24,25 @@ LAYER_NUM = 2 +def _make_pool(page_size): + return MHATokenToKVPool( + size=page_size * 7, + page_size=page_size, + dtype=torch.bfloat16, + head_num=HEAD_NUM, + head_dim=HEAD_DIM, + layer_num=LAYER_NUM, + device="cpu", + enable_memory_saver=False, + enable_alt_stream=False, + ) + + class TestMHAPoolFusedKVLayout(unittest.TestCase): def test_v_is_block_aligned_offset_from_k(self): for page_size in (16, 32, 64): with self.subTest(page_size=page_size): - pool = MHATokenToKVPool( - size=page_size * 7, - page_size=page_size, - dtype=torch.bfloat16, - head_num=HEAD_NUM, - head_dim=HEAD_DIM, - layer_num=LAYER_NUM, - device="cpu", - enable_memory_saver=False, - enable_alt_stream=False, - ) + pool = _make_pool(page_size) for layer_id in range(LAYER_NUM): k_buf, v_buf = pool.get_kv_buffer(layer_id) # Same paged view and block size the fmha_v2 path computes. @@ -47,6 +54,17 @@ def test_v_is_block_aligned_offset_from_k(self): self.assertEqual(rem, 0, f"layer {layer_id}") self.assertGreater(delta, 0, f"layer {layer_id}") + def test_rocm_keeps_k_and_v_in_separate_allocations(self): + with mock.patch.object(memory_pool, "_is_hip", True): + pool = _make_pool(page_size=16) + for layer_id in range(LAYER_NUM): + k_buf, v_buf = pool.get_kv_buffer(layer_id) + self.assertNotEqual( + k_buf.untyped_storage().data_ptr(), + v_buf.untyped_storage().data_ptr(), + f"layer {layer_id}", + ) + if __name__ == "__main__": unittest.main()