perf(disaggregation): cache static buffer metadata
This commit is contained in:
wxiwnd
@@ -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:
|
||||
|
||||
109
test/registered/unit/mem_cache/test_kv_pool_buf_infos.py
Normal file
109
test/registered/unit/mem_cache/test_kv_pool_buf_infos.py
Normal file
@@ -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()
|
||||
Reference in New Issue
Block a user