From 4c85172f3a05d7959a69f8179587b6ac92494d06 Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Thu, 17 Sep 2026 09:40:26 +0800 Subject: [PATCH] [HiCache] Back up MXFP8 KV scales in the host pool (#39089) --- python/sglang/srt/mem_cache/pool_host/mha.py | 21 +- .../srt/mem_cache/pool_host/mha_mxfp8.py | 266 ++++++++++++++++++ .../test_mxfp8_mha_pool_host_unit.py | 195 +++++++++++++ 3 files changed, 480 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/mem_cache/pool_host/mha_mxfp8.py create mode 100644 test/registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 8a3d488cc..f91e6d72a 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -28,7 +28,11 @@ from sglang.kernels.ops.kvcache.hicache import ( from sglang.kernels.ops.kvcache.hicache import ( transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla, ) -from sglang.srt.mem_cache.memory_pool import MHATokenToKOnlyPool, MHATokenToKVPool +from sglang.srt.mem_cache.memory_pool import ( + MHATokenToKOnlyPool, + MHATokenToKVPool, + MHATokenToKVPoolMXFP8, +) from sglang.srt.mem_cache.pool_host.base import ( _WRITE_BACK_STAGING_PAGE_CHUNK, HostKVCache, @@ -1415,9 +1419,22 @@ class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost): def get_mha_host_pool_cls(device_pool: MHATokenToKVPool) -> type: """Pick the right MHA host-pool class based on the device pool's K/V dims. - Returns ``AsymmetricMHATokenToKVPoolHost`` when ``head_dim != v_head_dim`` + Returns ``MHATokenToKVPoolMXFP8Host`` for the block-scaled MXFP8 pool (its + UE8M0 scales must travel with the payload), + ``AsymmetricMHATokenToKVPoolHost`` when ``head_dim != v_head_dim`` (e.g. MiMo-V2), else the default ``MHATokenToKVPoolHost``. """ + if isinstance(device_pool, MHATokenToKVPoolMXFP8): + if device_pool.head_dim != device_pool.v_head_dim: + raise NotImplementedError( + "MXFP8 HiCache does not support asymmetric K/V head dimensions yet." + ) + + from sglang.srt.mem_cache.pool_host.mha_mxfp8 import ( + MHATokenToKVPoolMXFP8Host, + ) + + return MHATokenToKVPoolMXFP8Host if device_pool.head_dim != device_pool.v_head_dim: return AsymmetricMHATokenToKVPoolHost return MHATokenToKVPoolHost diff --git a/python/sglang/srt/mem_cache/pool_host/mha_mxfp8.py b/python/sglang/srt/mem_cache/pool_host/mha_mxfp8.py new file mode 100644 index 000000000..1c4c0933c --- /dev/null +++ b/python/sglang/srt/mem_cache/pool_host/mha_mxfp8.py @@ -0,0 +1,266 @@ +"""Host (L2) pool for the MXFP8 block-scaled MHA KV cache. + +The fp8 payload rides the regular MHA host pool; this subclass adds the +per-page UE8M0 scale blocks, which the device pool keeps in the FA4 +interleaved layout (num_pages, head, 32, page_size // 32, sf_dim). A page's +scales are one contiguous block, so they move as whole-page rows of +`n_heads * page_size * sf_dim` bytes with the same staged JIT kernels the +payload uses, just at page granularity (page_size=1 in kernel terms). +""" + +from __future__ import annotations + +from typing import Sequence + +import torch + +from sglang.kernels.ops.kvcache.hicache import ( + can_use_write_back_jit_kernel, +) +from sglang.kernels.ops.kvcache.hicache import ( + transfer_hicache_all_layer_mla_staged_lf_pf as jit_transfer_hicache_all_layer_mla_staged_lf_pf, +) +from sglang.kernels.ops.kvcache.hicache import ( + transfer_hicache_one_layer_mla as jit_transfer_hicache_one_layer_mla, +) +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPoolMXFP8 +from sglang.srt.mem_cache.pool_host.common import ( + ALLOC_MEMORY_FUNCS, + _cuda_host_unregister, +) +from sglang.srt.mem_cache.pool_host.mha import ( + MHATokenToKVPoolHost, + _is_cuda, + _is_hip, +) + + +class MHATokenToKVPoolMXFP8Host(MHATokenToKVPoolHost): + device_pool: MHATokenToKVPoolMXFP8 + + def __init__( + self, + device_pool: MHATokenToKVPoolMXFP8, + host_to_device_ratio: float, + host_size: int, + page_size: int, + layout: str, + pin_memory: bool = True, + device: str = "cpu", + allocator_type: str = "default", + *, + mtp_draft_device_pools: Sequence = (), + pool_label: str = "kv", + ): + if layout != "page_first": + raise NotImplementedError( + f"MXFP8 KV host pool supports only the page_first layout, got {layout!r}." + ) + if mtp_draft_device_pools: + raise NotImplementedError( + "MXFP8 KV host pool does not pack MTP draft KV layers." + ) + if not device_pool.mxfp8_sf_interleaved: + raise NotImplementedError( + "MXFP8 KV host pool requires the interleaved (page_size=128) scale layout." + ) + # Bytes of UE8M0 scales one page holds per layer; needed by + # get_size_per_token before the base constructor sizes the pool. + self.k_sf_page_bytes = device_pool.k_scale_buffer[0][0].numel() + self.v_sf_page_bytes = device_pool.v_scale_buffer[0][0].numel() + self.k_scale_host: torch.Tensor | None = None + self.v_scale_host: torch.Tensor | None = None + super().__init__( + device_pool, + host_to_device_ratio, + host_size, + page_size, + layout, + pin_memory, + device, + allocator_type, + pool_label=pool_label, + ) + if self.page_size != device_pool.page_size: + raise ValueError( + "MXFP8 KV host pool moves scales per page, so the host page size " + f"({self.page_size}) must equal the device page size " + f"({device_pool.page_size})." + ) + if not self.can_use_write_back_jit or not all( + can_use_write_back_jit_kernel(element_size=size) + for size in (self.k_sf_page_bytes, self.v_sf_page_bytes) + ): + raise NotImplementedError( + "MXFP8 KV host pool needs the staged JIT write-back kernel " + "(io_backend='kernel', page_first layout, CUDA or HIP)." + ) + self._init_scale_buffers() + + def get_size_per_token(self): + payload = super().get_size_per_token() + scales = (self.k_sf_page_bytes + self.v_sf_page_bytes) // self.page_size + return payload + scales * self.layer_num + + def _init_scale_buffers(self): + alloc_func = ALLOC_MEMORY_FUNCS[self.device_pool.device] + # Host: page-first like the payload, one contiguous scale block per + # (page, layer) so a page's layers are a single memcpy span. + self.k_scale_host = alloc_func( + (self.page_num, self.layer_num, self.k_sf_page_bytes), + dtype=torch.uint8, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + self.v_scale_host = alloc_func( + (self.page_num, self.layer_num, self.v_sf_page_bytes), + dtype=torch.uint8, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) + # [page, layer, bytes] -> per-layer strided [page, bytes] views for H2D. + self.k_scale_host_layers = list(self.k_scale_host.transpose(0, 1)) + self.v_scale_host_layers = list(self.v_scale_host.transpose(0, 1)) + + # Device: each layer's interleaved scale tensor as [page, bytes] rows. + self.k_scale_device_rows = [ + buf.view(torch.uint8).reshape(buf.shape[0], -1) + for buf in self.device_pool.k_scale_buffer + ] + self.v_scale_device_rows = [ + buf.view(torch.uint8).reshape(buf.shape[0], -1) + for buf in self.device_pool.v_scale_buffer + ] + self.k_scale_device_ptrs = torch.tensor( + [rows.data_ptr() for rows in self.k_scale_device_rows], + dtype=torch.uint64, + device=self.device_pool.device, + ) + self.v_scale_device_ptrs = torch.tensor( + [rows.data_ptr() for rows in self.v_scale_device_rows], + dtype=torch.uint64, + device=self.device_pool.device, + ) + self.k_scale_staging = torch.empty( + (self.staging_page_capacity, self.layer_num, self.k_sf_page_bytes), + dtype=torch.uint8, + device=self.device_pool.device, + ) + self.v_scale_staging = torch.empty( + (self.staging_page_capacity, self.layer_num, self.v_sf_page_bytes), + dtype=torch.uint8, + device=self.device_pool.device, + ) + + def _page_ids(self, indices: torch.Tensor) -> torch.Tensor: + """Page-aligned token indices -> one page id per page, same device.""" + return indices[:: self.page_size] // self.page_size + + def backup_from_device_all_layer( + self, device_pool, host_indices, device_indices, io_backend + ): + super().backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend + ) + if io_backend != "kernel": + raise NotImplementedError( + f"MXFP8 KV host pool supports only io_backend='kernel', got {io_backend!r}." + ) + device_pages = self._page_ids(device_indices) + host_pages = self._page_ids(host_indices) + if host_pages.is_cuda: + host_pages = host_pages.cpu() + for ptr_src, staging, dst in ( + (self.k_scale_device_ptrs, self.k_scale_staging, self.k_scale_host), + (self.v_scale_device_ptrs, self.v_scale_staging, self.v_scale_host), + ): + jit_transfer_hicache_all_layer_mla_staged_lf_pf( + ptr_src=ptr_src, + src_indices=device_pages, + dst_indices=host_pages, + staging=staging, + dst=dst, + page_size=1, + ) + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + *, + is_draft: bool = False, + ): + super().load_to_device_per_layer( + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + is_draft=is_draft, + ) + if is_draft: + raise NotImplementedError("MXFP8 KV host pool has no draft layers.") + if io_backend != "kernel": + raise NotImplementedError( + f"MXFP8 KV host pool supports only io_backend='kernel', got {io_backend!r}." + ) + if not self._is_device_layer_owned(device_pool, layer_id): + return + host_layer_id = self._host_layer_index(layer_id) + device_pages = self._page_ids(device_indices) + host_pages = self._page_ids(host_indices) + for dst_rows, src_rows, sf_page_bytes in ( + ( + self.k_scale_device_rows[layer_id], + self.k_scale_host_layers[host_layer_id], + self.k_sf_page_bytes, + ), + ( + self.v_scale_device_rows[layer_id], + self.v_scale_host_layers[host_layer_id], + self.v_sf_page_bytes, + ), + ): + jit_transfer_hicache_one_layer_mla( + cache_dst=dst_rows, + indices_dst=device_pages, + cache_src=src_rows, + indices_src=host_pages, + element_dim=sf_page_bytes, + ) + + def destroy(self): + for buf in (self.k_scale_host, self.v_scale_host): + if buf is not None and self.pin_memory and (_is_cuda or _is_hip): + _cuda_host_unregister(buf) + self.k_scale_host = None + self.v_scale_host = None + super().destroy() + + def _storage_pages_unsupported(self) -> NotImplementedError: + return NotImplementedError( + "MXFP8 KV host pool does not expose flat storage (L3) pages yet: " + "a page's UE8M0 scales live outside kv_buffer." + ) + + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + raise self._storage_pages_unsupported() + + def get_dummy_flat_data_page(self) -> torch.Tensor: + raise self._storage_pages_unsupported() + + def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + raise self._storage_pages_unsupported() + + def get_page_buffer_meta(self, indices): + raise self._storage_pages_unsupported() + + def get_split_heads_page_buffer_meta( + self, indices: torch.Tensor, split_factor: int + ): + raise self._storage_pages_unsupported() diff --git a/test/registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py b/test/registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py new file mode 100644 index 000000000..f216c4f56 --- /dev/null +++ b/test/registered/unit/mem_cache/test_mxfp8_mha_pool_host_unit.py @@ -0,0 +1,195 @@ +"""Unit tests for the MXFP8 MHA host pool: UE8M0 scales must round-trip +through L2 alongside the fp8 payload.""" + +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPoolMXFP8 +from sglang.srt.mem_cache.pool_host.mha import ( + MHATokenToKVPoolHost, + get_mha_host_pool_cls, +) +from sglang.srt.mem_cache.pool_host.mha_mxfp8 import MHATokenToKVPoolMXFP8Host +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +MXFP8_MODULE = "sglang.srt.mem_cache.pool_host.mha_mxfp8" +MHA_MODULE = "sglang.srt.mem_cache.pool_host.mha" + +PAGE_SIZE = 4 +LAYER_NUM = 2 +HEAD_NUM = 1 +HEAD_DIM = 8 +SF_DIM = HEAD_DIM // 32 or 1 +SF_PAGE_BYTES = HEAD_NUM * PAGE_SIZE * SF_DIM + + +def _ptr_key(ptrs: torch.Tensor) -> tuple[int, ...]: + return tuple(int(ptr) for ptr in ptrs.cpu().tolist()) + + +def _cpu_mla_staged_lf_pf_copy( + src_registry, *, ptr_src, src_indices, dst_indices, dst, **_ +): + """CPU stand-in for the staged D2H kernel: per (row, layer) gather-scatter.""" + src_layers = src_registry[_ptr_key(ptr_src)] + for layer_id, src in enumerate(src_layers): + dst[dst_indices, layer_id] = src[src_indices] + + +def _cpu_one_layer_mla_copy(*, cache_dst, indices_dst, cache_src, indices_src, **_): + cache_dst[indices_dst] = cache_src[indices_src] + + +def _make_host(k_rows, v_rows): + """A page_first MXFP8 host pool over CPU tensors, built without the base + constructor so no pinned memory or device pool is needed.""" + host = MHATokenToKVPoolMXFP8Host.__new__(MHATokenToKVPoolMXFP8Host) + host.layout = "page_first" + host.page_size = PAGE_SIZE + host.layer_num = LAYER_NUM + host.page_num = 4 + host.k_sf_page_bytes = SF_PAGE_BYTES + host.v_sf_page_bytes = SF_PAGE_BYTES + host.device_pool = SimpleNamespace(layer_shard_enabled=False, layer_num=LAYER_NUM) + host.k_scale_host = torch.zeros( + host.page_num, LAYER_NUM, SF_PAGE_BYTES, dtype=torch.uint8 + ) + host.v_scale_host = torch.zeros( + host.page_num, LAYER_NUM, SF_PAGE_BYTES, dtype=torch.uint8 + ) + host.k_scale_host_layers = list(host.k_scale_host.transpose(0, 1)) + host.v_scale_host_layers = list(host.v_scale_host.transpose(0, 1)) + host.k_scale_device_rows = k_rows + host.v_scale_device_rows = v_rows + host.k_scale_device_ptrs = torch.tensor( + [r.data_ptr() for r in k_rows], dtype=torch.uint64 + ) + host.v_scale_device_ptrs = torch.tensor( + [r.data_ptr() for r in v_rows], dtype=torch.uint64 + ) + host.k_scale_staging = torch.empty(2, LAYER_NUM, SF_PAGE_BYTES, dtype=torch.uint8) + host.v_scale_staging = torch.empty(2, LAYER_NUM, SF_PAGE_BYTES, dtype=torch.uint8) + return host + + +class TestMXFP8MHATokenToKVPoolHost(CustomTestCase): + def test_factory_selects_mxfp8_host_pool(self): + mxfp8_pool = MHATokenToKVPoolMXFP8.__new__(MHATokenToKVPoolMXFP8) + mxfp8_pool.head_dim = mxfp8_pool.v_head_dim = 128 + plain_pool = SimpleNamespace(head_dim=4, v_head_dim=4) + + self.assertIs(get_mha_host_pool_cls(mxfp8_pool), MHATokenToKVPoolMXFP8Host) + self.assertIs(get_mha_host_pool_cls(plain_pool), MHATokenToKVPoolHost) + + def test_size_per_token_counts_scales(self): + host = MHATokenToKVPoolMXFP8Host.__new__(MHATokenToKVPoolMXFP8Host) + host.page_size = PAGE_SIZE + host.k_sf_page_bytes = SF_PAGE_BYTES + host.v_sf_page_bytes = SF_PAGE_BYTES + payload = 2 * LAYER_NUM * HEAD_NUM * HEAD_DIM + with mock.patch.object( + MHATokenToKVPoolHost, "get_size_per_token", return_value=payload + ): + host.layer_num = LAYER_NUM + self.assertEqual( + host.get_size_per_token(), + payload + 2 * HEAD_NUM * SF_DIM * LAYER_NUM, + ) + + def test_scales_round_trip_device_host_device(self): + num_device_pages = 4 + k_rows = [ + ( + torch.arange(num_device_pages * SF_PAGE_BYTES, dtype=torch.uint8) + + 10 * layer + ) + .reshape(num_device_pages, SF_PAGE_BYTES) + .clone() + for layer in range(LAYER_NUM) + ] + v_rows = [ + ( + torch.arange(num_device_pages * SF_PAGE_BYTES, dtype=torch.uint8) + + 100 + + 10 * layer + ) + .reshape(num_device_pages, SF_PAGE_BYTES) + .clone() + for layer in range(LAYER_NUM) + ] + host = _make_host(k_rows, v_rows) + # Two device pages (2, 3) back up into host pages (1, 0). + device_indices = torch.arange(2 * PAGE_SIZE, 4 * PAGE_SIZE, dtype=torch.int64) + host_indices = torch.cat( + [torch.arange(PAGE_SIZE, 2 * PAGE_SIZE), torch.arange(0, PAGE_SIZE)] + ).to(torch.int64) + expected_k = [rows[2:4].clone() for rows in k_rows] + expected_v = [rows[2:4].clone() for rows in v_rows] + registry = { + _ptr_key(host.k_scale_device_ptrs): k_rows, + _ptr_key(host.v_scale_device_ptrs): v_rows, + } + + with ( + mock.patch.object( + MHATokenToKVPoolHost, "backup_from_device_all_layer" + ) as payload_backup, + mock.patch.object( + MHATokenToKVPoolHost, "load_to_device_per_layer" + ) as payload_load, + mock.patch( + f"{MXFP8_MODULE}.jit_transfer_hicache_all_layer_mla_staged_lf_pf", + side_effect=lambda **kw: _cpu_mla_staged_lf_pf_copy(registry, **kw), + ) as staged, + mock.patch( + f"{MXFP8_MODULE}.jit_transfer_hicache_one_layer_mla", + side_effect=_cpu_one_layer_mla_copy, + ) as one_layer, + ): + host.backup_from_device_all_layer( + host.device_pool, host_indices, device_indices, io_backend="kernel" + ) + self.assertEqual(staged.call_count, 2) + # Device page 2 landed in host page 1, device page 3 in host page 0. + self.assertTrue( + torch.equal(host.k_scale_host[1], torch.stack([r[2] for r in k_rows])) + ) + self.assertTrue( + torch.equal(host.v_scale_host[0], torch.stack([r[3] for r in v_rows])) + ) + + for rows in k_rows + v_rows: + rows.zero_() + for layer_id in range(LAYER_NUM): + host.load_to_device_per_layer( + host.device_pool, + host_indices, + device_indices, + layer_id, + io_backend="kernel", + ) + + payload_backup.assert_called_once() + self.assertEqual(payload_load.call_count, LAYER_NUM) + self.assertEqual(one_layer.call_count, 2 * LAYER_NUM) + for layer in range(LAYER_NUM): + self.assertTrue(torch.equal(k_rows[layer][2:4], expected_k[layer])) + self.assertTrue(torch.equal(v_rows[layer][2:4], expected_v[layer])) + self.assertTrue(torch.all(k_rows[layer][:2] == 0)) + + def test_storage_pages_are_rejected(self): + host = MHATokenToKVPoolMXFP8Host.__new__(MHATokenToKVPoolMXFP8Host) + with self.assertRaises(NotImplementedError): + host.get_dummy_flat_data_page() + with self.assertRaises(NotImplementedError): + host.get_data_page(0) + + +if __name__ == "__main__": + unittest.main()