minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

feat(mem_cache): unified memory pool for hybrid Mamba / SWA models (#29678)

Browse Source 
Co-authored-by: lch1475369 <lch1475369@gmail.com>
This commit is contained in:
Cheng Wan
2026-07-01 13:21:59 -07:00
committed by GitHub
co-authored by lch1475369
parent 3adfd0f34b
commit 4a8e76805c
28 changed files with 8378 additions and 551 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_full_loc_fast_path.py
+206
View File
@@ -0,0 +1,206 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Routing tests for the composite write paths (`UnifiedSWAKVPool`,
`HybridLinearKVPool`).
All write-location info travels in the attention metadata (`KVWriteLoc`); the
pools hold none and never translate — the write loc reaching `set_kv_buffer` is
always PHYSICAL. Two routing contracts are pinned here:
1. Full-attention. The full-physical loc is carried in `KVWriteLoc.full_loc`
(from `ForwardBatch.out_cache_loc_full_physical`) and written directly.
`UnifiedSWAKVPool` asserts it's present (the unified memory pool always precomputes
it); `HybridLinearKVPool` falls back to `loc` for a static (non-shared) pool,
where `loc` is itself already physical.
2. SWA. The swa-physical loc rides the backend `swa_out_cache_loc` rail
(`KVWriteLoc.swa_loc`) and is written directly.
Pure dispatch tests: the inner sub-pools are recording stubs, so no GPU / real
buffers are needed (CPU CI).
python -m pytest test/registered/unit/mem_cache/test_full_loc_fast_path.py -v
"""
import types
import unittest
import torch
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _loc_info(virtual_loc, swa_phys=None, full_phys=None):
from sglang.srt.mem_cache.memory_pool import KVWriteLoc
return KVWriteLoc(virtual_loc, swa_phys, full_phys)
class _RecordingPool:
"""Stub sub-pool that records the `loc` and kwargs passed to `set_kv_buffer`."""
def __init__(self):
self.calls = []
def set_kv_buffer(self, layer, loc, cache_k, cache_v, *args, **kwargs):
self.calls.append((loc, kwargs))
class TestUnifiedSWARouting(unittest.TestCase):
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write the full-physical
`full_loc`; SWA layers write the swa-physical `swa_loc`. Both come from the
write metadata; the pool never translates."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
# Bypass the heavy __init__; set only the attributes set_kv_buffer reads.
pool = object.__new__(UnifiedSWAKVPool)
pool.full_kv_pool = _RecordingPool()
pool.swa_kv_pool = _RecordingPool()
# layer 0 -> full attention; layer 1 -> SWA. (pool_layer_id, is_swa)
pool.layers_mapping = {0: (0, False), 1: (0, True)}
return pool
def test_full_layer_writes_full_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
full_phys = torch.tensor([3, 4, 5], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0) # full layer
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, swa_phys, full_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
# Forward the full-physical tensor from the write metadata, NOT the
# virtual loc. No `already_physical` — the pool only ever gets physical.
self.assertIs(forwarded, full_phys)
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_full_layer_requires_full_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
# No full_loc precomputed -> fail loud (the unified memory pool must precompute
# out_cache_loc_full_physical) rather than write a virtual loc as physical.
with self.assertRaises(AssertionError):
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, swa_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
def test_swa_layer_writes_swa_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
swa_phys = torch.tensor([1, 2, 0], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=1) # SWA layer
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, swa_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.swa_kv_pool.calls), 1)
forwarded, kwargs = pool.swa_kv_pool.calls[0]
# SWA write rides the backend rail: forward the swa-physical loc directly.
self.assertIs(forwarded, swa_phys)
self.assertNotIn("already_physical", kwargs)
# Full pool untouched for an SWA layer.
self.assertEqual(len(pool.full_kv_pool.calls), 0)
def test_swa_layer_requires_swa_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=1) # SWA layer
# No swa_loc bundled -> the rail contract is violated; must assert
# rather than silently writing wrong (un-translated) locations.
with self.assertRaises(AssertionError):
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, None),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
class TestHybridLinearFullLocRouting(unittest.TestCase):
"""`HybridLinearKVPool.set_kv_buffer` (non-MLA) writes the full-physical
`full_loc` from the write metadata when present (unified memory pool), else the
already-physical `loc` (static pool). No translate, no `already_physical`."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
pool = object.__new__(HybridLinearKVPool)
pool.full_kv_pool = _RecordingPool()
pool.use_mla = False
pool.full_attention_layer_id_mapping = {0: 0}
return pool
def test_writes_full_loc_from_write_loc(self):
pool = self._make_bare_pool()
virtual_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
full_phys = torch.tensor([2, 3, 4], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(virtual_loc, full_phys=full_phys),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, full_phys)
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_falls_back_to_loc_when_absent(self):
# Static (non-shared) pool: no full_loc bundled; `loc` is already
# physical, so write it directly.
pool = self._make_bare_pool()
phys_loc = torch.tensor([7, 8, 9], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.set_kv_buffer(
layer,
_loc_info(phys_loc),
torch.zeros(3, 4, 8),
torch.zeros(3, 4, 8),
)
self.assertEqual(len(pool.full_kv_pool.calls), 1)
forwarded, kwargs = pool.full_kv_pool.calls[0]
self.assertIs(forwarded, phys_loc)
self.assertNotIn("already_physical", kwargs)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_layout_compat.py
+355
View File
@@ -0,0 +1,355 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Unit tests for the page-major layer-major byte layout.
Verifies that:
1. The new 4-D ``_build_mha_views`` output exposes correct byte addresses
for each (layer, page, tok_in_page, head, dim) — under both the
degenerate ``page_size=1`` case (byte-identical to the old per-token
envelope) and the new ``page_size>1`` layer-major case.
2. ``MHASubPoolSpec.layer_k_offset_in_page`` /
``layer_v_offset_in_page`` math matches the layout intent.
3. ``set_kv_buffer`` round-trips correctly for both page sizes.
4. Compaction (``move_kv_cache_native``) moves the right bytes for both
page sizes via the 4-D advanced indexing path.
CPU-only — no GPU / Triton needed.
python -m pytest test/registered/unit/mem_cache/test_layout_compat.py -v
"""
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
import unittest
import torch
from sglang.srt.mem_cache.memory_pool import move_kv_cache_native
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
UnifiedKVPool,
)
_DEV = "cpu"
def _make_mha_spec(name, grow, layer_num=2, head_num=2, head_dim=4):
return MHASubPoolSpec(
name=name,
layer_num=layer_num,
head_num=head_num,
head_dim=head_dim,
store_dtype=torch.float16,
grow_direction=grow,
)
def _make_mamba_spec(name, grow, layer_num=2):
return MambaSubPoolSpec(
name=name,
layer_num=layer_num,
conv_state_shapes=((4, 3),),
conv_dtype=torch.float32,
temporal_state_shape=(2, 2, 2),
temporal_dtype=torch.float32,
grow_direction=grow,
)
class TestMHASpecLayerOffsets(unittest.TestCase):
"""Verify ``layer_k_offset_in_page`` / ``layer_v_offset_in_page`` math."""
def test_offsets_at_page_size_1_match_envelope(self):
spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
# At ps=1, layer-major within a 1-token page IS envelope-per-token.
# Layer L's K offset = L * (k_row + v_row); V offset = +k_row.
k_row = spec.k_row_bytes()
v_row = spec.v_row_bytes()
for L in range(spec.layer_num):
self.assertEqual(
spec.layer_k_offset_in_page(L, page_size=1),
L * (k_row + v_row),
)
self.assertEqual(
spec.layer_v_offset_in_page(L, page_size=1),
L * (k_row + v_row) + k_row,
)
def test_offsets_at_page_size_gt_1(self):
spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
ps = 8
k_row = spec.k_row_bytes()
v_row = spec.v_row_bytes()
# Layer L's K block within the page starts at L * ps * (k_row+v_row).
# V block starts at +ps * k_row.
for L in range(spec.layer_num):
self.assertEqual(
spec.layer_k_offset_in_page(L, page_size=ps),
L * ps * (k_row + v_row),
)
self.assertEqual(
spec.layer_v_offset_in_page(L, page_size=ps),
L * ps * (k_row + v_row) + ps * k_row,
)
def test_page_bytes(self):
spec = _make_mha_spec("full", "up", layer_num=3, head_num=2, head_dim=4)
# page_bytes = page_size * entry_bytes (preserved invariant)
for ps in [1, 8, 64, 256]:
self.assertEqual(spec.page_bytes(ps), ps * spec.entry_bytes())
class TestBuildMHAViews(unittest.TestCase):
"""Verify the 4-D view shape + strides at both page sizes."""
def _build(self, page_size, layer_num=3, head_num=2, head_dim=4, n_full_slots=64):
full = _make_mha_spec(
"full", "up", layer_num=layer_num, head_num=head_num, head_dim=head_dim
)
swa = _make_mha_spec(
"swa", "down", layer_num=2, head_num=head_num, head_dim=head_dim
)
# Pad to ensure max_slots % page_size == 0 in both sub-pools.
# entry_bytes is fixed per spec; size accordingly.
total = full.entry_bytes() * n_full_slots + swa.entry_bytes() * n_full_slots
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[full, swa],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
return pool, full
def test_view_shape_is_4d(self):
for ps in [1, 8]:
pool, spec = self._build(page_size=ps)
k_views, v_views = pool.mha_views_for("full")
self.assertEqual(len(k_views), spec.layer_num)
max_slots = pool.max_slots("full")
for L in range(spec.layer_num):
self.assertEqual(k_views[L].ndim, 4)
self.assertEqual(
tuple(k_views[L].shape),
(max_slots // ps, ps, spec.head_num, spec.head_dim),
)
self.assertEqual(
tuple(v_views[L].shape),
(max_slots // ps, ps, spec.head_num, spec.v_head_dim),
)
def test_strides_at_page_size_1_match_envelope(self):
"""At ps=1, the 4-D view's stride[0] equals what today's 3-D view's
stride[0] would have been (= entry_bytes / itemsize)."""
pool, spec = self._build(page_size=1, layer_num=4, head_num=3, head_dim=8)
k_views, _ = pool.mha_views_for("full")
itemsize = spec.store_dtype.itemsize
for L in range(spec.layer_num):
# stride[0] = page_bytes/itemsize = entry_bytes/itemsize at ps=1
self.assertEqual(k_views[L].stride(0), spec.entry_bytes() // itemsize)
# stride[1] = k_row/itemsize (within-page token stride)
self.assertEqual(k_views[L].stride(1), spec.k_row_bytes() // itemsize)
# stride[2] = head_dim (head stride)
self.assertEqual(k_views[L].stride(2), spec.head_dim)
# stride[3] = 1 (innermost)
self.assertEqual(k_views[L].stride(3), 1)
def test_strides_at_page_size_gt_1(self):
pool, spec = self._build(page_size=8, layer_num=4, head_num=3, head_dim=8)
k_views, _ = pool.mha_views_for("full")
itemsize = spec.store_dtype.itemsize
for L in range(spec.layer_num):
# page_bytes = 8 * 4 * (k_row + v_row); stride[0] = that / itemsize
self.assertEqual(k_views[L].stride(0), spec.page_bytes(8) // itemsize)
# token stride within layer L's K block = k_row/itemsize
self.assertEqual(k_views[L].stride(1), spec.k_row_bytes() // itemsize)
self.assertEqual(k_views[L].stride(2), spec.head_dim)
self.assertEqual(k_views[L].stride(3), 1)
def test_distinct_layers_dont_alias_at_page_size_gt_1(self):
"""Writes to layer 0 must not affect layer 1's K/V values (under
layer-major within-page layout)."""
pool, spec = self._build(page_size=8, layer_num=3, head_num=2, head_dim=4)
k_views, v_views = pool.mha_views_for("full")
# Set page 0, token 3, layer 0 K to a distinct pattern.
target_val = 0.5
k_views[0][0, 3] = target_val
# Layer 1 K at the same (page, tok) should remain at default (0.0).
self.assertFalse(torch.all(k_views[1][0, 3] == target_val))
self.assertTrue(torch.all(k_views[1][0, 3] == 0.0))
# And layer 0 V at the same (page, tok) should remain at default.
self.assertFalse(torch.all(v_views[0][0, 3] == target_val))
self.assertTrue(torch.all(v_views[0][0, 3] == 0.0))
def test_distinct_pages_dont_alias_at_page_size_gt_1(self):
"""Writes to one page must not affect another page."""
pool, spec = self._build(page_size=8, layer_num=3, head_num=2, head_dim=4)
k_views, _ = pool.mha_views_for("full")
# Set page 0, token 3, layer 0 K to a distinct pattern.
k_views[0][0, 3] = 1.25
# Page 1, token 3, layer 0 K should remain at default.
self.assertTrue(torch.all(k_views[0][1, 3] == 0.0))
class TestMoveKVCacheNative4D(unittest.TestCase):
"""Verify ``move_kv_cache_native`` handles 4-D buffers at both
page_size=1 (degenerate envelope) and page_size>1 (layer-major)."""
def _build_buffer(
self, page_size, layer_num=2, head_num=2, head_dim=4, n_full_slots=64
):
full = _make_mha_spec(
"full", "up", layer_num=layer_num, head_num=head_num, head_dim=head_dim
)
swa = _make_mha_spec(
"swa", "down", layer_num=2, head_num=head_num, head_dim=head_dim
)
total = full.entry_bytes() * n_full_slots + swa.entry_bytes() * n_full_slots
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[full, swa],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
return pool
def test_move_kv_cache_page_size_1(self):
pool = self._build_buffer(page_size=1, layer_num=2, head_num=2, head_dim=4)
k_views, v_views = pool.mha_views_for("full")
# Write distinct markers at source slots 5, 6.
for L in range(2):
k_views[L][5, 0] = float(L + 1)
v_views[L][5, 0] = -float(L + 1)
k_views[L][6, 0] = float(L + 10)
v_views[L][6, 0] = -float(L + 10)
# Move 5 -> 8 and 6 -> 9.
move_kv_cache_native(
k_views,
v_views,
tgt_loc=torch.tensor([8, 9], dtype=torch.int64),
src_loc=torch.tensor([5, 6], dtype=torch.int64),
page_size=1,
)
for L in range(2):
self.assertTrue(torch.all(k_views[L][8, 0] == float(L + 1)))
self.assertTrue(torch.all(v_views[L][8, 0] == -float(L + 1)))
self.assertTrue(torch.all(k_views[L][9, 0] == float(L + 10)))
self.assertTrue(torch.all(v_views[L][9, 0] == -float(L + 10)))
def test_move_kv_cache_page_size_gt_1(self):
ps = 8
pool = self._build_buffer(page_size=ps, layer_num=2, head_num=2, head_dim=4)
k_views, v_views = pool.mha_views_for("full")
# Write markers at token ids 5 and 14 (different pages).
for L in range(2):
# token 5 = (page 0, tok 5)
k_views[L][0, 5] = float(L + 1)
v_views[L][0, 5] = -float(L + 1)
# token 14 = (page 1, tok 6)
k_views[L][1, 6] = float(L + 10)
v_views[L][1, 6] = -float(L + 10)
# Move token 5 -> token 23 (page 2, tok 7) and 14 -> 31 (page 3, tok 7).
move_kv_cache_native(
k_views,
v_views,
tgt_loc=torch.tensor([23, 31], dtype=torch.int64),
src_loc=torch.tensor([5, 14], dtype=torch.int64),
page_size=ps,
)
for L in range(2):
# 23 = page 2, tok 7
self.assertTrue(torch.all(k_views[L][2, 7] == float(L + 1)))
self.assertTrue(torch.all(v_views[L][2, 7] == -float(L + 1)))
# 31 = page 3, tok 7
self.assertTrue(torch.all(k_views[L][3, 7] == float(L + 10)))
self.assertTrue(torch.all(v_views[L][3, 7] == -float(L + 10)))
def test_move_kv_cache_3d_legacy_path_unchanged(self):
"""move_kv_cache_native(3-D, page_size=1) must take the legacy
else-branch and be byte-identical to today."""
k = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
v = [torch.zeros((32, 2, 4), dtype=torch.float16) for _ in range(2)]
for L in range(2):
k[L][5] = float(L + 1)
v[L][5] = -float(L + 1)
move_kv_cache_native(
k,
v,
tgt_loc=torch.tensor([7], dtype=torch.int64),
src_loc=torch.tensor([5], dtype=torch.int64),
page_size=1,
)
for L in range(2):
self.assertTrue(torch.all(k[L][7] == float(L + 1)))
self.assertTrue(torch.all(v[L][7] == -float(L + 1)))
class TestByteIdentityAtPageSize1(unittest.TestCase):
"""Verify that at page_size=1 the new 4-D view describes the SAME
physical bytes as the old 3-D view would have. The view
semantics differ (4-D vs 3-D shape) but the underlying byte layout is
identical — confirmed by manually computing expected byte offsets and
matching them against the 4-D view's strides + storage_offset.
"""
def test_byte_addresses_match_envelope(self):
spec = _make_mha_spec("full", "up", layer_num=4, head_num=2, head_dim=4)
ps = 1
# Build pool.
total = spec.entry_bytes() * 64 + spec.entry_bytes() * 32
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[
spec,
_make_mha_spec("swa", "down", layer_num=2),
],
device=_DEV,
enable_memory_saver=False,
page_size=ps,
)
k_views, v_views = pool.mha_views_for("full")
# For each (layer, slot), compute the expected byte address under
# the envelope layout and verify the 4-D view's data_ptr +
# advanced indexing agrees.
max_slots = pool.max_slots("full")
itemsize = spec.store_dtype.itemsize
base_addr = pool._raw.data_ptr()
for L in range(spec.layer_num):
for s in range(0, max_slots, max(1, max_slots // 4)):
# Envelope: bytes for slot s, layer L's K start at:
# s * entry_bytes + L * (k_row + v_row)
expected_k_byte_offset = s * spec.entry_bytes() + L * (
spec.k_row_bytes() + spec.v_row_bytes()
)
# 4-D view: k_views[L][page=s, tok=0, head=0, dim=0]
# storage_offset of the element [s, 0, 0, 0]:
view_offset_elems = (
k_views[L].storage_offset()
+ s * k_views[L].stride(0)
+ 0 * k_views[L].stride(1)
+ 0 * k_views[L].stride(2)
+ 0 * k_views[L].stride(3)
)
view_byte_offset = view_offset_elems * itemsize
# 4-D view sits over `_raw.view(spec.store_dtype)`, which
# has data_ptr == _raw.data_ptr() (same backing storage).
self.assertEqual(view_byte_offset, expected_k_byte_offset)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+2729
View File
File diff suppressed because it is too large Load Diff
test/registered/unit/mem_cache/test_store_cache_4d.py
+23 -115
View File
@@ -24,10 +24,10 @@ import torch
from sglang.test.ci.ci_register import register_cuda_ci
_HAS_CUDA = torch.cuda.is_available()
# The set_kv_buffer integration test needs SharedMHATokenToKVPool, which only
# exists once the shared-memory-pool feature lands; skip it where absent.
# The set_kv_buffer integration test needs UnifiedMHATokenToKVPool, which only
# exists once the shared-KV-pool feature lands; skip it where absent.
_HAS_SHARED_POOL = (
importlib.util.find_spec("sglang.srt.mem_cache.shared_memory_pool") is not None
importlib.util.find_spec("sglang.srt.mem_cache.unified_memory_pool") is not None
)
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
@@ -72,7 +72,7 @@ class TestStoreCache4D(unittest.TestCase):
seed: int = 0xC0FFEE,
):
torch.manual_seed(seed)
# The shared pool's views are 4-D `(num_pages, page_size, head_num,
# The unified memory pool's views are 4-D `(num_pages, page_size, head_num,
# head_dim)` with the trailing two dims contiguous. We allocate two
# independent contiguous buffers (one for the kernel-under-test,
# one as the legacy-path target) so we can compare them.
@@ -154,7 +154,7 @@ class TestStoreCache4D(unittest.TestCase):
def test_store_cache_4d_ps1_byte_identical(self):
"""At page_size=1 the kernel constexpr-folds to the slot-major
envelope view. Output must be byte-identical to advanced indexing.
This protects the Stage 1/2/3 green eval matrix from regression."""
This protects against byte-layout regression."""
self._check_parity(
num_pages=64,
page_size=1,
@@ -228,8 +228,7 @@ class TestStoreCache4D(unittest.TestCase):
def test_store_cache_4d_dtype_fp8_e5m2(self):
"""fp8_e5m2 is used for KV-cache quantization. Caller is responsible
for the cast (Phase 1); the kernel sees same-dtype source and
destination."""
for the cast; the kernel sees same-dtype source and destination."""
self._check_parity(
num_pages=16,
page_size=64,
@@ -317,31 +316,25 @@ class TestStoreCache4DAssertions(unittest.TestCase):
@unittest.skipUnless(
_HAS_CUDA and _HAS_SHARED_POOL,
"Triton kernels require CUDA; SharedMHATokenToKVPool required",
"Triton kernels require CUDA; UnifiedMHATokenToKVPool required",
)
class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
"""Integration parity test — exercises the kernel through the FULL
``SharedMHATokenToKVPool.set_kv_buffer`` path, including the
``_external_allocator`` v2p translation and the dtype cast. Confirms the
production code path produces bit-identical output to a PyTorch
advanced-indexing reference write.
``UnifiedMHATokenToKVPool.set_kv_buffer`` path (the direct PHYSICAL write +
the dtype cast; the pool no longer translates). Confirms it produces
bit-identical output to a PyTorch advanced-indexing reference write.
"""
def _build_pool_and_stub_alloc(self, page_size: int, v2p=None):
"""Build a small SharedMHATokenToKVPool wired to a stub allocator.
By default `virtual_to_physical` is identity (the kernel-vs-legacy
parity tests don't exercise virtual-id semantics). Pass an explicit
`v2p` tensor (sized `max_slots + 1`) to exercise a NON-identity
translation — used by the `set_full_loc` fast-path parity test, which
needs virtual != physical so the precomputed-physical fast path is
meaningfully different from the per-call gather."""
def _build_pool(self, page_size: int):
"""Build a small UnifiedMHATokenToKVPool. The pool writes PHYSICAL locs
directly (no allocator / v2p translate), so `set_kv_buffer` receives the
already-physical write location."""
import torch as _t
from sglang.srt.mem_cache.shared_memory_pool import (
from sglang.srt.mem_cache.unified_memory_pool import (
MHASubPoolSpec,
SharedMemoryPool,
SharedMHATokenToKVPool,
UnifiedKVPool,
UnifiedMHATokenToKVPool,
)
spec = MHASubPoolSpec(
@@ -362,15 +355,15 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
store_dtype=_t.bfloat16,
grow_direction="down",
)
pool = SharedMemoryPool(
pool = UnifiedKVPool(
total_bytes=total + peer.entry_bytes() * 16,
sub_pool_specs=[spec, peer],
device="cuda",
enable_memory_saver=False,
page_size=page_size,
)
kv_pool = SharedMHATokenToKVPool(
shared_buffer=pool,
kv_pool = UnifiedMHATokenToKVPool(
unified_buffer=pool,
sub_pool_name="full",
page_size=page_size,
start_layer=0,
@@ -378,21 +371,12 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
enable_alt_stream=False,
)
# Stub allocator with an identity (default) or caller-supplied v2p.
max_slots = pool.max_slots("full")
if v2p is None:
v2p = _t.arange(max_slots + 1, dtype=_t.int64, device="cuda")
class _StubAllocator:
virtual_to_physical = v2p
kv_pool.attach_allocator(_StubAllocator())
return kv_pool
def _run_set_kv_buffer_and_compare(self, page_size: int):
import torch as _t
kv_pool = self._build_pool_and_stub_alloc(page_size)
kv_pool = self._build_pool(page_size)
# A fake `layer` object with the minimum interface
# `set_kv_buffer` reads: `.layer_id`.
@@ -415,9 +399,8 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
k_kernel = kv_pool.k_buffer[0].clone()
v_kernel = kv_pool.v_buffer[0].clone()
# Reference: PyTorch advanced-indexing into a fresh view. The stub
# allocator's v2p is identity, so physical loc == virtual loc and no
# dtype cast happens (store_dtype == dtype), making this the exact
# Reference: PyTorch advanced-indexing into a fresh view at the same
# (physical) loc, with no dtype cast (store_dtype == dtype) — the exact
# write the kernel performs.
kv_pool.k_buffer[0].zero_()
kv_pool.v_buffer[0].zero_()
@@ -449,81 +432,6 @@ class TestStoreCache4DThroughSetKVBuffer(unittest.TestCase):
def test_integration_ps64(self):
self._run_set_kv_buffer_and_compare(page_size=64)
def _run_full_loc_fast_path_parity(self, page_size: int):
"""Stage 3.5 fast-path byte-identity: writing through the precomputed
full-physical loc (`set_loc` fast path) must produce a byte-identical
KV buffer to writing the virtual loc and letting `set_kv_buffer`
translate per call. Uses a NON-identity v2p so the two paths are
genuinely different code (fast path skips the gather)."""
import torch as _t
# Non-identity v2p: reverse-map the physical slot space so virtual i
# lands on a different physical slot. Keep slot 0 -> 0 (padding sink).
# Build the pool once to learn max_slots, then rebuild with the v2p.
probe = self._build_pool_and_stub_alloc(page_size)
max_slots = probe.k_buffer[0].shape[0] * page_size
v2p = _t.arange(max_slots + 1, dtype=_t.int64, device="cuda")
# Shuffle the interior [1, max_slots) so virtual != physical, leave
# 0 (sink) and the trailing sentinel (max_slots -> itself) alone.
interior = _t.randperm(max_slots - 1, device="cuda") + 1
v2p[1:max_slots] = interior
kv_pool = self._build_pool_and_stub_alloc(page_size, v2p=v2p)
class _FakeLayer:
layer_id = 0
layer = _FakeLayer()
head_num, head_dim = 4, 64
N = 16
num_pages = kv_pool.k_buffer[0].shape[0]
total = num_pages * page_size
# Draw virtual ids from [1, total) (avoid the padding sink at 0).
loc = (_t.randperm(total - 1, device="cuda")[:N] + 1).to(_t.int64)
cache_k = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
cache_v = _t.randn((N, head_num, head_dim), dtype=_t.bfloat16, device="cuda")
# SLOW path: no precompute pinned -> per-call v2p gather inside
# set_kv_buffer translates virtual -> physical.
kv_pool.set_loc(None)
kv_pool.set_kv_buffer(layer, loc, cache_k.clone(), cache_v.clone())
k_slow = kv_pool.k_buffer[0].clone()
v_slow = kv_pool.v_buffer[0].clone()
# FAST path: precompute the full-physical loc exactly as
# `set_kv_buffer`'s page math would, pin it via set_loc, and pass
# it as `loc` so the data-ptr fast path fires (no gather).
if page_size == 1:
phys = _t.clamp_min(v2p[loc], 0)
else:
virt_pages = loc // page_size
offsets = loc % page_size
phys = _t.clamp_min(v2p[virt_pages] * page_size + offsets, 0)
kv_pool.k_buffer[0].zero_()
kv_pool.v_buffer[0].zero_()
kv_pool.set_loc(phys)
try:
kv_pool.set_kv_buffer(layer, phys, cache_k.clone(), cache_v.clone())
k_fast = kv_pool.k_buffer[0].clone()
v_fast = kv_pool.v_buffer[0].clone()
finally:
kv_pool.set_loc(None)
self.assertTrue(
_t.equal(k_fast, k_slow),
f"K mismatch: full_loc fast path != per-call translate at ps={page_size}",
)
self.assertTrue(
_t.equal(v_fast, v_slow),
f"V mismatch: full_loc fast path != per-call translate at ps={page_size}",
)
def test_full_loc_fast_path_parity_ps1(self):
self._run_full_loc_fast_path_parity(page_size=1)
def test_full_loc_fast_path_parity_ps64(self):
self._run_full_loc_fast_path_parity(page_size=64)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_mamba_views.py
+291
View File
@@ -0,0 +1,291 @@
# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Round-trip correctness of ``UnifiedKVPool._build_mamba_views`` — the
envelope-strided conv/temporal (SSM) state views that back ``UnifiedMambaPool``.
This isolates the unified-memory-pool Mamba STATE layout from the full model. It guards
against a class of correctness defect where Falcon-H1 greedy decode is garbled
under the unified memory pool, isolated to the Mamba conv/temporal state path: a
stride/offset/alignment bug in the view construction (analogous to the fixed
`_extract_kv_strides` MHA bug).
Within one slot's envelope the bytes are
``[conv[0]·L0 | conv[0]·L1 | ... | conv[1]·L0 | ... | temporal·L0 | ...]`` and
across slots the layout is envelope (slot stride == entry_bytes). Each returned
view is ``(num_layers, max_slots, *inner_shape)``. The conv dtype (bf16, 2 B)
and temporal dtype (fp32, 4 B) DIFFER, so the temporal view's byte offset must
be a multiple of the temporal itemsize — an alignment hazard that
``_build_mamba_views`` now asserts.
These tests prove the views:
- round-trip every (tensor, layer, slot) element with the Falcon-like
bf16-conv / fp32-temporal dtype mix (catches stride/offset/alignment bugs);
- do NOT alias each other (conv[i] vs conv[j] vs temporal) or across
layers/slots (catches envelope-overlap);
- match a contiguous ``(num_layers, max_slots, *inner)`` reference exactly
(the shape `MambaPool.State.conv[i]` / `.temporal` expose);
- reject a deliberately mis-aligned spec via the alignment assert.
Skipped on CPU — these views back GPU kernels and we mirror the GPU path.
python -m pytest test/registered/unit/mem_cache/test_shared_mamba_views.py -v
"""
import unittest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
_HAS_CUDA = torch.cuda.is_available()
_DEV = "cuda" if _HAS_CUDA else "cpu"
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
def _make_pool(
*,
mamba_layer_num,
conv_state_shapes,
conv_dtype,
temporal_state_shape,
temporal_dtype,
want_slots=8,
device=_DEV,
):
"""Build a minimal 2-sub-pool ``UnifiedKVPool`` (a small MHA grow-up peer
+ the Mamba grow-down pool under test) sized to hold >= ``want_slots`` Mamba
slots, and return ``(pool, mamba_spec)``."""
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
UnifiedKVPool,
)
mamba_spec = MambaSubPoolSpec(
name="mamba",
layer_num=mamba_layer_num,
grow_direction="down",
conv_state_shapes=tuple(tuple(s) for s in conv_state_shapes),
conv_dtype=conv_dtype,
temporal_state_shape=tuple(temporal_state_shape),
temporal_dtype=temporal_dtype,
)
# Tiny full-attention peer (required: exactly one grow-up + one grow-down).
full_spec = MHASubPoolSpec(
name="full",
layer_num=1,
head_num=1,
head_dim=8,
store_dtype=torch.bfloat16,
grow_direction="up",
)
entry_mamba = mamba_spec.entry_bytes()
entry_full = full_spec.entry_bytes()
entry_max = max(entry_mamba, entry_full)
# Need max_slots("mamba") = total // entry_mamba >= want_slots, and total
# large enough that BOTH pools clear their min_slot_index. Add generous
# headroom, then round up to a multiple of 8 (covers bf16/fp32 .view()).
total_bytes = want_slots * entry_mamba + 8 * entry_max
total_bytes = ((total_bytes + 7) // 8) * 8
pool = UnifiedKVPool(
total_bytes=total_bytes,
sub_pool_specs=[full_spec, mamba_spec],
device=device,
enable_memory_saver=False,
)
return pool, mamba_spec
@unittest.skipUnless(_HAS_CUDA, "shared Mamba views back GPU kernels")
class TestUnifiedMambaViews(unittest.TestCase):
# Falcon-H1-like dims: even conv_dim, bf16 conv, fp32 temporal, several
# layers. (Mamba2 conv state is (conv_dim, kernel-1); temporal/SSM state is
# (nheads, head_dim, ssm_state_size).)
FALCON_KW = dict(
mamba_layer_num=5, # odd, to stress the temporal-offset alignment
conv_state_shapes=[(48, 3)], # conv_dim=48, kernel-1=3
conv_dtype=torch.bfloat16,
temporal_state_shape=(6, 8, 16), # nheads, head_dim, ssm_state
temporal_dtype=torch.float32,
)
def _fill_and_roundtrip(self, pool, mamba_spec):
"""Write a distinct random tensor to each conv view + the temporal view
(in their own dtypes), then read all back and assert exact equality.
Writing ALL views first and reading ALL after means any envelope overlap
(conv[i]/conv[j]/temporal aliasing) corrupts an earlier write → mismatch.
"""
conv_views, temporal_view = pool.mamba_views_for("mamba")
torch.manual_seed(0)
refs = []
for v in conv_views:
r = torch.randn(v.shape, device=v.device).to(v.dtype)
v.copy_(r)
refs.append(r)
rt = torch.randn(temporal_view.shape, device=temporal_view.device).to(
temporal_view.dtype
)
temporal_view.copy_(rt)
refs.append(rt)
# Read back AFTER all writes.
for i, v in enumerate(conv_views):
self.assertTrue(
torch.equal(v, refs[i]),
f"conv view[{i}] round-trip mismatch (stride/offset/overlap "
f"bug); shape={tuple(v.shape)} stride={v.stride()}",
)
self.assertTrue(
torch.equal(temporal_view, refs[-1]),
f"temporal view round-trip mismatch; shape={tuple(temporal_view.shape)} "
f"stride={temporal_view.stride()}",
)
def test_roundtrip_falcon_like(self):
pool, spec = _make_pool(**self.FALCON_KW)
self._fill_and_roundtrip(pool, spec)
def test_roundtrip_single_layer_single_slot_edges(self):
# 1 layer, multiple conv tensors, same-dtype conv/temporal.
pool, spec = _make_pool(
mamba_layer_num=1,
conv_state_shapes=[(16, 3), (8, 3)],
conv_dtype=torch.float32,
temporal_state_shape=(4, 8, 16),
temporal_dtype=torch.float32,
want_slots=4,
)
self._fill_and_roundtrip(pool, spec)
def test_roundtrip_multi_conv_tensors(self):
# Two conv tensors + bf16/fp32 mix — exercises the per-conv-tensor offset
# accumulation in _build_mamba_views.
pool, spec = _make_pool(
mamba_layer_num=3,
conv_state_shapes=[(32, 3), (16, 3)],
conv_dtype=torch.bfloat16,
temporal_state_shape=(8, 8, 16),
temporal_dtype=torch.float32,
want_slots=6,
)
self._fill_and_roundtrip(pool, spec)
def test_no_cross_region_overlap(self):
"""Zero buffer; write a sentinel to ONE view; every OTHER view must read
all-zero. Pinpoints conv[i]/conv[j]/temporal aliasing if present."""
pool, spec = _make_pool(**self.FALCON_KW)
conv_views, temporal_view = pool.mamba_views_for("mamba")
views = list(conv_views) + [temporal_view]
names = [f"conv[{i}]" for i in range(len(conv_views))] + ["temporal"]
for target in range(len(views)):
pool._raw.zero_()
views[target].fill_(7.0)
for other in range(len(views)):
if other == target:
self.assertTrue(
bool((views[other] == 7.0).all().item()),
f"write to {names[target]} did not fully land",
)
continue
self.assertTrue(
bool((views[other] == 0).all().item()),
f"writing {names[target]} CORRUPTED {names[other]} "
f"(envelope regions overlap)",
)
def test_per_layer_per_slot_addressing(self):
"""Distinct value per (layer, slot) on the temporal view; verify exact
addressing (no layer/slot aliasing). Uses small integers exactly
representable in the view dtype.
NB: ``temporal_view`` is a non-contiguous strided view, so we must NOT
``.reshape()`` it (that would COPY, breaking the alias) — we
broadcast-assign into the view in place and read back via basic
indexing (which keeps the view)."""
pool, spec = _make_pool(**self.FALCON_KW)
_, temporal_view = pool.mamba_views_for("mamba")
N, S = temporal_view.shape[0], temporal_view.shape[1]
inner_ndim = temporal_view.dim() - 2
# value = layer*S + slot (< N*S, small → exact in fp32)
base = (
torch.arange(N, device=temporal_view.device)[:, None] * S
+ torch.arange(S, device=temporal_view.device)[None, :]
).to(temporal_view.dtype)
# Broadcast (N, S) over the inner dims, in place into the strided view.
temporal_view[:] = base.view(N, S, *([1] * inner_ndim))
# Read back the first inner element of every (layer, slot) via basic
# indexing (stays a view).
readback = temporal_view[(slice(None), slice(None)) + (0,) * inner_ndim]
self.assertTrue(
torch.equal(readback, base),
"temporal (layer, slot) addressing wrong — layer/slot stride bug",
)
def test_matches_contiguous_reference(self):
"""The shared view must be a faithful relabeling of a contiguous
``(num_layers, max_slots, *inner)`` tensor: identical data written by the
same logical index reads back identically."""
pool, spec = _make_pool(**self.FALCON_KW)
conv_views, temporal_view = pool.mamba_views_for("mamba")
for v in conv_views + [temporal_view]:
ref = torch.randn(v.shape, device=v.device).to(v.dtype)
contig = ref.clone().contiguous()
v.copy_(ref)
self.assertEqual(tuple(v.shape), tuple(contig.shape))
self.assertTrue(
torch.equal(v.contiguous(), contig),
"shared view not equivalent to its contiguous counterpart",
)
def test_alignment_guard_fires_on_misaligned_spec(self):
"""A spec whose conv region (bf16) is an odd multiple of 2 B makes the
per-slot entry (= conv_region + N*temporal_row = 2 B + 4 B = 6 B) NOT a
multiple of the temporal itemsize (fp32, 4 B). The temporal/SSM-state
view's storage_offset is computed by integer-dividing a byte offset by
the temporal itemsize, so this would silently mis-offset the view.
``_build_mamba_views`` must reject it with a loud alignment assert.
NOTE: the ``entry_bytes % itemsize`` guard is what fires here, and it
subsumes the conv-region offset check (see the comment in
``_build_mamba_views``). We assert on the shared "misaligned" wording
rather than on which specific guard trips."""
with self.assertRaises(AssertionError) as cm:
_make_pool(
mamba_layer_num=1, # entry = 2 B conv + 4 B temporal = 6 B, not %4
conv_state_shapes=[(1, 1)],
conv_dtype=torch.bfloat16,
temporal_state_shape=(1,),
temporal_dtype=torch.float32,
want_slots=4,
)
self.assertIn("misalign", str(cm.exception).lower())
def test_alignment_ok_for_aligned_spec(self):
"""An aligned spec (conv region a multiple of the temporal itemsize)
must build and round-trip cleanly."""
# conv region = N * conv_dim*(k-1) * 2 ; with conv_dim=2 -> per-layer 2*3*2=12,
# times N=2 = 24, divisible by 4. Aligned.
pool, spec = _make_pool(
mamba_layer_num=2,
conv_state_shapes=[(2, 3)],
conv_dtype=torch.bfloat16,
temporal_state_shape=(2, 4, 4),
temporal_dtype=torch.float32,
want_slots=4,
)
self._fill_and_roundtrip(pool, spec)
if __name__ == "__main__":
unittest.main()