feat(attention): add architecture-owned SM12x FA4 kernels (#32991)

python/sglang/kernels/ops/attention/fa4_sm120/__init__.py

@@ -0,0 +1 @@
+"""SGLang-owned FlashAttention-4 kernels and launch policy for SM120."""

python/sglang/kernels/ops/attention/fa4_sm120/dispatch.py

@@ -0,0 +1,58 @@
+# Copyright (c) 2026, SGLang Team.
+"""Lightweight dispatch bridge for SGLang-owned SM120 FA4 kernels.
+
+The vendored FA4 interface dispatches through this module so the SM120
+implementation and its launch state remain outside ``flash_attn/cute``.
+"""
+
+from functools import lru_cache
+
+
+@lru_cache(maxsize=None)
+def get_forward_host(arch: int):
+    """Return the optional architecture-owned forward host."""
+    if arch // 10 == 12:
+        from sglang.kernels.ops.attention.fa4_sm120.runtime import (
+            sm120_forward_host,
+        )
+
+        return sm120_forward_host
+    return None
+
+
+def resolve_runtime_policy(
+    *,
+    device_capability: tuple[int, int],
+    deterministic: bool,
+) -> tuple[int, int, bool]:
+    """Resolve generic and architecture-owned SplitKV launch policy."""
+    arch = device_capability[0] * 10 + device_capability[1]
+    uses_arch_decode_policy = get_forward_host(arch) is not None
+    no_splitkv = device_capability < (9, 0) or uses_arch_decode_policy
+    num_splits = 1 if deterministic or no_splitkv else 0
+    decode_num_splits = (
+        0 if uses_arch_decode_policy and not deterministic else num_splits
+    )
+    return num_splits, decode_num_splits, uses_arch_decode_policy
+
+
+@lru_cache(maxsize=None)
+def get_forward_arch(device) -> int | None:
+    """Return the device arch when it has an architecture-owned forward host."""
+    import torch
+
+    major, minor = torch.cuda.get_device_capability(device)
+    arch = major * 10 + minor
+    return arch if get_forward_host(arch) is not None else None
+
+
+def try_cached_paged_decode(*, arch: int, **kwargs):
+    """Try an architecture-owned paged-decode launch plan."""
+    host = get_forward_host(arch)
+    return None if host is None else host.try_paged_decode(arch=arch, **kwargs)
+
+
+def try_cached_varlen(*, arch: int, **kwargs):
+    """Try an architecture-owned varlen launch plan."""
+    host = get_forward_host(arch)
+    return None if host is None else host.try_varlen(arch=arch, **kwargs)

python/sglang/kernels/ops/attention/fa4_sm120/flash_fwd_decode.py

@@ -0,0 +1,711 @@
+# Copyright (c) 2026, SGLang Team.
+"""End-to-end transposed SM120 paged-decode specialization.
+
+This path keeps the packed query axis on the N=8 dimension of warp MMA:
+
+    scores.T = K @ Q.T        # (64, 8)
+    output.T = V.T @ P.T      # (256, 8)
+
+The dataflow is isolated from the general SM120 kernel because its page-TMA
+transport, column-wise online softmax, and transposed epilogue form one compile
+specialization.
+"""
+
+import math
+from typing import Optional
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Float32, Int32, const_expr
+from cutlass.cute.nvgpu import warp
+from cutlass.pipeline import PipelineAsync, PipelineState
+from quack import layout_utils
+
+from sglang.kernels.ops.attention.fa4_sm120.flash_fwd import (
+    FlashAttentionForwardSm120,
+)
+from sglang.kernels.ops.attention.flash_attn.cute import utils
+from sglang.kernels.ops.attention.flash_attn.cute.block_info import BlockInfo
+from sglang.kernels.ops.attention.flash_attn.cute.named_barrier import NamedBarrierFwd
+from sglang.kernels.ops.attention.flash_attn.cute.pack_gqa import PackGQA
+from sglang.kernels.ops.attention.flash_attn.cute.seqlen_info import SeqlenInfoQK
+from sglang.kernels.ops.attention.flash_attn.cute.utils import AuxData
+
+
+class FlashAttentionForwardSm120DecodeTranspose(FlashAttentionForwardSm120):
+    """M64N8 QK and transposed PV for qualified packed single-token decode."""
+
+    # Paged TMA is part of this kernel's dataflow, not a runtime tuning knob.
+    # Keeping it on the distinct class identity prevents a gather-compiled
+    # specialization from being reused for the TMA tensor layout.
+    paged_tma = True
+    query_mma_n = 8
+    query_in_regs = True
+
+    def _uses_n_distributed_qk(self) -> bool:
+        # Reuse the base kernel's four-consumer-warp dispatch. All four warps
+        # participate in both transposed MMA phases.
+        return True
+
+    def _uses_split_pv_warps(self) -> bool:
+        # Experimental mixed-HDV channel slices retain the same shared P
+        # handoff even though they are outside the base kernel's qualified
+        # (HDQ, HDV) configuration table.
+        return True
+
+    def _setup_attributes(self):
+        super()._setup_attributes()
+        # The tiled MMA is deliberately warp-local. The base kernel normally
+        # derives its cooperative-group size from TiledMma.size, which would
+        # expose only one consumer warp here. Four physical consumer warps
+        # instead operate on disjoint K rows.
+        self.num_qk_threads = self.num_threads
+        self.num_mma_threads = self.num_threads
+        self.num_Q_load_threads = self.num_threads
+        self.num_epilogue_threads = self.num_threads
+
+    def _get_tiled_mma(self):
+        mma_op = warp.MmaF16BF16Op(self.dtype, Float32, (16, 8, 16))
+        tiled_mma_qk = cute.make_tiled_mma(
+            mma_op,
+            (1, 1, 1),
+            permutation_mnk=(16, self.query_mma_n, 16),
+        )
+        # Each warp covers 64 value rows; the four disjoint SMEM views cover
+        # HDV=256 without a CTA-level M permutation.
+        tiled_mma_pv = cute.make_tiled_mma(
+            mma_op,
+            (1, 1, 1),
+            permutation_mnk=(64, self.query_mma_n, 16),
+        )
+        return tiled_mma_qk, tiled_mma_pv
+
+    @cute.jit
+    def _gemm_n8(
+        self,
+        tiled_mma: cute.TiledMma,
+        acc: cute.Tensor,
+        tCrA: cute.Tensor,
+        tCrB: cute.Tensor,
+        tCsA: cute.Tensor,
+        tCsB: cute.Tensor,
+        smem_thr_copy_A: cute.TiledCopy,
+        smem_thr_copy_B: cute.TiledCopy,
+        B_in_regs: cutlass.Constexpr[bool] = False,
+    ):
+        """Issue an N=8 warp-MMA mainloop through the underlying MMA atom.
+
+        CuTe DSL's tiled-MMA verifier rejects m16n8k16 when logical N is
+        exactly eight because it compares the raw A value mode against the C
+        value mode. The hardware atom itself has the correct native fragment
+        contract, so keep tiling for partitioning/copies and issue GEMM through
+        that atom.
+        """
+        mma_atom = cute.make_mma_atom(tiled_mma.op)
+        tCrA_copy_view = smem_thr_copy_A.retile(tCrA)
+        tCrB_copy_view = smem_thr_copy_B.retile(tCrB)
+        cute.copy(
+            smem_thr_copy_A,
+            tCsA[None, None, 0],
+            tCrA_copy_view[None, None, 0],
+        )
+        if const_expr(not B_in_regs):
+            cute.copy(
+                smem_thr_copy_B,
+                tCsB[None, None, 0],
+                tCrB_copy_view[None, None, 0],
+            )
+        for k in cutlass.range_constexpr(cute.size(tCsA.shape[2])):
+            if k < cute.size(tCsA.shape[2]) - 1:
+                cute.copy(
+                    smem_thr_copy_A,
+                    tCsA[None, None, k + 1],
+                    tCrA_copy_view[None, None, k + 1],
+                )
+                if const_expr(not B_in_regs):
+                    cute.copy(
+                        smem_thr_copy_B,
+                        tCsB[None, None, k + 1],
+                        tCrB_copy_view[None, None, k + 1],
+                    )
+            cute.gemm(
+                mma_atom,
+                acc,
+                tCrA[None, None, k],
+                tCrB[None, None, k],
+                acc,
+            )
+
+    @cute.jit
+    def _rescale_transposed_o(
+        self,
+        acc_O: cute.Tensor,
+        tiled_mma_pv: cute.TiledMma,
+        tidx: Int32,
+        sRowScale: cute.Tensor,
+        scale_row: cutlass.Constexpr[int],
+    ):
+        lane_idx = tidx % cute.arch.WARP_SIZE
+        thr_mma_pv = tiled_mma_pv.get_slice(lane_idx)
+        acc_O_qd = layout_utils.reshape_acc_to_mn(acc_O, transpose=True)
+        cO = cute.make_identity_tensor((64, self.query_mma_n))
+        tOcO_qd = layout_utils.reshape_acc_to_mn(
+            thr_mma_pv.partition_C(cO), transpose=True
+        )
+        for r in cutlass.range(cute.size(acc_O_qd, mode=[0]), unroll_full=True):
+            query_row = tOcO_qd[r, 0][1]
+            acc_O_qd[r, None].store(
+                acc_O_qd[r, None].load() * sRowScale[scale_row, query_row]
+            )
+
+    @cute.jit
+    def _compute_one_n_block_transposed(
+        self,
+        n_block: Int32,
+        consumer_state: PipelineState,
+        acc_O: cute.Tensor,
+        sQ: cute.Tensor,
+        sK: cute.Tensor,
+        sV: cute.Tensor,
+        sP: cute.Tensor,
+        sRowScale: cute.Tensor,
+        pipeline_k: PipelineAsync,
+        pipeline_v: PipelineAsync,
+        tiled_mma_qk: cute.TiledMma,
+        tiled_mma_pv: cute.TiledMma,
+        smem_thr_copy_K: cute.TiledCopy,
+        smem_thr_copy_Q: cute.TiledCopy,
+        smem_thr_copy_V: cute.TiledCopy,
+        smem_thr_copy_P: cute.TiledCopy,
+        tSrK: cute.Tensor,
+        tSrQ: cute.Tensor,
+        tOrV: cute.Tensor,
+        tOrP: cute.Tensor,
+        tKsK: cute.Tensor,
+        tQsQ: cute.Tensor,
+        tVsV: cute.Tensor,
+        tPsP: cute.Tensor,
+        tidx: Int32,
+        softmax_scale_log2: Float32,
+        seqlen: SeqlenInfoQK,
+        window_size_left: Optional[Int32],
+        is_first_n_block: cutlass.Constexpr[bool] = False,
+    ):
+        num_qk_warps = const_expr(4)
+        local_sum_base = const_expr(num_qk_warps)
+        global_max_row = const_expr(2 * num_qk_warps)
+        global_sum_row = const_expr(global_max_row + 1)
+        old_o_scale_row = const_expr(global_max_row + 2)
+        warp_scale_base = const_expr(global_max_row + 3)
+        p_stage = consumer_state.index
+        warp_idx = tidx // cute.arch.WARP_SIZE
+        lane_idx = tidx % cute.arch.WARP_SIZE
+        key_row_base = warp_idx * const_expr(16)
+
+        k_wait_token = pipeline_k.consumer_try_wait(consumer_state)
+        pipeline_k.consumer_wait(consumer_state, k_wait_token)
+
+        thr_mma_qk = tiled_mma_qk.get_slice(lane_idx)
+        acc_shape_S = thr_mma_qk.partition_shape_C((16, self.query_mma_n))
+        acc_S = cute.make_rmem_tensor(acc_shape_S, Float32)
+        acc_S.fill(0.0)
+        self._gemm_n8(
+            tiled_mma_qk,
+            acc_S,
+            tSrK,
+            tSrQ,
+            tKsK[None, None, None, p_stage],
+            tQsQ,
+            smem_thr_copy_K,
+            smem_thr_copy_Q,
+            B_in_regs=self.query_in_regs,
+        )
+        pipeline_k.consumer_release(consumer_state)
+
+        acc_S_qk = layout_utils.reshape_acc_to_mn(acc_S, transpose=True)
+        cS = cute.make_identity_tensor((16, self.query_mma_n))
+        tScS_qk = layout_utils.reshape_acc_to_mn(
+            thr_mma_qk.partition_C(cS), transpose=True
+        )
+        num_query_rows = cute.size(acc_S_qk, mode=[0])
+        row_max_local = cute.make_rmem_tensor(num_query_rows, Float32)
+        row_sum_local = cute.make_rmem_tensor(num_query_rows, Float32)
+        for r in cutlass.range(num_query_rows, unroll_full=True):
+            # The final paged tile can be only partially populated. The page
+            # gather leaves invalid SMEM rows untouched. Keep the test outside
+            # the unrolled fragment loop so complete 64-token tiles pay no
+            # per-element coordinate/comparison cost.
+            tile_start = n_block * self.tile_n
+            local_window_start = (
+                cutlass.max(
+                    seqlen.seqlen_k - 1 - window_size_left,
+                    0,
+                )
+                if const_expr(self.is_local and window_size_left is not None)
+                else Int32(0)
+            )
+            if tile_start + self.tile_n > seqlen.seqlen_k or (
+                const_expr(self.is_local and window_size_left is not None)
+                and tile_start < local_window_start
+            ):
+                for c in cutlass.range(cute.size(acc_S_qk, mode=[1]), unroll_full=True):
+                    key_row = tile_start + key_row_base + tScS_qk[r, c][0]
+                    if key_row >= seqlen.seqlen_k or (
+                        const_expr(self.is_local and window_size_left is not None)
+                        and key_row < local_window_start
+                    ):
+                        acc_S_qk[r, c] = -Float32.inf
+
+            acc_S_row = acc_S_qk[r, None].load()
+            row_max = utils.fmax_reduce(acc_S_row)
+            # In an m16n8 accumulator, lanes with the same low two lane
+            # bits own the same pair of N/query columns. Reducing across M
+            # (keys) therefore uses the strided lane group
+            # {lane, lane^4, lane^8, lane^16}, not a contiguous width-4
+            # group as in the ordinary row-wise QK layout.
+            for offset in cutlass.range_constexpr(2, 5):
+                row_max = utils.fmax(
+                    row_max,
+                    cute.arch.shuffle_sync_bfly(row_max, offset=1 << offset),
+                )
+            row_max_safe = 0.0 if row_max == -Float32.inf else row_max
+            acc_S_row_exp = cute.math.exp2(
+                (acc_S_row - row_max_safe) * softmax_scale_log2,
+                fastmath=True,
+            )
+            row_sum = utils.fadd_reduce(acc_S_row_exp)
+            for offset in cutlass.range_constexpr(2, 5):
+                row_sum += cute.arch.shuffle_sync_bfly(row_sum, offset=1 << offset)
+            row_max_local[r] = row_max
+            row_sum_local[r] = row_sum
+            acc_S_qk[r, None].store(acc_S_row_exp)
+
+        keys_per_warp = const_expr(self.tile_n // num_qk_warps)
+        if tScS_qk[0, 0][0] % keys_per_warp == 0:
+            for r in cutlass.range(num_query_rows, unroll_full=True):
+                query_row = tScS_qk[r, 0][1]
+                sRowScale[warp_idx, query_row] = row_max_local[r]
+                sRowScale[local_sum_base + warp_idx, query_row] = row_sum_local[r]
+        cute.arch.fence_view_async_shared()
+
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.PFull),
+            number_of_threads=self.num_mma_threads,
+        )
+        if tidx < self.query_mma_n:
+            query_row = tidx
+            row_max = sRowScale[0, query_row]
+            for warp_idx_it in cutlass.range_constexpr(1, num_qk_warps):
+                row_max = utils.fmax(row_max, sRowScale[warp_idx_it, query_row])
+            row_max_prev = (
+                row_max
+                if const_expr(is_first_n_block)
+                else sRowScale[global_max_row, query_row]
+            )
+            row_max_new = (
+                row_max
+                if const_expr(is_first_n_block)
+                else utils.fmax(row_max_prev, row_max)
+            )
+            row_max_new_safe = 0.0 if row_max_new == -Float32.inf else row_max_new
+            old_o_scale = (
+                1.0
+                if const_expr(is_first_n_block)
+                else cute.math.exp2(
+                    (row_max_prev - row_max_new_safe) * softmax_scale_log2,
+                    fastmath=True,
+                )
+            )
+            row_sum_new = (
+                0.0
+                if const_expr(is_first_n_block)
+                else sRowScale[global_sum_row, query_row] * old_o_scale
+            )
+            for warp_idx_it in cutlass.range_constexpr(num_qk_warps):
+                warp_scale = cute.math.exp2(
+                    (sRowScale[warp_idx_it, query_row] - row_max_new_safe)
+                    * softmax_scale_log2,
+                    fastmath=True,
+                )
+                sRowScale[warp_scale_base + warp_idx_it, query_row] = warp_scale
+                row_sum_new += (
+                    sRowScale[local_sum_base + warp_idx_it, query_row] * warp_scale
+                )
+            sRowScale[global_max_row, query_row] = row_max_new
+            sRowScale[global_sum_row, query_row] = row_sum_new
+            if const_expr(not is_first_n_block):
+                sRowScale[old_o_scale_row, query_row] = old_o_scale
+        cute.arch.fence_view_async_shared()
+
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.PEmpty),
+            number_of_threads=self.num_mma_threads,
+        )
+        for r in cutlass.range(num_query_rows, unroll_full=True):
+            query_row = tScS_qk[r, 0][1]
+            warp_scale = sRowScale[warp_scale_base + warp_idx, query_row]
+            acc_S_qk[r, None].store(acc_S_qk[r, None].load() * warp_scale)
+        for r in cutlass.range(num_query_rows, unroll_full=True):
+            query_row = tScS_qk[r, 0][1]
+            for c in cutlass.range(cute.size(acc_S_qk, mode=[1]), unroll_full=True):
+                key_row = key_row_base + tScS_qk[r, c][0]
+                sP[query_row, key_row, p_stage] = self.dtype(acc_S_qk[r, c])
+        cute.arch.fence_view_async_shared()
+
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.PFull),
+            number_of_threads=self.num_mma_threads,
+        )
+        # acc_O is zero-initialized immediately before the first KV tile, so
+        # multiplying it by that tile's compile-time unit scale is pure
+        # overhead. Later tiles still rescale accumulated output before PV.
+        if const_expr(not is_first_n_block):
+            self._rescale_transposed_o(
+                acc_O,
+                tiled_mma_pv,
+                tidx,
+                sRowScale,
+                old_o_scale_row,
+            )
+
+        v_wait_token = pipeline_v.consumer_try_wait(consumer_state)
+        pipeline_v.consumer_wait(consumer_state, v_wait_token)
+
+        self._gemm_n8(
+            tiled_mma_pv,
+            acc_O,
+            tOrV,
+            tOrP,
+            tVsV[None, None, None, p_stage],
+            tPsP[None, None, None, p_stage],
+            smem_thr_copy_V,
+            smem_thr_copy_P,
+        )
+        pipeline_v.consumer_release(consumer_state)
+
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.PEmpty),
+            number_of_threads=self.num_mma_threads,
+        )
+        consumer_state.advance()
+        return consumer_state
+
+    @cute.jit
+    def _store_transposed_output(
+        self,
+        acc_O: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        sLSE: cute.Tensor,
+        tiled_mma_pv: cute.TiledMma,
+        tidx: Int32,
+        seqlen: SeqlenInfoQK,
+        m_block: Int32,
+        head_idx: Int32,
+        batch_idx: Int32,
+        split_idx: Int32,
+    ):
+        row_limit = seqlen.seqlen_q * self.qhead_per_kvhead
+        warp_idx = tidx // cute.arch.WARP_SIZE
+        lane_idx = tidx % cute.arch.WARP_SIZE
+        thr_mma_pv = tiled_mma_pv.get_slice(lane_idx)
+        acc_O_qd = layout_utils.reshape_acc_to_mn(acc_O, transpose=True)
+        cO = cute.make_identity_tensor((64, self.query_mma_n))
+        tOcO_qd = layout_utils.reshape_acc_to_mn(
+            thr_mma_pv.partition_C(cO), transpose=True
+        )
+        mO_cur = (
+            seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx, split_idx]
+            if const_expr(self.is_split_kv)
+            else seqlen.offset_batch_Q(mO, batch_idx, dim=3)[None, None, head_idx]
+        )
+        for r in cutlass.range(cute.size(acc_O_qd, mode=[0]), unroll_full=True):
+            query_row_local = tOcO_qd[r, 0][1]
+            query_row = m_block * self.tile_m + query_row_local
+            if query_row < row_limit:
+                for c in cutlass.range(cute.size(acc_O_qd, mode=[1]), unroll_full=True):
+                    value_col = warp_idx * const_expr(64) + tOcO_qd[r, c][0]
+                    mO_cur[query_row, value_col] = self.dtype(acc_O_qd[r, c])
+
+        if const_expr(mLSE is not None):
+            if tidx < self.query_mma_n:
+                query_row = m_block * self.tile_m + tidx
+                if query_row < row_limit:
+                    mLSE_cur = (
+                        seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[
+                            None, head_idx, split_idx
+                        ]
+                        if const_expr(self.is_split_kv)
+                        else seqlen.offset_batch_Q(mLSE, batch_idx, dim=2)[
+                            None, head_idx
+                        ]
+                    )
+                    mLSE_cur[query_row] = sLSE[tidx]
+
+    @cute.jit
+    def mma(
+        self,
+        mQ: cute.Tensor,
+        mO: cute.Tensor,
+        mLSE: Optional[cute.Tensor],
+        sQ: cute.Tensor,
+        sK: cute.Tensor,
+        sV: cute.Tensor,
+        sO: cute.Tensor,
+        sP: Optional[cute.Tensor],
+        sRowScale: Optional[cute.Tensor],
+        sLSE: Optional[cute.Tensor],
+        learnable_sink: Optional[cute.Tensor],
+        pipeline_k: PipelineAsync,
+        pipeline_v: PipelineAsync,
+        gmem_tiled_copy_Q: cute.TiledCopy,
+        gmem_tiled_copy_O: cute.TiledCopy,
+        tiled_mma_qk: cute.TiledMma,
+        tiled_mma_pv: cute.TiledMma,
+        tidx: Int32,
+        softmax_scale_log2: Float32,
+        softmax_scale: Optional[Float32],
+        consumer_state: PipelineState,
+        block_info: BlockInfo,
+        seqlen: SeqlenInfoQK,
+        n_block_min: Int32,
+        n_block_max: Int32,
+        m_block: Int32,
+        head_idx: Int32,
+        batch_idx: Int32,
+        split_idx: Int32,
+        is_qk_owner: cutlass.Constexpr[bool],
+        aux_data: AuxData = AuxData(),
+        fastdiv_mods=None,
+    ):
+        assert self.paged_kv
+        assert self.pack_gqa
+        assert self.is_causal or self.is_local
+        assert self.score_mod is None
+        assert self.mask_mod is None
+        assert self.tile_m == 16 and self.tile_n == 64
+        assert self.tile_hdim == 256
+        assert self.tile_hdimv in (64, 256)
+        assert self.qhead_per_kvhead <= self.query_mma_n
+        assert is_qk_owner
+
+        mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
+        warp_idx = tidx // cute.arch.WARP_SIZE
+        lane_idx = tidx % cute.arch.WARP_SIZE
+        thr_mma_qk = tiled_mma_qk.get_slice(lane_idx)
+        thr_mma_pv = tiled_mma_pv.get_slice(lane_idx)
+        acc_shape_O = thr_mma_pv.partition_shape_C((64, self.query_mma_n))
+        acc_O = cute.make_rmem_tensor(acc_shape_O, Float32)
+        acc_O.fill(0.0)
+
+        smem_copy_atom_k_major = cute.make_copy_atom(
+            warp.LdMatrix8x8x16bOp(transpose=False, num_matrices=4),
+            self.dtype,
+        )
+        smem_copy_atom_v_major = cute.make_copy_atom(
+            warp.LdMatrix8x8x16bOp(transpose=True, num_matrices=4),
+            self.dtype,
+        )
+        smem_thr_copy_K = utils.make_tiled_copy_A(
+            smem_copy_atom_k_major, tiled_mma_qk
+        ).get_slice(lane_idx)
+        smem_thr_copy_Q = utils.make_tiled_copy_B(
+            smem_copy_atom_k_major, tiled_mma_qk
+        ).get_slice(lane_idx)
+        smem_thr_copy_V = utils.make_tiled_copy_A(
+            smem_copy_atom_v_major, tiled_mma_pv
+        ).get_slice(lane_idx)
+        smem_thr_copy_P = utils.make_tiled_copy_B(
+            smem_copy_atom_k_major, tiled_mma_pv
+        ).get_slice(lane_idx)
+
+        sK_warp = cute.local_tile(
+            sK,
+            (16, self.tile_hdim, self._num_k_stages()),
+            (warp_idx, 0, 0),
+        )
+        sQ_query = cute.local_tile(sQ, (self.query_mma_n, self.tile_hdim), (0, 0))
+        sP_query = cute.local_tile(
+            sP,
+            (self.query_mma_n, self.tile_n, self._num_p_stages()),
+            (0, 0, 0),
+        )
+        tSrK = thr_mma_qk.make_fragment_A(
+            thr_mma_qk.partition_A(sK_warp[None, None, 0])
+        )
+        tSrQ = thr_mma_qk.make_fragment_B(thr_mma_qk.partition_B(sQ_query))
+        sV_warp = cute.local_tile(
+            sV,
+            (64, self.tile_n, self._num_v_stages()),
+            (warp_idx, 0, 0),
+        )
+        tOrV = thr_mma_pv.make_fragment_A(
+            thr_mma_pv.partition_A(sV_warp[None, None, 0])
+        )
+        tOrP = thr_mma_pv.make_fragment_B(
+            thr_mma_pv.partition_B(sP_query[None, None, 0])
+        )
+        tVsV = smem_thr_copy_V.partition_S(sV_warp)
+        tKsK = smem_thr_copy_K.partition_S(sK_warp)
+        tQsQ = smem_thr_copy_Q.partition_S(sQ_query)
+        tPsP = smem_thr_copy_P.partition_S(sP_query)
+
+        PackGQA(
+            self.tile_m,
+            self.tile_hdim,
+            self.check_hdim_oob,
+            self.qhead_per_kvhead,
+        ).load_Q(
+            mQ_cur,
+            sQ,
+            gmem_tiled_copy_Q,
+            tidx,
+            m_block,
+            seqlen.seqlen_q,
+        )
+        cute.arch.cp_async_commit_group()
+        cute.arch.cp_async_wait_group(0)
+        cute.arch.barrier(
+            barrier_id=1,
+            number_of_threads=self.num_Q_load_threads,
+        )
+
+        if const_expr(self.query_in_regs):
+            # Q.T is invariant across all KV tiles and its N=8 fragment is
+            # small. Load it while the first K TMA is in flight, then reuse it
+            # instead of reloading Q from SMEM per tile.
+            tSrQ_copy_view = smem_thr_copy_Q.retile(tSrQ)
+            for k in cutlass.range_constexpr(cute.size(tQsQ.shape[2])):
+                cute.copy(
+                    smem_thr_copy_Q,
+                    tQsQ[None, None, k],
+                    tSrQ_copy_view[None, None, k],
+                )
+
+        n_block = cutlass.max(n_block_max - 1, n_block_min)
+        consumer_state = self._compute_one_n_block_transposed(
+            n_block,
+            consumer_state,
+            acc_O,
+            sQ,
+            sK,
+            sV,
+            sP,
+            sRowScale,
+            pipeline_k,
+            pipeline_v,
+            tiled_mma_qk,
+            tiled_mma_pv,
+            smem_thr_copy_K,
+            smem_thr_copy_Q,
+            smem_thr_copy_V,
+            smem_thr_copy_P,
+            tSrK,
+            tSrQ,
+            tOrV,
+            tOrP,
+            tKsK,
+            tQsQ,
+            tVsV,
+            tPsP,
+            tidx,
+            softmax_scale_log2,
+            seqlen,
+            block_info.window_size_left,
+            is_first_n_block=True,
+        )
+        for n_tile in cutlass.range(n_block - n_block_min, unroll=1):
+            consumer_state = self._compute_one_n_block_transposed(
+                n_block - n_tile - 1,
+                consumer_state,
+                acc_O,
+                sQ,
+                sK,
+                sV,
+                sP,
+                sRowScale,
+                pipeline_k,
+                pipeline_v,
+                tiled_mma_qk,
+                tiled_mma_pv,
+                smem_thr_copy_K,
+                smem_thr_copy_Q,
+                smem_thr_copy_V,
+                smem_thr_copy_P,
+                tSrK,
+                tSrQ,
+                tOrV,
+                tOrP,
+                tKsK,
+                tQsQ,
+                tVsV,
+                tPsP,
+                tidx,
+                softmax_scale_log2,
+                seqlen,
+                block_info.window_size_left,
+            )
+
+        global_max_row = const_expr(8)
+        global_sum_row = const_expr(9)
+        final_scale_row = const_expr(0)
+        row_limit = seqlen.seqlen_q * self.qhead_per_kvhead
+        if tidx < self.query_mma_n:
+            query_row = tidx
+            row_max = sRowScale[global_max_row, query_row]
+            row_sum = sRowScale[global_sum_row, query_row]
+            if query_row < row_limit:
+                if const_expr(learnable_sink is not None):
+                    if const_expr(not self.is_split_kv) or split_idx == 0:
+                        q_head_idx = query_row + head_idx * self.qhead_per_kvhead
+                        sink_val = Float32(learnable_sink[q_head_idx])
+                        log2_e = math.log2(math.e)
+                        if row_max == -Float32.inf:
+                            row_max = sink_val * (log2_e / softmax_scale_log2)
+                            row_sum = 1.0
+                        else:
+                            row_sum += cute.math.exp2(
+                                sink_val * log2_e - row_max * softmax_scale_log2,
+                                fastmath=True,
+                            )
+                row_sum_is_zero_or_nan = row_sum == 0.0 or row_sum != row_sum
+                sRowScale[final_scale_row, query_row] = cute.arch.rcp_approx(
+                    row_sum if not row_sum_is_zero_or_nan else 1.0
+                )
+                sLSE[query_row] = (
+                    (
+                        row_max * softmax_scale_log2
+                        + cute.math.log2(row_sum, fastmath=True)
+                    )
+                    * math.log(2.0)
+                    if not row_sum_is_zero_or_nan
+                    else -Float32.inf
+                )
+            else:
+                sRowScale[final_scale_row, query_row] = 1.0
+                sLSE[query_row] = -Float32.inf
+        cute.arch.fence_view_async_shared()
+        cute.arch.barrier(
+            barrier_id=int(NamedBarrierFwd.PFull),
+            number_of_threads=self.num_mma_threads,
+        )
+        self._rescale_transposed_o(
+            acc_O,
+            tiled_mma_pv,
+            tidx,
+            sRowScale,
+            final_scale_row,
+        )
+        self._store_transposed_output(
+            acc_O,
+            mO,
+            mLSE,
+            sLSE,
+            tiled_mma_pv,
+            tidx,
+            seqlen,
+            m_block,
+            head_idx,
+            batch_idx,
+            split_idx,
+        )

