diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 8a61142bf..86179c82f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -91,6 +91,22 @@ def get_tensor_size_bytes(t: Union[torch.Tensor, List[torch.Tensor]]): return np.prod(t.shape) * t.dtype.itemsize +def _copy_buf_infos( + buf_infos: Tuple[List[int], List[int], List[int]], +) -> Tuple[List[int], List[int], List[int]]: + return buf_infos[0].copy(), buf_infos[1].copy(), buf_infos[2].copy() + + +def _build_single_pool_buf_infos( + buffers: List[torch.Tensor], page_size: int +) -> Tuple[List[int], List[int], List[int]]: + return ( + [buf.data_ptr() for buf in buffers], + [buf.nbytes for buf in buffers], + [buf[0].nbytes * page_size for buf in buffers], + ) + + def _set_kv_buffer_impl( k: torch.Tensor, v: torch.Tensor, @@ -335,6 +351,7 @@ class MambaPool: ) self.mem_usage = self.mamba_cache.mem_usage_bytes() / GB self.num_mamba_layers = num_mamba_layers + self._contiguous_buf_infos = None def get_speculative_mamba2_params_all_layers(self) -> SpeculativeState: assert isinstance(self.mamba_cache, self.SpeculativeState) @@ -400,28 +417,34 @@ class MambaPool: Only returns conv and temporal state buffers, excluding intermediate buffers used for speculative decoding (intermediate_ssm, intermediate_conv_window). """ - state_tensors = [] - for field in vars(self.mamba_cache): - # Skip intermediate buffers used only for speculative decoding - # These buffers have different size (spec_state_size + 1) and should not be transferred - if field in ("intermediate_ssm", "intermediate_conv_window"): - continue - value = getattr(self.mamba_cache, field) - if isinstance(value, list): - state_tensors.extend(value) - else: - state_tensors.append(value) - data_ptrs, data_lens, item_lens = [], [], [] + if self._contiguous_buf_infos is None: + state_tensors = [] + for field in vars(self.mamba_cache): + # Skip intermediate buffers used only for speculative decoding. + # These buffers have different size (spec_state_size + 1) and should not be transferred. + if field in ("intermediate_ssm", "intermediate_conv_window"): + continue + value = getattr(self.mamba_cache, field) + if isinstance(value, list): + state_tensors.extend(value) + else: + state_tensors.append(value) - for _, state_tensor in enumerate(state_tensors): - data_ptrs += [ - state_tensor[i].data_ptr() for i in range(self.num_mamba_layers) - ] - data_lens += [state_tensor[i].nbytes for i in range(self.num_mamba_layers)] - item_lens += [ - state_tensor[i][0].nbytes for i in range(self.num_mamba_layers) - ] - return data_ptrs, data_lens, item_lens + data_ptrs, data_lens, item_lens = [], [], [] + for state_tensor in state_tensors: + data_ptrs.extend( + state_tensor[i].data_ptr() for i in range(self.num_mamba_layers) + ) + data_lens.extend( + state_tensor[i].nbytes for i in range(self.num_mamba_layers) + ) + item_lens.extend( + state_tensor[i][0].nbytes for i in range(self.num_mamba_layers) + ) + + self._contiguous_buf_infos = (data_ptrs, data_lens, item_lens) + + return _copy_buf_infos(self._contiguous_buf_infos) def get_state_dim_per_tensor(self): """Get the sliceable dimension size for each state tensor. @@ -784,6 +807,7 @@ class MHATokenToKVPool(KVCache): else: self._kv_copy_config = None + self._contiguous_buf_infos = None self._finalize_allocation_log(size) # for store_cache JIT kernel @@ -885,6 +909,7 @@ class MHATokenToKVPool(KVCache): def _clear_buffers(self): del self.k_buffer del self.v_buffer + self._contiguous_buf_infos = None def get_kv_size_bytes(self): assert hasattr(self, "k_buffer") @@ -899,30 +924,16 @@ class MHATokenToKVPool(KVCache): # for disagg def get_contiguous_buf_infos(self): - # layer_num x [seq_len, head_num, head_dim] - # layer_num x [page_num, page_size, head_num, head_dim] - kv_data_ptrs = [ - self._get_key_buffer(i).data_ptr() - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self._get_value_buffer(i).data_ptr() - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - kv_data_lens = [ - self._get_key_buffer(i).nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self._get_value_buffer(i).nbytes - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - kv_item_lens = [ - self._get_key_buffer(i)[0].nbytes * self.page_size - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] + [ - self._get_value_buffer(i)[0].nbytes * self.page_size - for i in range(self.start_layer, self.start_layer + self.layer_num) - ] - return kv_data_ptrs, kv_data_lens, kv_item_lens + if self._contiguous_buf_infos is None: + start = self.start_layer + end = self.start_layer + self.layer_num + key_buffers = [self._get_key_buffer(i) for i in range(start, end)] + value_buffers = [self._get_value_buffer(i) for i in range(start, end)] + self._contiguous_buf_infos = _build_single_pool_buf_infos( + key_buffers + value_buffers, self.page_size + ) + + return _copy_buf_infos(self._contiguous_buf_infos) def get_cpu_copy(self, indices): torch.cuda.synchronize() @@ -1131,6 +1142,35 @@ class MHATokenToKVPoolFP4(MHATokenToKVPool): del self.v_buffer del self.k_scale_buffer del self.v_scale_buffer + self._contiguous_buf_infos = None + + def get_contiguous_buf_infos(self): + # Keep the pre-bucket-2 behavior for FP4 MHA pools. + # _get_key_buffer/_get_value_buffer materialize dequantized temporaries, so + # this path must not memoize metadata across calls until FP4 disagg buffer + # registration is specialized around stable packed storage. + kv_data_ptrs = [ + self._get_key_buffer(i).data_ptr() + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self._get_value_buffer(i).data_ptr() + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + kv_data_lens = [ + self._get_key_buffer(i).nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self._get_value_buffer(i).nbytes + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + kv_item_lens = [ + self._get_key_buffer(i)[0].nbytes * self.page_size + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + [ + self._get_value_buffer(i)[0].nbytes * self.page_size + for i in range(self.start_layer, self.start_layer + self.layer_num) + ] + return kv_data_ptrs, kv_data_lens, kv_item_lens def _get_key_buffer(self, layer_id: int): # for internal use of referencing @@ -1460,6 +1500,7 @@ class MLATokenToKVPool(KVCache): dtype=torch.uint64, device=self.device, ) + self._contiguous_buf_infos = None if not use_nsa: # NSA will allocate indexer KV cache later and then log the total size self._finalize_allocation_log(size) @@ -1483,6 +1524,7 @@ class MLATokenToKVPool(KVCache): def _clear_buffers(self): del self.kv_buffer + self._contiguous_buf_infos = None def get_kv_size_bytes(self): assert hasattr(self, "kv_buffer") @@ -1493,13 +1535,13 @@ class MLATokenToKVPool(KVCache): # for disagg def get_contiguous_buf_infos(self): - # MLA has only one kv_buffer, so only the information of this buffer needs to be returned. - kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in range(self.layer_num)] - kv_data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)] - kv_item_lens = [ - self.kv_buffer[i][0].nbytes * self.page_size for i in range(self.layer_num) - ] - return kv_data_ptrs, kv_data_lens, kv_item_lens + if self._contiguous_buf_infos is None: + # MLA has only one kv_buffer, so only the information of this buffer needs to be returned. + self._contiguous_buf_infos = _build_single_pool_buf_infos( + self.kv_buffer, self.page_size + ) + + return _copy_buf_infos(self._contiguous_buf_infos) def get_key_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: @@ -1671,6 +1713,7 @@ class MLATokenToKVPoolFP4(MLATokenToKVPool): def _clear_buffers(self): del self.kv_buffer del self.kv_scale_buffer + self._contiguous_buf_infos = None def get_key_buffer(self, layer_id: int): if self.layer_transfer_counter is not None: diff --git a/test/registered/unit/mem_cache/test_kv_pool_buf_infos.py b/test/registered/unit/mem_cache/test_kv_pool_buf_infos.py new file mode 100644 index 000000000..8bd49117f --- /dev/null +++ b/test/registered/unit/mem_cache/test_kv_pool_buf_infos.py @@ -0,0 +1,109 @@ +"""Unit tests for static disagg KV buffer metadata caches.""" + +import unittest +from unittest.mock import MagicMock, patch + +import torch + +from sglang.srt.mem_cache.memory_pool import ( + MHATokenToKVPool, + MHATokenToKVPoolFP4, + MLATokenToKVPool, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=6, suite="stage-a-test-cpu") + + +class _BrokenBuffer: + def data_ptr(self): + raise AssertionError("buffer infos were recomputed") + + @property + def nbytes(self): + raise AssertionError("buffer infos were recomputed") + + def __getitem__(self, _key): + raise AssertionError("buffer infos were recomputed") + + +class TestKVPoolBufInfos(CustomTestCase): + def test_mha_contiguous_buf_infos_are_cached_but_return_copies(self): + pool = MHATokenToKVPool( + size=16, + page_size=1, + dtype=torch.bfloat16, + head_num=2, + head_dim=8, + layer_num=2, + device="cpu", + enable_memory_saver=False, + ) + + first = pool.get_contiguous_buf_infos() + first[0][0] = -1 + + with ( + patch.object( + pool, + "_get_key_buffer", + side_effect=AssertionError("key buffer infos were recomputed"), + ), + patch.object( + pool, + "_get_value_buffer", + side_effect=AssertionError("value buffer infos were recomputed"), + ), + ): + second = pool.get_contiguous_buf_infos() + + self.assertNotEqual(second[0][0], -1) + self.assertIsNot(first[0], second[0]) + self.assertIsNot(first[1], second[1]) + self.assertIsNot(first[2], second[2]) + + def test_mla_contiguous_buf_infos_are_cached_but_return_copies(self): + pool = MLATokenToKVPool( + size=16, + page_size=1, + dtype=torch.bfloat16, + kv_lora_rank=8, + qk_rope_head_dim=4, + layer_num=2, + device="cpu", + enable_memory_saver=False, + ) + + first = pool.get_contiguous_buf_infos() + first[0][0] = -1 + pool.kv_buffer = [_BrokenBuffer() for _ in range(pool.layer_num)] + + second = pool.get_contiguous_buf_infos() + + self.assertNotEqual(second[0][0], -1) + self.assertIsNot(first[0], second[0]) + self.assertIsNot(first[1], second[1]) + self.assertIsNot(first[2], second[2]) + + def test_mha_fp4_contiguous_buf_infos_do_not_cache_temporary_views(self): + pool = object.__new__(MHATokenToKVPoolFP4) + pool.start_layer = 0 + pool.layer_num = 1 + pool.page_size = 1 + + key_buf = torch.zeros((2, 2), dtype=torch.bfloat16) + value_buf = torch.zeros((2, 2), dtype=torch.bfloat16) + pool._get_key_buffer = MagicMock(return_value=key_buf) + pool._get_value_buffer = MagicMock(return_value=value_buf) + + first = pool.get_contiguous_buf_infos() + second = pool.get_contiguous_buf_infos() + + self.assertEqual(pool._get_key_buffer.call_count, 6) + self.assertEqual(pool._get_value_buffer.call_count, 6) + self.assertEqual(first, second) + + +if __name__ == "__main__": + unittest.main()