Support CPU offload for mxfp8 KV cache (#35888)

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
+    def _read_scales(self, idx, loc):
-    # instead of silently corrupting dequantization.
+        """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.

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)
+            # A chunk is whatever the sub-pool produced: k/v, plus the block
-            v_cpu = torch.cat([chunk[1] for chunk in layer_chunks], dim=0)
+            # scales for a quantized pool. Filter every tensor it carries.
-            k_cpu = k_cpu[row_mask]
+            num_tensors = len(layer_chunks[0])
-            v_cpu = v_cpu[row_mask]
+            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):
+            for i in range(0, len(tensors[0]), chunk_size):
-                filtered_layer.append(
+                filtered_layer.append([t[i : i + chunk_size] for t in tensors])
-                    [k_cpu[i : i + chunk_size], v_cpu[i : i + chunk_size]]
-                )
             filtered.append(filtered_layer)
         return filtered
 

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