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