diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 06cca265d..74ddba1ab 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_mem_pool_host.py b/test/registered/unit/mem_cache/test_mem_pool_host.py index 4a7179536..3d9f39c06 100644 --- a/test/registered/unit/mem_cache/test_mem_pool_host.py +++ b/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)