python/sglang/kernels/ops/attention/fa4_sm120/paged_kv.py

@@ -0,0 +1,216 @@
+import math
+from dataclasses import dataclass
+from typing import Type
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32, const_expr
+from cutlass.cute import FastDivmodDivisor
+from cutlass.cute.nvgpu import cpasync
+from quack.cute_dsl_utils import ParamsBase
+
+from sglang.kernels.ops.attention.flash_attn.cute import utils
+
+
+@dataclass
+class Sm120PagedKVManager(ParamsBase):
+    """SM120 paged-KV loader for the stage-sliced FA4 pipeline."""
+
+    mPageTable: cute.Tensor
+    mK_paged: cute.Tensor
+    mV_paged: cute.Tensor
+    thread_idx: Int32
+
+    page_size_divmod: FastDivmodDivisor
+    seqlen_k: Int32
+    leftpad_k: Int32
+    n_block_size: cutlass.Constexpr[Int32]
+    num_threads: cutlass.Constexpr[Int32]
+    head_dim_padded: cutlass.Constexpr[Int32]
+    head_dim_v_padded: cutlass.Constexpr[Int32]
+
+    gmem_threads_per_row: cutlass.Constexpr[Int32]
+    page_entry_per_thread: cutlass.Constexpr[Int32]
+    async_copy_elems: cutlass.Constexpr[Int32]
+
+    gmem_tiled_copy_KV: cute.TiledCopy
+    gmem_thr_copy_KV: cute.TiledCopy
+    tPrPage: cute.Tensor
+    tPrPageOffset: cute.Tensor
+
+    @staticmethod
+    def create(
+        mPageTable: cute.Tensor,
+        mK_paged: cute.Tensor,
+        mV_paged: cute.Tensor,
+        page_size_divmod: FastDivmodDivisor,
+        bidb: Int32,
+        bidh: Int32,
+        thread_idx: Int32,
+        seqlen_k: Int32,
+        leftpad_k: Int32,
+        n_block_size: cutlass.Constexpr[Int32],
+        head_dim_padded: cutlass.Constexpr[Int32],
+        head_dim_v_padded: cutlass.Constexpr[Int32],
+        num_threads: cutlass.Constexpr[Int32],
+        dtype: Type[cutlass.Numeric],
+    ):
+        universal_copy_bits = 128
+        async_copy_elems = universal_copy_bits // dtype.width
+        dtype_bytes = dtype.width // 8
+        gmem_k_block_size = math.gcd(
+            head_dim_padded,
+            head_dim_v_padded,
+            128 // dtype_bytes,
+        )
+        assert gmem_k_block_size % async_copy_elems == 0
+        gmem_threads_per_row = gmem_k_block_size // async_copy_elems
+        assert cute.arch.WARP_SIZE % gmem_threads_per_row == 0
+
+        atom_async_copy = cute.make_copy_atom(
+            cpasync.CopyG2SOp(cache_mode=cpasync.LoadCacheMode.GLOBAL),
+            dtype,
+            num_bits_per_copy=universal_copy_bits,
+        )
+        thr_layout = cute.make_ordered_layout(
+            (num_threads // gmem_threads_per_row, gmem_threads_per_row),
+            order=(1, 0),
+        )
+        val_layout = cute.make_layout((1, async_copy_elems))
+        gmem_tiled_copy_KV = cute.make_tiled_copy_tv(
+            atom_async_copy, thr_layout, val_layout
+        )
+        gmem_thr_copy_KV = gmem_tiled_copy_KV.get_slice(thread_idx)
+
+        # SM120 decode tiles can have fewer rows than DMA threads. Keep one
+        # register entry per thread so those shapes do not create zero-sized
+        # register tensors.
+        page_entry_per_thread = max(1, (n_block_size + num_threads - 1) // num_threads)
+        tPrPage = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
+        tPrPageOffset = cute.make_rmem_tensor((page_entry_per_thread,), Int32)
+
+        return Sm120PagedKVManager(
+            mPageTable[bidb, None],
+            mK_paged[None, None, bidh, None],
+            mV_paged[None, None, bidh, None],
+            thread_idx,
+            page_size_divmod,
+            seqlen_k,
+            leftpad_k,
+            n_block_size,
+            num_threads,
+            head_dim_padded,
+            head_dim_v_padded,
+            gmem_threads_per_row,
+            page_entry_per_thread,
+            async_copy_elems,
+            gmem_tiled_copy_KV,
+            gmem_thr_copy_KV,
+            tPrPage,
+            tPrPageOffset,
+        )
+
+    @cute.jit
+    def _load_page_table_entry(self, i: Int32, n_block: Int32):
+        row = (
+            i * self.num_threads
+            + (self.thread_idx % self.gmem_threads_per_row)
+            * (self.num_threads // self.gmem_threads_per_row)
+            + (self.thread_idx // self.gmem_threads_per_row)
+        )
+        row_idx = n_block * self.n_block_size + row
+        page_idx, page_offset = divmod(row_idx + self.leftpad_k, self.page_size_divmod)
+        is_valid = (
+            (i + 1) * self.num_threads <= self.n_block_size or row < self.n_block_size
+        ) and row_idx < self.seqlen_k
+        page = self.mPageTable[page_idx] if is_valid else 0
+        self.tPrPage[i] = page
+        self.tPrPageOffset[i] = page_offset
+
+    @cute.jit
+    def load_page_table(self, n_block: Int32):
+        # The entry count is a specialization constant for SM120. Expanding
+        # this small loop removes a measurable dynamic-loop cost in decode.
+        for i in cutlass.range_constexpr(self.page_entry_per_thread):
+            self._load_page_table_entry(i, n_block)
+
+    @cute.jit
+    def compute_X_ptr(self, K_or_V: str):
+        tPrXPtr = cute.make_rmem_tensor((self.page_entry_per_thread,), cutlass.Int64)
+        mX = self.mK_paged if const_expr(K_or_V == "K") else self.mV_paged
+        for i in cutlass.range_constexpr(self.page_entry_per_thread):
+            page = self.tPrPage[i]
+            page_offset = self.tPrPageOffset[i]
+            # SGLang stores both paged K and paged V as
+            # (page_size, head_dim, num_pages).
+            tPrXPtr[i] = utils.elem_pointer(mX, (page_offset, 0, page)).toint()
+        return tPrXPtr
+
+    @cute.jit
+    def _copy_row_async(
+        self,
+        tXsX: cute.Tensor,
+        tXcX: cute.Tensor,
+        mX_paged_cur_copy: cute.Tensor,
+        m: Int32,
+        should_load: cute.Tensor,
+    ):
+        for k in cutlass.range_constexpr(cute.size(tXsX, mode=[2])):
+            ki = tXcX[0, 0, k][1] // self.async_copy_elems
+            mX_paged_cur_copy_ki = mX_paged_cur_copy[None, ki]
+            tXsX_k = tXsX[None, m, k]
+            mX_paged_cur_copy_ki = cute.make_tensor(
+                mX_paged_cur_copy_ki.iterator, tXsX_k.layout
+            )
+            cute.copy(
+                self.gmem_tiled_copy_KV,
+                mX_paged_cur_copy_ki,
+                tXsX_k,
+                pred=should_load,
+            )
+
+    @cute.jit
+    def load_KV(self, n_block: Int32, sX: cute.Tensor, K_or_V: str):
+        assert K_or_V in ("K", "V")
+
+        tPrXPtr = self.compute_X_ptr(K_or_V)
+
+        # The SM120 pipeline passes one stage at a time. V has already been
+        # transposed by the caller's shared-memory view.
+        sX_pi = cute.group_modes(sX, 0, 1)
+        head_dim = (
+            self.head_dim_v_padded
+            if const_expr(K_or_V == "V")
+            else self.head_dim_padded
+        )
+        cX = cute.make_identity_tensor((self.n_block_size, head_dim))
+        tXsX = self.gmem_thr_copy_KV.partition_D(sX_pi)
+        tXcX = self.gmem_thr_copy_KV.partition_S(cX)
+        tXc0X = self.gmem_thr_copy_KV.get_slice(0).partition_S(cX)
+
+        seqlenk_row_limit = (
+            self.seqlen_k - n_block * self.n_block_size - tXcX[0][0]
+            if n_block >= 0
+            else 0
+        )
+        for m in cutlass.range_constexpr(cute.size(tXsX, mode=[1])):
+            row_valid = tXc0X[0, m, 0][0] < seqlenk_row_limit
+            should_load = cute.make_fragment_like(tXsX[(0, None), m, 0], cute.Boolean)
+            should_load.fill(row_valid)
+
+            x_ptr_i64 = utils.shuffle_sync(
+                tPrXPtr[m // self.gmem_threads_per_row],
+                m % self.gmem_threads_per_row,
+                width=self.gmem_threads_per_row,
+            )
+            x_gmem_ptr = cute.make_ptr(
+                self.mK_paged.element_type,
+                x_ptr_i64,
+                cute.AddressSpace.gmem,
+                assumed_align=16,
+            )
+            mX_paged_cur = cute.make_tensor(x_gmem_ptr, cute.make_layout((head_dim,)))
+            mX_paged_cur_copy = cute.tiled_divide(
+                mX_paged_cur, (self.async_copy_elems,)
+            )
+            self._copy_row_async(tXsX, tXcX, mX_paged_cur_copy, m, should_load)

python/sglang/kernels/ops/attention/fa4_sm120/policy.py

@@ -0,0 +1,90 @@
+# Copyright (c) 2026, SGLang Team.
+"""Pure workload qualification shared by the SM120 FA4 host and kernel."""
+
+from typing import Optional
+
+LOW_HD_DECODE_SHAPES = frozenset({(64, 64), (128, 128)})
+LOW_HD_DECODE_TILE_N = 64
+LOW_HD_DECODE_MIN_VISIBLE_K = 256
+LOW_HD_DECODE_SHORT_VISIBLE_K = 512
+LOW_HD_DECODE_MAX_SPLITS = 8
+LOW_HD_DECODE_SHORT_MIN_TILES_PER_CTA = 2
+LOW_HD_DECODE_LONG_MIN_TILES_PER_CTA = 8
+
+
+def visible_decode_seqlen_k(
+    max_seqlen_k: int,
+    *,
+    is_local: bool,
+    window_size_left: Optional[int],
+    window_size_right: Optional[int],
+) -> int:
+    """Return the exact KV span visible to a single query position."""
+    if not is_local:
+        return max(0, max_seqlen_k)
+    left = max_seqlen_k if window_size_left is None else max(0, window_size_left)
+    right = max_seqlen_k if window_size_right is None else max(0, window_size_right)
+    return max(0, min(max_seqlen_k, left + 1 + right))
+
+
+def low_hd_paged_decode_tile_m(
+    *,
+    head_dim: int,
+    head_dim_v: int,
+    paged_kv: bool,
+    seqlen_q: Optional[int],
+    visible_seqlen_k: Optional[int],
+    qhead_per_kvhead: Optional[int],
+    num_sms: Optional[int] = None,
+    total_mblocks: Optional[int] = None,
+) -> Optional[int]:
+    """Return the qualified low-HD decode M tile, or ``None`` for fallback."""
+    if (
+        not paged_kv
+        or seqlen_q != 1
+        or visible_seqlen_k is None
+        or visible_seqlen_k <= LOW_HD_DECODE_MIN_VISIBLE_K
+        or (head_dim, head_dim_v) not in LOW_HD_DECODE_SHAPES
+    ):
+        return None
+    qhead_ratio = 1 if qhead_per_kvhead is None else qhead_per_kvhead
+    if (
+        (head_dim, head_dim_v) == (64, 64)
+        and qhead_ratio >= 8
+        and visible_seqlen_k <= LOW_HD_DECODE_SHORT_VISIBLE_K
+    ):
+        if num_sms is not None and total_mblocks is not None:
+            num_n_blocks = (
+                visible_seqlen_k + LOW_HD_DECODE_TILE_N - 1
+            ) // LOW_HD_DECODE_TILE_N
+            max_short_splits = min(
+                LOW_HD_DECODE_MAX_SPLITS,
+                num_n_blocks // LOW_HD_DECODE_SHORT_MIN_TILES_PER_CTA,
+            )
+            # Use the one-warp M16 CTA only when bounded SplitKV can fill an
+            # SM wave without reducing each partition below two KV tiles.
+            if total_mblocks * max_short_splits >= num_sms:
+                return 16
+        return None
+    if qhead_ratio >= 8 and visible_seqlen_k <= LOW_HD_DECODE_SHORT_VISIBLE_K:
+        return 32
+    return 16
+
+
+def is_low_hd_paged_decode_tile(
+    *,
+    head_dim: int,
+    head_dim_v: int,
+    paged_kv: bool,
+    seqlen_q: int,
+    tile_m: int,
+    tile_n: int,
+) -> bool:
+    """Return whether a selected tile belongs to the qualified low-HD path."""
+    return (
+        paged_kv
+        and seqlen_q == 1
+        and (head_dim, head_dim_v) in LOW_HD_DECODE_SHAPES
+        and tile_m in (16, 32)
+        and tile_n == LOW_HD_DECODE_TILE_N
+    )

python/sglang/kernels/ops/attention/fa4_sm120/scheduler.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2026, SGLang Team.
+"""Schedulers owned by the SGLang SM120 FA4 implementation."""
+
+from dataclasses import dataclass
+from typing import Tuple
+
+import cutlass
+import cutlass.cute as cute
+from cutlass import Int32
+from quack.cute_dsl_utils import ParamsBase
+
+from sglang.kernels.ops.attention.flash_attn.cute.tile_scheduler import (
+    SchedulingMode,
+    TileSchedulerArguments,
+    WorkTileInfo,
+)
+
+
+class Sm120UniformBatchScheduler:
+    """Map uniform varlen batches without the generic prefix-sum walk.
+
+    SM120 paged decode uses one equally sized query segment per request. Its
+    compile-time dispatch proves that invariant before selecting this scheduler,
+    so each CTA can recover ``(block, head, batch)`` arithmetically.
+    """
+
+    @dataclass
+    class Params(ParamsBase):
+        num_head: Int32
+        num_batch: Int32
+        total_q: Int32
+        num_splits: Int32
+        tile_shape_mn: cutlass.Constexpr[Tuple[int, int]]
+        is_split_kv: cutlass.Constexpr[bool] = False
+
+        @staticmethod
+        @cute.jit
+        def create(
+            args: TileSchedulerArguments, *, loc=None, ip=None
+        ) -> "Sm120UniformBatchScheduler.Params":
+            assert args.cluster_shape_mn == (
+                1,
+                1,
+            ), "SM120 uniform-batch scheduling requires a 1x1 cluster"
+            return Sm120UniformBatchScheduler.Params(
+                num_head=args.num_head,
+                num_batch=args.num_batch,
+                total_q=args.total_q,
+                num_splits=args.num_splits,
+                tile_shape_mn=args.tile_shape_mn,
+                is_split_kv=args.is_split_kv,
+            )
+
+    def __init__(
+        self,
+        params: Params,
+        tile_idx: Int32,
+        split_idx: Int32,
+        *,
+        loc=None,
+        ip=None,
+    ):
+        self.params = params
+        self._tile_idx = tile_idx
+        self._split_idx = split_idx
+        self._is_first_block = True
+        self._loc = loc
+        self._ip = ip
+
+    @staticmethod
+    def to_underlying_arguments(
+        args: TileSchedulerArguments,
+        *,
+        scheduling_mode: SchedulingMode = SchedulingMode.STATIC,
+        loc=None,
+        ip=None,
+    ) -> Params:
+        assert (
+            scheduling_mode == SchedulingMode.STATIC
+        ), f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}"
+        return Sm120UniformBatchScheduler.Params.create(args, loc=loc, ip=ip)
+
+    @staticmethod
+    @cute.jit
+    def create(
+        params: Params, clc=None, *, loc=None, ip=None
+    ) -> "Sm120UniformBatchScheduler":
+        tile_idx, split_idx, _ = cute.arch.block_idx()
+        return Sm120UniformBatchScheduler(params, tile_idx, split_idx, loc=loc, ip=ip)
+
+    @staticmethod
+    @cute.jit
+    def get_grid_shape(
+        params: Params, *, loc=None, ip=None
+    ) -> Tuple[Int32, Int32, Int32]:
+        rows_per_batch = params.total_q // params.num_batch
+        num_m_blocks = cute.ceil_div(rows_per_batch, params.tile_shape_mn[0])
+        return (
+            num_m_blocks * params.num_head * params.num_batch,
+            params.num_splits,
+            Int32(1),
+        )
+
+    @cute.jit
+    def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        params = self.params
+        rows_per_batch = params.total_q // params.num_batch
+        num_m_blocks = cute.ceil_div(rows_per_batch, params.tile_shape_mn[0])
+        mh_blocks_per_batch = num_m_blocks * params.num_head
+        batch_idx = self._tile_idx // mh_blocks_per_batch
+        mh_block = self._tile_idx - batch_idx * mh_blocks_per_batch
+        block = mh_block // params.num_head
+        head_idx = mh_block - block * params.num_head
+        split_idx = (
+            self._split_idx if cutlass.const_expr(params.is_split_kv) else Int32(0)
+        )
+        return WorkTileInfo(
+            (Int32(block), Int32(head_idx), Int32(batch_idx), split_idx),
+            self._is_first_block,
+        )
+
+    def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo:
+        return self.get_current_work(loc=loc, ip=ip)
+
+    def prefetch_next_work(self, *, loc=None, ip=None):
+        pass
+
+    def advance_to_next_work(self, *, loc=None, ip=None) -> WorkTileInfo:
+        self._is_first_block = False
+        return self.get_current_work(loc=loc, ip=ip)
+
+    def producer_tail(self, *, loc=None, ip=None):
+        pass
+
+    def __extract_mlir_values__(self):
+        values, self._values_pos = [], []
+        for obj in (self.params, self._tile_idx, self._split_idx):
+            obj_values = cutlass.extract_mlir_values(obj)
+            values += obj_values
+            self._values_pos.append(len(obj_values))
+        return values
+
+    def __new_from_mlir_values__(self, values):
+        objects = []
+        for obj, n_items in zip(
+            (self.params, self._tile_idx, self._split_idx), self._values_pos
+        ):
+            objects.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
+            values = values[n_items:]
+        return Sm120UniformBatchScheduler(*objects, loc=self._loc)

python/sglang/kernels/ops/attention/flash_attention.py

@@ -48,6 +48,7 @@ def flash_attn_with_kvcache(
     rel_bias_prep_cache=None,
     ver=3,
     out=None,
+    max_seqlen_k: Optional[int] = None,
 ):
     """
     If k and v are not None, k_cache and v_cache will be updated *inplace* with the new values from
@@ -194,6 +195,7 @@ def flash_attn_with_kvcache(
             page_table=page_table,
             cu_seqlens_q=cu_seqlens_q,
             max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
             rotary_seqlens=rotary_seqlens,
             q_descale=q_descale,
             k_descale=k_descale,
@@ -213,6 +215,7 @@ def flash_attn_with_kvcache(
             rel_bias=rel_bias,
             rel_bias_prep_cache=rel_bias_prep_cache,
             return_softmax_lse=return_softmax_lse,
+            out=out,
         )
     else:
         raise RuntimeError(f"Unknown flash attention version {ver}")
@@ -319,6 +322,7 @@ def flash_attn_varlen_func(
             rel_bias=rel_bias,
             rel_bias_prep_cache=rel_bias_prep_cache,
             return_softmax_lse=return_softmax_lse,
+            out=out,
         )
     else:
         raise RuntimeError(f"Unknown flash attention version {ver}")

python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py

@@ -0,0 +1,339 @@
+# Copyright (c) 2026, SGLang Team.
+"""SGLang-facing FlashAttention-4 APIs specialized for SM12x."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+
+from sglang.kernel_api_logging import debug_kernel_api
+from sglang.kernels.ops.attention.flash_attention_v4 import (
+    _flash_attn_import_error,
+    _flash_attn_varlen_func,
+    _maybe_contiguous,
+    _pad_mla_q_heads,
+    _unpad_mla_result,
+)
+
+if os.environ.get("SGLANG_INKLING_FA4_USE_PIP") == "1":
+    # The pip escape hatch deliberately bypasses SGLang-owned SM12x kernels.
+    get_forward_arch = None
+    resolve_runtime_policy = None
+    try_cached_paged_decode = None
+else:
+    from sglang.kernels.ops.attention.fa4_sm120.dispatch import (
+        get_forward_arch,
+        resolve_runtime_policy,
+        try_cached_paged_decode,
+    )
+
+
+@dataclass(frozen=True)
+class FlashAttentionV4SM120RuntimePolicy:
+    num_splits: int
+    decode_num_splits: int
+    decode_uses_static_max_seqlen_k: bool
+
+
+def get_flash_attention_v4_sm120_runtime_policy(
+    *,
+    device_capability: tuple[int, int],
+    deterministic: bool,
+) -> FlashAttentionV4SM120RuntimePolicy:
+    """Resolve the SM12x FA4 launch policy exposed to SGLang."""
+    if resolve_runtime_policy is None:
+        num_splits = 1 if deterministic or device_capability < (9, 0) else 0
+        return FlashAttentionV4SM120RuntimePolicy(
+            num_splits=num_splits,
+            decode_num_splits=num_splits,
+            decode_uses_static_max_seqlen_k=False,
+        )
+    num_splits, decode_num_splits, decode_uses_static_max_seqlen_k = (
+        resolve_runtime_policy(
+            device_capability=device_capability,
+            deterministic=deterministic,
+        )
+    )
+    return FlashAttentionV4SM120RuntimePolicy(
+        num_splits=num_splits,
+        decode_num_splits=decode_num_splits,
+        decode_uses_static_max_seqlen_k=decode_uses_static_max_seqlen_k,
+    )
+
+
+def _validate_out_contract(out: Optional[torch.Tensor]) -> None:
+    if out is None:
+        return
+    if out.requires_grad:
+        raise ValueError("out must not require gradients")
+    if out.stride(-1) != 1:
+        raise ValueError("out must have stride 1 in the last dimension")
+
+
+@debug_kernel_api
+def flash_attn_varlen_func(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k: Optional[torch.Tensor] = None,
+    qv: Optional[torch.Tensor] = None,
+    seqused_q: Optional[torch.Tensor] = None,
+    seqused_k: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    max_seqlen_k: Optional[int] = None,
+    page_table: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    softcap: Optional[float] = None,
+    window_size: Tuple[Optional[int], Optional[int]] = (-1, -1),
+    learnable_sink: Optional[torch.Tensor] = None,
+    sinks: Optional[torch.Tensor] = None,
+    num_splits: int = 1,
+    pack_gqa: Optional[bool] = None,
+    score_mod: Optional[Callable] = None,
+    aux_tensors: Optional[list] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    sfq: Optional[torch.Tensor] = None,
+    sfk: Optional[torch.Tensor] = None,
+    sfv: Optional[torch.Tensor] = None,
+    rel_bias: Optional[torch.Tensor] = None,
+    rel_bias_prep_cache: Optional[dict] = None,
+    return_softmax_lse: bool = False,
+    out: Optional[torch.Tensor] = None,
+    **_: object,
+):
+    if _flash_attn_varlen_func is None:  # pragma: no cover
+        raise ImportError(
+            "FlashAttention-4 CUTE is not available. Install flash-attn-4 with "
+            "its CUDA/CUTE dependencies, or run from a source tree where the "
+            "vendored FA4 package is importable."
+        ) from _flash_attn_import_error
+
+    _validate_out_contract(out)
+    q, k, v, qv = [_maybe_contiguous(t) for t in (q, k, v, qv)]
+    if qv is None and q.shape[-1] == 256 and k.shape[-1] == 256 and v.shape[-1] == 256:
+        # The vendored hd256 kernel assumes dense Q/K/V strides.
+        q, k, v = [t.contiguous() for t in (q, k, v)]
+    q, qv, mla_head_padding = _pad_mla_q_heads(q, qv, v, pack_gqa)
+    if qv is not None and num_splits < 1:
+        # FA4 MLA does not implement split-KV; auto mode must use one split.
+        num_splits = 1
+    cu_seqlens_q, cu_seqlens_k = [
+        _maybe_contiguous(t) for t in (cu_seqlens_q, cu_seqlens_k)
+    ]
+    seqused_q, seqused_k = [_maybe_contiguous(t) for t in (seqused_q, seqused_k)]
+    page_table = _maybe_contiguous(page_table)
+
+    if learnable_sink is None and sinks is not None:
+        learnable_sink = sinks
+    if window_size == (-1, -1):
+        window_size = (None, None)
+
+    sf_kwargs = {}
+    if sfq is not None:
+        sf_kwargs["sfq"] = sfq
+    if sfk is not None:
+        sf_kwargs["sfk"] = sfk
+    if sfv is not None:
+        sf_kwargs["sfv"] = sfv
+
+    descale_kwargs = {}
+    if q_descale is not None:
+        descale_kwargs["q_descale"] = q_descale
+    if k_descale is not None:
+        descale_kwargs["k_descale"] = k_descale
+    if v_descale is not None:
+        descale_kwargs["v_descale"] = v_descale
+
+    rel_bias_kwargs = {}
+    if rel_bias is not None:
+        rel_bias_kwargs["rel_bias"] = rel_bias
+    if rel_bias_prep_cache is not None:
+        rel_bias_kwargs["rel_bias_prep_cache"] = rel_bias_prep_cache
+
+    result = _flash_attn_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        qv=qv,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        seqused_q=seqused_q,
+        seqused_k=seqused_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        softcap=softcap,
+        window_size=window_size,
+        learnable_sink=learnable_sink,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        score_mod=score_mod,
+        aux_tensors=aux_tensors,
+        return_lse=return_softmax_lse,
+        out=out,
+        **sf_kwargs,
+        **descale_kwargs,
+        **rel_bias_kwargs,
+    )
+    result = _unpad_mla_result(result, mla_head_padding)
+
+    if return_softmax_lse:
+        return result
+    if isinstance(result, tuple):
+        return result[0]
+    return result
+
+
+@debug_kernel_api
+def flash_attn_with_kvcache(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    k: Optional[torch.Tensor] = None,
+    v: Optional[torch.Tensor] = None,
+    qv: Optional[torch.Tensor] = None,
+    rotary_cos: Optional[torch.Tensor] = None,
+    rotary_sin: Optional[torch.Tensor] = None,
+    cache_seqlens: Optional[Union[int, torch.Tensor]] = None,
+    cache_batch_idx: Optional[torch.Tensor] = None,
+    cache_leftpad: Optional[torch.Tensor] = None,
+    page_table: Optional[torch.Tensor] = None,
+    cu_seqlens_q: Optional[torch.Tensor] = None,
+    cu_seqlens_k_new: Optional[torch.Tensor] = None,
+    max_seqlen_q: Optional[int] = None,
+    rotary_seqlens: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    window_size: Tuple[int, int] = (-1, -1),
+    attention_chunk: Optional[int] = None,
+    softcap: float = 0.0,
+    rotary_interleaved: bool = True,
+    scheduler_metadata=None,
+    num_splits: int = 0,
+    pack_gqa: Optional[bool] = None,
+    sm_margin: int = 0,
+    sinks: Optional[torch.Tensor] = None,
+    score_mod: Optional[Callable] = None,
+    aux_tensors: Optional[list] = None,
+    sfq: Optional[torch.Tensor] = None,
+    sfk: Optional[torch.Tensor] = None,
+    sfv: Optional[torch.Tensor] = None,
+    rel_bias: Optional[torch.Tensor] = None,
+    rel_bias_prep_cache: Optional[dict] = None,
+    return_softmax_lse: bool = False,
+    out: Optional[torch.Tensor] = None,
+    max_seqlen_k: Optional[int] = None,
+    **_: object,
+):
+    _validate_out_contract(out)
+    if k is not None or v is not None:
+        raise NotImplementedError("FA4 does not support updating KV cache in-place.")
+    if rotary_cos is not None or rotary_sin is not None or rotary_seqlens is not None:
+        raise NotImplementedError("FA4 path does not support rotary embedding.")
+    if cache_batch_idx is not None or cache_leftpad is not None:
+        raise NotImplementedError(
+            "FA4 path does not support non-consecutive batch indices or left padding."
+        )
+    if isinstance(cache_seqlens, int):
+        cache_seqlens = torch.full(
+            (k_cache.shape[0],), cache_seqlens, dtype=torch.int32, device=k_cache.device
+        )
+
+    forward_arch = get_forward_arch(q.device) if get_forward_arch is not None else None
+    if (
+        forward_arch is not None
+        and not return_softmax_lse
+        and softcap in (None, 0.0)
+        and all(
+            value is None
+            for value in (
+                qv,
+                score_mod,
+                aux_tensors,
+                q_descale,
+                k_descale,
+                v_descale,
+                sfq,
+                sfk,
+                sfv,
+                rel_bias,
+                rel_bias_prep_cache,
+            )
+        )
+    ):
+        q, k_cache, v_cache = [_maybe_contiguous(t) for t in (q, k_cache, v_cache)]
+        cu_seqlens_q, cache_seqlens, page_table = [
+            _maybe_contiguous(t) for t in (cu_seqlens_q, cache_seqlens, page_table)
+        ]
+        fast_result = try_cached_paged_decode(
+            arch=forward_arch,
+            q=q,
+            k=k_cache,
+            v=v_cache,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=None,
+            seqused_q=None,
+            seqused_k=cache_seqlens,
+            page_table=page_table,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size_left=window_size[0],
+            window_size_right=window_size[1],
+            learnable_sink=sinks,
+            requested_num_splits=num_splits,
+            pack_gqa=pack_gqa,
+            out=out,
+        )
+        if fast_result is not None:
+            return fast_result[0]
+
+    result = flash_attn_varlen_func(
+        q=q,
+        k=k_cache,
+        v=v_cache,
+        qv=qv,
+        cu_seqlens_q=cu_seqlens_q,
+        seqused_k=cache_seqlens,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        page_table=page_table,
+        softmax_scale=softmax_scale,
+        causal=causal,
+        softcap=softcap if softcap != 0.0 else None,
+        window_size=window_size,
+        num_splits=num_splits,
+        pack_gqa=pack_gqa,
+        learnable_sink=sinks,
+        score_mod=score_mod,
+        aux_tensors=aux_tensors,
+        q_descale=q_descale,
+        k_descale=k_descale,
+        v_descale=v_descale,
+        sfq=sfq,
+        sfk=sfk,
+        sfv=sfv,
+        rel_bias=rel_bias,
+        rel_bias_prep_cache=rel_bias_prep_cache,
+        return_softmax_lse=return_softmax_lse if forward_arch is not None else True,
+        out=out,
+    )
+
+    if return_softmax_lse:
+        return result
+    if isinstance(result, tuple):
+        return result[0]
+    return result

python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_sm120.py

@@ -1,63 +0,0 @@
-# Copyright (c) 2025, Jay Shah, Ganesh Bikshandi, Ying Zhang, Vijay Thakkar, Pradeep Ramani, Tri Dao.
-# SM120 (Blackwell GeForce / DGX Spark) forward pass.
-#
-# SM120 uses the same SM80-era MMA instructions (mma.sync.aligned.m16n8k16) but has
-# a smaller shared memory capacity (99 KB vs 163 KB on SM80). This module subclasses
-# FlashAttentionForwardSm80 and overrides the SMEM capacity check accordingly.
-
-import cutlass
-import cutlass.utils as utils_basic
-
-from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd import (
-    FlashAttentionForwardSm80,
-)
-
-
-class FlashAttentionForwardSm120(FlashAttentionForwardSm80):
-    # Keep arch = 80 to use CpAsync code paths (no TMA for output).
-    # The compilation target is determined by the GPU at compile time, not this field.
-    arch = 80
-
-    @staticmethod
-    def can_implement(
-        dtype,
-        head_dim,
-        head_dim_v,
-        tile_m,
-        tile_n,
-        num_stages,
-        num_threads,
-        is_causal,
-        Q_in_regs=False,
-    ) -> bool:
-        """Check if the kernel can be implemented on SM120.
-
-        Same logic as SM80 but uses SM120's shared memory capacity (99 KB).
-        """
-        if dtype not in [cutlass.Float16, cutlass.BFloat16]:
-            return False
-        if head_dim % 8 != 0:
-            return False
-        if head_dim_v % 8 != 0:
-            return False
-        if tile_n % 16 != 0:
-            return False
-        if num_threads % 32 != 0:
-            return False
-        # Shared memory usage: Q tile + (K tile + V tile)
-        smem_usage_Q = tile_m * head_dim * 2
-        smem_usage_K = tile_n * head_dim * num_stages * 2
-        smem_usage_V = tile_n * head_dim_v * num_stages * 2
-        smem_usage_QV = (
-            (smem_usage_Q + smem_usage_V)
-            if not Q_in_regs
-            else max(smem_usage_Q, smem_usage_V)
-        )
-        smem_usage = smem_usage_QV + smem_usage_K
-        # SM120 has 99 KB shared memory (vs 163 KB on SM80)
-        smem_capacity = utils_basic.get_smem_capacity_in_bytes("sm_120")
-        if smem_usage > smem_capacity:
-            return False
-        if (tile_m * 2) % num_threads != 0:
-            return False
-        return True

python/sglang/kernels/ops/attention/flash_attn/cute/interface.py

@@ -58,8 +58,10 @@ from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm100 import (
     DescaleTensors,
     FlashAttentionForwardSm100,
 )
-from sglang.kernels.ops.attention.flash_attn.cute.flash_fwd_sm120 import (
-    FlashAttentionForwardSm120,
+from sglang.kernels.ops.attention.fa4_sm120.dispatch import (
+    get_forward_host,
+    try_cached_paged_decode,
+    try_cached_varlen,
 )
 from sglang.kernels.ops.attention.flash_attn.cute.shearing_bias import ShearingBias
 
@@ -85,9 +87,8 @@ def _parse_arch_str(arch_str):
 def _get_device_arch():
     """Cached device arch check.
 
-    Override with FLASH_ATTENTION_ARCH (e.g. 'sm_80' or '80') to select which
-    kernel path to use (SM80/SM90/SM100/SM120) independently of the compilation
-    target (CUTE_DSL_ARCH).
+    Override with FLASH_ATTENTION_ARCH (e.g. 'sm_80' or '80') to select the
+    kernel path independently of the compilation target (CUTE_DSL_ARCH).
 
     For CPU-only compilation (no GPU), set both:
       FLASH_ATTENTION_ARCH=sm_80  (kernel selection)
@@ -100,6 +101,12 @@ def _get_device_arch():
     return major * 10 + int(minor)
 
 
+@lru_cache(maxsize=None)
+def _get_device_num_sms(device: torch.device) -> int:
+    """Return the stable SM count without querying CUDA on every launch."""
+    return torch.cuda.get_device_properties(device).multi_processor_count
+
+
 def _validate_head_dims(
     head_dim: int, head_dim_v: int, compute_capability: int, alignment: int
 ) -> None:
@@ -342,6 +349,60 @@ def _flash_attn_fwd(
         aux_tensors: Some score_mods will want to read from global aux_tensors. This is how we thread them through to the inner kernel.
         aux_scalars: Runtime scalar captures used by score_mod or mask_mod.
     """
+    fake_mode = is_fake_mode()
+    arch = _get_device_arch() if _arch is None else _arch
+    arch_forward_host = get_forward_host(arch)
+    requested_num_splits = num_splits
+    if (
+        not fake_mode
+        and arch_forward_host is not None
+        and not return_lse
+        and lse is None
+        and softcap in (None, 0.0)
+        and all(
+            value is None
+            for value in (
+                qv,
+                gather_kv_indices,
+                score_mod,
+                mask_mod,
+                block_sparse_tensors,
+                aux_tensors,
+                aux_scalars,
+                q_descale,
+                k_descale,
+                v_descale,
+                rel_bias,
+                sfq,
+                sfk,
+                sfv,
+            )
+        )
+    ):
+        fast_result = try_cached_paged_decode(
+            arch=arch,
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            seqused_q=seqused_q,
+            seqused_k=seqused_k,
+            page_table=page_table,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+            learnable_sink=learnable_sink,
+            requested_num_splits=requested_num_splits,
+            pack_gqa=pack_gqa,
+            out=out,
+        )
+        if fast_result is not None:
+            return fast_result
+
     aux_scalars = tuple(aux_scalars) if aux_scalars else None
     q, k, v, qv = [maybe_contiguous(t) for t in (q, k, v, qv)]
     assert q is not None or qv is not None
@@ -478,7 +539,7 @@ def _flash_attn_fwd(
         assert learnable_sink.shape == (num_head,)
         assert learnable_sink.dtype == torch.bfloat16, "learnable_sink must be bfloat16"
 
-    if not is_fake_mode():
+    if not fake_mode:
         assert all(
             t is None or t.is_cuda
             for t in (
@@ -497,17 +558,13 @@ def _flash_attn_fwd(
                 learnable_sink,
             )
         ), "inputs must be on CUDA device"
-    arch = _get_device_arch() if _arch is None else _arch
-    assert arch // 10 in [
-        8,
-        9,
-        10,
-        11,
-        12,
-    ], "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, 12.x"
+    assert arch // 10 in [8, 9, 10, 11] or arch_forward_host is not None, (
+        "Unsupported compute capability. Supported: 8.x, 9.x, 10.x, 11.x, "
+        "and architectures registered through the forward-host bridge"
+    )
     assert num_head % num_head_kv == 0, "num_head must be divisible by num_head_kv"
     alignment = 16 // v.element_size()
-    if arch // 10 not in [8, 12]:
+    if arch // 10 != 8 and arch_forward_host is None:
         _validate_head_dims(head_dim, head_dim_v, arch // 10, alignment)
     if softmax_scale is None:
         softmax_scale = (
@@ -565,6 +622,10 @@ def _flash_attn_fwd(
             device=device,
         )
     else:
+        if out.requires_grad:
+            raise ValueError("out must not require gradients")
+        if out.stride(-1) != 1:
+            raise ValueError("out must have stride 1 in the last dimension")
         _validate_tensor(
             out,
             "out",
@@ -619,21 +680,41 @@ def _flash_attn_fwd(
 
     current_stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
 
-    # SM80/SM120: uses SM80 MMA, 128 threads (4 warps)
-    if arch // 10 in [8, 12]:
+    if arch // 10 == 8:
         num_threads = 128
+    num_SMs = 132 if fake_mode else _get_device_num_sms(device)
 
     fwd_cfg = FwdConfig(128, 128, True, True)  # default
-    if tile_mn is None:
-        if arch // 10 == 12:
-            # SM120 tile sizes tuned for 99 KB SMEM capacity:
-            # D<=64:  128x128 → 48 KB (good occupancy)
-            # D>64:   128x64  → 64 KB (128x128 would use 96 KB, hurting occupancy)
-            if head_dim <= 64:
-                fwd_cfg = FwdConfig(128, 128, True, True)
-            else:
-                fwd_cfg = FwdConfig(128, 64, True, True)
-        elif arch // 10 == 8:
+    arch_forward_config = None
+    if arch_forward_host is not None:
+        arch_forward_config = arch_forward_host.select_config(
+            head_dim=head_dim,
+            head_dim_v=head_dim_v,
+            tile_mn=tile_mn,
+            has_bias=rel_bias is not None,
+            total_q_rows=total_q * num_head,
+            num_sms=None if fake_mode else num_SMs,
+            num_batch=batch_size,
+            seqlen_q=(max_seqlen_q if max_seqlen_q is not None else seqlen_q),
+            seqlen_k=(max_seqlen_k if max_seqlen_k is not None else seqlen_k),
+            num_head_kv=num_head_kv,
+            qhead_per_kvhead=qhead_per_kvhead,
+            is_causal=causal,
+            is_local=local,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+            pack_gqa=pack_gqa,
+            paged_kv=page_table is not None,
+        )
+        fwd_cfg = FwdConfig(
+            arch_forward_config.tile_m,
+            arch_forward_config.tile_n,
+            True,
+            True,
+        )
+        num_threads = arch_forward_config.num_threads
+    elif tile_mn is None:
+        if arch // 10 == 8:
             fwd_cfg = FwdConfig(128, 64, True, True)  # SM80, should tune
         elif arch // 10 == 9:
             sparse_q = get_sparse_q_block_size(block_sparse_tensors, seqlen_q)
@@ -683,13 +764,58 @@ def _flash_attn_fwd(
     ) // m_block_size_effective
     total_mblocks = batch_size * num_head_kv * num_m_blocks
     num_n_blocks = (seqlen_k_loaded + tile_n - 1) // tile_n
-    num_SMs = (
-        132
-        if is_fake_mode()
-        else torch.cuda.get_device_properties(device).multi_processor_count
-    )
-    if num_splits < 1:
-        num_splits = num_splits_heuristic(total_mblocks, num_SMs, num_n_blocks, 128)
+    arch_forward_plan = None
+    if arch_forward_host is not None:
+        arch_forward_plan = arch_forward_host.resolve_plan(
+            requested_num_splits=num_splits,
+            generic_num_n_blocks=num_n_blocks,
+            head_dim=head_dim,
+            head_dim_v=head_dim_v,
+            batch_size=batch_size,
+            num_head_kv=num_head_kv,
+            paged_kv=page_table is not None,
+            page_size=page_size,
+            k=k,
+            v=v,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            pack_gqa=pack_gqa,
+            element_size=q.element_size(),
+            packed_q_rows=seqlen_q_packgqa,
+            tile_m=tile_m,
+            tile_n=tile_n,
+            num_m_blocks=num_m_blocks,
+            total_mblocks=total_mblocks,
+            num_sms=num_SMs,
+            total_q=total_q,
+            has_cu_seqlens_q=cu_seqlens_q is not None,
+            has_seqused_q=seqused_q is not None,
+            has_seqused_k=seqused_k is not None,
+            is_causal=causal,
+            is_local=local,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+            has_score_or_mask_mod=(
+                softcap is not None
+                or score_mod is not None
+                or mask_mod is not None
+                or rel_bias is not None
+            ),
+            is_stream_capturing=(
+                not fake_mode and torch.cuda.is_current_stream_capturing()
+            ),
+            device=device,
+            fake_mode=fake_mode,
+            generic_heuristic=num_splits_heuristic,
+        )
+        num_splits = arch_forward_plan.num_splits
+    elif num_splits < 1:
+        num_splits = num_splits_heuristic(
+            total_mblocks,
+            num_SMs,
+            num_n_blocks,
+            128,
+        )
 
     # SplitKV uses float32 partial output, which doubles the O buffer size
     # in shared memory, causing OOM for diff-headdim (192, 128)
@@ -859,8 +985,8 @@ def _flash_attn_fwd(
         disable_sparse_kv_bitmask = None
         p = row_max = None
 
-    # rel_bias -> sheared bias (Inkling relative attention). Produces `bias`, the column-aligned
-    # bias the SM100 kernel adds to pre-softmax scores via its dedicated TMA pipeline.
+    # Inkling relative attention. Shear the relative rows into the column-aligned
+    # tiles consumed by the architecture-specific attention mainloop.
     rel_extent = 0
     rel_extent_padded = 0
     bias = None
@@ -868,16 +994,17 @@ def _flash_attn_fwd(
     cu_total_m_blocks_bias = None
     blocks_to_batch_idx = None
     if rel_bias is not None:
-        assert arch // 10 in [
-            9,
-            10,
-            11,
-        ], "rel_bias (sheared bias) is only supported on SM9x/10x"
+        assert arch // 10 in [9, 10, 11] or arch_forward_host is not None, (
+            "rel_bias requires SM9x/10x/11x or an architecture-owned "
+            "forward implementation"
+        )
         qhead_per_kvhead_packgqa = qhead_per_kvhead if pack_gqa else 1
         rel_extent = rel_bias.shape[-1]
         rel_extent_padded = rel_extent + 256
         assert rel_extent % 128 == 0
-        assert tile_m == 128 and tile_n == 128
+        assert tile_n == 128
+        if arch_forward_host is None:
+            assert tile_m == 128
         assert (
             causal
             or window_size_left is None
@@ -975,7 +1102,7 @@ def _flash_attn_fwd(
                         current_stream,
                         options="--enable-tvm-ffi",
                     )
-                if not is_fake_mode():
+                if not fake_mode:
                     _flash_attn_fwd.compile_cache_prepare_shear_bias[
                         compile_key_prepare
                     ](
@@ -1003,6 +1130,7 @@ def _flash_attn_fwd(
             qhead_per_kvhead,
             rows_per_cta,
             group_tile_bias,
+            tile_m,
             max_m_blocks_leq_one,
             cu_total_m_blocks_bias is not None,
             blocks_to_batch_idx is not None,
@@ -1039,6 +1167,7 @@ def _flash_attn_fwd(
                     qhead_per_kvhead=qhead_per_kvhead,
                     rows_per_cta=rows_per_cta,
                     tile_m=group_tile_bias,
+                    attention_tile_m=tile_m,
                     max_m_blocks_leq_one=max_m_blocks_leq_one,
                     use_pdl=use_pdl,
                 ),
@@ -1057,7 +1186,7 @@ def _flash_attn_fwd(
                 current_stream,
                 options="--enable-tvm-ffi",
             )
-        if not is_fake_mode():
+        if not fake_mode:
             _flash_attn_fwd.compile_cache_shear_bias[shear_compile_key](
                 rel_bias,
                 bias,
@@ -1111,6 +1240,8 @@ def _flash_attn_fwd(
         is_split_kv,
         pack_gqa,
         arch,
+        arch_forward_config.compile_key if arch_forward_config is not None else None,
+        arch_forward_plan.compile_key if arch_forward_plan is not None else None,
         page_size not in [None, tile_n],  # paged KV non-TMA
         use_2cta_instrs,
         q_subtile_factor,
@@ -1364,31 +1495,37 @@ def _flash_attn_fwd(
                         )
                     ),
                 )
-        elif arch // 10 == 12:
-            # SM120 (Blackwell GeForce / DGX Spark): uses SM80 MMA with SM120 SMEM capacity
-            assert not use_block_sparsity, "Block sparsity not supported on SM 12.0"
-            assert page_table is None, "Paged KV not supported on SM 12.0 in this PR"
-            assert not is_split_kv, "SplitKV not supported on SM 12.0 in this PR"
-            fa_fwd = FlashAttentionForwardSm120(
-                dtype,
-                head_dim,
-                head_dim_v,
-                qhead_per_kvhead,
+        elif arch_forward_host is not None:
+            assert not use_block_sparsity, (
+                "Block sparsity is not supported by the architecture-owned "
+                "forward implementation"
+            )
+            assert arch_forward_config is not None
+            assert arch_forward_plan is not None
+            fa_fwd = arch_forward_host.make_kernel(
+                dtype=dtype,
+                head_dim=head_dim,
+                head_dim_v=head_dim_v,
+                qhead_per_kvhead=qhead_per_kvhead,
                 is_causal=causal,
                 is_local=local,
                 pack_gqa=pack_gqa,
-                tile_m=tile_m,
-                tile_n=tile_n,
-                num_stages=1,
-                num_threads=num_threads,
-                Q_in_regs=False,
+                config=arch_forward_config,
+                paged_kv=page_table is not None,
                 score_mod=score_mod,
                 mask_mod=mask_mod,
                 has_aux_tensors=aux_tensors is not None,
+                is_split_kv=is_split_kv,
+                has_bias=bias is not None,
+                bias_block_size=tile_bias,
+                rel_extent_padded=rel_extent_padded,
+                plan=arch_forward_plan,
             )
         else:
             raise ValueError(
-                f"Unsupported compute capability: {arch}. Supported: 8.x, 9.x, 10.x, 11.x, 12.x"
+                f"Unsupported compute capability: {arch}. Supported: 8.x, 9.x, "
+                "10.x, 11.x, and architectures registered through the "
+                "forward-host bridge"
             )
         # TODO: check @can_implement
         if qv is not None:
@@ -1440,7 +1577,7 @@ def _flash_attn_fwd(
                     AuxData(cute_aux_tensors, aux_scalars),
                 ]
             )
-            if arch // 10 in [9, 10, 11]:
+            if arch // 10 in [9, 10, 11] or arch_forward_host is not None:
                 compile_args.append(bias_tensor)  # mBias
             if arch // 10 in [10, 11]:
                 if not use_dedicated_hd256_kernel:
@@ -1453,12 +1590,17 @@ def _flash_attn_fwd(
                             v_sf_vec_size,
                         ]
                     )
+            if arch_forward_host is not None:
+                assert arch_forward_plan is not None
+                compile_args.extend(
+                    arch_forward_host.compile_arguments(arch_forward_plan)
+                )
             compile_args.append(current_stream)
             _flash_attn_fwd.compile_cache[compile_key] = cute.compile(
                 *compile_args, options="--enable-tvm-ffi"
             )
 
-    if not is_fake_mode():
+    if not fake_mode:
         q_call, k_call, v_call, qv_call = [
             t.detach() if t is not None else None for t in (q, k, v, qv)
         ]
@@ -1543,7 +1685,7 @@ def _flash_attn_fwd(
                     AuxData(aux_tensors, aux_scalars),
                 ]
             )
-            if arch // 10 in [9, 10, 11]:
+            if arch // 10 in [9, 10, 11] or arch_forward_host is not None:
                 call_args.append(bias)  # mBias
             if arch // 10 in [10, 11]:
                 if not use_dedicated_hd256_kernel:
@@ -1556,7 +1698,52 @@ def _flash_attn_fwd(
                             sfv_call,  # mSFV (None unless v_blockscaled)
                         ]
                     )
-            _flash_attn_fwd.compile_cache[compile_key](*call_args)
+            if arch_forward_host is not None:
+                assert arch_forward_plan is not None
+                call_args.extend(arch_forward_host.runtime_arguments(arch_forward_plan))
+            compiled_fwd = _flash_attn_fwd.compile_cache[compile_key]
+            if (
+                arch_forward_host is not None
+                and cu_seqlens_q is not None
+                and cu_seqlens_k is not None
+                and seqused_q is None
+                and seqused_k is None
+                and not is_split_kv
+                and not requires_grad
+                and (learnable_sink is None or not learnable_sink.requires_grad)
+                and lse is None
+                and softcap is None
+                and score_mod is None
+                and mask_mod is None
+                and block_sparse_tensors is None
+                and aux_tensors is None
+                and aux_scalars is None
+                and q_descale is None
+                and k_descale is None
+                and v_descale is None
+                and rel_bias is None
+                and sfq is None
+                and sfk is None
+                and sfv is None
+            ):
+                arch_forward_host.register_varlen(
+                    arch=arch,
+                    compiled_fn=compiled_fwd,
+                    compile_key=compile_key,
+                    q=q,
+                    k=k,
+                    v=v,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_k=cu_seqlens_k,
+                    max_seqlen_q=max_seqlen_q,
+                    max_seqlen_k=max_seqlen_k,
+                    causal=causal,
+                    window_size_left=window_size_left,
+                    window_size_right=window_size_right,
+                    learnable_sink=learnable_sink,
+                    pack_gqa=pack_gqa,
+                )
+            compiled_fwd(*call_args)
     if is_split_kv:
         _flash_attn_fwd_combine(
             out_partial,
@@ -1566,6 +1753,72 @@ def _flash_attn_fwd(
             cu_seqlens_q,
             seqused_q,
         )
+    if (
+        not fake_mode
+        and not torch.cuda.is_current_stream_capturing()
+        and arch_forward_host is not None
+        and q is not None
+        and k is not None
+        and qv is None
+        and cu_seqlens_q is not None
+        and cu_seqlens_k is None
+        and seqused_q is None
+        and seqused_k is not None
+        and page_table is not None
+        and not requires_grad
+        and (learnable_sink is None or not learnable_sink.requires_grad)
+        and lse is None
+        and softcap is None
+        and score_mod is None
+        and mask_mod is None
+        and block_sparse_tensors is None
+        and aux_tensors is None
+        and aux_scalars is None
+        and q_descale is None
+        and k_descale is None
+        and v_descale is None
+        and rel_bias is None
+        and sfq is None
+        and sfk is None
+        and sfv is None
+    ):
+        compiled_combine = None
+        if is_split_kv:
+            lse_partial_transposed = lse_partial.transpose(-1, -2)
+            combine_key = _fwd_combine_compile_key(
+                out_partial,
+                out,
+                None,
+                cu_seqlens_q,
+                None,
+                None,
+            )
+            compiled_combine = _flash_attn_fwd_combine.compile_cache[combine_key]
+        arch_forward_host.register_paged_decode(
+            arch=arch,
+            compiled_fn=compiled_fwd,
+            compile_key=compile_key,
+            compiled_combine=compiled_combine,
+            actual_num_splits=num_splits,
+            tile_m=tile_m,
+            tile_n=tile_n,
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens_q,
+            seqused_k=seqused_k,
+            page_table=page_table,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            causal=causal,
+            window_size_left=window_size_left,
+            window_size_right=window_size_right,
+            learnable_sink=learnable_sink,
+            pack_gqa=pack_gqa,
+            requested_num_splits=requested_num_splits,
+            out_partial=out_partial if is_split_kv else None,
+            lse_partial=lse_partial if is_split_kv else None,
+        )
     return out, lse
 
 
@@ -1689,6 +1942,7 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
         v_sf_vec_size: Optional[int] = None,
         rel_bias_prep_cache: Optional[dict] = None,
         return_lse: bool = False,
+        out: Optional[torch.Tensor] = None,
     ):
         aux_scalars = tuple(aux_scalars) if aux_scalars else None
         shared_kv = k is v
@@ -1736,32 +1990,34 @@ class FlashAttnVarlenFunc(torch.autograd.Function):
             qk_sf_vec_size=qk_sf_vec_size,
             v_sf_vec_size=v_sf_vec_size,
             rel_bias_prep_cache=rel_bias_prep_cache,
+            out=out,
         )
-        ctx.save_for_backward(
-            q,
-            k,
-            v,
-            out,
-            lse,
-            cu_seqlens_q,
-            cu_seqlens_k,
-            seqused_q,
-            seqused_k,
-            *(aux_tensors or ()),
-        )
-        ctx.softmax_scale = softmax_scale
-        ctx.causal = causal
-        ctx.window_size = window_size
-        ctx.softcap = softcap
-        ctx.deterministic = deterministic
-        ctx.max_seqlen_q = max_seqlen_q
-        ctx.max_seqlen_k = max_seqlen_k
-        ctx.return_lse = return_lse
-        ctx.score_mod = score_mod
-        ctx.score_mod_bwd = score_mod_bwd
-        ctx.mask_mod = mask_mod
-        ctx.aux_scalars = aux_scalars
-        ctx.set_materialize_grads(False)
+        if ctx is not None:
+            ctx.save_for_backward(
+                q,
+                k,
+                v,
+                out,
+                lse,
+                cu_seqlens_q,
+                cu_seqlens_k,
+                seqused_q,
+                seqused_k,
+                *(aux_tensors or ()),
+            )
+            ctx.softmax_scale = softmax_scale
+            ctx.causal = causal
+            ctx.window_size = window_size
+            ctx.softcap = softcap
+            ctx.deterministic = deterministic
+            ctx.max_seqlen_q = max_seqlen_q
+            ctx.max_seqlen_k = max_seqlen_k
+            ctx.return_lse = return_lse
+            ctx.score_mod = score_mod
+            ctx.score_mod_bwd = score_mod_bwd
+            ctx.mask_mod = mask_mod
+            ctx.aux_scalars = aux_scalars
+            ctx.set_materialize_grads(False)
         return out, lse
 
 
@@ -1852,6 +2108,7 @@ def flash_attn_varlen_func(
     v_sf_vec_size: Optional[int] = None,
     rel_bias_prep_cache: Optional[dict] = None,
     return_lse: bool = False,
+    out: Optional[torch.Tensor] = None,
 ):
     """
     Tensor arguments:
@@ -1891,7 +2148,63 @@ def flash_attn_varlen_func(
         qk_sf_vec_size = 32
     if v_sf_vec_size is None and sfv is not None and sfv.dtype == torch.float8_e8m0fnu:
         v_sf_vec_size = 32
-    return FlashAttnVarlenFunc.apply(
+    if out is not None and out.requires_grad:
+        raise ValueError("out must not require gradients")
+    if out is not None and out.stride(-1) != 1:
+        raise ValueError("out must have stride 1 in the last dimension")
+    runtime_arch = _get_device_arch()
+    forward_host = None if is_fake_mode() else get_forward_host(runtime_arch)
+    if (
+        forward_host is not None
+        and q is not None
+        and k is not None
+        and cu_seqlens_q is not None
+        and cu_seqlens_k is not None
+        and num_splits == 1
+        and softcap in (None, 0.0)
+        and not return_lse
+        and all(
+            value is None
+            for value in (
+                qv,
+                seqused_q,
+                seqused_k,
+                gather_kv_indices,
+                page_table,
+                score_mod,
+                mask_mod,
+                block_sparse_tensors,
+                aux_tensors,
+                aux_scalars,
+                q_descale,
+                k_descale,
+                v_descale,
+                rel_bias,
+                sfq,
+                sfk,
+                sfv,
+            )
+        )
+    ):
+        fast_result = try_cached_varlen(
+            arch=runtime_arch,
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_q,
+            max_seqlen_k=max_seqlen_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            learnable_sink=learnable_sink,
+            pack_gqa=pack_gqa,
+            out=out,
+        )
+        if fast_result is not None:
+            return fast_result
+    autograd_args = (
         q,
         k,
         v,
@@ -1930,7 +2243,34 @@ def flash_attn_varlen_func(
         v_sf_vec_size,
         rel_bias_prep_cache,
         return_lse,
+        out,
     )
+    needs_autograd = False
+    if forward_host is not None or out is not None:
+        differentiable_tensors = (
+            q,
+            k,
+            v,
+            qv,
+            learnable_sink,
+            q_descale,
+            k_descale,
+            v_descale,
+            rel_bias,
+            sfq,
+            sfk,
+            sfv,
+            *(aux_tensors or ()),
+        )
+        needs_autograd = torch.is_grad_enabled() and any(
+            tensor is not None and tensor.requires_grad
+            for tensor in differentiable_tensors
+        )
+    if needs_autograd and out is not None:
+        raise ValueError("out is only supported for forward-only inference")
+    if not needs_autograd and forward_host is not None:
+        return FlashAttnVarlenFunc.forward(None, *autograd_args)
+    return FlashAttnVarlenFunc.apply(*autograd_args)
 
 
 def _compile_fwd_combine(
@@ -2039,6 +2379,35 @@ def _compile_fwd_combine(
     )
 
 
+def _fwd_combine_compile_key(
+    out_partial: torch.Tensor,
+    out: torch.Tensor,
+    lse: Optional[torch.Tensor],
+    cu_seqlens: Optional[torch.Tensor],
+    seqused: Optional[torch.Tensor],
+    varlen_batch_idx: Optional[torch.Tensor],
+) -> tuple:
+    head_dim = out_partial.shape[-1]
+    num_splits = out_partial.shape[0]
+    k_block_size = 64 if head_dim <= 64 else 128
+    tile_m = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32)
+    log_max_splits = max(math.ceil(math.log2(num_splits)), 4)
+    if tile_m == 8:
+        log_max_splits = max(log_max_splits, 5)
+    return (
+        torch2cute_dtype_map[out.dtype],
+        torch2cute_dtype_map[out_partial.dtype],
+        head_dim,
+        tile_m,
+        k_block_size,
+        log_max_splits,
+        cu_seqlens is not None,
+        seqused is not None,
+        lse is not None,
+        varlen_batch_idx is not None,
+    )
+
+
 def _flash_attn_fwd_combine(
     out_partial: torch.Tensor,
     lse_partial: torch.Tensor,
@@ -2092,38 +2461,18 @@ def _flash_attn_fwd_combine(
             if not is_fake_mode():
                 assert t.is_cuda, f"{name} must be on CUDA device"
             assert t.is_contiguous(), f"{name} must be contiguous"
-    head_dim = out_partial.shape[-1]
     num_splits = out_partial.shape[0]
     assert num_splits <= 256
-    # If hdim is 96 or 192, it's faster to round them to 128 or 256 respectively
-    # so that kBlockM is smaller and we have more parallelism.
-    k_block_size = 64 if head_dim <= 64 else 128
-    # We want kBlockM to be as small as possible to maximize parallelism.
-    # E.g., if hdim is 64, we want kBlockM to be 16 so that we can use 256 threads, each reading 4 elements (floats).
-    tile_m = 8 if k_block_size % 128 == 0 else (16 if k_block_size % 64 == 0 else 32)
-    log_max_splits = max(math.ceil(math.log2(num_splits)), 4)
-    if tile_m == 8:
-        # If kBlockM == 8 then the minimum number of splits is 32.
-        # TODO: we can deal w this by using 128 threads instead
-        log_max_splits = max(log_max_splits, 5)
-
-    # Create combine kernel configuration
-    dtype = torch2cute_dtype_map[out.dtype]
-    dtype_partial = torch2cute_dtype_map[out_partial.dtype]
     # Device architecture is invariant for the lifetime of this server/JIT
     # cache, so PDL does not belong in the compile key.
     use_pdl = is_arch_support_pdl()
-    compile_key = (
-        dtype,
-        dtype_partial,
-        head_dim,
-        tile_m,
-        k_block_size,
-        log_max_splits,
-        cu_seqlens is not None,
-        seqused is not None,
-        lse is not None,
-        varlen_batch_idx is not None,
+    compile_key = _fwd_combine_compile_key(
+        out_partial,
+        out,
+        lse,
+        cu_seqlens,
+        seqused,
+        varlen_batch_idx,
     )
     if compile_key not in _flash_attn_fwd_combine.compile_cache:
         _flash_attn_fwd_combine.compile_cache[compile_key] = _compile_fwd_combine(

python/sglang/kernels/ops/attention/flash_attn/cute/mask.py

@@ -168,6 +168,9 @@ class AttentionMask:
         1  # only pass in if we're doing PackGQA
     )
     swap_AB: cutlass.Constexpr[bool] = False
+    # R2P assumes the canonical row ownership of the QK accumulator. Kernels
+    # with a different accumulator mapping must use the direct predicate path.
+    enable_r2p_optimization: cutlass.Constexpr[bool] = True
 
     @property
     def seqlen_q(self) -> Int32:
@@ -330,7 +333,7 @@ class AttentionMask:
                 )
                 if const_expr(mask_causal):
                     r2p = const_expr(
-                        not self.swap_AB
+                        not self.swap_AB and self.enable_r2p_optimization
                     )  # R2P trick, see apply_mask_sm100
                     for r in cutlass.range(
                         cute.size(tScS_mn.shape[0]), unroll_full=True
@@ -375,7 +378,9 @@ class AttentionMask:
                         if const_expr(self.window_size_left is not None)
                         else None
                     )
-                    r2p_local = const_expr(not self.swap_AB)
+                    r2p_local = const_expr(
+                        not self.swap_AB and self.enable_r2p_optimization
+                    )
                     for r in cutlass.range(
                         cute.size(tScS_mn.shape[0]), unroll_full=True
                     ):

python/sglang/kernels/ops/attention/flash_attn/cute/shearing_bias.py

@@ -35,6 +35,7 @@ class ShearingBias:
         qhead_per_kvhead: cutlass.Constexpr[int] = 1,
         rows_per_cta: int = 4,
         tile_m: int = 128,
+        attention_tile_m: int = 128,
         max_m_blocks_leq_one: bool = False,
         use_pdl: bool = False,
         clamp_subtiles: bool = True,
@@ -67,6 +68,12 @@ class ShearingBias:
 
         # only used with block packed scheduling
         self.tile_m = tile_m
+        # The output columns are aligned to the rightmost N tile visible to the
+        # attention CTA. Keep this distinct from ``tile_m`` above: the shear
+        # scheduler may group rows in larger blocks than the attention kernel
+        # consumes per CTA.
+        assert attention_tile_m % self.rows_per_cta == 0
+        self.attention_tile_m = attention_tile_m
         # Shrink the subtile grid dim to the rows a block can actually hold
         # (decode blocks hold qhead_per_kvhead*seqlen_q rows, not tile_m).
         self.clamp_subtiles = clamp_subtiles
@@ -335,7 +342,7 @@ class ShearingBias:
             )
 
             block_info = BlockInfo(
-                128,
+                self.attention_tile_m,
                 128,
                 self.is_causal,
                 self.is_local,
@@ -386,7 +393,7 @@ class ShearingBias:
 
             # Convention: inclusive min, exclusive max
             m_idx = m_block * self.rows_per_cta + warp_idx
-            attn_m_block = m_idx // 128
+            attn_m_block = m_idx // self.attention_tile_m
 
             _, attn_n_block_max = block_info.get_n_block_min_max(
                 seqlen_info,

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -253,6 +253,7 @@ class FlashAttentionBackend(AttentionBackend):
 
         # Select version
         self.fa_impl_ver = fa_impl_ver
+        device_capability = get_device_capability()
         if self.fa_impl_ver == 3:
             from sgl_kernel.flash_attn import (
                 flash_attn_varlen_func,
@@ -261,11 +262,25 @@ class FlashAttentionBackend(AttentionBackend):
             )
 
             self._get_scheduler_metadata = get_scheduler_metadata
+            self._get_fa_runtime_policy = None
         elif self.fa_impl_ver == 4:
-            from sglang.kernels.ops.attention.flash_attention_v4 import (
-                flash_attn_varlen_func,
-                flash_attn_with_kvcache,
-            )
+            if device_capability[0] == 12:
+                from sglang.kernels.ops.attention.flash_attention_v4_sm120 import (
+                    flash_attn_varlen_func,
+                    flash_attn_with_kvcache,
+                    get_flash_attention_v4_sm120_runtime_policy,
+                )
+
+                self._get_fa_runtime_policy = (
+                    get_flash_attention_v4_sm120_runtime_policy
+                )
+            else:
+                from sglang.kernels.ops.attention.flash_attention_v4 import (
+                    flash_attn_varlen_func,
+                    flash_attn_with_kvcache,
+                )
+
+                self._get_fa_runtime_policy = None
 
             self._get_scheduler_metadata = None
             if model_runner.server_args.enable_deterministic_inference:
@@ -295,15 +310,22 @@ class FlashAttentionBackend(AttentionBackend):
         )
         self.has_softcap = _softcapping is not None and _softcapping > 0.0
 
-        # If num_splits == 0, we use a heuristic to automatically determine the number of splits.
-        # We set nums splits to 1 if deterministic inference is enabled.
-        # See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ for more details.
-        fa4_no_splitkv = self.fa_impl_ver == 4 and get_device_capability() < (9, 0)
-        self.num_splits = (
-            1
-            if model_runner.server_args.enable_deterministic_inference or fa4_no_splitkv
-            else 0
-        )
+        # num_splits == 0 delegates SplitKV sizing to the selected FA runtime.
+        deterministic = model_runner.server_args.enable_deterministic_inference
+        if self._get_fa_runtime_policy is None:
+            self.num_splits = 1 if deterministic else 0
+            self.decode_num_splits = self.num_splits
+            self._decode_uses_static_max_seqlen_k = False
+        else:
+            runtime_policy = self._get_fa_runtime_policy(
+                device_capability=device_capability,
+                deterministic=deterministic,
+            )
+            self.num_splits = runtime_policy.num_splits
+            self.decode_num_splits = runtime_policy.decode_num_splits
+            self._decode_uses_static_max_seqlen_k = (
+                runtime_policy.decode_uses_static_max_seqlen_k
+            )
         # Set (never getattr'd) so forward_extend can identity-check "is this the
         # full-CG prefill metadata?" to disable the pointer-keyed shear-bias
         # block-schedule cache (see forward_extend rel_bias handling).
@@ -506,6 +528,12 @@ class FlashAttentionBackend(AttentionBackend):
                 # Local attention and scheduler metadata require capture-time slice sizing.
                 # Both depend on data already filled by replay above.
                 metadata = self.decode_cuda_graph_metadata[bs]
+                if self._decode_uses_static_max_seqlen_k:
+                    # FA4 bakes its N-tile grid and SplitKV specialization into
+                    # the graph. Capture against the full replay bound, not the
+                    # padded seq-len fill value (1), or a later long-context
+                    # replay would leave K/V tiles uncovered.
+                    metadata.max_seq_len_k = self.max_context_len
                 self._maybe_update_local_attn_metadata_for_capture(metadata, bs)
                 if self._sched_meta_buf is not None:
                     sched = self._compute_scheduler_metadata(
@@ -1860,6 +1888,8 @@ class FlashAttentionBackend(AttentionBackend):
 
             if layer.is_cross_attention:
                 # Always use non-chunked logic for cross-attention
+                if self._decode_uses_static_max_seqlen_k:
+                    kwargs["max_seqlen_k"] = metadata.encoder_max_seq_len_k
                 o = flash_attn_with_kvcache(
                     q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                     k_cache=key_cache,
@@ -1879,6 +1909,8 @@ class FlashAttentionBackend(AttentionBackend):
                 )
             elif use_local_attn:
                 # Use chunked (local) attention batching for self-attention
+                if self._decode_uses_static_max_seqlen_k:
+                    kwargs["max_seqlen_k"] = local_attn_metadata.local_max_seq_len
                 o = flash_attn_with_kvcache(
                     q=q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
                     k_cache=key_cache,
@@ -1932,6 +1964,8 @@ class FlashAttentionBackend(AttentionBackend):
                     and not pa_swa_active
                 ):
                     sched_meta = metadata.scheduler_metadata
+                if self._decode_uses_static_max_seqlen_k:
+                    kwargs["max_seqlen_k"] = metadata.max_seq_len_k
                 result = flash_attn_with_kvcache(
                     q=q_reshaped,
                     k_cache=key_cache,
@@ -1945,7 +1979,13 @@ class FlashAttentionBackend(AttentionBackend):
                     window_size=window_size,
                     softcap=layer.logit_cap,
                     return_softmax_lse=use_cascade_attn,
-                    num_splits=self.num_splits,
+                    num_splits=(
+                        self.decode_num_splits
+                        if not is_swa_layer
+                        and not use_cascade_attn
+                        and not pa_swa_active
+                        else self.num_splits
+                    ),
                     out=_fa_out,
                     ver=self.fa_impl_ver,
                     scheduler_metadata=sched_meta,
@@ -1953,6 +1993,10 @@ class FlashAttentionBackend(AttentionBackend):
                 )
                 if use_cascade_attn:
                     o, softmax_lse, *rest = result
+                    if self._decode_uses_static_max_seqlen_k:
+                        kwargs["max_seqlen_k"] = (
+                            self.forward_metadata_spec_decode_expand.max_seq_len_k
+                        )
                     o_expand, softmax_lse_expand, *rest_expand = (
                         flash_attn_with_kvcache(
                             q=q_reshaped,