From 73a2c117c619eaabdae41b12dac92561899f3c26 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Fri, 21 Aug 2026 13:05:06 +0800 Subject: [PATCH] Support mxfp8 KV cache in PD transfer (#35718) --- python/sglang/srt/disaggregation/base/conn.py | 5 ++ python/sglang/srt/disaggregation/decode.py | 10 ++-- .../srt/disaggregation/mooncake/conn.py | 2 + python/sglang/srt/disaggregation/prefill.py | 10 ++-- python/sglang/srt/disaggregation/utils.py | 26 +++++++++ python/sglang/srt/mem_cache/memory_pool.py | 23 ++++++-- .../sglang/srt/mem_cache/swa_memory_pool.py | 6 +++ .../test_mxfp8_scale_transfer_buffers.py | 54 +++++++++++++++++++ 8 files changed, 121 insertions(+), 15 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_mxfp8_scale_transfer_buffers.py diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index f88eda509..88c956f4b 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index 402b830c1..e2c74d376 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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 ( diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index dba26671f..ab158e8be 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/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: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 5c44bb1eb..daf27ee38 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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(), diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 96eaa1e5a..b7962a2d8 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1bbc997d9..05bfc1259 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/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( diff --git a/python/sglang/srt/mem_cache/swa_memory_pool.py b/python/sglang/srt/mem_cache/swa_memory_pool.py index 1ec9d3aa8..f52c90c01 100644 --- a/python/sglang/srt/mem_cache/swa_memory_pool.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_mxfp8_scale_transfer_buffers.py b/test/registered/unit/mem_cache/test_mxfp8_scale_transfer_buffers.py new file mode 100644 index 000000000..9ed3689c1 --- /dev/null +++ b/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()