[Fix] [AMD] memory_pool: Keep K and V in separate allocations on ROCm

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

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()