Check KV4 compatibility with attention backends and add KV4 support to the attention_backend doc (#14467)
Signed-off-by: Ho-Ren (Jack) Chuang <horenchuang@bytedance.com>
This commit is contained in:
Ho-Ren (Jack) Chuang
@@ -15,34 +15,34 @@ The support matrix is split into two parts: MHA (standard attention) and MLA (mu
|
|||||||
|
|
||||||
### MHA Backends
|
### MHA Backends
|
||||||
|
|
||||||
| **Backend** | **Page Size > 1 (native)** | **FP8 KV Cache** | **Spec topk=1** | **Spec topk>1** | **Sliding Window** | **MultiModal** |
|
| **Backend** | **Page Size > 1 (native)** | **FP8 KV Cache** | **FP4 KV Cache** | **Spec topk=1** | **Spec topk>1** | **Sliding Window** | **MultiModal** |
|
||||||
|---------------------------------|-----------------------------|------------------|-----------------|-----------------|--------------------|----------------|
|
|---------------------------------|-----------------------------|------------------|-----------------|-----------------|-----------------|--------------------|----------------|
|
||||||
| **FlashInfer** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
| **FlashInfer** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ |
|
||||||
| **FA3 (FlashAttention 3)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
| **FA3 (FlashAttention 3)** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
|
||||||
| **FA4 (FlashAttention 4)** | 128 | ❌ | ❌ | ❌ | ❌ | ❌ |
|
| **FA4 (FlashAttention 4)** | 128 | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||||
| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ |
|
| **Triton** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||||
| **Torch Native (SDPA)** | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
|
| **Torch Native (SDPA)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ |
|
||||||
| **FlexAttention (PyTorch)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
| **FlexAttention (PyTorch)** | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||||
| **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ❌ | ✅ | ❌ |
|
| **TRTLLM MHA** | 16, 32 or 64 | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||||
| **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
| **Dual Chunk FlashAttention** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||||
| **AITER (ROCm)** | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ |
|
| **AITER (ROCm)** | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ |
|
||||||
| **Wave (ROCm)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
| **Wave (ROCm)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||||
| **Ascend (NPU)** | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ |
|
| **Ascend (NPU)** | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ |
|
||||||
| **Intel XPU** | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
| **Intel XPU** | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||||
| **Intel AMX (CPU)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
| **Intel AMX (CPU)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||||
|
|
||||||
### MLA Backends
|
### MLA Backends
|
||||||
|
|
||||||
| **Backend** | **Native Page Sizes** | **FP8 KV Cache** | **Chunked Prefix Cache** | **Spec topk=1** | **Spec topk>1** |
|
| **Backend** | **Native Page Sizes** | **FP8 KV Cache** | **FP4 KV Cache** | **Chunked Prefix Cache** | **Spec topk=1** | **Spec topk>1** |
|
||||||
|----------------------------|---------------------------|------------------|--------------------------|-----------------|-----------------|
|
|----------------------------|---------------------------|------------------|------------------|--------------------------|-----------------|-----------------|
|
||||||
| **FlashInfer MLA** | 1 | ❌ | ✅ | ✅ | ❌ |
|
| **FlashInfer MLA** | 1 | ❌ | ✅ | ✅ | ✅ | ❌ |
|
||||||
| **FlashMLA** | 64 | ✅ | ✅ | ✅ | ❌ |
|
| **FlashMLA** | 64 | ✅ | ❌ | ✅ | ✅ | ❌ |
|
||||||
| **Cutlass MLA** | 128 | ✅ | ✅ | ✅ | ❌ |
|
| **Cutlass MLA** | 128 | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||||
| **TRTLLM MLA (Blackwell)** | 32 or 64 | ✅ | ✅ | ✅ | ❌ |
|
| **TRTLLM MLA (Blackwell)** | 32 or 64 | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||||
| **FA3 (FlashAttention 3)** | n/a | ❌ | ✅ | ✅ | ⚠️ (page_size=1 only) |
|
| **FA3 (FlashAttention 3)** | n/a | ❌ | ❌ | ✅ | ✅ | ⚠️ (page_size=1 only) |
|
||||||
| **Triton** | n/a | ❌ | ❌ | ✅ | ⚠️ (page_size=1 only) |
|
| **Triton** | n/a | ❌ | ❌ | ❌ | ✅ | ⚠️ (page_size=1 only) |
|
||||||
| **FA4** | 1 | ❌ | ❌ | ❌ | ❌ |
|
| **FA4** | 1 | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||||
| **Ascend MLA (NPU)** | 128 | ❌ | ❌ | ❌ | ❌ |
|
| **Ascend MLA (NPU)** | 128 | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||||
|
|
||||||
```{note}
|
```{note}
|
||||||
Multimodal attention is selected by `--mm-attention-backend`. The "MultiModal" column indicates whether a corresponding multimodal implementation exists for that backend family.
|
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/).
|
- 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}
|
```{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.
|
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.
|
||||||
```
|
```
|
||||||
|
|||||||
@@ -656,6 +656,7 @@ class ServerArgs:
|
|||||||
# Set kernel backends.
|
# Set kernel backends.
|
||||||
self._handle_sampling_backend()
|
self._handle_sampling_backend()
|
||||||
self._handle_attention_backend_compatibility()
|
self._handle_attention_backend_compatibility()
|
||||||
|
self._handle_kv4_compatibility()
|
||||||
self._handle_page_size()
|
self._handle_page_size()
|
||||||
self._handle_amd_specifics()
|
self._handle_amd_specifics()
|
||||||
self._handle_grammar_backend()
|
self._handle_grammar_backend()
|
||||||
@@ -1555,6 +1556,84 @@ class ServerArgs:
|
|||||||
self.enable_mixed_chunk = False
|
self.enable_mixed_chunk = False
|
||||||
self.disable_radix_cache = True
|
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):
|
def _handle_page_size(self):
|
||||||
if self.page_size is None:
|
if self.page_size is None:
|
||||||
self.page_size = 1
|
self.page_size = 1
|
||||||
|
|||||||
Reference in New Issue
Block a user