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 = []
-        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: