From 33c27d8e7f4f1c6bb479c702cde92567c67f8391 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Fri, 31 Jul 2026 01:32:08 -0700 Subject: [PATCH] [unified-memory] Let Kimi-Linear use the paged MLA attention backends (#32972) --- .../sglang/kernels/ops/kvcache/kv_indices.py | 16 +- .../attention/flashinfer_mla_backend.py | 76 ++++ .../layers/attention/trtllm_mla_backend.py | 67 +++- python/sglang/srt/mem_cache/memory_pool.py | 11 +- .../srt/mem_cache/multi_ended_allocator.py | 9 + python/sglang/srt/server_args.py | 28 +- .../test_kimi_linear_unified_memory.py | 77 ++++ .../test_unified_mla_dense_block_table.py | 375 ++++++++++++++++++ .../test_page_major_backend_allowlist.py | 111 ++++++ 9 files changed, 759 insertions(+), 11 deletions(-) create mode 100644 test/registered/models_e2e/test_kimi_linear_unified_memory.py create mode 100644 test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py create mode 100644 test/registered/unit/server_args/test_page_major_backend_allowlist.py diff --git a/python/sglang/kernels/ops/kvcache/kv_indices.py b/python/sglang/kernels/ops/kvcache/kv_indices.py index 46bb37e7d..1ee8804e6 100644 --- a/python/sglang/kernels/ops/kvcache/kv_indices.py +++ b/python/sglang/kernels/ops/kvcache/kv_indices.py @@ -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, ) diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 0d68b303c..4093911c0 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 1c1d07401..267e2502e 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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: diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index ba582b52f..4ac0aeae9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -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) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index 61f17bff4..ec2358bc5 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1f874fa5d..5e92a5e0a 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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. diff --git a/test/registered/models_e2e/test_kimi_linear_unified_memory.py b/test/registered/models_e2e/test_kimi_linear_unified_memory.py new file mode 100644 index 000000000..6583ccd50 --- /dev/null +++ b/test/registered/models_e2e/test_kimi_linear_unified_memory.py @@ -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() diff --git a/test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py b/test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py new file mode 100644 index 000000000..68252ca8f --- /dev/null +++ b/test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py @@ -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() diff --git a/test/registered/unit/server_args/test_page_major_backend_allowlist.py b/test/registered/unit/server_args/test_page_major_backend_allowlist.py new file mode 100644 index 000000000..d6d9287ec --- /dev/null +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -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()