diff --git a/python/sglang/kernels/ops/attention/metadata.py b/python/sglang/kernels/ops/attention/metadata.py index 858042733..0f236a2a1 100644 --- a/python/sglang/kernels/ops/attention/metadata.py +++ b/python/sglang/kernels/ops/attention/metadata.py @@ -193,6 +193,11 @@ def _fused_metadata_kernel_general( use_swa: tl.constexpr, SHIFT: tl.constexpr, BLOCK_COLS: tl.constexpr, + # Unified-memory dense-view path (page-major envelope shared with the mamba + # sub-pool). Both default to the identity for the statically-partitioned + # pool, where req_to_token already holds physical ids. + v2p_ptr=None, + PAGE_MULT: tl.constexpr = 1, ): pid_b = tl.program_id(0) # batch index pid_c = tl.program_id(1) # column chunk index @@ -251,6 +256,13 @@ def _fused_metadata_kernel_general( else: page_table_val = page_index >> SHIFT + # Unified memory: virtual page -> physical page -> that layer's dense block. + # Derived from page_table_val, NOT page_index, which the SWA branch below + # still needs in virtual space. Masked so padded lanes never index the table. + if v2p_ptr is not None: + page_table_val = tl.load(v2p_ptr + page_table_val, mask=mask, other=0) + page_table_val = page_table_val * PAGE_MULT + # Store to page_table pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1 tl.store(page_table + pt_offsets, page_table_val, mask=mask, cache_modifier=".cg") @@ -295,6 +307,9 @@ def _fused_metadata_kernel_ps1_no_swa( max_seq_pages, seq_len_delta: tl.constexpr, BLOCK_COLS: tl.constexpr, + # Unified-memory dense-view path; identity defaults for the static pool. + v2p_ptr=None, + PAGE_MULT: tl.constexpr = 1, ): pid_b = tl.program_id(0) # batch index pid_c = tl.program_id(1) # column chunk index @@ -338,6 +353,10 @@ def _fused_metadata_kernel_ps1_no_swa( ) # page_table = page_index // 1 = page_index + # Unified memory: at page_size 1 the virtual token id IS the virtual page id. + if v2p_ptr is not None: + page_index = tl.load(v2p_ptr + page_index, mask=mask, other=0) + page_index = page_index * PAGE_MULT pt_offsets = i * page_table_stride_0 + col_offsets * page_table_stride_1 tl.store(page_table + pt_offsets, page_index, mask=mask, cache_modifier=".cg") @@ -565,6 +584,8 @@ def normal_decode_set_metadata( page_size: int, swa_page_table: Optional[torch.Tensor] = None, token_to_kv_pool: Optional["SWAKVPool"] = None, + v2p_page_table: Optional[torch.Tensor] = None, + kernel_page_multiplier: int = 1, ): """ Fused Triton implementation that replaces 4-5 sequential CUDA kernels with 1-2 kernels: @@ -572,8 +593,14 @@ def normal_decode_set_metadata( 2. cu_seqlens_k = cumsum(cache_seqlens) (prefix-sum) 3. page_indices = req_to_token[pool_idx, stride_idx] (2-D gather) 4. page_table = page_indices // page_size (floor-divide) + 4b. (unified memory) page_table = v2p_page_table[page] * kernel_page_multiplier 5. (optional) swa_page_table for sliding window attention + Step 4b is folded in rather than applied afterwards so the capture-stable + page_table is written already translated: no separate pass a caller could + forget, and no temporary to keep pointer-stable across cuda-graph replays. + Identity (None / 1) for the statically-partitioned pool. + Achieves ~5.2x speedup on H200 hardware for typical decode workloads. Contract: only the live prefix (cdiv(cache_seqlens, page_size) pages) of each @@ -633,6 +660,8 @@ def normal_decode_set_metadata( max_seq_pages, seq_len_delta, BLOCK_COLS=BLOCK_COLS, + v2p_ptr=v2p_page_table, + PAGE_MULT=kernel_page_multiplier, num_warps=8, num_stages=3, ) @@ -696,6 +725,8 @@ def normal_decode_set_metadata( use_swa, shift, BLOCK_COLS=BLOCK_COLS, + v2p_ptr=v2p_page_table, + PAGE_MULT=kernel_page_multiplier, num_warps=4, num_stages=3, ) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 7448e8d3c..7535a8d8e 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -18,6 +18,7 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import ( ) from sglang.srt.configs.model_config import AttentionArch from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy from sglang.srt.layers.cp.utils import is_cp_v2_active @@ -184,6 +185,11 @@ class FlashAttentionBackend(AttentionBackend): # seq_lens_cpu / seq_lens_sum D2H sync is ever needed. self.needs_cpu_seq_lens = False self.use_mla = model_runner.model_config.attention_arch == AttentionArch.MLA + # Unified pool: req_to_token holds VIRTUAL ids but the MLA per-layer views + # are DENSE, so every page_table needs remapping. MLA-only -- the MHA/SWA + # sub-pools keep the strided envelope layout FA3 cannot read at all. + self._unified_hooks = unified_mla_hooks(model_runner.token_to_kv_pool_allocator) + self._unified_dense = self._unified_hooks.enabled and self.use_mla self.skip_prefill = skip_prefill self.attn_cp_size = model_runner.ps.attn_cp_size self._verify_mask = None @@ -1040,6 +1046,26 @@ class FlashAttentionBackend(AttentionBackend): ) ) + # Unified pool: one remap for every eager branch above, which all filled + # page_table with VIRTUAL token ids. Rebinding is safe here because those + # branches each produced a fresh tensor; the captured path instead folds + # the remap into normal_decode_set_metadata, which must write in place. + # + # Placed BEFORE the `// page_size` reduction, in token space: since + # dense(t) = phys_page * (ps * L) + t % ps, dense(page_start) // ps is + # phys_page * L, the dense page id the kernel wants. One site then serves + # both page sizes, and it inherits translate_kv_loc_dense's tombstone + # clamp so an unwritten req_to_token slot lands in the page-0 sink. + if self._unified_dense and metadata.page_table is not None: + # Flattened: the page_size == 1 translate path uses index_select, + # which rejects a 2-D index. + pt = metadata.page_table + metadata.page_table = ( + self._unified_hooks.translate_kv_loc_dense(pt.reshape(-1)) + .to(torch.int32) + .view(pt.shape) + ) + # Convert the page table to a strided format which is needed by FA3 API if self.page_size > 1: self.strided_indices = torch.arange( @@ -2632,6 +2658,12 @@ class FlashAttentionBackend(AttentionBackend): if self.use_sliding_window_kv_pool else None ), + v2p_page_table=( + self._unified_hooks.v2p_page_table + if self._unified_dense + else None + ), + kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier, ) else: @@ -2748,6 +2780,12 @@ class FlashAttentionBackend(AttentionBackend): if self.use_sliding_window_kv_pool else None ), + v2p_page_table=( + self._unified_hooks.v2p_page_table + if self._unified_dense + else None + ), + kernel_page_multiplier=self._unified_hooks.kernel_page_multiplier, ) self._maybe_update_local_attn_metadata_for_replay( diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 4093911c0..ebb5b40ef 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -23,6 +23,7 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.flashinfer_backend import ( create_flashinfer_kv_indices_triton, ) +from sglang.srt.layers.attention.unified_mem_hooks import unified_mla_hooks from sglang.srt.layers.dcp import ( DecodeContextParallelMetadata, update_local_kv_lens_for_dcp, @@ -66,51 +67,6 @@ 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 diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 267e2502e..bb66ddc58 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -34,8 +34,8 @@ 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.unified_mem_hooks import 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 from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( diff --git a/python/sglang/srt/layers/attention/unified_mem_hooks.py b/python/sglang/srt/layers/attention/unified_mem_hooks.py new file mode 100644 index 000000000..c15f39283 --- /dev/null +++ b/python/sglang/srt/layers/attention/unified_mem_hooks.py @@ -0,0 +1,70 @@ +# 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. +# ============================================================================== +"""Allocator hooks the paged MLA attention backends need under the unified +memory pool. + +Lives in its own module because three unrelated backend families consume it +(fa3, flashinfer_mla, and trtllm_mla with its cutedsl_mla / tokenspeed_mla +subclasses) and none of them should have to import another's module to get it. +""" + +from __future__ import annotations + +from typing import Callable, Optional + +import msgspec +import torch + + +class UnifiedMLAHooks(msgspec.Struct, frozen=True): + """Dense-view hooks for one KV allocator. + + All-``None``/1/``False`` for the statically-partitioned pool, where + ``req_to_token`` already holds physical ids and no translation is needed. + """ + + # Page-level virtual->physical table, gathered through by block-table kernels. + v2p_page_table: Optional[torch.Tensor] + # Virtual token id -> DENSE kernel-facing id (tombstones clamped to the sink). + translate_kv_loc_dense: Optional[Callable[..., torch.Tensor]] + # Dense page stride scale (= number of full-attention MLA layers). + kernel_page_multiplier: int + enabled: bool + + +_STATIC_POOL = UnifiedMLAHooks( + v2p_page_table=None, + translate_kv_loc_dense=None, + kernel_page_multiplier=1, + enabled=False, +) + + +def unified_mla_hooks(allocator) -> UnifiedMLAHooks: + """Probe ``allocator`` for the unified-pool dense-view hooks. + + Detection keys on the v2p table, NOT on ``kernel_page_multiplier > 1``: a + rank owning exactly ONE full-attention layer has multiplier 1 while its + ``req_to_token`` is still virtual. There the dense id collapses onto the + physical id, so the v2p gather alone is the whole translation. + """ + v2p = getattr(allocator, "full_v2p_page_table", None) + if v2p is None: + return _STATIC_POOL + 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, + ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 5e92a5e0a..9905744fd 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -7712,12 +7712,14 @@ class ServerArgs: # 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. + # TRTLLMMLABackend and inherit its dense read/write path; "fa3" remaps its + # page_table (in-kernel for captured decode, one funnel for eager). # 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", + "fa3", "trtllm_mla", "flashinfer", "cutedsl_mla", diff --git a/test/registered/models_e2e/test_kimi_linear_unified_memory.py b/test/registered/models_e2e/test_kimi_linear_unified_memory.py index 6583ccd50..efec93bba 100644 --- a/test/registered/models_e2e/test_kimi_linear_unified_memory.py +++ b/test/registered/models_e2e/test_kimi_linear_unified_memory.py @@ -1,28 +1,26 @@ """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. +Under `--enable-unified-memory` the MLA full side is exposed as DENSE per-layer +views and every loc the kernels see is a translated virtual id, so the whole +read/write path differs from the static pool. The unit tests pin that pool in +isolation; this is the end-to-end guard. `test_prefix_cache_branching` carries +most of the weight: a radix hit replays virtual locs whose physical pages may +have moved under 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. +No `--attention-backend` is pinned on purpose -- the test runs whatever the host +resolves to (`fa3` on this suite's H100 runner, also the H200 default). Both +defects found in review on #32972 were reachable only under a resolved default, +which a pinned test hides by construction. -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. +Reference GSM8K, all with `--enable-unified-memory`: + - 2x H200 TP2, resolved default (fa3): 0.917 @400, vs 0.915 static (1 sigma + ~= 0.015). This file as written scores 0.920 @200. + - 2x H200 TP2, `--attention-backend triton`: 0.900 @400. + - 1x B300 TP1: 0.915 flashinfer, 0.900 trtllm_mla, @200. + +Nightly-only: a second full 48B server launch is too much for 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 """ @@ -45,7 +43,7 @@ class TestKimiLinearUnifiedMemory( 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). + # cost accuracy (measured 0.917 vs 0.915 static, see the module docstring). gsm8k_score_threshold = 0.88 other_args = [ "--trust-remote-code", @@ -54,22 +52,6 @@ class TestKimiLinearUnifiedMemory( "--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", ] 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 index 68252ca8f..196f265f1 100644 --- 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 @@ -20,9 +20,12 @@ 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. +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_dense` over the + token-level kv_indices; + - `normal_decode_set_metadata` in-kernel, for fa3's captured-decode page table. Covered here: - kernel identity: `v2p_ptr=None, PAGE_MULT=1` is byte-identical to main; @@ -30,7 +33,10 @@ Covered here: 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. + page-level block table the trtllm path builds; + - 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. python -m pytest test/registered/unit/mem_cache/test_unified_mla_dense_block_table.py -v """ @@ -199,6 +205,105 @@ 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. + """ + + def _run(self, page_size, *, v2p, mult, bs=5, max_ctx=2048): + from sglang.kernels.ops.attention.metadata import normal_decode_set_metadata + + 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, + ) + torch.cuda.synchronize() + want = _reference(rt, rpi, sl, page_size, v2p=v2p_arg, mult=mult) + return page_table, want, sl + + def _assert_live_prefix(self, got, want, sl, page_size): + """The kernel contract only (re)writes each row's live page prefix; the + tail keeps stale values that consumers bound by cache_seqlens.""" + 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(), want[r, :n_pages]), + f"row {r} (page_size={page_size}):\n" + f"got ={got[r, :n_pages]}\nwant={want[r, :n_pages]}", + ) + + def test_identity_when_hooks_absent(self): + """Static pool: no v2p, multiplier 1 -> byte-identical to pre-change.""" + for page_size in (1, 64): + got, want, sl = self._run(page_size, v2p=False, mult=1) + self._assert_live_prefix(got, want, sl, page_size) + + def test_dense_mapping_ps1_fast_path(self): + got, want, sl = self._run(1, v2p=True, mult=_LAYERS) + self._assert_live_prefix(got, want, sl, 1) + + def test_dense_mapping_general_path(self): + got, want, sl = self._run(64, v2p=True, mult=_LAYERS) + self._assert_live_prefix(got, want, sl, 64) + + def test_single_full_attention_layer(self): + """multiplier 1 with a real v2p: the gather alone is the translation.""" + for page_size in (1, 64): + got, want, sl = self._run(page_size, v2p=True, mult=1) + self._assert_live_prefix(got, want, sl, page_size) + virtual = _reference( + *TestDenseBlockTable._make_batch(self, page_size)[:3], + 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_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 @@ -207,7 +312,7 @@ class TestUnifiedMLAHookDetection(unittest.TestCase): @staticmethod def _probe(**attrs): - from sglang.srt.layers.attention.flashinfer_mla_backend import ( + from sglang.srt.layers.attention.unified_mem_hooks import ( unified_mla_hooks, ) 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 index d6d9287ec..2550e70b1 100644 --- a/test/registered/unit/server_args/test_page_major_backend_allowlist.py +++ b/test/registered/unit/server_args/test_page_major_backend_allowlist.py @@ -16,12 +16,15 @@ 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. +backends can read it directly once their kv_indices / block tables are remapped +to dense ids -- `fa3`, `flashinfer`'s MLA backend, and `trtllm_mla` with its +`cutedsl_mla` / `tokenspeed_mla` subclasses. 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. +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. python -m pytest test/registered/unit/server_args/test_page_major_backend_allowlist.py -v """ @@ -66,9 +69,15 @@ def _accepts(backend: str, *, use_mla: bool, unified: bool = True) -> bool: 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") + DENSE_MLA_BACKENDS = ( + "fa3", + "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") + UNWIRED_BACKENDS = ("flashmla", "cutlass_mla", "trtllm_mha", "aiter") def test_triton_always_allowed(self): for use_mla in (True, False):