From 5662c0336320ba42be3983b2a8c8de6c30fba743 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sat, 22 Aug 2026 10:19:06 +0800 Subject: [PATCH] Support CPU offload for mxfp8 KV cache (#35888) --- python/sglang/srt/mem_cache/memory_pool.py | 68 +++++++++++++++++-- .../sglang/srt/mem_cache/swa_memory_pool.py | 17 ++--- .../mem_cache/test_swa_cpu_copy_filter.py | 61 +++++++++++++++++ 3 files changed, 134 insertions(+), 12 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 05bfc1259..da331decf 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -3547,13 +3547,73 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool): self.k_scale_buffer[idx][tgt_loc] = self.k_scale_buffer[idx][src_loc] self.v_scale_buffer[idx][tgt_loc] = self.v_scale_buffer[idx][src_loc] - # These paths copy k/v buffers without the scale buffers; fail loudly - # instead of silently corrupting dequantization. + def _read_scales(self, idx, loc): + """Per-token UE8M0 K/V scales at ``loc``, inverse of ``_write_scales``.""" + if self.mxfp8_sf_interleaved: + return ( + self._read_sf_interleaved(self.k_scale_buffer[idx], loc), + self._read_sf_interleaved(self.v_scale_buffer[idx], loc), + ) + return self.k_scale_buffer[idx][loc], self.v_scale_buffer[idx][loc] + def get_cpu_copy(self, indices, mamba_indices=None): - raise NotImplementedError("CPU offloading is unsupported for MXFP8 KV cache.") + # The scales travel with their fp8 payload; a restored slot dequantizes + # against mismatched exponents without them. + assert not self.use_hnd, ( + "CPU KV offload indexes by slot (NHD); HND KV cache " + "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet." + ) + current_platform.synchronize() + kv_cache_cpu = [] + chunk_size = self.cpu_offloading_chunk_size + for layer_id in range(self.layer_num): + kv_cache_cpu.append([]) + for i in range(0, len(indices), chunk_size): + chunk_indices = indices[i : i + chunk_size] + k_scale, v_scale = self._read_scales(layer_id, chunk_indices) + kv_cache_cpu[-1].append( + [ + self.k_buffer[layer_id][chunk_indices].to( + "cpu", non_blocking=True + ), + self.v_buffer[layer_id][chunk_indices].to( + "cpu", non_blocking=True + ), + k_scale.to("cpu", non_blocking=True), + v_scale.to("cpu", non_blocking=True), + ] + ) + current_platform.synchronize() + return kv_cache_cpu def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None): - raise NotImplementedError("CPU offloading is unsupported for MXFP8 KV cache.") + assert not self.use_hnd, ( + "CPU KV offload indexes by slot (NHD); HND KV cache " + "(SGLANG_USE_HND_KVCACHE) is not supported with CPU offload yet." + ) + current_platform.synchronize() + device = self.k_buffer[0].device + chunk_size = self.cpu_offloading_chunk_size + for layer_id in range(self.layer_num): + for i in range(0, len(indices), chunk_size): + chunk_indices = indices[i : i + chunk_size] + k_cpu, v_cpu, k_scale_cpu, v_scale_cpu = kv_cache_cpu[layer_id][ + i // chunk_size + ] + assert k_cpu.shape[0] == v_cpu.shape[0] == len(chunk_indices) + self.k_buffer[layer_id][chunk_indices] = k_cpu.to( + device, non_blocking=True + ) + self.v_buffer[layer_id][chunk_indices] = v_cpu.to( + device, non_blocking=True + ) + self._write_scales( + layer_id, + chunk_indices, + k_scale_cpu.to(device, non_blocking=True), + v_scale_cpu.to(device, non_blocking=True), + ) + current_platform.synchronize() def get_kv_scale_buf_infos(self): """(ptrs, lens, item_lens) for the UE8M0 scale buffers, k then v. diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index f52c90c01..8b7cf4821 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/python/sglang/srt/mem_cache/swa_memory_pool.py @@ -245,16 +245,17 @@ class SWAKVPool(BaseSWAKVPool): filtered.append([]) continue - k_cpu = torch.cat([chunk[0] for chunk in layer_chunks], dim=0) - v_cpu = torch.cat([chunk[1] for chunk in layer_chunks], dim=0) - k_cpu = k_cpu[row_mask] - v_cpu = v_cpu[row_mask] + # A chunk is whatever the sub-pool produced: k/v, plus the block + # scales for a quantized pool. Filter every tensor it carries. + num_tensors = len(layer_chunks[0]) + tensors = [ + torch.cat([chunk[t] for chunk in layer_chunks], dim=0)[row_mask] + for t in range(num_tensors) + ] filtered_layer = [] - for i in range(0, len(k_cpu), chunk_size): - filtered_layer.append( - [k_cpu[i : i + chunk_size], v_cpu[i : i + chunk_size]] - ) + for i in range(0, len(tensors[0]), chunk_size): + filtered_layer.append([t[i : i + chunk_size] for t in tensors]) filtered.append(filtered_layer) return filtered diff --git a/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py b/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py new file mode 100644 index 000000000..fa5c1a6d8 --- /dev/null +++ b/test/registered/unit/mem_cache/test_swa_cpu_copy_filter.py @@ -0,0 +1,61 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +ROWS = 4 + + +def _pool(chunk_size: int = ROWS) -> SWAKVPool: + pool = object.__new__(SWAKVPool) + pool.swa_kv_pool = SimpleNamespace(cpu_offloading_chunk_size=chunk_size) + return pool + + +def _rows(t: int): + return torch.arange(ROWS * 2).reshape(ROWS, 2) + t * 100 + + +def _chunk(num_tensors: int): + """swa_kv_cpu is layers -> chunks -> tensors; one layer, one chunk here.""" + return [[[_rows(t) for t in range(num_tensors)]]] + + +class TestSWACpuCopyFilter(unittest.TestCase): + def test_keeps_every_tensor_a_chunk_carries(self): + """A quantized sub-pool puts the block scales in the same chunk as K/V. + Trimming rows must not drop them: the load side unpacks whatever the get + side produced, and a short chunk resumes against the wrong exponents.""" + row_mask = torch.tensor([True, False, True, False]) + + filtered = _pool()._filter_swa_cpu_copy(_chunk(4), row_mask) + + self.assertEqual(len(filtered[0][0]), 4) + for t, tensor in enumerate(filtered[0][0]): + expected = _rows(t)[row_mask] + self.assertTrue(torch.equal(tensor, expected)) + + def test_unquantized_chunk_is_unchanged(self): + row_mask = torch.tensor([False, True, True, False]) + + filtered = _pool()._filter_swa_cpu_copy(_chunk(2), row_mask) + + self.assertEqual(len(filtered[0][0]), 2) + self.assertEqual(filtered[0][0][0].shape[0], 2) + + def test_all_rows_kept_returns_the_input(self): + original = _chunk(4) + + self.assertIs( + _pool()._filter_swa_cpu_copy(original, torch.tensor([True] * ROWS)), + original, + ) + + +if __name__ == "__main__": + unittest.main()