diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md index 54b2dc989..12649c305 100644 --- a/docs/advanced_features/attention_backend.md +++ b/docs/advanced_features/attention_backend.md @@ -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)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 7556df36c..7c6e1be96 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -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` | diff --git a/python/sglang/srt/layers/attention/vision.py b/python/sglang/srt/layers/attention/vision.py index 9dd35e5e3..79d9cd3c5 100644 --- a/python/sglang/srt/layers/attention/vision.py +++ b/python/sglang/srt/layers/attention/vision.py @@ -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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 2dc11cd2e..8aca7047e 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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.", )