[HiCache] Size MHA host pools from device row width (#40304)

python/sglang/srt/mem_cache/pool_host/mha.py

@@ -102,7 +102,7 @@ class MHATokenToKVPoolHost(HostKVCache):
             allocator_type,
             pool_label=pool_label,
         )
-        self.element_dim = self.device_pool.head_num * self.device_pool.head_dim
+        self.element_dim = self.head_num * self.head_dim
         # The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
         # helpers in hicache.cuh are guarded by USE_ROCM and the staged
         # write-back kernel has a ROCm path, so enable them on HIP too. This
@@ -156,7 +156,8 @@ class MHATokenToKVPoolHost(HostKVCache):
         self._init_write_back_staging_buffers()
 
     def get_size_per_token(self):
-        self.head_num = self.device_pool.head_num
+        # One allocator token may hold multiple attention rows.
+        self.head_num = self.device_pool.row_dim // self.device_pool.head_dim
         self.head_dim = self.device_pool.head_dim
         self.layer_num = self.target_layer_num + len(self.mtp_draft_device_pools)
         return self.head_dim * self.head_num * self.layer_num * self.dtype.itemsize * 2

test/registered/unit/mem_cache/test_mem_pool_host.py

@@ -47,6 +47,47 @@ class TestHostKVCache(CustomTestCase):
             allocator_type="default",
         )
 
+    def test_multiple_attention_rows_per_token(self):
+        for rows_per_token in (1, 3):
+            device_pool = MHATokenToKVPool(
+                size=4,
+                page_size=self.page_size,
+                dtype=torch.float16,
+                head_num=2 * rows_per_token,
+                head_dim=4,
+                layer_num=2,
+                device="cpu",
+                enable_memory_saver=False,
+            )
+            # Report logical heads while retaining the wider physical rows.
+            device_pool.head_num = 2
+            for layout in (
+                "layer_first",
+                "page_first",
+                "page_first_direct",
+                "page_head",
+            ):
+                with self.subTest(rows_per_token=rows_per_token, layout=layout):
+                    host_pool = MHATokenToKVPoolHost(
+                        device_pool=device_pool,
+                        host_to_device_ratio=2.0,
+                        host_size=0,
+                        page_size=self.page_size,
+                        layout=layout,
+                        pin_memory=False,
+                    )
+                    device_row = device_pool.k_buffer[0][0]
+                    row_bytes = device_row.numel() * device_row.element_size()
+                    self.assertEqual(host_pool.element_dim, device_row.numel())
+                    self.assertEqual(host_pool.token_stride_size, row_bytes)
+                    self.assertEqual(
+                        host_pool.size_per_token, 2 * device_pool.layer_num * row_bytes
+                    )
+                    self.assertEqual(
+                        host_pool.kv_buffer.nbytes,
+                        host_pool.size * host_pool.size_per_token,
+                    )
+
     def test_double_alloc(self):
         indices = self.host_pool.alloc(4)
         self.assertEqual(len(indices), 4)