Support mxfp8 KV cache in PD transfer (#35718)

python/sglang/srt/disaggregation/base/conn.py

@@ -24,6 +24,11 @@ class StateType(str, enum.Enum):
     SWA_RING = "swa_ring"
     # DeepSeek-V4 online C128 request-scoped state.
     C128_STATE = "c128_state"
+    # A block-scaled KV dtype keeps its per-block scales in buffers parallel to
+    # K/V, one component per sub-pool so each carries the index payload of the
+    # KV it describes (whole sequence for full attention, window for SWA).
+    BLOCK_SCALE = "block_scale"
+    BLOCK_SCALE_SWA = "block_scale_swa"
 
 
 @dataclasses.dataclass

python/sglang/srt/disaggregation/decode.py

@@ -1379,7 +1379,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 )
                 return kv_to_page_indices(window_kv_indices_swa, page_size)
 
-            def _dsa_payload():
+            def _full_kv_pages_payload():
                 kv_indices_full = self.req_to_token_pool.req_to_token[
                     decode_req.req.req_pool_idx, :seq_len
                 ]
@@ -1415,15 +1415,15 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 )
                 if clear_c128_state is not None:
                     clear_c128_state(int(decode_req.req.req_pool_idx))
-            # MINIMAX_INDEX_K reuses _dsa_payload: index rows live at the same loc
-            # as main KV on the same page_size.
             payloads = {
                 StateType.MAMBA: _mamba_payload,
                 StateType.SWA: _swa_payload,
-                StateType.DSA: _dsa_payload,
-                StateType.MINIMAX_INDEX_K: _dsa_payload,
+                StateType.DSA: _full_kv_pages_payload,
+                StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
                 StateType.SWA_RING: _swa_ring_payload,
                 StateType.C128_STATE: _c128_state_payload,
+                StateType.BLOCK_SCALE: _full_kv_pages_payload,
+                StateType.BLOCK_SCALE_SWA: _swa_payload,
             }
             if _is_npu and isinstance(self.token_to_kv_pool, DeepSeekV4TokenToKVPool):
                 from sglang.srt.hardware_backend.npu.dsv4.dsv4_common_hooks import (

python/sglang/srt/disaggregation/mooncake/conn.py

@@ -1255,6 +1255,8 @@ class MooncakeKVManager(CommonKVManager):
             StateType.DSA,
             StateType.SWA_RING,
             StateType.C128_STATE,
+            StateType.BLOCK_SCALE,
+            StateType.BLOCK_SCALE_SWA,
         )
 
     def _requires_exact_state_index_match(self, st: StateType) -> bool:

python/sglang/srt/disaggregation/prefill.py

@@ -1212,7 +1212,7 @@ class SchedulerDisaggregationPrefillMixin:
                 )
                 return kv_to_page_indices(window_kv_indices_swa, page_size)
 
-            def _dsa_payload():
+            def _full_kv_pages_payload():
                 kv_indices_full = self.req_to_token_pool.req_to_token[
                     req.req_pool_idx, :seq_len
                 ]
@@ -1250,15 +1250,15 @@ class SchedulerDisaggregationPrefillMixin:
             state_types = (
                 self.disagg_prefill_bootstrap_queue.kv_manager.kv_args.state_types
             )
