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

[unified-memory] Let Kimi-Linear use the paged MLA attention backends (#32972)

Browse Source
This commit is contained in:
Cheng Wan
2026-07-31 01:32:08 -07:00
committed by GitHub
parent 937c77cf50
commit 33c27d8e7f
9 changed files with 759 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/models_e2e/test_kimi_linear_unified_memory.py
+77
View File
@@ -0,0 +1,77 @@
"""Kimi-Linear (MLA full attention + KDA linear attention) served from the
unified memory pool.
`--enable-unified-memory` replaces the statically-partitioned hybrid pools with
one byte buffer split dynamically between the full-attention KV sub-pool and the
per-request KDA state sub-pool. For an MLA model the full side is exposed as
DENSE per-layer views (`build_dense_mla_views`) and every loc the kernels see is
a translated virtual id, so the whole read/write path differs from the static
pool: `translate_kv_loc_dense` for kv_indices and the cuda-graph write loc,
`HybridLinearKVPool._full_translate` for the model-level MLA entry points, and
page-envelope relocation on allocator compaction.
None of that is covered by the CPU/GPU unit tests, which pin the pool in
isolation. This is the end-to-end guard: accuracy must match the static-pool
baseline, and the prefix-cache branching case must still hit, since a radix hit
replays virtual locs whose physical pages may have moved under compaction.
Reference numbers on 2x H200 TP2, GSM8K 400 examples (2026-07-30):
static pools 0.915, `--enable-unified-memory` 0.900 (1 sigma ~= 0.015) -- both with
the attention backend pinned to triton, as this test runs it. For reference the
paged MLA kernels land in the same band on a single B300 TP1, GSM8K 200, unified
(2026-07-31): 0.915 with flashinfer prefill+decode, 0.900 with trtllm_mla. Those
are not exercised here (see the comment on `other_args`).
Nightly-only: it needs a second full 48B server launch, which is too much to add
to per-PR CI on top of the existing Kimi-Linear e2e coverage.
python -m pytest test/registered/models_e2e/test_kimi_linear_unified_memory.py -v
"""
import unittest
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
register_cuda_ci(est_time=600, suite="nightly-4-gpu", nightly=True)
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
class TestKimiLinearUnifiedMemory(
GSM8KMixin, PrefixCacheBranchingMixin, DefaultServerBase
):
model = KIMI_LINEAR_MODEL
cache_chunk_size = 64
# Same bar as the static-pool Kimi-Linear e2e test: unified memory must not
# cost accuracy (measured 0.900 vs 0.915 static, see the module docstring).
gsm8k_score_threshold = 0.88
other_args = [
"--trust-remote-code",
"--tp-size",
"2",
"--chunked-prefill-size",
"2048",
"--enable-unified-memory",
# Pinned because the resolved default is not portable: on pre-Blackwell
# (this suite's runner is H100) an unspecified backend resolves to `fa3`,
# which cannot read the dense views at all, so the un-pinned form fails at
# startup with the page-major allowlist assertion. Unified memory on such a
# host currently REQUIRES an explicit compatible --attention-backend; that
# is a real usability gap, tracked separately, not something this test can
# paper over.
#
# Consequence to keep in mind: pinning triton means this test does NOT
# cover the paged MLA backends (trtllm_mla / flashinfer / cutedsl_mla /
# tokenspeed_mla), which is where dense-id translation bugs live -- a
# captured flashinfer decode reading untranslated virtual ids scored GSM8K
# 0.000 on a healthy server. Those paths are covered by the unit tests plus
# manual B300 runs; an sm100-gated case here would close the gap.
"--attention-backend",
"triton",
]
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py
+375
View File
@@ -0,0 +1,375 @@
# 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.
# ==============================================================================
"""DENSE block table / kv_indices for the paged MLA backends under the unified
memory pool (Kimi-Linear).
`req_to_token` holds VIRTUAL token ids, while the per-layer MLA views are dense
(`build_dense_mla_views`). The paged MLA backends therefore need their page-level
block table filled with DENSE page ids:
dense_page(virtual_page) = v2p[virtual_page] * layer_num
`create_flashmla_kv_indices_triton` does that in-kernel via `v2p_ptr` / `PAGE_MULT`
(trtllm_mla / cutedsl_mla / tokenspeed_mla), and the flashinfer_mla updaters do it
by post-gathering `translate_kv_loc_dense` over the token-level kv_indices.
Covered here:
- kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main;
- kernel dense mapping against the python reference, for several page sizes,
ragged sequence lengths and a non-identity v2p permutation;
- padded block-table lanes never index the v2p table out of bounds;
- the token-level dense translate the flashinfer updaters apply agrees with the
page-level block table the trtllm path builds.
python -m pytest test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py -v
"""
import unittest
import torch
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=30, stage="base-b", runner_config="1-gpu-small")
_HAS_CUDA = torch.cuda.is_available()
_DEV = "cuda"
_LAYERS = 24 # K3 MLA full-attention layer count
def _fill_block_table(
req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult
):
from sglang.kernels.ops.kvcache.kv_indices import (
create_flashmla_kv_indices_triton,
get_num_kv_index_blocks_flashmla,
)
bs = req_pool_indices.shape[0]
max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size
out = torch.full((bs, max_blocks), -1, dtype=torch.int32, device=_DEV)
grid = (bs, get_num_kv_index_blocks_flashmla(max_blocks, page_size))
create_flashmla_kv_indices_triton[grid](
req_to_token,
req_pool_indices,
seq_lens,
None,
out,
req_to_token.stride(0),
max_blocks,
PAGED_SIZE=page_size,
v2p_ptr=v2p,
PAGE_MULT=mult,
)
return out
def _reference(req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult):
"""Python reference: virtual token -> virtual page -> physical page -> dense."""
bs = req_pool_indices.shape[0]
max_blocks = (int(seq_lens.max().item()) + page_size - 1) // page_size
ref = torch.full((bs, max_blocks), -1, dtype=torch.int64, device=_DEV)
for r in range(bs):
n_pages = (int(seq_lens[r].item()) + page_size - 1) // page_size
row = req_to_token[int(req_pool_indices[r].item())]
virt_pages = row[: n_pages * page_size : page_size] // page_size
pages = v2p[virt_pages.long()] if v2p is not None else virt_pages.long()
ref[r, :n_pages] = pages * mult
return ref
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestDenseBlockTable(unittest.TestCase):
def _make_batch(self, page_size, bs=5, max_ctx=2048, n_pages=512):
"""Ragged batch with a non-identity virtual->physical page permutation."""
g = torch.Generator(device="cpu").manual_seed(97 + page_size)
seq_lens = torch.tensor(
[page_size, page_size + 1, 3 * page_size, 7 * page_size - 3, 1],
dtype=torch.int32,
device=_DEV,
)[:bs]
req_to_token = torch.zeros((bs, max_ctx), dtype=torch.int32, device=_DEV)
# Every request gets a distinct, non-monotonic run of virtual pages.
virt_page_perm = torch.randperm(n_pages - 1, generator=g)[: bs * 8] + 1
for r in range(bs):
n = int(seq_lens[r].item())
pages = virt_page_perm[r * 8 : r * 8 + 8].to(_DEV)
toks = (
pages[:, None] * page_size
+ torch.arange(page_size, device=_DEV)[None, :]
).reshape(-1)
req_to_token[r, :n] = toks[:n].to(torch.int32)
req_pool_indices = torch.arange(bs, dtype=torch.int32, device=_DEV)
# Non-identity page-level v2p, with a tombstone that no request references.
v2p = torch.randperm(n_pages, generator=g).to(_DEV).to(torch.int64)
v2p[0] = 0 # page 0 is the reserved sink
return req_to_token, req_pool_indices, seq_lens, v2p
def test_identity_when_hooks_absent(self):
"""v2p_ptr=None / PAGE_MULT=1 must reproduce the pre-change behaviour."""
for page_size in (1, 32, 64):
rt, rpi, sl, _ = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=None, mult=1)
want = _reference(rt, rpi, sl, page_size, v2p=None, mult=1)
self.assertTrue(
torch.equal(got.long(), want), f"page_size={page_size}: {got} != {want}"
)
def test_dense_block_table_matches_reference(self):
for page_size in (1, 32, 64):
rt, rpi, sl, v2p = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
want = _reference(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
self.assertTrue(
torch.equal(got.long(), want),
f"page_size={page_size}:\ngot ={got}\nwant={want}",
)
def test_single_full_attention_layer_still_maps_v2p(self):
"""A config with exactly ONE full-attention layer (e.g. a PP rank owning a
single MLA layer) has `kernel_page_multiplier == 1`, but its req_to_token
still holds VIRTUAL ids. The dense id collapses onto the physical id, so
the v2p gather alone IS the whole translation -- it must not be skipped.
Regression guard for detecting the unified pool via `multiplier > 1`:
that predicate treats this config as a static pool and leaves the block
table in virtual id space.
"""
for page_size in (1, 64):
rt, rpi, sl, v2p = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=1).long()
want = _reference(rt, rpi, sl, page_size, v2p=v2p, mult=1)
self.assertTrue(torch.equal(got, want), f"page_size={page_size}")
# ... and the v2p permutation is non-trivial here, so a skipped
# translation would be visibly different rather than accidentally equal.
virtual = _reference(rt, rpi, sl, page_size, v2p=None, mult=1)
self.assertFalse(
torch.equal(want, virtual),
"test batch degenerated: v2p is the identity on the pages used",
)
def test_padded_lanes_stay_untouched(self):
"""Lanes past a request's page count keep the -1 fill: the masked v2p load
must not write a translated value (nor read out of bounds)."""
page_size = 64
rt, rpi, sl, v2p = self._make_batch(page_size)
got = _fill_block_table(rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS)
for r in range(got.shape[0]):
n_pages = (int(sl[r].item()) + page_size - 1) // page_size
self.assertTrue(
torch.all(got[r, n_pages:] == -1),
f"row {r} padded lanes were written: {got[r]}",
)
def test_agrees_with_token_level_dense_translate(self):
"""The flashinfer updaters translate TOKEN ids with
`translate_kv_loc_dense`; the trtllm path builds PAGE ids in-kernel. Both
must address the same dense page block."""
page_size = 64
rt, rpi, sl, v2p = self._make_batch(page_size)
block_table = _fill_block_table(
rt, rpi, sl, page_size, v2p=v2p, mult=_LAYERS
).long()
for r in range(rt.shape[0]):
n = int(sl[r].item())
virt_tokens = rt[r, :n].long()
# translate_kv_loc_dense's formula, applied to token ids.
dense_tokens = (
v2p[virt_tokens // page_size] * (page_size * _LAYERS)
+ virt_tokens % page_size
)
# The block-table entry scaled by page_size must be the dense id of
# each page's first token.
first_of_page = dense_tokens[::page_size]
n_pages = (n + page_size - 1) // page_size
self.assertTrue(
torch.equal(block_table[r, :n_pages] * page_size, first_of_page),
f"row {r}: block table and token translate disagree",
)
class TestUnifiedMLAHookDetection(unittest.TestCase):
"""`unified_mla_hooks` decides whether the paged MLA backends translate at
all. Getting the predicate wrong is silent: the block table and KV write loc
stay in virtual id space and address the wrong pages once virtual and
physical diverge (e.g. after compaction)."""
@staticmethod
def _probe(**attrs):
from sglang.srt.layers.attention.flashinfer_mla_backend import (
unified_mla_hooks,
)
class _Alloc:
pass
alloc = _Alloc()
for k, v in attrs.items():
setattr(alloc, k, v)
return unified_mla_hooks(alloc)
def test_static_pool_disables_every_hook(self):
"""No v2p table -> statically-partitioned pool; req_to_token is already
physical, so all hooks must stay off (byte-identical to pre-change)."""
hooks = self._probe()
self.assertFalse(hooks.enabled)
self.assertIsNone(hooks.v2p_page_table)
self.assertIsNone(hooks.translate_kv_loc_dense)
self.assertEqual(hooks.kernel_page_multiplier, 1)
def test_multi_layer_unified_pool(self):
table = torch.arange(8)
hooks = self._probe(
full_v2p_page_table=table,
translate_kv_loc_dense=lambda x, **kw: x,
kernel_page_multiplier=_LAYERS,
)
self.assertTrue(hooks.enabled)
self.assertIs(hooks.v2p_page_table, table)
self.assertIsNotNone(hooks.translate_kv_loc_dense)
self.assertEqual(hooks.kernel_page_multiplier, _LAYERS)
def test_single_full_attention_layer_pool_is_still_unified(self):
"""REGRESSION: `kernel_page_multiplier == 1` does NOT mean static.
A hybrid MLA config with exactly one full-attention layer (e.g. a
pipeline-parallel rank owning a single MLA layer) has multiplier 1, yet
its locs are still virtual. Detecting on `multiplier > 1` would disable
the v2p gather here and corrupt reads/writes after compaction.
"""
table = torch.arange(8)
hooks = self._probe(
full_v2p_page_table=table,
translate_kv_loc_dense=lambda x, **kw: x,
kernel_page_multiplier=1,
)
self.assertTrue(hooks.enabled, "single-layer unified pool read as static")
self.assertIs(hooks.v2p_page_table, table)
self.assertIsNotNone(hooks.translate_kv_loc_dense)
# Multiplier stays 1: dense id == physical id, so the v2p gather alone is
# the whole translation and PAGE_MULT must not scale it.
self.assertEqual(hooks.kernel_page_multiplier, 1)
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestInPlaceKvIndicesTranslate(unittest.TestCase):
"""The flashinfer decode updater must translate kv_indices IN PLACE.
Under cuda-graph replay the `kv_indices` it is handed IS the capture-stable
buffer the captured wrapper reads (`fast_decode_kwargs["kv_indices"]`), and
`fast_mla_decode_plan` ignores its `kv_indices` argument -- so rebinding the
local name to a fresh translated tensor leaves the graph reading VIRTUAL ids.
These pin the write-back contract that fix relies on.
"""
def _allocator(self, page_size=1, n_full_tokens=4096):
from sglang.srt.mem_cache.multi_ended_allocator import MultiEndedAllocator
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
)
full = MLASubPoolSpec(
name="full",
layer_num=_LAYERS,
kv_lora_rank=512,
qk_rope_head_dim=64,
store_dtype=torch.bfloat16,
grow_direction="down",
)
mamba = MambaSubPoolSpec(
name="mamba",
layer_num=2,
conv_state_shapes=((8, 16),),
conv_dtype=torch.bfloat16,
temporal_state_shape=(4, 8, 8),
temporal_dtype=torch.float32,
grow_direction="up",
)
pool = UnifiedKVPool(
total_bytes=full.entry_bytes() * n_full_tokens + mamba.entry_bytes() * 16,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=page_size,
view_tail_pad_bytes=page_size * full.entry_bytes(),
)
class _Stub:
def move_kv_cache(self, dst, src):
pass
full_alloc = MultiEndedAllocator(
kvcache=_Stub(),
unified_buffer=pool,
sub_pool_name="full",
device=_DEV,
is_id_owner=True,
page_size=page_size,
kernel_page_multiplier=_LAYERS,
)
mamba_alloc = MultiEndedAllocator(
kvcache=_Stub(),
unified_buffer=pool,
sub_pool_name="mamba",
device=_DEV,
is_id_owner=True,
)
full_alloc.bind_peer(mamba_alloc)
mamba_alloc.bind_peer(full_alloc)
return full_alloc
def test_int32_buffer_prefix_translated_tail_untouched(self):
"""Mirrors the updater: an int32 capture-stable buffer holding VIRTUAL
ids in [:n] gets the dense ids written back in place, narrowed to int32,
with the stale tail left alone (it must never index the v2p table)."""
alloc = self._allocator()
virt = alloc.alloc(64)
self.assertIsNotNone(virt)
n = virt.numel()
# Capture-stable int32 buffer: [:n] freshly filled with virtual ids by
# create_flashinfer_kv_indices_triton, tail = stale junk from a bigger replay.
buf = torch.full((n * 3,), 2**30, dtype=torch.int32, device=_DEV)
buf[:n] = virt.to(torch.int32)
tail_before = buf[n:].clone()
valid = buf[:n]
valid.copy_(alloc.translate_kv_loc_dense(valid))
expected = alloc.translate_kv_loc_dense(virt)
self.assertEqual(buf.dtype, torch.int32)
self.assertTrue(
torch.equal(buf[:n].long(), expected),
"in-place translate did not land dense ids in the stable buffer",
)
self.assertTrue(
torch.equal(buf[n:], tail_before),
"stale tail was modified -- it can hold ids outside the v2p table",
)
def test_dense_ids_differ_from_virtual(self):
"""Guard the guard: if dense == virtual the in-place test proves nothing."""
alloc = self._allocator()
virt = alloc.alloc(64)
self.assertIsNotNone(virt)
self.assertFalse(
torch.equal(alloc.translate_kv_loc_dense(virt), virt),
"dense ids coincide with virtual ids; pick a different allocation",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_page_major_backend_allowlist.py
+111
View File
@@ -0,0 +1,111 @@
# 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.
# ==============================================================================
"""`--enable-page-major-kv-layout` full-attention backend allowlist.
The page-major envelope K/V views are strided, which only the Triton attention
kernels read. The one exception is the unified-memory MLA pool: it exposes each
layer as a DENSE contiguous view (`build_dense_mla_views`), so the paged MLA
backends (`trtllm_mla` and its `cutedsl_mla` / `tokenspeed_mla` subclasses, plus
`flashinfer`'s MLA backend) can read it directly once their kv_indices / block
tables are remapped to dense ids.
Pinned here so the exception cannot silently widen to a backend that has no
dense-id remapping (`fa3`, `flashmla`, ...) or leak into the MHA path.
python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
"""
import unittest
from sglang.srt.server_args import ServerArgs
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool:
"""Run just `_handle_page_major_kv_layout` against a minimal stand-in.
ServerArgs' real constructor pulls in a model config; this exercises the
single handler under test with the fields it reads.
"""
sa = ServerArgs.__new__(ServerArgs)
for name, value in {
"enable_unified_memory": unified,
# The unified pool sets this itself; without it the flag must be explicit
# or the handler returns before reaching the allowlist.
"enable_page_major_kv_layout": not unified,
"attention_backend": backend,
"prefill_attention_backend": None,
"decode_attention_backend": None,
"linear_attn_backend": "triton",
"linear_attn_decode_backend": None,
"linear_attn_prefill_backend": None,
"mamba_backend": "triton",
}.items():
object.__setattr__(sa, name, value)
sa.use_mla_backend = lambda: use_mla
sa._resolved_attention_backends = lambda: [backend]
try:
ServerArgs._handle_page_major_kv_layout(sa)
return True
except AssertionError:
return False
class TestPageMajorBackendAllowlist(unittest.TestCase):
# Wired for the dense per-layer MLA views (see the module docstring).
DENSE_MLA_BACKENDS = ("trtllm_mla", "flashinfer", "cutedsl_mla", "tokenspeed_mla")
# No dense-id remapping: must stay rejected until they get one.
UNWIRED_BACKENDS = ("fa3", "flashmla", "cutlass_mla", "trtllm_mha", "aiter")
def test_triton_always_allowed(self):
for use_mla in (True, False):
self.assertTrue(_accepts("triton", use_mla=use_mla))
def test_dense_mla_backends_allowed_under_unified_mla(self):
for backend in self.DENSE_MLA_BACKENDS:
self.assertTrue(
_accepts(backend, use_mla=True),
f"{backend} should be allowed with the unified-memory MLA pool",
)
def test_dense_mla_backends_rejected_for_mha(self):
"""The dense-view exception is MLA-only -- MHA sub-pools stay strided."""
for backend in self.DENSE_MLA_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=False),
f"{backend} must stay rejected for a non-MLA model",
)
def test_dense_mla_backends_rejected_without_unified_memory(self):
"""Plain --enable-page-major-kv-layout (no unified pool) keeps the
strided views, so only Triton can read them."""
for backend in self.DENSE_MLA_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=True, unified=False),
f"{backend} must stay rejected without --enable-unified-memory",
)
def test_unwired_backends_always_rejected(self):
for backend in self.UNWIRED_BACKENDS:
for use_mla in (True, False):
self.assertFalse(
_accepts(backend, use_mla=use_mla),
f"{backend} has no dense-id remapping and must be rejected",
)
if __name__ == "__main__":
unittest.main()