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
python/sglang/kernels/ops/kvcache/kv_indices.py
+15 -1
View File
@@ -105,6 +105,15 @@ def create_flashmla_kv_indices_triton(
req_to_token_ptr_stride: tl.constexpr,
kv_indices_ptr_stride: tl.constexpr,
PAGED_SIZE: tl.constexpr = 64,
# Unified-memory dense-view path (page-major envelope shared with the mamba
# sub-pool). req_to_token holds VIRTUAL token ids; the block table the MLA
# kernel consumes must hold DENSE page ids. When v2p_ptr is given, map each
# virtual page through it to the physical page, then scale by PAGE_MULT
# (= num MLA layers) so the entry addresses the layer's dense per-page block
# in the (num_pages*L, page_size, kv_cache_dim) reshaped view. Both default
# to the identity (v2p_ptr None, PAGE_MULT 1) for the static pool.
v2p_ptr=None,
PAGE_MULT: tl.constexpr = 1,
):
NUM_PAGE_PER_BLOCK: tl.constexpr = (
FLASHMLA_CREATE_KV_BLOCK_SIZE_TRITON // PAGED_SIZE
@@ -145,8 +154,13 @@ def create_flashmla_kv_indices_triton(
+ paged_offset,
mask=mask,
)
page = data // PAGED_SIZE
if v2p_ptr is not None:
# virtual page -> physical page (page-level v2p); masked so padded
# lanes never index the table out of bounds.
page = tl.load(v2p_ptr + page, mask=mask_out, other=0)
tl.store(
kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
data // PAGED_SIZE,
page * PAGE_MULT,
mask=mask_out,
)
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
+76
View File
@@ -66,6 +66,51 @@ if is_flashinfer_available():
)
@dataclass(frozen=True)
class UnifiedMLAHooks:
"""Allocator hooks the paged MLA backends need under the unified memory pool.
All-``None``/1/``False`` for the statically-partitioned pool, where
``req_to_token`` already holds physical ids.
"""
# Page-level virtual->physical table, gathered through by the block-table kernel.
v2p_page_table: Optional[torch.Tensor]
# Virtual token id -> DENSE kernel-facing id.
translate_kv_loc_dense: Optional[Callable[..., torch.Tensor]]
# Dense page stride scale (= number of full-attention MLA layers).
kernel_page_multiplier: int
enabled: bool
def unified_mla_hooks(allocator) -> UnifiedMLAHooks:
"""Probe ``allocator`` for the unified-pool dense-view hooks.
Detection keys on the page-level v2p table, NOT on
``kernel_page_multiplier > 1``: a configuration with exactly ONE
full-attention layer (e.g. a pipeline-parallel rank that owns a single MLA
layer) has multiplier 1 while its ``req_to_token`` still holds VIRTUAL ids.
With multiplier 1 the dense id collapses onto the physical id, so the v2p
gather alone is the whole translation -- skipping it would leave the block
table and the KV write loc in virtual space and silently address the wrong
pages once virtual and physical diverge (e.g. after compaction).
"""
v2p = getattr(allocator, "full_v2p_page_table", None)
if v2p is None:
return UnifiedMLAHooks(
v2p_page_table=None,
translate_kv_loc_dense=None,
kernel_page_multiplier=1,
enabled=False,
)
return UnifiedMLAHooks(
v2p_page_table=v2p,
translate_kv_loc_dense=getattr(allocator, "translate_kv_loc_dense", None),
kernel_page_multiplier=getattr(allocator, "kernel_page_multiplier", 1),
enabled=True,
)
@dataclass
class DecodeMetadata:
decode_wrapper: BatchMLAPagedAttentionWrapper
@@ -675,6 +720,10 @@ class FlashInferMLAIndicesUpdaterDecode:
self.kv_indptr = attn_backend.kv_indptr
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.q_indptr = attn_backend.q_indptr_decode
# Unified dense MLA pool: VIRTUAL -> DENSE kv_indices (see prefill updater).
self._translate_kv_loc_dense = unified_mla_hooks(
model_runner.token_to_kv_pool_allocator
).translate_kv_loc_dense
def update(
self,
@@ -732,6 +781,21 @@ class FlashInferMLAIndicesUpdaterDecode:
kv_indices,
self.req_to_token.shape[1],
)
# Unified pool: VIRTUAL -> DENSE, written back IN PLACE.
#
# On the cuda-graph replay path `kv_indices` IS the capture-stable
# buffer (fast_decode_kwargs["kv_indices"] == cuda_graph_kv_indices)
# that the captured wrapper reads, and `fast_mla_decode_plan` ignores
# the kv_indices argument entirely -- rebinding the local name to a
# fresh tensor would leave the graph reading VIRTUAL ids. Only the
# [:paged_kernel_lens_sum] prefix the index kernel just filled is
# translated; the stale tail is left alone so it can never index the
# v2p table out of bounds. The int64 translate result narrows back to
# the buffer's int32 on copy_ (flashinfer requires int32; dense ids
# fit comfortably).
if self._translate_kv_loc_dense is not None:
valid = kv_indices[:paged_kernel_lens_sum]
valid.copy_(self._translate_kv_loc_dense(valid))
if get_parallel().dcp_enabled:
plan_dcp_decode_metadata(
@@ -797,6 +861,12 @@ class FlashInferMLAIndicesUpdaterPrefill:
self.qo_indptr = attn_backend.qo_indptr
self.req_to_token = model_runner.req_to_token_pool.req_to_token
self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged
# Unified dense MLA pool: kv_indices built from req_to_token are VIRTUAL;
# the paged wrapper reads the dense per-layer view, so remap them to DENSE
# token ids. None (identity) unless the unified MLA pool is active.
self._translate_kv_loc_dense = unified_mla_hooks(
model_runner.token_to_kv_pool_allocator
).translate_kv_loc_dense
def update(
self,
@@ -867,6 +937,12 @@ class FlashInferMLAIndicesUpdaterPrefill:
kv_indices,
self.req_to_token.shape[1],
)
# Unified pool: VIRTUAL -> DENSE token ids for the paged wrapper.
# Prefill is not cuda-graph captured under unified memory, so an eager
# gather is safe. Dense ids fit int32 (max = full_slots*num_layers ~
# 1e7 << 2^31); the flashinfer wrapper requires int32.
if self._translate_kv_loc_dense is not None:
kv_indices = self._translate_kv_loc_dense(kv_indices).to(torch.int32)
qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
qo_indptr = qo_indptr[: bs + 1]
custom_mask = None
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+64 -3
View File
@@ -34,6 +34,7 @@ from sglang.srt.environ import envs
from sglang.srt.layers.attention.flashinfer_mla_backend import (
FlashInferMLAAttnBackend,
FlashInferMLAMultiStepDraftBackend,
unified_mla_hooks,
)
from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -246,6 +247,23 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# Tree-mask scratch is fetched from the target backend only.
self.is_draft_runner = model_runner.is_draft_worker
# Unified-memory dense-view hooks (None on the static pool). req_to_token
# holds VIRTUAL token ids; the block table needs DENSE page ids, so the
# kv-index kernels gather virtual->physical page through `_v2p_page_table`
# then scale by `_kernel_page_multiplier` (= num MLA layers). See
# build_dense_mla_views / create_flashmla_kv_indices_triton.
_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator)
self._v2p_page_table = _hooks.v2p_page_table
self._kernel_page_multiplier = _hooks.kernel_page_multiplier
self._unified_mla = _hooks.enabled
# virtual token id -> DENSE kernel-facing id, for the KV write loc.
self._translate_kv_loc_dense = _hooks.translate_kv_loc_dense
# Per-forward dense write loc ([:n] view of a capture-stable buffer),
# set by the cuda-graph out-graph hook; None on the eager path (where the
# write translates through the pool's _full_translate hook instead).
self._decode_dense_loc: Optional[torch.Tensor] = None
self.cuda_graph_out_cache_loc_dense: Optional[torch.Tensor] = None
def _calc_padded_blocks(self, max_seq_len: int) -> int:
"""
Calculate padded block count that satisfies both TRT-LLM and Triton constraints.
@@ -308,6 +326,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
max_blocks,
PAGED_SIZE=self.page_size,
v2p_ptr=self._v2p_page_table,
PAGE_MULT=self._kernel_page_multiplier,
)
return block_kv_indices
@@ -325,6 +345,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.decode_cuda_graph_kv_indices = torch.full(
(max_bs, max_blocks_per_seq), -1, dtype=torch.int32, device=self.device
)
# Unified pool: capture-stable buffer for the DENSE KV write loc, filled
# out-of-graph in init_forward_metadata_out_graph so the in-graph
# set_mla_kv_buffer captures no translate.
if self._unified_mla:
self.cuda_graph_out_cache_loc_dense = torch.zeros(
max_num_tokens, dtype=torch.int64, device=self.device
)
num_tokens_per_req = max_num_tokens // max_bs
if is_float4_e2m1fn_x2(self.data_type):
@@ -464,6 +491,8 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
self.req_to_token.stride(0),
metadata.block_kv_indices.shape[1],
PAGED_SIZE=self.page_size,
v2p_ptr=self._v2p_page_table,
PAGE_MULT=self._kernel_page_multiplier,
)
def get_cuda_graph_seq_len_fill_value(self) -> int:
@@ -522,8 +551,32 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
forward_mode=forward_mode,
)
# Unified pool: precompute the DENSE KV write loc into the capture-stable
# buffer (both capture and each replay-prep run this out of the graph),
# so the in-graph set_mla_kv_buffer writes a dense loc without capturing a
# translate. Only decode writes KV under unified (spec is gated off).
if self._unified_mla and forward_mode.is_decode_or_idle():
out_cache_loc = forward_batch.out_cache_loc
n = out_cache_loc.shape[0]
dst = self.cuda_graph_out_cache_loc_dense[:n]
self._translate_kv_loc_dense(out_cache_loc, out=dst)
# Replay-prep receives the RAW (unpadded) out_cache_loc
# (build_replay_fb_view), but the captured write kernel consumes the
# full captured tier of this buffer. Zero the tail so pad rows write
# to the dense sink (row 0) instead of stale dense locs left by
# earlier larger replays — a stale tail scatters pad-row garbage into
# live KV pages. Mirrors the runner's PaddingPolicy.ZERO on its own
# out_cache_loc slot.
self.cuda_graph_out_cache_loc_dense[n:].zero_()
self._decode_dense_loc = dst
else:
self._decode_dense_loc = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Initialize the metadata for a forward pass."""
# Eager path: no capture-stable dense write loc; the pool's _full_translate
# hook translates the write loc (safe out of a cuda graph).
self._decode_dense_loc = None
# Delegate to parent for non-decode modes.
if (
forward_batch.forward_mode.is_extend()
@@ -803,9 +856,17 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
assert (
k is not None and k_rope is not None
), "For populating trtllm_mla kv cache, both k_nope and k_rope should be not None."
self.token_to_kv_pool.set_mla_kv_buffer(
layer, forward_batch.out_cache_loc, k, k_rope
)
if self._decode_dense_loc is not None:
# cuda-graph path: dense write loc precomputed out-of-graph, so
# the in-graph write captures no translate allocation.
self.token_to_kv_pool.set_mla_kv_buffer(
layer, self._decode_dense_loc, k, k_rope, loc_is_dense=True
)
else:
# eager (or static pool): the pool's _full_translate handles it.
self.token_to_kv_pool.set_mla_kv_buffer(
layer, forward_batch.out_cache_loc, k, k_rope
)
# Prepare query tensor inline
if merge_query:
python/sglang/srt/mem_cache/memory_pool.py
+8 -3
View File
@@ -3836,12 +3836,17 @@ class HybridLinearKVPool(KVCache):
loc: torch.Tensor,
cache_k_nope: torch.Tensor,
cache_k_rope: torch.Tensor,
loc_is_dense: bool = False,
):
assert self.use_mla, "set_mla_kv_buffer called when use_mla is False"
# Model-level MLA entry point: `loc` is a VIRTUAL loc under the unified
# pool (eager prefill only; the decode write goes through set_kv_buffer's
# pre-translated `full_loc`), so translate to the dense id space here.
loc = self._full_translate(loc)
# pool, so translate to the dense id space here.
#
# `loc_is_dense`: the caller already translated `loc` (the unified-pool
# cuda-graph decode precomputes it out-of-graph into a capture-stable
# buffer, so the in-graph write does not capture a translate allocation).
if not loc_is_dense:
loc = self._full_translate(loc)
with self._transfer_id_context(layer):
self.full_kv_pool.set_mla_kv_buffer(layer, loc, cache_k_nope, cache_k_rope)
python/sglang/srt/mem_cache/multi_ended_allocator.py
+9
View File
@@ -1898,6 +1898,15 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def kernel_page_multiplier(self) -> int:
return self.full_attn_allocator.kernel_page_multiplier
@property
def full_v2p_page_table(self) -> torch.Tensor:
"""Page-level virtual->physical table of the full sub-pool. Kernels that
build the MLA block table directly from req_to_token (e.g. trtllm_mla,
flashmla) gather through this to turn a VIRTUAL page into a physical one,
then scale by `kernel_page_multiplier` to reach the dense per-page block.
"""
return self.full_attn_allocator.virtual_to_physical
def translate_kv_loc_dense(
self,
loc: torch.Tensor,
python/sglang/srt/server_args.py
+24 -4
View File
@@ -7705,13 +7705,33 @@ class ServerArgs:
if not self.enable_page_major_kv_layout:
return
# Only the Triton attention kernels read the strided 4-D envelope K/V
# views; FA3 / FlashInfer do not.
# views; FA3 / FlashInfer do not. EXCEPTION: the unified-memory MLA pool
# exposes each layer as a DENSE contiguous per-layer view
# (build_dense_mla_views), which the paged MLA kernels consume directly,
# with their kv_indices / block tables remapped to dense ids. Names below
# are the RESOLVED ids from _resolved_attention_backends: "flashinfer" is
# FlashInferMLAAttnBackend for an MLA model, "trtllm_mla" the trtllm
# decode kernel; "cutedsl_mla" and "tokenspeed_mla" subclass
# TRTLLMMLABackend and inherit its dense read/write path.
# flashmla / cutlass_mla share the create_flashmla block-table path and
# can be added the same way once exercised.
if self.enable_unified_memory and self.use_mla_backend():
allowed_full = {
"triton",
"trtllm_mla",
"flashinfer",
"cutedsl_mla",
"tokenspeed_mla",
}
else:
allowed_full = {"triton"}
backends = set(self._resolved_attention_backends())
backends.discard(None)
assert backends <= {"triton"}, (
assert backends <= allowed_full, (
"--enable-page-major-kv-layout requires the Triton attention backend "
f"for the full-attention layers; got {sorted(backends)}. Pass "
"--attention-backend triton."
"for the full-attention layers (unified-memory MLA also allows the "
f"paged MLA backends); got {sorted(backends)}, allowed "
f"{sorted(allowed_full)}. Pass a compatible --attention-backend."
)
# The Mamba state is stored in envelope-strided views; only the
# stride-aware Triton causal-conv / SSM kernels read them correctly.
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()