-            # MINIMAX_INDEX_K reuses _dsa_payload: index rows live at the same loc
-            # as main KV on the same page_size.
             payloads = {
                 StateType.MAMBA: _mamba_payload,
                 StateType.SWA: _swa_payload,
-                StateType.DSA: _dsa_payload,
-                StateType.MINIMAX_INDEX_K: _dsa_payload,
+                StateType.DSA: _full_kv_pages_payload,
+                StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
                 StateType.SWA_RING: _swa_ring_payload,
                 StateType.C128_STATE: _c128_state_payload,
+                StateType.BLOCK_SCALE: _full_kv_pages_payload,
+                StateType.BLOCK_SCALE_SWA: _swa_payload,
             }
             if _is_npu and isinstance(
                 self.token_to_kv_pool_allocator.get_kvcache(),

python/sglang/srt/disaggregation/utils.py

@@ -1051,8 +1051,10 @@ def setup_state_kv_args(
     from sglang.srt.mem_cache.memory_pool import (
         DSATokenToKVPool,
         HybridLinearKVPool,
+        MHATokenToKVPoolMXFP8,
         MiniMaxSparseKVPool,
     )
+    from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 
     kv_args.state_types = []
     kv_args.state_data_ptrs = []
@@ -1064,6 +1066,13 @@ def setup_state_kv_args(
     kv_args.is_hybrid_mla_backend = False
     kv_args.state_conv_shard_groups = []
 
+    if isinstance(token_to_kv_pool, MHATokenToKVPoolMXFP8):
+        append_state_component(
+            kv_args,
+            StateType.BLOCK_SCALE,
+            *token_to_kv_pool.get_kv_scale_buf_infos(),
+        )
+
     if isinstance(token_to_kv_pool, MiniMaxSparseKVPool):
         if token_to_kv_pool.index_kv_pool is not None:
             raise NotImplementedError(
@@ -1082,6 +1091,23 @@ def setup_state_kv_args(
             append_state_component(
                 kv_args, StateType.SWA, data_ptrs, data_lens, item_lens
             )
+            # MXFP8 KV: each sub-pool's block scales ride as their own component
+            # so they inherit the index payload of the KV they describe.
+            # Only the concrete SWAKVPool owns a full sub-pool; other
+            # BaseSWAKVPool implementations describe their state per entry.
+            if isinstance(token_to_kv_pool, SWAKVPool) and isinstance(
+                token_to_kv_pool.full_kv_pool, MHATokenToKVPoolMXFP8
+            ):
+                append_state_component(
+                    kv_args,
+                    StateType.BLOCK_SCALE,
+                    *token_to_kv_pool.get_kv_scale_buf_infos(),
+                )
+                append_state_component(
+                    kv_args,
+                    StateType.BLOCK_SCALE_SWA,
+                    *token_to_kv_pool.get_swa_kv_scale_buf_infos(),
+                )
             # unified_kv: the SWA ring lives in the unified buffers (no separate
             # swa_kv_pool) and is addressed per-row, so ship it as SWA_RING.
             if getattr(token_to_kv_pool, "_unified_kv", False) and hasattr(

python/sglang/srt/mem_cache/memory_pool.py

@@ -3403,6 +3403,9 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool):
             ],
             device=self.device,
         )
+        # This override replaces the base allocation, so the PD-transfer
+        # descriptors for the packed data buffers are built here too.
+        self._kv_buffer_descs = self._build_kv_buffer_descs()
 
     def _clear_buffers(self):
         del self.k_buffer
@@ -3552,11 +3555,21 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool):
     def load_cpu_copy(self, kv_cache_cpu, indices, mamba_indices=None):
         raise NotImplementedError("CPU offloading is unsupported for MXFP8 KV cache.")
 
-    def get_contiguous_buf_infos(self):
-        raise NotImplementedError(
-            "KV transfer / disaggregation is unsupported for MXFP8 KV cache "
-            "(scale buffers are not exposed)."
-        )
+    def get_kv_scale_buf_infos(self):
+        """(ptrs, lens, item_lens) for the UE8M0 scale buffers, k then v.
+
+        The interleaved layout puts pages on the leading axis, so a page's
+        scales are one contiguous row; the flat layout is per slot.
+        """
+        tensors = self.k_scale_buffer + self.v_scale_buffer
+        ptrs = [t.data_ptr() for t in tensors]
+        lens = [t.nbytes for t in tensors]
+        row_bytes = [t[0].nbytes for t in tensors]
+        if self.mxfp8_sf_interleaved:
+            item_lens = row_bytes
+        else:
+            item_lens = [rb * self.page_size for rb in row_bytes]
+        return ptrs, lens, item_lens
 
     def set_kv_buffer_prefix_valid(self, *args, **kwargs):
         raise NotImplementedError(

python/sglang/srt/mem_cache/swa_memory_pool.py

@@ -137,6 +137,12 @@ class SWAKVPool(BaseSWAKVPool):
             full_kv_item_lens,
         )
 
+    def get_kv_scale_buf_infos(self):
+        return self.full_kv_pool.get_kv_scale_buf_infos()
+
+    def get_swa_kv_scale_buf_infos(self):
+        return self.swa_kv_pool.get_kv_scale_buf_infos()
+
     def get_state_buf_infos(self):
         swa_kv_data_ptrs, swa_kv_data_lens, swa_kv_item_lens = (
             self.swa_kv_pool.get_contiguous_buf_infos()

test/registered/unit/mem_cache/test_mxfp8_scale_transfer_buffers.py

@@ -0,0 +1,54 @@
+import unittest
+
+import torch
+
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPoolMXFP8
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+LAYERS = 2
+PAGE_SIZE = 128
+
+
+def _pool(interleaved: bool) -> MHATokenToKVPoolMXFP8:
+    """A pool stub carrying only what the scale accessor reads."""
+    pool = object.__new__(MHATokenToKVPoolMXFP8)
+    pool.page_size = PAGE_SIZE
+    pool.mxfp8_sf_interleaved = interleaved
+    # Interleaved keeps pages on the leading axis; flat keeps slots.
+    shape = (4, 3, 32, 4, 2) if interleaved else (4 * PAGE_SIZE, 3, 2)
+    pool.k_scale_buffer = [
+        torch.zeros(shape, dtype=torch.float8_e8m0fnu) for _ in range(LAYERS)
+    ]
+    pool.v_scale_buffer = [
+        torch.zeros(shape, dtype=torch.float8_e8m0fnu) for _ in range(LAYERS)
+    ]
+    return pool
+
+
+class TestMXFP8ScaleTransferBuffers(unittest.TestCase):
+    def test_interleaved_item_len_is_one_page_row(self):
+        pool = _pool(interleaved=True)
+
+        ptrs, lens, item_lens = pool.get_kv_scale_buf_infos()
+
+        self.assertEqual(len(ptrs), 2 * LAYERS)
+        self.assertNotIn(0, lens)
+        row = pool.k_scale_buffer[0][0].nbytes
+        self.assertEqual(item_lens[0], row)
+        self.assertEqual(lens[0] // item_lens[0], pool.k_scale_buffer[0].shape[0])
+
+    def test_flat_item_len_covers_a_whole_page(self):
+        """A flat buffer is indexed per slot, so a page's worth is page_size rows."""
+        pool = _pool(interleaved=False)
+
+        _, lens, item_lens = pool.get_kv_scale_buf_infos()
+
+        row = pool.k_scale_buffer[0][0].nbytes
+        self.assertEqual(item_lens[0], row * PAGE_SIZE)
+        self.assertEqual(lens[0] // item_lens[0], 4)
+
+
+if __name__ == "__main__":
+    unittest.main()