chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

add the fa4 mm backend and varlen func (#13539)

Browse Source 
Signed-off-by: vincentzed <207368749+vincentzed@users.noreply.github.com>
Co-authored-by: Brayden Zhong <b8zhong@uwaterloo.ca>
This commit is contained in:
Yi Zhong
2026-01-23 10:12:06 -05:00
committed by GitHub
parent 2fb328109f
commit 08fcda2f63
4 changed files with 58 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/advanced_features/attention_backend.md
View File

@@ -19,7 +19,7 @@ The support matrix is split into two parts: MHA (standard attention) and MLA (mu
|---------------------------------|-----------------------------|------------------|-----------------|-----------------|-----------------|--------------------|----------------|
| **FlashInfer** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ |
| **FA3 (FlashAttention 3)** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
| **FA4 (FlashAttention 4)** | 128 | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FA4 (FlashAttention 4)** | 128 | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Torch Native (SDPA)** | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ |
| **FlexAttention (PyTorch)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |

2
docs/advanced_features/server_arguments.md
View File

@@ -266,7 +266,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
| `--decode-attention-backend` | Choose the kernels for decode attention layers (have priority over --attention-backend). | `None` | `triton`, `torch_native`, `flex_attention`, `nsa`, `cutlass_mla`, `fa3`, `fa4`, `flashinfer`, `flashmla`, `trtllm_mla`, `trtllm_mha`, `dual_chunk_flash_attn`, `aiter`, `wave`, `intel_amx`, `ascend` |
| `--sampling-backend` | Choose the kernels for sampling layers. | `None` | `flashinfer`, `pytorch`, `ascend` |
| `--grammar-backend` | Choose the backend for grammar-guided decoding. | `None` | `xgrammar`, `outlines`, `llguidance`, `none` |
| `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `triton_attn`, `ascend_attn`, `aiter_attn` |
| `--mm-attention-backend` | Set multimodal attention backend. | `None` | `sdpa`, `fa3`, `fa4`, `triton_attn`, `ascend_attn`, `aiter_attn` |
| `--nsa-prefill-backend` | Choose the NSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek NSA-style attention). | `flashmla_sparse` | `flashmla_sparse`, `flashmla_kv`, `flashmla_auto`, `fa3`, `tilelang`, `aiter` |
| `--nsa-decode-backend` | Choose the NSA backend for the decode stage when running DeepSeek NSA-style attention. Overrides `--attention-backend` for decoding. | `fa3` | `flashmla_sparse`, `flashmla_kv`, `fa3`, `tilelang`, `aiter` |
| `--fp8-gemm-backend` | Choose the runner backend for Blockwise FP8 GEMM operations. Options: 'auto' (default, auto-selects based on hardware), 'deep_gemm' (JIT-compiled; enabled by default on NVIDIA Hopper (SM90) and Blackwell (SM100) when DeepGEMM is installed), 'flashinfer_trtllm' (optimal for Blackwell and low-latency), 'cutlass' (optimal for Hopper/Blackwell GPUs and high-throughput), 'triton' (fallback, widely compatible), 'aiter' (ROCm only). **NOTE**: This replaces the deprecated environment variables SGLANG_ENABLE_FLASHINFER_FP8_GEMM and SGLANG_SUPPORT_CUTLASS_BLOCK_FP8. | `auto` | `auto`, `deep_gemm`, `flashinfer_trtllm`, `cutlass`, `triton`, `aiter` |

57
python/sglang/srt/layers/attention/vision.py
View File

@@ -18,7 +18,7 @@ from sglang.srt.models.utils import apply_qk_norm
from sglang.srt.utils import (
get_bool_env_var,
get_device_capability,
is_blackwell,
is_blackwell_supported,
is_cuda,
is_hip,
is_npu,
@@ -400,6 +400,58 @@ class VisionFlash3Attention(nn.Module):
return output
class VisionFlash4Attention(nn.Module):
def __init__(
self,
**kwargs,
):
if not _is_cuda:
raise Exception("VisionFlash4Attention is only available for cuda")
super().__init__()
def forward(
self,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
cu_seqlens: torch.Tensor | SingletonCache | None,
bsz: int,
seq_len: int,
**kwargs,
) -> torch.Tensor:
r"""
Args:
cu_seqlens: [b]
Returns:
[b * s, h, head_size]
"""
if cu_seqlens is None:
cu_seqlens = _get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
elif isinstance(cu_seqlens, SingletonCache):
if cu_seqlens.empty():
cu_seqlens.set_data(
_get_cu_seqlens_for_shape(bsz, seq_len, device=q.device)
)
cu_seqlens = cu_seqlens.get_data()
cu_seqlens = cu_seqlens.to(dtype=torch.int32).to(q.device)
seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
max_seqlen = seq_lens.max().item()
output = flash_attn_varlen_func(
q,
k,
v,
cu_seqlens_q=cu_seqlens,
cu_seqlens_k=cu_seqlens,
max_seqlen_q=max_seqlen,
max_seqlen_k=max_seqlen,
ver=4,
)
return output
class VisionAiterAttention(nn.Module):
def __init__(
self,
@@ -499,6 +551,7 @@ QKV_BACKEND_IMPL = {
"triton_attn": VisionTritonAttention,
"sdpa": VisionSdpaAttention,
"fa3": VisionFlash3Attention,
"fa4": VisionFlash4Attention,
"ascend_attn": VisionAscendAttention,
"aiter_attn": VisionAiterAttention,
}
@@ -671,7 +724,7 @@ class VisionAttention(nn.Module):
backend = "triton_attn"
else:
backend = "sdpa"
if backend == "fa3" and is_blackwell():
if backend == "fa3" and is_blackwell_supported():
raise ValueError("The 'fa3' backend is not supported on Blackwell GPUs")
return backend

2
python/sglang/srt/server_args.py
View File

@@ -3620,7 +3620,7 @@ class ServerArgs:
parser.add_argument(
"--mm-attention-backend",
type=str,
choices=["sdpa", "fa3", "triton_attn", "ascend_attn", "aiter_attn"],
choices=["sdpa", "fa3", "fa4", "triton_attn", "ascend_attn", "aiter_attn"],
default=ServerArgs.mm_attention_backend,
help="Set multimodal attention backend.",
)