diff --git a/docs/advanced_features/attention_backend.md b/docs/advanced_features/attention_backend.md index b290c3278..3010b7463 100644 --- a/docs/advanced_features/attention_backend.md +++ b/docs/advanced_features/attention_backend.md @@ -15,34 +15,34 @@ The support matrix is split into two parts: MHA (standard attention) and MLA (mu ### MHA Backends -| **Backend** | **Page Size > 1 (native)** | **FP8 KV Cache** | **Spec topk=1** | **Spec topk>1** | **Sliding Window** | **MultiModal** | -|---------------------------------|-----------------------------|------------------|-----------------|-----------------|--------------------|----------------| -| **FlashInfer** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | -| **FA3 (FlashAttention 3)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | -| **FA4 (FlashAttention 4)** | 128 | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | -| **Torch Native (SDPA)** | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | -| **FlexAttention (PyTorch)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ❌ | ✅ | ❌ | -| **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **AITER (ROCm)** | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | -| **Wave (ROCm)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | -| **Ascend (NPU)** | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | -| **Intel XPU** | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ | -| **Intel AMX (CPU)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Backend** | **Page Size > 1 (native)** | **FP8 KV Cache** | **FP4 KV Cache** | **Spec topk=1** | **Spec topk>1** | **Sliding Window** | **MultiModal** | +|---------------------------------|-----------------------------|------------------|-----------------|-----------------|-----------------|--------------------|----------------| +| **FlashInfer** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | +| **FA3 (FlashAttention 3)** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | +| **FA4 (FlashAttention 4)** | 128 | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | +| **Torch Native (SDPA)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | +| **FlexAttention (PyTorch)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | +| **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | +| **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **AITER (ROCm)** | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | +| **Wave (ROCm)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| **Ascend (NPU)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | +| **Intel XPU** | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | +| **Intel AMX (CPU)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ### MLA Backends -| **Backend** | **Native Page Sizes** | **FP8 KV Cache** | **Chunked Prefix Cache** | **Spec topk=1** | **Spec topk>1** | -|----------------------------|---------------------------|------------------|--------------------------|-----------------|-----------------| -| **FlashInfer MLA** | 1 | ❌ | ✅ | ✅ | ❌ | -| **FlashMLA** | 64 | ✅ | ✅ | ✅ | ❌ | -| **Cutlass MLA** | 128 | ✅ | ✅ | ✅ | ❌ | -| **TRTLLM MLA (Blackwell)** | 32 or 64 | ✅ | ✅ | ✅ | ❌ | -| **FA3 (FlashAttention 3)** | n/a | ❌ | ✅ | ✅ | ⚠️ (page_size=1 only) | -| **Triton** | n/a | ❌ | ❌ | ✅ | ⚠️ (page_size=1 only) | -| **FA4** | 1 | ❌ | ❌ | ❌ | ❌ | -| **Ascend MLA (NPU)** | 128 | ❌ | ❌ | ❌ | ❌ | +| **Backend** | **Native Page Sizes** | **FP8 KV Cache** | **FP4 KV Cache** | **Chunked Prefix Cache** | **Spec topk=1** | **Spec topk>1** | +|----------------------------|---------------------------|------------------|------------------|--------------------------|-----------------|-----------------| +| **FlashInfer MLA** | 1 | ❌ | ✅ | ✅ | ✅ | ❌ | +| **FlashMLA** | 64 | ✅ | ❌ | ✅ | ✅ | ❌ | +| **Cutlass MLA** | 128 | ✅ | ✅ | ✅ | ✅ | ❌ | +| **TRTLLM MLA (Blackwell)** | 32 or 64 | ✅ | ✅ | ✅ | ✅ | ❌ | +| **FA3 (FlashAttention 3)** | n/a | ❌ | ❌ | ✅ | ✅ | ⚠️ (page_size=1 only) | +| **Triton** | n/a | ❌ | ❌ | ❌ | ✅ | ⚠️ (page_size=1 only) | +| **FA4** | 1 | ❌ | ✅ | ❌ | ❌ | ❌ | +| **Ascend MLA (NPU)** | 128 | ❌ | ❌ | ❌ | ❌ | ❌ | ```{note} Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" column indicates whether a corresponding multimodal implementation exists for that backend family. @@ -53,6 +53,10 @@ Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" c - NSA is specifically designed for [DeepSeek V3.2 DSA](https://lmsys.org/blog/2025-09-29-deepseek-V32/). ``` +```{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. +``` + ```{tip} Speculative decoding topk: `topk` is the number of draft tokens sampled per step from the draft model. `topk = 1` follows classic EAGLE; `topk > 1` explores multiple branches and requires backend support in both draft and verification paths. ``` diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ea1470300..b3e166403 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -656,6 +656,7 @@ class ServerArgs: # Set kernel backends. self._handle_sampling_backend() self._handle_attention_backend_compatibility() + self._handle_kv4_compatibility() self._handle_page_size() self._handle_amd_specifics() self._handle_grammar_backend() @@ -1555,6 +1556,84 @@ class ServerArgs: self.enable_mixed_chunk = False self.disable_radix_cache = True + def _handle_kv4_compatibility(self): + """Check FP4 KV cache compatibility with the attention backend""" + if self.kv_cache_dtype != "fp4_e2m1": + return + + use_mla_backend = self.use_mla_backend() + # self.attention_backend didn't overwrite self.prefill/decode_attention_backend yet + self.prefill_attention_backend_str, self.decode_attention_backend_str = ( + self.get_attention_backends() + ) + + if is_cuda(): + if ( + self.prefill_attention_backend_str != self.decode_attention_backend_str + and self.prefill_attention_backend_str != "fa4" + ): # Take care of prefill=fa4 later + logger.warning( + f"Attention: Using KV4 with PREFILL = {self.prefill_attention_backend_str} " + f"and DECODE = {self.decode_attention_backend_str}. " + f"Compatibility issues are unlikely, but may occur in rare edge cases." + ) + else: + if self.prefill_attention_backend_str == "fa4": + if use_mla_backend: # FA4 + MLA + KV4_FA4_MLA_BACKEND_CHOICES = [ + "cutlass_mla", + "flashinfer", + "trtllm_mla", + ] + assert ( + self.decode_attention_backend_str + in KV4_FA4_MLA_BACKEND_CHOICES + ), ( + f"KV4 FA4 MLA expects decode_attention_backend to be one of " + f"{KV4_FA4_MLA_BACKEND_CHOICES}, but got {self.decode_attention_backend_str}" + ) + else: # FA4 + MHA + KV4_FA4_MHA_BACKEND_CHOICES = [ + "triton", + "torch_native", + "flex_attention", + ] + assert ( + self.decode_attention_backend_str + in KV4_FA4_MHA_BACKEND_CHOICES + ), ( + f"KV4 FA4 MHA expects decode_attention_backend to be one of " + f"{KV4_FA4_MHA_BACKEND_CHOICES}, but got {self.decode_attention_backend_str}" + ) + else: + if use_mla_backend: # !FA4 + MLA + KV4_ATTENTION_MLA_BACKEND_CHOICES = [ + "cutlass_mla", + "flashinfer", + "trtllm_mla", + ] + assert ( + self.attention_backend in KV4_ATTENTION_MLA_BACKEND_CHOICES + ), ( + f"KV4 MLA expects attention_backend to be one of " + f"{KV4_ATTENTION_MLA_BACKEND_CHOICES}, but got {self.attention_backend}" + ) + else: # !FA4 + MHA + KV4_ATTENTION_MHA_BACKEND_CHOICES = [ + "triton", + "torch_native", + "flex_attention", + "trtllm_mha", + ] + assert ( + self.attention_backend in KV4_ATTENTION_MHA_BACKEND_CHOICES + ), ( + f"KV4 MHA expects attention_backend to be one of " + f"{KV4_ATTENTION_MHA_BACKEND_CHOICES}, but got {self.attention_backend}" + ) + else: + raise RuntimeError("KV4 is not tested on non-CUDA platforms.") + def _handle_page_size(self): if self.page_size is None: self.page_size = 1