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

[Feature] Unified memory: support decode context parallelism for Kimi-Linear (#36890)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-01 12:44:26 -07:00
committed by GitHub
parent 3315356cc0
commit 0b1ce3d140
20 changed files with 671 additions and 105 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_full_loc_fast_path.py
+59
View File
@@ -176,7 +176,10 @@ class TestUnifiedSWATombstoneClamp(unittest.TestCase):
# A real sub-allocator (not a stand-in): the translation reads its v2p
# table, and the pool reaches it through the allocator's own method.
swa_allocator = object.__new__(MultiEndedAllocator)
# `page_size` is the WIDENED (DCP) surface, `pool_page_size` the physical
# rows per page; equal at dcp_size == 1, which is what this fixture is.
swa_allocator.page_size = page_size
swa_allocator.pool_page_size = page_size
swa_allocator.virtual_to_physical = v2p
swa_allocator.kernel_page_multiplier = multiplier
pool = object.__new__(UnifiedSWAKVPool)
@@ -352,5 +355,61 @@ class TestHybridLinearMLARouting(unittest.TestCase):
self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
class TestMlaWriteDoorsUnderDcp(unittest.TestCase):
"""Which MLA write door is DCP-aware, and which refuses.
`set_mla_kv_buffer` resolves the owner rule inside its kernel, so it owns
the DCP write. `set_kv_buffer` (the combined latent+rope row) cannot: the
two backends that could reach it disagree on the loc space -- flashinfer's
`k_rope is None` branch passes a WIDENED loc, the Triton backend one it
already collapsed -- so there is no single correct translation. It used to
select `loc % dcp_size == dcp_rank` and then write WITHOUT dividing, i.e.
widened ids straight into a rank-local buffer. Refusing is the contract;
a re-added masked-but-undivided write is what this guards."""
def _bare_mla_pool(self):
from sglang.srt.mem_cache.memory_pool import MLATokenToKVPool
pool = object.__new__(MLATokenToKVPool)
pool.size = 64
pool.page_size = 1
pool.kernel_page_blocks = 1
pool.start_layer = 0
pool.dtype = torch.float16
pool.store_dtype = torch.float16
pool.dsa_kv_cache_store_fp8 = False
pool.kv_buffer = [torch.zeros((65, 1, 8), dtype=torch.float16)]
return pool
def test_set_kv_buffer_refuses_under_dcp(self):
from sglang.srt.runtime_context import get_parallel
pool = self._bare_mla_pool()
layer = types.SimpleNamespace(layer_id=0)
loc = torch.tensor([0, 1, 2, 3], dtype=torch.int64)
cache_k = torch.ones((4, 1, 8), dtype=torch.float16)
with get_parallel().override(
dcp_enabled=True, attn_dcp_size=2, attn_dcp_rank=1
):
with self.assertRaises(AssertionError) as cm:
pool.set_kv_buffer(layer, _loc_info(loc), cache_k, None)
self.assertIn("set_mla_kv_buffer", str(cm.exception))
# Nothing was written on the way to refusing.
self.assertTrue(bool((pool.kv_buffer[0] == 0).all()))
def test_set_kv_buffer_still_writes_without_dcp(self):
pool = self._bare_mla_pool()
layer = types.SimpleNamespace(layer_id=0)
loc = torch.tensor([3, 5], dtype=torch.int64)
cache_k = torch.ones((2, 1, 8), dtype=torch.float16)
pool.set_kv_buffer(layer, _loc_info(loc), cache_k, None)
self.assertTrue(bool((pool.kv_buffer[0][3] == 1).all()))
self.assertTrue(bool((pool.kv_buffer[0][5] == 1).all()))
self.assertTrue(bool((pool.kv_buffer[0][4] == 0).all()))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+196
View File
@@ -24,6 +24,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
import contextlib
import random
import unittest
@@ -41,6 +42,7 @@ from sglang.srt.mem_cache.unified_memory_pool import (
MLASubPoolSpec,
UnifiedKVPool,
)
from sglang.srt.runtime_context import get_parallel
_DEV = "cpu"
@@ -3299,5 +3301,199 @@ class TestFloatMultiEndedAllocator(unittest.TestCase):
fla.bind_peer(sa)
class TestDcpWidening(unittest.TestCase):
"""`dcp_size > 1`: the alloc surface speaks a widened virtual id space while
the pool keeps storing one row per `dcp_size` logical ids."""
@contextlib.contextmanager
def _dcp(self, dcp_size, dcp_rank=0):
"""The width comes from the parallel context, not a constructor
argument, so one scope has to hold construction and every read."""
with get_parallel().override(
dcp_enabled=dcp_size > 1,
attn_dcp_size=dcp_size,
attn_dcp_rank=dcp_rank,
):
yield
def _build_pair(self, *, page_size, n_full_slots=64):
"""(full, mamba) as the composite wires them: only full shards."""
full = _make_mha_spec("full", "up", layer_num=2)
mamba = _make_mamba_spec("mamba", "down", layer_num=2)
pool = UnifiedKVPool(
total_bytes=full.entry_bytes() * n_full_slots + mamba.entry_bytes() * 16,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
alloc = MultiEndedAllocator(
kvcache=_FakeKVCache(pool.max_slots("full")),
unified_buffer=pool,
sub_pool_name="full",
device=_DEV,
is_id_owner=True,
page_size=page_size,
shards_under_dcp=True,
)
# The peer stays slot-granular: mamba state is replicated, not sharded.
mamba = MultiEndedAllocator(
kvcache=_FakeKVCache(pool.max_slots("mamba")),
unified_buffer=pool,
sub_pool_name="mamba",
device=_DEV,
is_id_owner=True,
)
alloc.bind_peer(mamba)
return alloc, mamba
def _build(self, *, page_size, n_full_slots=64):
return self._build_pair(page_size=page_size, n_full_slots=n_full_slots)[0]
def test_replicated_peer_stays_slot_granular_under_dcp(self):
"""BUG REGRESSION. Widening every sub-allocator off the process DCP
width, rather than only the sharding one, makes the Mamba page size
dcp_size; state is allocated one slot per request, so the first request
fails the page-multiple check."""
with self._dcp(4):
_, mamba = self._build_pair(page_size=2)
self.assertEqual(mamba.page_size, 1)
self.assertEqual(mamba.page_size, mamba.pool_page_size)
self.assertIsNotNone(mamba.alloc(1))
def test_capacity_scales_but_physical_pages_do_not(self):
for page_size in (1, 8):
with self._dcp(1):
base = self._build(page_size=page_size)
base_pages = base.num_pages
base_page_bytes = base.entry_bytes_per_page
base_avail = base.available_size()
for dcp_size in (2, 4):
with self._dcp(dcp_size):
a = self._build(page_size=page_size)
self.assertEqual(a.page_size, page_size * dcp_size)
self.assertEqual(a.pool_page_size, page_size)
# Same rows, same bytes per page; only the id space grows.
self.assertEqual(a.num_pages, base_pages)
self.assertEqual(a.entry_bytes_per_page, base_page_bytes)
self.assertEqual(a.available_size(), base_avail * dcp_size)
def test_alloc_returns_whole_widened_pages(self):
with self._dcp(2, 1):
a = self._build(page_size=4)
ids = a.alloc(3 * 8) # 3 widened pages of 4*2 ids
self.assertIsNotNone(ids)
pages = ids.view(3, 8)
self.assertTrue(
torch.equal(pages[:, 1:] - pages[:, :-1], torch.ones(3, 7).long())
)
self.assertTrue(bool((pages[:, 0] % 8 == 0).all()))
# Freeing the widened ids releases exactly the pages they came from.
before = a.available_size()
a.free(ids)
self.assertEqual(a.available_size(), before + 3 * 8)
def test_every_rank_maps_a_widened_page_to_one_physical_page(self):
"""The DCP ranks must agree on the physical page a widened page uses;
only the row WITHIN it differs, by `(loc % dcp) -> loc // dcp`."""
dcp_size = 4
with self._dcp(dcp_size):
allocs = [self._build(page_size=2) for _ in range(dcp_size)]
ids = [a.alloc(2 * dcp_size * 2) for a in allocs]
for i in ids:
self.assertIsNotNone(i)
# Same allocation order -> same widened ids on every rank.
for i in ids[1:]:
self.assertTrue(torch.equal(i, ids[0]))
for rank, (a, i) in enumerate(zip(allocs, ids)):
with self._dcp(dcp_size, rank):
owned = (i % dcp_size) == rank
self.assertEqual(int(owned.sum()), i.numel() // dcp_size)
phys = a.translate_kv_loc(i[owned] // dcp_size)
# Collapsed ids land inside this rank's physical rows,
# contiguously within each page, never on the reserved sink.
self.assertTrue(bool((phys > 0).all()))
self.assertTrue(bool((phys < a.max_slots).all()))
self.assertEqual(len(set(phys.tolist())), phys.numel())
def test_write_translate_tombstones_unowned_ids(self):
dcp_size = 2
for rank in range(dcp_size):
with self._dcp(dcp_size, rank):
a = self._build(page_size=2)
ids = a.alloc(2 * dcp_size * 3)
written = a.translate_write_loc_for_kernel(ids)
owned = (ids % dcp_size) == rank
# Owned ids agree with the read translate of the collapsed id...
self.assertTrue(
torch.equal(
written[owned],
a.translate_kv_loc_for_kernel(ids[owned] // dcp_size),
)
)
# ...and the rest go to the sink the write kernels skip.
self.assertTrue(bool((written[~owned] == 0).all()))
self.assertTrue(bool((written[owned] > 0).all()))
def _build_composite(self, *, page_size):
from sglang.srt.mem_cache.multi_ended_allocator import (
UnifiedMambaTokenToKVPoolAllocator,
)
full_spec = _make_mha_spec("full", "up", layer_num=2)
mamba_spec = _make_mamba_spec("mamba", "down", layer_num=2)
pool = UnifiedKVPool(
total_bytes=16 * page_size * full_spec.entry_bytes()
+ 8 * mamba_spec.entry_bytes(),
sub_pool_specs=[full_spec, mamba_spec],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
)
full_kv = _FakeKVCache(pool.max_slots("full"))
mamba_kv = _FakeKVCache(pool.max_slots("mamba"))
mamba_kv._copy_from_physical = lambda src, dst: None
class _FakeHybridLinearKVPool:
full_kv_pool = full_kv
mamba_pool = mamba_kv
return UnifiedMambaTokenToKVPoolAllocator(
unified_buffer=pool,
kvcache=_FakeHybridLinearKVPool(),
device=_DEV,
page_size=page_size,
need_sort=False,
forward_stream=None,
)
def test_mamba_slot_cost_is_in_the_same_units_as_available_size(self):
"""The planner charges `mamba_slot_full_token_cost()` against a budget
fed by `available_size()`. Both are bytes/entry_bytes conversions, so
both carry `dcp_size`; if only the budget widens, every Mamba state is
under-reserved by that factor and a batch is admitted whose later
allocations cross the shared byte frontier."""
for page_size in (1, 8):
with self._dcp(1):
base = self._build_composite(page_size=page_size)
base_cost = base.mamba_slot_full_token_cost()
base_avail = base.available_size()
self.assertGreater(base_cost, 0)
for dcp_size in (2, 4):
with self._dcp(dcp_size):
a = self._build_composite(page_size=page_size)
self.assertEqual(a.available_size(), base_avail * dcp_size)
mamba_bytes = a.mamba_allocator.entry_bytes_per_page
full_entry = a.full_attn_allocator.entry_bytes
cost = a.mamba_slot_full_token_cost()
# A widened token is `full_entry / dcp_size` bytes, so the
# reservation covers the slot...
self.assertGreaterEqual(cost * full_entry, mamba_bytes * dcp_size)
# ...and stays tight (rounds up by less than one token).
self.assertLess((cost - 1) * full_entry, mamba_bytes * dcp_size)
# The un-scaled cost -- the bug -- would not have covered it.
self.assertLess(base_cost * full_entry, mamba_bytes * dcp_size)
if __name__ == "__main__":
unittest.main()