feat: add FA4 SM90 paged KV decode support & update attention docs (#18442)

Co-authored-by: Zeyu Wang <zeyu.wang@yahooinc.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>

docs/advanced_features/attention_backend.md

@@ -49,10 +49,14 @@ Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" c
 ```
 
 ```{note}
-- FlashAttention 4 is prefill-only for now.
+- FlashAttention 4 supports both prefill and decode on SM90 (Hopper) and SM100 (Blackwell). On SM90, `page_size` must be 128.
 - NSA is specifically designed for [DeepSeek V3.2 DSA](https://lmsys.org/blog/2025-09-29-deepseek-V32/).
 ```
 
+```{warning}
+**FA4 on Hopper (SM90):** FA4 decode speed decreases as sequence length grows due to lack of SplitKV support. At batch=1 compared to FA3 on H100: ~-10% at 2K tokens, ~-18% at 4K, ~-31% at 8K, ~-49% at 16K. Larger batch sizes reduce the gap (e.g., batch=8: ~-2% at 2K, ~-8% at 4K). Blackwell (SM100) is not affected.
+```
+
 ```{note}
 For the KV4 FA4 scenario, FA4 requires using a different --decode-attention-backend to run. Except for trtllm_mha being incompatible with FA4, all other decode backends behave as shown in the table.
 ```
@@ -204,6 +208,13 @@ python3 -m sglang.launch_server \
 
 - FlashAttention 4 (MHA & MLA)
 ```bash
+# FA4 for both prefill and decode on SM90/SM100
+python3 -m sglang.launch_server \
+  --model-path Qwen/Qwen3-30B-A3B-Instruct-2507-FP8 \
+  --attention-backend fa4 \
+  --page-size 128 \
+  --trust-remote-code
+
 python3 -m sglang.launch_server \
   --tp 8 \
   --model deepseek-ai/DeepSeek-R1 \

python/sglang/jit_kernel/flash_attention/cute/flash_fwd.py

@@ -69,6 +69,7 @@ class FlashAttentionForwardBase:
         score_mod: Optional[cutlass.Constexpr] = None,
         mask_mod: Optional[cutlass.Constexpr] = None,
         has_aux_tensors: bool = False,
+        page_size: Optional[int] = None,
     ):
         """Initializes the configuration for a flash attention kernel.
 
@@ -111,6 +112,7 @@ class FlashAttentionForwardBase:
         self.score_mod = score_mod
         self.mask_mod = mask_mod
         self.qk_acc_dtype = Float32
+        self.page_size = page_size
         if const_expr(has_aux_tensors):
             self.vec_size: cutlass.Constexpr = 1
         else:
@@ -1332,6 +1334,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
         self.num_producer_threads = 32
         self.num_Q_load_threads = self.num_mma_threads  # If not TMA_Q, MMA threads load Q
         self.num_epilogue_threads = self.num_mma_threads
+        self.tiles_per_page = self.page_size // self.tile_n if cutlass.const_expr(self.page_size is not None) else None
         self.num_mma_regs = (
             256
             if self.num_mma_warp_groups == 1
@@ -1533,6 +1536,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
             mCuSeqlensK,
             mSeqUsedQ,
             mSeqUsedK,
+            mPageTable,
             tma_atom_Q,
             tma_atom_K,
             tma_atom_V,
@@ -1579,6 +1583,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
         mCuSeqlensK: Optional[cute.Tensor],
         mSeqUsedQ: Optional[cute.Tensor],
         mSeqUsedK: Optional[cute.Tensor],
+        mPageTable: Optional[cute.Tensor],
         tma_atom_Q: Optional[cute.CopyAtom],
         tma_atom_K: Optional[cute.CopyAtom],
         tma_atom_V: Optional[cute.CopyAtom],
@@ -1682,7 +1687,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
         SeqlenInfoCls = partial(
             SeqlenInfoQK.create,
             seqlen_q_static=mQ.shape[0] if const_expr(not self.pack_gqa) else mQ.shape[0][1],
-            seqlen_k_static=mK.shape[0],
+            seqlen_k_static=mK.shape[0] if const_expr(mPageTable is None) else mK.shape[0] * mPageTable.shape[1],
             mCuSeqlensQ=mCuSeqlensQ,
             mCuSeqlensK=mCuSeqlensK,
             mSeqUsedQ=mSeqUsedQ,
@@ -1707,6 +1712,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
                 sQ,
                 sK,
                 sV,
+                mPageTable,
                 tma_atom_Q,
                 tma_atom_K,
                 tma_atom_V,
@@ -1766,6 +1772,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
         sQ: cute.Tensor,
         sK: cute.Tensor,
         sV: cute.Tensor,
+        mPageTable: Optional[cute.Tensor],
         tma_atom_Q: cute.CopyAtom,
         tma_atom_K: cute.CopyAtom,
         tma_atom_V: cute.CopyAtom,
@@ -1790,13 +1797,21 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
                 m_block, head_idx, batch_idx, _ = work_tile.tile_idx
                 seqlen = SeqlenInfoCls(batch_idx)
                 mQ_cur = seqlen.offset_batch_Q(mQ, batch_idx, dim=3)[None, None, head_idx]
-                head_idx_kv = (
-                    head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx
-                )
-                mK_cur = seqlen.offset_batch_K(mK, batch_idx, dim=3)[None, None, head_idx_kv]
-                mV_cur = seqlen.offset_batch_K(mV, batch_idx, dim=3)[None, None, head_idx_kv]
-                gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0))
-                gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0))
+                head_idx_kv = head_idx // self.qhead_per_kvhead if const_expr(not self.pack_gqa) else head_idx
+                if const_expr(mPageTable is None):
+                    if const_expr(not seqlen.has_cu_seqlens_k):
+                        mK_cur, mV_cur = [t[None, None, head_idx_kv, batch_idx] for t in (mK, mV)]
+                    else:
+                        mK_cur = cute.domain_offset((seqlen.offset_k, 0), mK[None, None, head_idx_kv])
+                        mV_cur = cute.domain_offset((seqlen.offset_k, 0), mV[None, None, head_idx_kv])
+                    gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0))
+                    gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0))
+                else:
+                    mK_cur, mV_cur = [t[None, None, head_idx_kv, None] for t in (mK, mV)]
+                    gK = cute.local_tile(mK_cur, (self.tile_n, self.tile_hdim), (None, 0, None))
+                    gV = cute.local_tile(mV_cur, (self.tile_n, self.tile_hdimv), (None, 0, None))
+                    gK = cute.group_modes(gK, 2, 4)
+                    gV = cute.group_modes(gV, 2, 4)
                 if const_expr(self.use_tma_Q):
                     gQ = cute.local_tile(mQ_cur, (self.tile_m, self.tile_hdim), (m_block, 0))
                     load_Q, _, _ = copy_utils.tma_get_copy_fn(
@@ -1818,7 +1833,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
                     # if cute.arch.thread_idx()[0] == 0:
                     #     cute.printf("m_block = %d, n_block_min: %d, n_block_max: %d", m_block, n_block_min, n_block_max)
                     # First iteration: load both Q & K with the same mbarrier
-                    n_block = n_block_max - 1
+                    n_block = self.get_n_block(batch_idx, n_block_max - 1, mPageTable)
                     pipeline_k.producer_acquire(
                         kv_producer_state,
                         extra_tx_count=self.tma_copy_bytes["Q"]
@@ -1834,7 +1849,7 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
                         load_V(src_idx=n_block, producer_state=kv_producer_state)
                         kv_producer_state.advance()
                         for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1):
-                            n_block = n_block_max - 1 - i - 1
+                            n_block = self.get_n_block(batch_idx, n_block_max - 1 - i - 1, mPageTable)
                             pipeline_k.producer_acquire(kv_producer_state)
                             load_K(src_idx=n_block, producer_state=kv_producer_state)
                             pipeline_v.producer_acquire(kv_producer_state)
@@ -1842,15 +1857,15 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
                             kv_producer_state.advance()
                     else:
                         for i in cutlass.range(n_block_max - 1 - n_block_min, unroll=1):
-                            n_block_prev = n_block_max - i - 1
-                            n_block = n_block_prev - 1
+                            n_block_prev = self.get_n_block(batch_idx, n_block_max - i - 1, mPageTable)
+                            n_block = self.get_n_block(batch_idx, n_block_max - i - 2, mPageTable)
                             kv_producer_state_prev = kv_producer_state.clone()
                             kv_producer_state.advance()
                             pipeline_k.producer_acquire(kv_producer_state)
                             load_K(src_idx=n_block, producer_state=kv_producer_state)
                             pipeline_v.producer_acquire(kv_producer_state_prev)
                             load_V(src_idx=n_block_prev, producer_state=kv_producer_state_prev)
-                        n_block = n_block_min
+                        n_block = self.get_n_block(batch_idx, n_block_min, mPageTable)
                         pipeline_v.producer_acquire(kv_producer_state)
                         load_V(src_idx=n_block, producer_state=kv_producer_state)
                         kv_producer_state.advance()
@@ -1877,6 +1892,19 @@ class FlashAttentionForwardSm90(FlashAttentionForwardBase):
                 work_tile = tile_scheduler.get_current_work()
                 # End of persistent scheduler loop
 
+    @cute.jit
+    def get_n_block(
+        self,
+        batch_idx: int,
+        n_block: int,
+        mPageTable: Optional[cute.Tensor],
+    ):
+        if cutlass.const_expr(mPageTable is not None):
+            page_idx = mPageTable[batch_idx, n_block // self.tiles_per_page]
+            residue = n_block % self.tiles_per_page
+            n_block = page_idx * self.tiles_per_page + residue
+        return n_block
+
     @cute.jit
     def mma(
         self,

python/sglang/jit_kernel/flash_attention/cute/interface.py

@@ -378,7 +378,7 @@ def _flash_attn_fwd(
         cu_seqlens_k is None,
         seqused_q is None,
         seqused_k is None,
-        page_table is not None,
+        page_table is not None, page_size,
         window_size_left is not None,
         window_size_right is not None,
         learnable_sink is not None,
@@ -428,7 +428,7 @@ def _flash_attn_fwd(
             cute_aux_tensors = [to_cute_tensor(buf, assumed_align=None, fully_dynamic=True) for buf in aux_tensors]
 
         if compute_capability == 9:
-            assert page_table is None, "paged KV not supported on SM 9.0"
+            assert page_size is None or page_size % n_block_size == 0, f"Only page_size values that are multiples of {n_block_size} are supported for paged KV on SM 9.0"
             assert not is_split_kv, "SplitKV not supported on SM 9.0"
             # fa_fwd = FlashAttentionForwardSm80(
             fa_fwd = FlashAttentionForwardSm90(
@@ -450,6 +450,7 @@ def _flash_attn_fwd(
                 mask_mod=mask_mod,
                 score_mod=score_mod,
                 has_aux_tensors=aux_tensors is not None,
+                page_size=page_size,
             )
         elif compute_capability in [10, 11]:
             fa_fwd = FlashAttentionForwardSm100(

python/sglang/srt/server_args.py

@@ -1963,7 +1963,11 @@ class ServerArgs:
             )
             self.attention_backend = "triton"
 
-        if self.prefill_attention_backend == "fa4" and not self.use_mla_backend():
+        if (
+            self.prefill_attention_backend == "fa4"
+            and not self.use_mla_backend()
+            and is_sm100_supported()
+        ):
             logger.warning(
                 f"FA4 backend only supports page size 128 for non-MLA model architectures, changing page_size from {self.page_size} to 128."
             )