[unified-memory] Support fa3, the default MLA backend on pre-Blackwell hosts (#33046)

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,
         )

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(

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

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 (

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,
+    )

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",

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",
     ]
 
 

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,
         )
 

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):