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

feat(unified-memory): read unified pool from attention backends fa3/flashinfer/trtllm_mha/flashmla (#34613)

Browse Source 
Co-authored-by: Caihua Li <caihua.li@bytedance.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: Cheng Wan <cheng.wan@radixark.ai>
This commit is contained in:
caihuali95
2026-08-30 23:58:24 -07:00
committed by GitHub
co-authored by Caihua Li Claude Fable 5 Cheng Wan
parent 29578d5578
commit 8bb776dc48
31 changed files with 1182 additions and 757 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/attention/test_normal_decode_set_metadata.py
+5 -6
View File
@@ -31,7 +31,6 @@ def reference_normal_decode_set_metadata(
page_table: torch.Tensor,
req_to_token: torch.Tensor,
req_pool_indices: torch.Tensor,
strided_indices: torch.Tensor,
max_seq_pages: int,
seq_lens: torch.Tensor,
seq_len_delta: int,
@@ -45,6 +44,11 @@ def reference_normal_decode_set_metadata(
"""
cache_seqlens_int32.copy_(seq_lens + seq_len_delta)
cu_seqlens_k[1:].copy_(torch.cumsum(cache_seqlens_int32, dim=0, dtype=torch.int32))
# Page-start columns, derived internally (the wrapper's dead
# strided_indices parameter was removed alongside its v2p args).
strided_indices = torch.arange(
0, req_to_token.shape[1], page_size, device=req_to_token.device
)
page_indices = req_to_token[
req_pool_indices[:, None],
strided_indices[:max_seq_pages][None, :],
@@ -213,7 +217,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
ref_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
test_data["seq_len_delta"],
@@ -229,7 +232,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
test_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
test_data["seq_len_delta"],
@@ -351,7 +353,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
test_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
test_data["seq_len_delta"],
@@ -408,7 +409,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
ref_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
0,
@@ -423,7 +423,6 @@ class TestNormalDecodeSetMetadata(CustomTestCase):
test_data["page_table"],
test_data["req_to_token"],
test_data["req_pool_indices"],
test_data["strided_indices"],
test_data["max_seq_pages"],
test_data["seq_lens"],
0,
test/registered/attention/test_trtllm_mha_graph_metadata.py
+72
View File
@@ -30,6 +30,8 @@ PAGE_SIZE = 128
def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
backend = TRTLLMHAAttnBackend.__new__(TRTLLMHAAttnBackend)
backend.device = torch.device("cpu")
backend.max_context_len = 1024
@@ -45,6 +47,15 @@ def _make_backend_for_hook_test(speculative_num_draft_tokens=None):
backend.decode_cuda_graph_metadata = {}
backend.target_verify_metadata = {}
backend.draft_extend_metadata = {}
# Passthrough source (static pool): every unified-arm branch stays off,
# matching the real __init__'s parent binding.
backend.kv_index_translator = KVIndexTranslator(
req_to_token=backend.req_to_token,
token_to_kv_pool_allocator=SimpleNamespace(),
token_to_kv_pool=SimpleNamespace(),
page_size=PAGE_SIZE,
device="cpu",
)
backend.init_cuda_graph_state(max_bs=4, max_num_tokens=16)
return backend
@@ -543,6 +554,67 @@ def test_metadata_correctness(bs, seqlen_offset, q_mode, with_swa, static_width)
torch.testing.assert_close(swa_out_cache_loc, out_ref, rtol=0, atol=0)
@pytest.mark.parametrize("pass_tables", [False, True])
def test_skip_page_table_updates_seqlens_only(pass_tables):
"""The unified-memory arm: skip_page_table=True must still rebuild the
seqlen metadata in-graph but leave every page-table byte alone -- the bound
tables are capture-stable read tables the translator refreshes
out-of-graph, and an in-graph write would clobber them with virtual-derived
pages. Covers both call shapes: page_table=None (what the backend passes)
and a real sentinel-filled table (pins that the writes are compiled out,
not just unpassed)."""
if not torch.cuda.is_available():
pytest.skip("CUDA required")
bs, seqlen_offset, seed = 5, 1, 4242
pool_size, max_num_pages = 64, 16
seq_max = (max_num_pages - 2) * PAGE_SIZE
(
req_to_token,
req_pool_indices,
seq_lens,
_stride,
_cap,
) = _build_inputs(bs, pool_size, max_num_pages, None, seq_max, seed)
cache_seqlens = torch.zeros(bs, dtype=torch.int32, device=DEVICE)
cu_seqlens_k = torch.zeros(bs + 1, dtype=torch.int32, device=DEVICE)
sentinel_pt = None
sentinel_swa = None
if pass_tables:
sentinel_pt = torch.full(
(bs, max_num_pages), 777, dtype=torch.int32, device=DEVICE
)
sentinel_swa = torch.full(
(bs, max_num_pages), 888, dtype=torch.int32, device=DEVICE
)
update_trtllm_mha_graph_metadata(
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
req_to_token=req_to_token,
cache_seqlens=cache_seqlens,
cu_seqlens_k=cu_seqlens_k,
page_table=sentinel_pt,
bs=bs,
seqlen_offset=seqlen_offset,
max_seq_pages=max_num_pages,
page_size=PAGE_SIZE,
swa_page_table=sentinel_swa,
skip_page_table=True,
)
torch.cuda.synchronize()
cache_seqlens_ref = _ref_cache_seqlens(seq_lens, seqlen_offset)
torch.testing.assert_close(cache_seqlens, cache_seqlens_ref, rtol=0, atol=0)
cu_k_ref = torch.zeros(bs + 1, dtype=torch.int32, device=DEVICE)
cu_k_ref[1:] = torch.cumsum(cache_seqlens_ref, dim=0, dtype=torch.int32)
torch.testing.assert_close(cu_seqlens_k, cu_k_ref, rtol=0, atol=0)
if pass_tables:
assert bool((sentinel_pt == 777).all()), "page_table written despite skip"
assert bool((sentinel_swa == 888).all()), "swa_page_table written despite skip"
def test_bs_zero_noop():
if not torch.cuda.is_available():
pytest.skip("CUDA required")
test/registered/models_e2e/test_kimi_linear_unified_memory.py
+15 -1
View File
@@ -32,7 +32,7 @@ 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=570, stage="nightly", runner_config="4-gpu-h100")
register_cuda_ci(est_time=1200, stage="nightly", runner_config="4-gpu-h100")
KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
@@ -55,5 +55,19 @@ class TestKimiLinearUnifiedMemory(
]
class TestKimiLinearUnifiedMemoryFlashMLA(TestKimiLinearUnifiedMemory):
"""flashmla at its ps=64 snap: the canonical block-table route
(KVIndexTranslator.build_into into flashmla's padded tables) plus the ps=64
sub-pool sizing (64-token sink floor, dense-view tail pad) end to end.
Hopper-only, like the rest of this nightly suite."""
other_args = TestKimiLinearUnifiedMemory.other_args + [
"--attention-backend",
"flashmla",
"--page-size",
"64",
]
if __name__ == "__main__":
unittest.main()
test/registered/page_major/test_page_major_gpt_oss.py
+21 -10
View File
@@ -1,15 +1,12 @@
"""
End-to-end accuracy test for the unified memory pool on a hybrid-SWA MoE model.
"""Unified memory pool on a hybrid-SWA MoE model, across the backend matrix.
Launches gpt-oss-20b with ``--enable-unified-memory`` on the Triton attention
backend and checks that GSM8K accuracy holds. This exercises the SWA +
full-attention KV sub-pools stored as per-layer views in the unified
page-major envelope.
gpt-oss-20b is uniform-row hybrid-SWA, so its MHA and SWA sub-pools are
per-layer views and the fa3 cell reads them through the translator's read
tables. The resolved-default cell pins the no-pin path, since a pinned
backend hides default-resolution breakage by construction. flashinfer is
absent on purpose: gpt-oss uses attention sinks, which it does not support.
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
Usage:
python3 -m unittest test_page_major_gpt_oss
"""
import unittest
@@ -20,7 +17,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST_MXFP4_WITH_MOE
register_cuda_ci(est_time=420, stage="extra-a", runner_config="1-gpu-large")
register_cuda_ci(est_time=1500, stage="extra-a", runner_config="1-gpu-large")
_UNIFIED_COMMON_ARGS = [
"--enable-unified-memory",
@@ -64,5 +61,19 @@ class TestUnifiedGptOssTriton(DefaultServerBase):
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
class TestUnifiedGptOssFa3(TestUnifiedGptOssTriton):
"""fa3 pinned: the per-layer views read through the translator's read
tables (eager direct-bind + captured fused copy)."""
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
class TestUnifiedGptOssResolvedDefault(TestUnifiedGptOssTriton):
"""No backend pin: whatever the host resolves must be in the allow-list,
or the server fails to boot under its own defaults."""
other_args = _UNIFIED_COMMON_ARGS
if __name__ == "__main__":
unittest.main()
test/registered/page_major/test_page_major_qwen_hybrid.py
+29 -12
View File
@@ -1,17 +1,14 @@
"""
End-to-end accuracy test for the unified memory pool on a GDN-hybrid model.
"""Unified memory pool on a GDN-hybrid model, across the backend matrix.
Launches Qwen3.5-4B (a gated-delta-net / linear-attention hybrid) with
``--enable-unified-memory`` on the Triton attention + linear-attn + Mamba
backends and checks that GSM8K accuracy holds. This exercises the unified
envelope's most bug-prone path: the Mamba conv/SSM state stored as a strided
envelope view, plus the full-attention KV stored as per-layer views,
both read/written by the GDN prefill and decode kernels.
Qwen3.5-4B is a gated-delta-net / linear-attention hybrid, which exercises the
path most prone to subtle bugs: the Mamba conv/SSM state stays a strided
envelope view (its kernels are stride-aware by design) while the
full-attention KV is per-layer views, which the fa3 / flashinfer cells read
through the translator's read tables. The resolved-default cell pins the
no-pin path, since a pinned backend hides default-resolution breakage by
construction.
Registered to the label-gated ``run-ci-extra`` suite (opt-in, not per-commit).
Usage:
python3 -m unittest test_page_major_qwen_hybrid
"""
import unittest
@@ -22,7 +19,7 @@ from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.server_fixtures.default_fixture import DefaultServerBase
from sglang.test.test_utils import DEFAULT_HYBRID_GDN_SMALL_MODEL_NAME_FOR_TEST
register_cuda_ci(est_time=300, stage="extra-a", runner_config="1-gpu-large")
register_cuda_ci(est_time=1600, stage="extra-a", runner_config="1-gpu-large")
_UNIFIED_COMMON_ARGS = [
"--trust-remote-code",
@@ -74,5 +71,25 @@ class TestUnifiedQwenHybridTriton(DefaultServerBase):
self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
class TestUnifiedQwenHybridFa3(TestUnifiedQwenHybridTriton):
"""fa3 pinned: read tables, eager direct-bind + captured fused copy."""
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "fa3"]
class TestUnifiedQwenHybridFlashinfer(TestUnifiedQwenHybridTriton):
"""flashinfer pinned: token ids reconstructed from the read table by the
ENTRY_PAGE_SIZE CSR builder."""
other_args = _UNIFIED_COMMON_ARGS + ["--attention-backend", "flashinfer"]
class TestUnifiedQwenHybridResolvedDefault(TestUnifiedQwenHybridTriton):
"""No backend pin: whatever the host resolves must be in the allow-list,
or the server fails to boot under its own defaults."""
other_args = _UNIFIED_COMMON_ARGS
if __name__ == "__main__":
unittest.main()
test/registered/unit/layers/attention/test_flashattention_graph_metadata.py
+10
View File
@@ -4,6 +4,7 @@ from types import SimpleNamespace
import torch
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -46,6 +47,15 @@ class TestFlashAttentionGraphMetadata(CustomTestCase):
backend.req_to_token_pool = SimpleNamespace(
req_to_token=torch.zeros((1, 16), dtype=torch.int32)
)
# A real source over the stub pool: the probe disables it, giving the
# backend the strict passthrough view it now reads its tables from.
backend.kv_index_translator = KVIndexTranslator(
req_to_token=backend.req_to_token_pool.req_to_token,
token_to_kv_pool_allocator=SimpleNamespace(),
token_to_kv_pool=SimpleNamespace(),
page_size=1,
device="cpu",
)
backend.is_prefill_aware_swa = False
backend.has_swa = False
backend.use_sliding_window_kv_pool = False
test/registered/unit/layers/attention/test_flashattention_pa_swa_prefill_lens_size.py
+14 -4
View File
@@ -24,6 +24,7 @@ import torch
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
from sglang.srt.mem_cache.kv_index_translator import KVIndexTranslator
from sglang.srt.runtime_context import get_context
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.test_utils import CustomTestCase
@@ -62,15 +63,24 @@ def _make_prefill_aware_swa_runner(*, pool_size: int, max_context_len: int = 64)
enable_prefill_cp=False,
enable_dp_attention=False,
)
token_to_kv_pool = object()
token_to_kv_pool_allocator = object()
return SimpleNamespace(
sliding_window_size=None,
model_config=model_config,
device=device,
req_to_token_pool=req_to_token_pool,
token_to_kv_pool=object(), # not a SWAKVPool instance -> use_sliding_window_kv_pool=False
# getattr(..., "full_v2p_page_table", None) is None -> unified_mla_hooks
# falls back to the static (disabled) hook set.
token_to_kv_pool_allocator=object(),
token_to_kv_pool=token_to_kv_pool, # not a SWAKVPool -> use_sliding_window_kv_pool=False
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
# A real KVIndexTranslator over this non-unified pool: the probe finds no
# unified composite, so it is the strict passthrough the backend reads.
kv_index_translator=KVIndexTranslator(
req_to_token=req_to_token_pool.req_to_token,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
token_to_kv_pool=token_to_kv_pool,
page_size=1,
device=device,
),
kv_cache_dtype=torch.float16,
kv_cache_dtype_str="auto",
page_size=1,
test/registered/unit/layers/attention/test_kv_translate_ownership.py
+83
View File
@@ -0,0 +1,83 @@
"""Nothing under layers/attention may translate KV ids for itself.
Ownership is exactly two places: `KVIndexTranslator` for READS (indices are
born kernel-facing, backends consume its tables) and the ForwardBatch rebind
(`rebind_write_loc`) for WRITES. Virtual and physical ids share a value range,
so a backend that forgets a translate -- or does one twice -- reads the wrong
rows and nothing crashes. This scan makes both unrepresentable.
Out of scope, deliberately: the allocator-internal implementations
(`multi_ended_allocator` / `unified_memory_pool`), which ARE the mechanism the
translator calls; the PD transfer plane's `translate_kv_indices_for_transfer`,
which stages for RDMA outside the forward path; and the STATIC SWA pool's
legacy full->swa slot map, a different mapping kind with no virtual/physical
ambiguity -- its call sites are count-pinned below so new ones are added
consciously.
python3 -m pytest test/registered/unit/layers/attention/test_kv_translate_ownership.py -v
"""
import os
import re
import unittest
from sglang.srt.layers.attention import triton_backend as _anchor_module
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
# The attention package is a namespace package (no __init__), so anchor the
# scan on a concrete module inside it.
_ATTN_DIR = os.path.dirname(os.path.abspath(_anchor_module.__file__))
def _iter_sources():
for root, _dirs, files in os.walk(_ATTN_DIR):
for name in sorted(files):
if not name.endswith(".py"):
continue
path = os.path.join(root, name)
with open(path, encoding="utf-8") as fh:
yield os.path.relpath(path, _ATTN_DIR), fh.read()
class TestUnifiedTranslateBanned(CustomTestCase):
def test_no_unified_translate_calls(self):
"""No backend calls the unified translate surfaces. A hit here means
a backend re-grew its own id-space transition -- the design whose two
failure modes (forgotten translate, duplicated translate) this scan
exists to prevent. Route reads through KVIndexTranslator views and
writes through the ForwardBatch rebind instead."""
banned = re.compile(r"\.translate_kv_loc(_kernel_id)?\(")
hits = [
f"{rel}: {m.group(0)}"
for rel, src in _iter_sources()
for m in banned.finditer(src)
]
self.assertEqual(hits, [])
def test_no_translate_capability_probing(self):
"""No backend probes an allocator for translate capability -- the
getattr-hook pattern is how per-backend translation grew the first
time."""
probing = re.compile(r"""getattr\([^)]*['"]translate_kv_loc""")
hits = [rel for rel, src in _iter_sources() if probing.search(src)]
self.assertEqual(hits, [])
def test_hooks_module_deleted_and_unimported(self):
"""The per-backend hooks module (the previous owner of backend-side
v2p knowledge) stays deleted, and nothing imports it."""
self.assertFalse(
os.path.exists(os.path.join(_ATTN_DIR, "unified_mem_hooks.py"))
)
hits = [
rel
for rel, src in _iter_sources()
if "unified_mem_hooks" in src or "unified_mla_hooks" in src
]
self.assertEqual(hits, [])
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_full_loc_fast_path.py
+38 -38
View File
@@ -59,9 +59,11 @@ class _RecordingPool:
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."""
"""`UnifiedSWAKVPool.set_kv_buffer` routing: full layers write `full_loc`
when present (triton's capture-stable buffer), else the rebound generic
`loc` -- the same id space once the loc is rebound; SWA layers write the swa-physical
`swa_loc`, which has no fallback (a different id space). The pool never
translates."""
def _make_bare_pool(self):
from sglang.srt.mem_cache.unified_memory_pool import UnifiedSWAKVPool
@@ -96,21 +98,29 @@ class TestUnifiedSWARouting(unittest.TestCase):
self.assertIsNot(forwarded, virtual_loc)
self.assertNotIn("already_physical", kwargs)
def test_full_layer_requires_full_loc(self):
def test_full_layer_falls_back_to_generic_loc(self):
"""Bug regression: fa3 x unified-SWA crashed at gpt-oss
cuda-graph capture because every backend except triton bundles the
2-arg KVWriteLoc(loc, swa) and the full-layer door demanded an explicit
full_loc. Once the loc is rebound the generic `loc` IS the full-side kernel-facing id
(rebind_write_loc runs at ForwardBatch construction),
so the door must fall back to it -- the pool still never translates."""
pool = self._make_bare_pool()
virtual_loc = torch.tensor([10, 11, 12], dtype=torch.int64)
rebound_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),
)
pool.set_kv_buffer(
layer,
_loc_info(rebound_loc, swa_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, rebound_loc)
self.assertNotIn("already_physical", kwargs)
def test_swa_layer_writes_swa_loc(self):
pool = self._make_bare_pool()
@@ -264,19 +274,17 @@ class TestHybridLinearMLARouting(unittest.TestCase):
- `set_kv_buffer` (MLA branch) mirrors the MHA branch — write the
pre-translated `KVWriteLoc.full_loc` when present (unified pool, where it
carries the DENSE loc), else the raw `loc` (static pool, already physical).
- `set_mla_kv_buffer` forwards `loc` untouched (kernel-facing since the
ForwardBatch rebind); `get_mla_kv_buffer` applies `_full_translate`
exactly once (its indices are req_to_token-produced, virtual under the
unified pool)."""
- `set_mla_kv_buffer` / `get_mla_kv_buffer` forward `loc` untouched:
writes are kernel-facing since the ForwardBatch rebind, and read
indices are translated at their production sites."""
def _make_bare_pool(self, translate=None):
def _make_bare_pool(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
pool = object.__new__(HybridLinearKVPool)
pool.full_kv_pool = _RecordingMLAPool()
pool.use_mla = True
pool.full_attention_layer_id_mapping = {0: 0}
pool._full_translate = translate if translate is not None else (lambda ids: ids)
return pool
def test_mla_writes_full_loc_from_write_loc(self):
@@ -328,28 +336,20 @@ class TestHybridLinearMLARouting(unittest.TestCase):
self.assertEqual(len(pool.full_kv_pool.mla_set_calls), 1)
self.assertIs(pool.full_kv_pool.mla_set_calls[0], loc)
def test_get_mla_kv_buffer_translates_exactly_once(self):
"""READ door: `loc` is produced from req_to_token (VIRTUAL under the
unified pool), so the get side still translates here — exactly once.
The WRITE door (case above) never translates: the split is the write
flip's contract."""
calls = []
def translate(ids):
calls.append(ids)
return ids + 100
pool = self._make_bare_pool(translate=translate)
virtual_loc = torch.tensor([4, 5], dtype=torch.int64)
def test_get_mla_kv_buffer_door_never_translates(self):
"""Kernel-facing contract, read side: `loc` is a read-index tensor
already translated at its production site
(fetch_mha_one_shot_kv_indices / prepare_chunked_kv_indices); the
door forwards it UNTOUCHED -- a re-added door translate would
double-translate every unified MLA prefix read."""
pool = self._make_bare_pool()
loc = torch.tensor([104, 105], dtype=torch.int64)
layer = types.SimpleNamespace(layer_id=0)
pool.get_mla_kv_buffer(layer, virtual_loc)
pool.get_mla_kv_buffer(layer, loc)
self.assertEqual(len(calls), 1)
self.assertEqual(len(pool.full_kv_pool.mla_get_calls), 1)
self.assertTrue(
torch.all(pool.full_kv_pool.mla_get_calls[0] == virtual_loc + 100)
)
self.assertIs(pool.full_kv_pool.mla_get_calls[0], loc)
if __name__ == "__main__":
test/registered/unit/mem_cache/test_multi_ended_allocator.py
+76
View File
@@ -31,11 +31,13 @@ import torch
from sglang.srt.mem_cache.multi_ended_allocator import (
MultiEndedAllocator,
UnifiedMambaTokenToKVPoolAllocator,
UnifiedSWATokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.unified_memory_pool import (
MambaSubPoolSpec,
MHASubPoolSpec,
MLASubPoolSpec,
UnifiedKVPool,
)
@@ -2746,5 +2748,79 @@ class TestSWACompositeDenseSurface(unittest.TestCase):
self.assertTrue(bool((got[in_tomb] == 0).all().item()))
class TestPs64MLACompositeFeasibility(unittest.TestCase):
"""The Kimi/flashmla shape: MLA + mamba composite at page_size=64 (the
flashmla arg snap). Large pages stress every sizing derivation at once —
the 64-token sink-page floor, the ps*entry_bytes per-layer-view tail pad, and
the page-granular alloc — so this pins that the factory-shaped
construction stays FEASIBLE and the dense surface stays on-formula when
the page size jumps from the usual 1..4 to 64."""
PS = 64
LAYERS = 3
def _build(self):
full = MLASubPoolSpec(
name="full",
layer_num=self.LAYERS,
kv_lora_rank=64,
qk_rope_head_dim=16,
store_dtype=torch.float16,
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",
)
n_full = 8 * self.PS # 8 pages incl. the sink page
total = n_full * full.entry_bytes() + 16 * mamba.entry_bytes()
pool = UnifiedKVPool(
total_bytes=total,
sub_pool_specs=[full, mamba],
device=_DEV,
enable_memory_saver=False,
page_size=self.PS,
)
full_kv = _FakeKVCache(pool.max_slots("full"))
full_kv.attach_allocator = lambda allocator: None
mamba_kv = _FakeKVCache(pool.max_slots("mamba"))
mamba_kv.attach_allocator = lambda allocator: None
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=self.PS,
need_sort=False,
forward_stream=None,
)
def test_construction_alloc_and_dense_formula(self):
a = self._build()
# MLA: one latent row per layer, so the spec reports LAYERS blocks.
self.assertEqual(a.kernel_page_multiplier, self.LAYERS)
v = a.alloc(2 * self.PS)
self.assertIsNotNone(v, "2-page alloc infeasible at ps=64")
# Page-aligned virtual run (page-granular allocator invariant).
self.assertEqual(int(v[0].item()) % self.PS, 0)
# Dense translate follows the affine formula at ps=64, and every id
# fits int32 (the canonical narrows on store).
v2p = a.full_v2p_page_table
want = v2p[v // self.PS] * (self.PS * self.LAYERS) + v % self.PS
got = a.translate_kv_loc_for_kernel(v)
self.assertTrue(torch.equal(got, want), "kernel-facing formula broke at ps=64")
self.assertTrue(bool((got < 2**31).all().item()))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_mla_block_table.py
+106 -227
View File
@@ -20,20 +20,26 @@ block table filled with kernel-facing page ids:
dense_page(virtual_page) = v2p[virtual_page] * layer_num
Three backend families reach that same formula by different routes:
- `create_flashmla_kv_indices_triton` in-kernel via `v2p_ptr` / `PAGE_MULT`
(trtllm_mla / cutedsl_mla / tokenspeed_mla);
- the flashinfer_mla updaters, post-gathering `translate_kv_loc_for_kernel` over the
token-level kv_indices;
- `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table.
Since the read-path translator, ONE builder computes that formula for every
family — `build_kv_read_table` (the canonical) — and the backends only
differ in how they consume it:
- trtllm_mla / cutedsl_mla / tokenspeed_mla / flashmla: rows filled straight
into their padded block tables (`KVIndexTranslator.build_into`, prefix-only so
the backends' own -1 / stale tail sentinels survive);
- the flashinfer updaters: token ids reconstructed from the canonical by
`create_flashinfer_kv_indices_triton[ENTRY_PAGE_SIZE=ps]`;
- fa3's captured decode: `normal_decode_set_metadata` copies the canonical
rows' live prefixes (src_is_read_table=True).
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 kernel-facing translate the flashinfer updaters apply agrees with the
page-level block table the trtllm path builds;
- the static `create_flashmla_kv_indices_triton` (no id-space knowledge left)
still matches the plain token//ps reference;
- the canonical route against the python dense reference, for several page
sizes, ragged sequence lengths and a non-identity v2p permutation;
- lanes past a row's live prefix keep the backend's -1 sentinel (prefix-only
discipline — the trtllm/flashmla tail contract);
- the token-level kernel-facing translate the flashinfer updaters used to apply
agrees with the canonical page table (page-affinity of the id space);
- fa3's fused metadata kernels agree with the same reference, on both the
page_size == 1 fast path (which is what Kimi-Linear takes: fa3 imposes no
page-size constraint) and the general path.
@@ -57,6 +63,28 @@ _LAYERS = 24 # K3 MLA full-attention layer count
def _fill_block_table(
req_to_token, req_pool_indices, seq_lens, page_size, *, v2p, mult
):
"""The unified route: canonical builder into a -1-filled block table
(exactly what KVIndexTranslator.build_into does for trtllm_mla/flashmla)."""
from sglang.kernels.ops.kvcache.kv_read_table import build_kv_read_table
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)
build_kv_read_table(
req_to_token=req_to_token,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens.to(torch.int64),
v2p=v2p,
multiplier=mult,
page_size=page_size,
max_pages=max_blocks,
out=out,
)
return out
def _fill_block_table_static(req_to_token, req_pool_indices, seq_lens, page_size):
"""The static-pool route: the stripped flashmla kernel, token//ps verbatim."""
from sglang.kernels.ops.kvcache.kv_indices import (
create_flashmla_kv_indices_triton,
get_num_kv_index_blocks_flashmla,
@@ -75,8 +103,6 @@ def _fill_block_table(
req_to_token.stride(0),
max_blocks,
PAGED_SIZE=page_size,
v2p_ptr=v2p,
PAGE_MULT=mult,
)
return out
@@ -122,11 +148,12 @@ class TestDenseBlockTable(unittest.TestCase):
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."""
def test_static_kernel_matches_reference(self):
"""The stripped (id-space-free) flashmla kernel is byte-identical to the
plain token//ps reference -- guards the v2p-arg removal itself."""
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)
got = _fill_block_table_static(rt, rpi, sl, page_size)
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}"
@@ -166,8 +193,9 @@ class TestDenseBlockTable(unittest.TestCase):
)
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)."""
"""Lanes past a request's page count keep the -1 fill: the prefix-only
canonical build must never write a backend's tail sentinel (the
trtllm/flashmla block-table contract)."""
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)
@@ -207,43 +235,77 @@ class TestDenseBlockTable(unittest.TestCase):
@unittest.skipUnless(_HAS_CUDA, "requires CUDA")
class TestFa3MetadataDenseBlockTable(unittest.TestCase):
"""fa3 folds the unified remap into `normal_decode_set_metadata`, the fused
gather that writes its captured-decode page table, so the kernel itself has
to get the mapping right. Two kernels back it: a page_size == 1 / no-SWA fast
path (what Kimi-Linear takes, since fa3 imposes no page-size constraint) and
a general one.
"""fa3's captured-decode page table is written by `normal_decode_set_metadata`
fed with the translator's read table kernel page table
(src_is_read_table=True): the fused kernel copies the canonical
rows' live prefixes into the capture-stable buffer. Pinned END-TO-END:
build_kv_read_table -> wrapper -> page_table must equal the python
reference of the kernel-facing formula, on both the page_size == 1 / no-SWA fast
path (what Kimi-Linear takes) and the general kernel. The static call
(no source flag) stays byte-identical to the pre-translator kernel.
"""
def _run(self, page_size, *, v2p, mult, bs=5, max_ctx=2048):
from sglang.kernels.ops.attention.metadata import normal_decode_set_metadata
from sglang.kernels.ops.kvcache.kv_read_table import (
build_kv_read_table,
)
maker = TestDenseBlockTable._make_batch
rt, rpi, sl, v2p_full = maker(self, page_size, bs=bs, max_ctx=max_ctx)
v2p_arg = v2p_full if v2p else None
max_pages = (max_ctx + page_size - 1) // page_size
page_table = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
cache_seqlens = torch.zeros((bs,), dtype=torch.int32, device=_DEV)
cu_seqlens_k = torch.zeros((bs + 1,), dtype=torch.int32, device=_DEV)
strided = torch.arange(0, max_ctx, page_size, device=_DEV)
max_seq_pages = (int(sl.max().item()) + page_size - 1) // page_size
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
rt,
rpi,
strided,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
v2p_page_table=v2p_arg,
kernel_page_multiplier=mult,
)
if v2p:
# The translator's canonical, then the wrapper copies its rows.
canonical = torch.zeros((bs, max_pages), dtype=torch.int32, device=_DEV)
build_kv_read_table(
req_to_token=rt,
req_pool_indices=rpi,
seq_lens=sl.to(torch.int64),
v2p=v2p_full,
multiplier=mult,
page_size=page_size,
max_pages=max_pages,
out=canonical,
)
rows = torch.arange(bs, dtype=torch.int64, device=_DEV)
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
canonical,
rows,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
None,
None,
src_is_read_table=True,
)
else:
normal_decode_set_metadata(
cache_seqlens,
cu_seqlens_k,
page_table,
rt,
rpi,
max_seq_pages,
sl.to(torch.int64),
0,
page_size,
None,
None,
)
torch.cuda.synchronize()
want = _reference(rt, rpi, sl, page_size, v2p=v2p_arg, mult=mult)
want = _reference(
rt, rpi, sl, page_size, v2p=(v2p_full if v2p else None), mult=mult
)
return page_table, want, sl
def _assert_live_prefix(self, got, want, sl, page_size):
@@ -287,192 +349,9 @@ class TestFa3MetadataDenseBlockTable(unittest.TestCase):
"test batch degenerated: v2p is the identity on the pages used",
)
def test_agrees_with_flashmla_block_table(self):
"""fa3 and trtllm_mla build the same table two different ways; a
disagreement means one family is addressing the wrong pages."""
for page_size in (1, 64):
got, _, sl = self._run(page_size, v2p=True, mult=_LAYERS)
rt, rpi, sl2, v2p = TestDenseBlockTable._make_batch(self, page_size)
other = _fill_block_table(
rt, rpi, sl2, page_size, v2p=v2p, mult=_LAYERS
).long()
for r in range(got.shape[0]):
n_pages = (int(sl[r].item()) + page_size - 1) // page_size
self.assertTrue(
torch.equal(got[r, :n_pages].long(), other[r, :n_pages]),
f"fa3 and flashmla block tables disagree (row {r}, ps={page_size})",
)
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.unified_mem_hooks 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_for_kernel)
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_for_kernel=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_for_kernel)
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_for_kernel=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_for_kernel)
# Multiplier stays 1: kernel-facing 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,
)
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 kernel-facing 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_for_kernel(valid))
expected = alloc.translate_kv_loc_for_kernel(virt)
self.assertEqual(buf.dtype, torch.int32)
self.assertTrue(
torch.equal(buf[:n].long(), expected),
"in-place translate did not land kernel-facing 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_for_kernel(virt), virt),
"kernel-facing ids coincide with virtual ids; pick a different allocation",
)
# (The old fa3<->flashmla agreement case is gone: both families now
# consume the SAME canonical builder, so cross-family agreement holds by
# construction and the per-family cases above cover the two consumers.)
if __name__ == "__main__":
test/registered/unit/model_executor/test_unified_out_cache_loc_rebind.py
+100
View File
@@ -27,6 +27,7 @@ import inspect
import textwrap
import unittest
from types import SimpleNamespace
from unittest.mock import create_autospec
import torch
@@ -177,5 +178,104 @@ class TestPadComposesWithDerivation(CustomTestCase):
self.assertEqual(src._swa_write_loc_unified(fb.out_cache_loc).numel(), 0)
class TestReadRailTranslatesAtProduction(CustomTestCase):
"""The model-door READ indices (req_to_token-derived, VIRTUAL under the
unified pool) are translated at their PRODUCTION site -- the cache then
holds the kernel-facing result and the pool door never translates."""
def _fb_for_one_shot(self):
fb = _make_fb(torch.tensor([1, 2], dtype=torch.int64))
fb.batch_size = 2
fb.seq_lens = torch.tensor([2, 3], dtype=torch.int64)
fb.seq_lens_cpu = torch.tensor([2, 3], dtype=torch.int32)
fb.req_pool_indices = torch.tensor([0, 1], dtype=torch.int64)
return fb
def test_one_shot_indices_translated_once_and_cached(self):
from unittest.mock import patch
from sglang.srt.model_executor import forward_batch_deepseek_mha_mixin as mix
calls = []
sentinel = torch.arange(5, dtype=torch.int64) + 5000
def translate(t):
calls.append(t)
return sentinel
fb = self._fb_for_one_shot()
fake_pool = SimpleNamespace(
req_to_token=torch.zeros((4, 16), dtype=torch.int32)
)
# autospec, not a bare namespace: setting a name the translator does
# not have raises, so renaming the method breaks this test loudly.
fake_translator = create_autospec(KVIndexTranslator, instance=True)
fake_translator.translate_full_attn_ids = translate
fake_backend = SimpleNamespace(kv_index_translator=fake_translator)
with (
patch.object(mix, "get_req_to_token_pool", return_value=fake_pool),
patch.object(mix, "get_attn_backend", return_value=fake_backend),
patch.object(mix, "create_flashinfer_kv_indices_triton"),
):
r1 = fb.fetch_mha_one_shot_kv_indices()
r2 = fb.fetch_mha_one_shot_kv_indices()
self.assertIs(r1, sentinel) # production site translated
self.assertIs(r2, sentinel) # cache holds the TRANSLATED result
self.assertEqual(len(calls), 1) # translated exactly once
self.assertEqual(calls[0].dtype, torch.int32) # raw producer output
def test_one_shot_indices_noop_on_unmigrated_backend(self):
from unittest.mock import patch
from sglang.srt.model_executor import forward_batch_deepseek_mha_mixin as mix
fb = self._fb_for_one_shot()
fake_pool = SimpleNamespace(
req_to_token=torch.zeros((4, 16), dtype=torch.int32)
)
# A backend that never set the attribute inherits the base-class None.
fake_backend = SimpleNamespace(kv_index_translator=None)
with (
patch.object(mix, "get_req_to_token_pool", return_value=fake_pool),
patch.object(mix, "get_attn_backend", return_value=fake_backend),
patch.object(mix, "create_flashinfer_kv_indices_triton"),
):
r = fb.fetch_mha_one_shot_kv_indices()
# The raw int32 producer output passes through untouched.
self.assertEqual(r.dtype, torch.int32)
def test_get_mla_kv_buffer_door_passes_loc_untranslated(self):
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool
recorded = {}
class _RecordingLeafPool:
def get_mla_kv_buffer(self, layer, loc, dst_dtype):
recorded["loc"] = loc
return None, None
def get_kv_size_bytes(self):
return 0
pool = HybridLinearKVPool(
size=16,
dtype=torch.float16,
page_size=1,
head_num=1,
head_dim=8,
full_attention_layer_ids=[0],
device=_DEV,
mamba_pool=SimpleNamespace(get_size_per_token=lambda: 0),
enable_memory_saver=False,
use_mla=True,
start_layer=0,
full_kv_pool=_RecordingLeafPool(),
)
loc = torch.tensor([9, 10], dtype=torch.int64)
pool.get_mla_kv_buffer(SimpleNamespace(layer_id=0), loc, torch.float16)
self.assertIs(recorded["loc"], loc)
if __name__ == "__main__":
unittest.main()
test/registered/unit/server_args/test_page_major_backend_allowlist.py
+57 -27
View File
@@ -13,18 +13,28 @@
# ==============================================================================
"""`--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 contiguous view (`build_mla_views`), so the paged MLA
backends can read it directly once their kv_indices / block tables are remapped
to kernel-facing ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its
`cutedsl_mla` / `tokenspeed_mla` subclasses.
Two-way gate (see `_handle_page_major_kv_layout`), because the unified pool
exposes per-layer views and nothing else:
* unified-memory MLA models (`build_mla_views`) allow the whole wired
paged MLA family -- `fa3`, `flashinfer`'s MLA backend, `trtllm_mla` with
its `cutedsl_mla` / `tokenspeed_mla` subclasses, and `flashmla` (ps=64
snap);
* unified-memory MHA/SWA models (`build_mha_views`) allow `fa3` /
`fa4` / `flashinfer` / `trtllm_mha` alongside Triton;
* plain `--enable-page-major-kv-layout` without the unified pool keeps the
envelope-strided 4-D views only the stride-aware Triton kernels read.
Pinned here so the exception cannot silently widen to a backend that has no
dense-id remapping (`flashmla`, `cutlass_mla`, ...) or leak into the MHA path.
`fa3` matters most: it is the resolved default on pre-Blackwell hosts, so it is
the one entry whose absence used to make `--enable-unified-memory` fail to boot
under its own default configuration.
The same handler also screens the pool itself: the dense MHA/SWA views need
uniform K/V rows, so an asymmetric-K/V model (MiMoV2: head_dim 192 !=
v_head_dim 128) cannot run `--enable-unified-memory` at all and is rejected on
EVERY backend, Triton included. MLA models are exempt -- their sub-pool keeps
one latent row per layer, and several MLA configs (Kimi-Linear: head_dim 72,
v_head_dim 128) report asymmetric dims while running the unified pool today.
Pinned here so no arm silently widens to an unwired backend (`cutlass_mla`,
`aiter`) and no arm silently narrows: `fa3` is the resolved default on
pre-Blackwell hosts, so its absence from an arm makes `--enable-unified-memory`
fail to boot under its own default configuration.
python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v
"""
@@ -96,13 +106,23 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
"flashinfer",
"cutedsl_mla",
"tokenspeed_mla",
"flashmla",
)
# No dense-id remapping: must stay rejected until they get one.
UNWIRED_BACKENDS = ("flashmla", "cutlass_mla", "trtllm_mha", "aiter")
# Wired for the dense per-layer MHA/SWA views (uniform-row models).
DENSE_MHA_BACKENDS = ("fa3", "fa4", "flashinfer", "trtllm_mha")
# MLA-family kernels that must never leak into the MHA arm.
MLA_ONLY_BACKENDS = ("trtllm_mla", "cutedsl_mla", "tokenspeed_mla", "flashmla")
# No dense-id wiring anywhere: must stay rejected until they get one.
UNWIRED_BACKENDS = ("cutlass_mla", "aiter")
def test_triton_always_allowed(self):
def test_triton_allowed_on_every_arm(self):
"""Triton reads both view families, so it is the one backend neither
the MLA nor the MHA arm can narrow away."""
for use_mla in (True, False):
self.assertTrue(_accepts("triton", use_mla=use_mla))
# The uniform-row screen is a property of the model, not of the
# backend, so it rejects even Triton.
self.assertFalse(_accepts("triton", use_mla=False, has_asymmetric_kv=True))
def test_dense_mla_backends_allowed_under_unified_mla(self):
for backend in self.DENSE_MLA_BACKENDS:
@@ -111,21 +131,18 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
f"{backend} should be allowed with the unified-memory MLA pool",
)
def test_dense_mla_backends_rejected_for_mha(self):
"""The per-layer-view exception is MLA-only -- MHA sub-pools stay strided."""
for backend in self.DENSE_MLA_BACKENDS:
self.assertFalse(
def test_dense_mha_backends_allowed_for_uniform_row_models(self):
for backend in self.DENSE_MHA_BACKENDS:
self.assertTrue(
_accepts(backend, use_mla=False),
f"{backend} must stay rejected for a non-MLA model",
f"{backend} should be allowed for a uniform-row MHA 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:
def test_mla_only_backends_rejected_for_mha(self):
for backend in self.MLA_ONLY_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=True, unified=False),
f"{backend} must stay rejected without --enable-unified-memory",
_accepts(backend, use_mla=False),
f"{backend} is an MLA kernel and must stay out of the MHA arm",
)
def test_plain_page_major_arm_is_gated_at_boot(self):
@@ -143,7 +160,7 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
"""head_dim != v_head_dim (MiMoV2): no uniform rows, so no per-layer views
and no unified pool. The rejection is the POOL's, not a backend's, so
it must fire on every backend -- Triton included."""
for backend in ("triton",) + self.DENSE_MLA_BACKENDS:
for backend in ("triton",) + self.DENSE_MHA_BACKENDS:
self.assertFalse(
_accepts(backend, use_mla=False, has_asymmetric_kv=True),
f"--enable-unified-memory + {backend} must be rejected for an "
@@ -162,12 +179,25 @@ class TestPageMajorBackendAllowlist(unittest.TestCase):
"K/V head dims",
)
def test_page_major_rejected_without_unified_memory(self):
"""--enable-page-major-kv-layout without the unified pool is rejected
outright, Triton included: the static page-major arm went away with
the strided views and awaits its per-layer-view reimplementation."""
for backend in ("triton",) + tuple(
set(self.DENSE_MLA_BACKENDS + self.DENSE_MHA_BACKENDS)
):
for use_mla in (True, False):
self.assertFalse(
_accepts(backend, use_mla=use_mla, 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",
f"{backend} has no dense-id wiring and must be rejected",
)
def test_helion_linear_attention_is_kda_only(self):