perf(disaggregation): cache static buffer metadata

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 = []
+        if self._contiguous_buf_infos is None:
-        for field in vars(self.mamba_cache):
+            state_tensors = []
-            # Skip intermediate buffers used only for speculative decoding
+            for field in vars(self.mamba_cache):
-            # These buffers have different size (spec_state_size + 1) and should not be transferred
+                # Skip intermediate buffers used only for speculative decoding.
-            if field in ("intermediate_ssm", "intermediate_conv_window"):
+                # These buffers have different size (spec_state_size + 1) and should not be transferred.
-                continue
+                if field in ("intermediate_ssm", "intermediate_conv_window"):
-            value = getattr(self.mamba_cache, field)
+                    continue
-            if isinstance(value, list):
+                value = getattr(self.mamba_cache, field)
-                state_tensors.extend(value)
+                if isinstance(value, list):
-            else:
+                    state_tensors.extend(value)
-                state_tensors.append(value)
+                else:
-        data_ptrs, data_lens, item_lens = [], [], []
+                    state_tensors.append(value)
 
-        for _, state_tensor in enumerate(state_tensors):
+            data_ptrs, data_lens, item_lens = [], [], []
-            data_ptrs += [
+            for state_tensor in state_tensors:
-                state_tensor[i].data_ptr() for i in range(self.num_mamba_layers)
+                data_ptrs.extend(
-            ]
+                    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 += [
+                data_lens.extend(
-                state_tensor[i][0].nbytes for i in range(self.num_mamba_layers)
+                    state_tensor[i].nbytes for i in range(self.num_mamba_layers)
-            ]
+                )
-        return data_ptrs, data_lens, item_lens
+                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]
+        if self._contiguous_buf_infos is None:
-        # layer_num x [page_num, page_size, head_num, head_dim]
+            start = self.start_layer
-        kv_data_ptrs = [
+            end = self.start_layer + self.layer_num
-            self._get_key_buffer(i).data_ptr()
+            key_buffers = [self._get_key_buffer(i) for i in range(start, end)]
-            for i in range(self.start_layer, self.start_layer + self.layer_num)
+            value_buffers = [self._get_value_buffer(i) for i in range(start, end)]
-        ] + [
+            self._contiguous_buf_infos = _build_single_pool_buf_infos(
-            self._get_value_buffer(i).data_ptr()
+                key_buffers + value_buffers, self.page_size
-            for i in range(self.start_layer, self.start_layer + self.layer_num)
+            )
-        ]
+
-        kv_data_lens = [
+        return _copy_buf_infos(self._contiguous_buf_infos)
-            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_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.
+        if self._contiguous_buf_infos is None:
-        kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in range(self.layer_num)]
+            # MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
-        kv_data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)]
+            self._contiguous_buf_infos = _build_single_pool_buf_infos(
-        kv_item_lens = [
+                self.kv_buffer, self.page_size
-            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
+        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:

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