A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 4, 16:15 UTC
[Attention] use diff kv backend for mimo v2 flash (#40045)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Apr 24, 2026 at 19:25 UTC
e8ee2a78dbc08d398d5e798a149657b8aa821850
8 files changed
+112
-24
docs/design/attention_backends.md
+1
-1
@@ -172,7 +172,7 @@ Priority is **1 = highest** (tried first).
172
| `FLASHINFER` | TRTLLM† | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32, 64 | 64, 128, 256 | ✅ | ❌ | ✅ | Decoder | 10.x |
173
| `FLASH_ATTN` | FA2* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥8.0 |
174
| `FLASH_ATTN` | FA3* | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | %16 | Any | ✅ | ❌ | ✅ | All | 9.x |
175
-| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ❌ | ❌ | ✅ | All | ≥10.0 |
175
+| `FLASH_ATTN` | FA4* | fp16, bf16 | `auto`, `float16`, `bfloat16` | %16 | Any | ✅ | ❌ | ✅ | All | ≥10.0 |
176
| `FLASH_ATTN_DIFFKV` | | fp16, bf16 | `auto` | Any | Any | ❌ | ❌ | ✅ | Decoder | Any |
177
| `FLEX_ATTENTION` | | fp16, bf16, fp32 | `auto`, `float16`, `bfloat16` | Any | Any | ❌ | ✅ | ❌ | Decoder, Encoder Only | Any |
178
| `ROCM_AITER_FA` | | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3`, `fp8_e5m2` | 16, 32 | 64, 128, 256 | ❌ | ❌ | ❌ | Decoder | N/A |
tools/pre_commit/generate_attention_backend_docs.py
+41
-8
@@ -634,9 +634,10 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
634
except Exception:
635
return {}
636
637
- # Analyze the functions to determine FA3-specific features
637
+ # Analyze the functions to determine FA3/FA4-specific features
638
fa3_supports_fp8 = False
639
fa3_supports_sinks = False
640
+ fa4_supports_sinks = False
641
fa3_compute_cap: str | None = None
642
fa4_compute_cap: str | None = None
643
@@ -656,17 +657,49 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
657
fa3_supports_fp8 = True
658
break
659
659
- # Check flash_attn_supports_sinks - looks for `get_flash_attn_version() == 3`
660
+ # Check flash_attn_supports_sinks - looks for `fa_version == 3/4`
661
+ # or `get_flash_attn_version() == 3/4` (also accepts `in (3, 4)`)
662
if node.name == "flash_attn_supports_sinks":
663
for n in ast.walk(node):
664
if (
665
isinstance(n, ast.Compare)
664
- and isinstance(n.left, ast.Call)
665
- and isinstance(n.left.func, ast.Name)
666
- and n.left.func.id == "get_flash_attn_version"
666
+ and len(n.ops) == 1
667
+ and isinstance(n.ops[0], ast.Eq)
668
+ and isinstance(n.comparators[0], ast.Constant)
669
):
668
- fa3_supports_sinks = True
669
- break
670
+ is_version_compare = (
671
+ isinstance(n.left, ast.Name) and n.left.id == "fa_version"
672
+ ) or (
673
+ isinstance(n.left, ast.Call)
674
+ and isinstance(n.left.func, ast.Name)
675
+ and n.left.func.id == "get_flash_attn_version"
676
+ )
677
+ if is_version_compare:
678
+ val = n.comparators[0].value
679
+ if val == 3:
680
+ fa3_supports_sinks = True
681
+ elif val == 4:
682
+ fa4_supports_sinks = True
683
+ elif (
684
+ isinstance(n, ast.Compare)
685
+ and len(n.ops) == 1
686
+ and isinstance(n.ops[0], ast.In)
687
+ and isinstance(n.comparators[0], (ast.Tuple, ast.List, ast.Set))
688
+ ):
689
+ is_version_compare = (
690
+ isinstance(n.left, ast.Name) and n.left.id == "fa_version"
691
+ ) or (
692
+ isinstance(n.left, ast.Call)
693
+ and isinstance(n.left.func, ast.Name)
694
+ and n.left.func.id == "get_flash_attn_version"
695
+ )
696
+ if is_version_compare:
697
+ for elt in n.comparators[0].elts:
698
+ if isinstance(elt, ast.Constant):
699
+ if elt.value == 3:
700
+ fa3_supports_sinks = True
701
+ elif elt.value == 4:
702
+ fa4_supports_sinks = True
703
704
# Check get_flash_attn_version for FA3/FA4 compute capability
705
if node.name == "get_flash_attn_version":
@@ -731,7 +764,7 @@ def parse_flash_attn_features() -> dict[str, dict[str, Any]]:
764
"fa4": {
765
"compute_capability": fa4_compute_cap,
766
"supports_fp8": False,
734
- "supports_sink": False,
767
+ "supports_sink": fa4_supports_sinks,
768
},
769
}
770
vllm/model_executor/layers/attention/attention.py
+1
@@ -597,6 +597,7 @@ class Attention(nn.Module, AttentionLayerBase):
597
block_size=block_size,
598
num_kv_heads=self.num_kv_heads,
599
head_size=self.head_size,
600
+ head_size_v=self.head_size_v,
601
dtype=self.kv_cache_torch_dtype,
602
kv_quant_mode=quant_mode,
603
sliding_window=self.sliding_window,
vllm/model_executor/models/mimo_v2_flash.py
+14
-8
@@ -46,6 +46,9 @@ from vllm.model_executor.model_loader.weight_utils import (
46
from vllm.model_executor.models.utils import sequence_parallel_chunk
47
from vllm.sequence import IntermediateTensors
48
from vllm.v1.attention.backend import AttentionType
49
+from vllm.v1.attention.backends.flash_attn_diffkv import (
50
+ FlashAttentionDiffKVBackend,
51
+)
52
53
from .interfaces import MixtureOfExperts, SupportsPP
54
from .utils import (
@@ -287,6 +290,15 @@ class MiMoV2Attention(nn.Module):
290
)
291
292
sliding_window = sliding_window_size if sliding_window_size > -1 else None
293
+
294
+ # Use DiffKV backend when V has a different head dim than K
295
+ if self.v_head_dim != self.head_dim:
296
+ FlashAttentionDiffKVBackend.set_head_size_v(self.v_head_dim)
297
+ attn_backend = FlashAttentionDiffKVBackend
298
+ logger.info_once("Using FlashAttentionDiffKVBackend for attention.")
299
+ else:
300
+ attn_backend = None
301
+
302
self.attn = Attention(
303
self.num_heads,
304
self.head_dim,
@@ -298,6 +310,8 @@ class MiMoV2Attention(nn.Module):
310
attn_type=AttentionType.DECODER,
311
prefix=f"{prefix}.attn",
312
sinks=self.attention_sink_bias,
313
+ attn_backend=attn_backend,
314
+ head_size_v=self.v_head_dim,
315
)
316
317
def forward(
@@ -313,16 +327,8 @@ class MiMoV2Attention(nn.Module):
327
if self.v_scale is not None:
328
v = v * self.v_scale
329
316
- v = v.view(-1, self.num_kv_heads, self.v_head_dim)
317
- v = torch.nn.functional.pad(v, [0, self.head_dim - self.v_head_dim], value=0)
318
- v = v.view(-1, self.num_kv_heads * self.head_dim)
319
-
330
attn_output = self.attn(q, k, v)
331
322
- attn_output = attn_output.view(-1, self.num_heads, self.head_dim)[
323
- ..., : self.v_head_dim
324
- ].reshape(-1, self.num_heads * self.v_head_dim)
325
-
332
output, _ = self.o_proj(attn_output)
333
return output
334
vllm/v1/attention/backends/fa_utils.py
+22
-3
@@ -54,7 +54,10 @@ elif current_platform.is_rocm():
54
55
56
def get_flash_attn_version(
57
- requires_alibi: bool = False, head_size: int | None = None
57
+ requires_alibi: bool = False,
58
+ head_size: int | None = None,
59
+ head_size_v: int | None = None,
60
+ has_sinks: bool = False,
61
) -> int | None:
62
if current_platform.is_xpu():
63
return 2
@@ -112,6 +115,23 @@ def get_flash_attn_version(
115
)
116
fa_version = 2
117
118
+ # The FA3 kernel rejects s_aux (sinks) when hdim != hdim_v; upgrade to
119
+ # FA4 on SM90 when available.
120
+ if (
121
+ fa_version == 3
122
+ and has_sinks
123
+ and head_size is not None
124
+ and head_size_v is not None
125
+ and head_size != head_size_v
126
+ and device_capability.major == 9
127
+ and is_fa_version_supported(4)
128
+ ):
129
+ logger.info_once(
130
+ "Diff-KV with sinks: upgrading FlashAttention 3 -> 4",
131
+ scope="local",
132
+ )
133
+ fa_version = 4
134
+
135
# FA4 currently uses batch-shape-dependent scheduling
136
# heuristics on SM100+, which breaks batch invariance.
137
if envs.VLLM_BATCH_INVARIANT and fa_version == 4:
@@ -180,8 +200,7 @@ def flash_attn_supports_quant_query_input() -> bool:
200
def flash_attn_supports_sinks() -> bool:
201
if current_platform.is_xpu():
202
return True
183
- else:
184
- return get_flash_attn_version() == 3
203
+ return get_flash_attn_version() in (3, 4)
204
205
206
def flash_attn_supports_mla():
vllm/v1/attention/backends/flash_attn_diffkv.py
+18
-4
@@ -6,14 +6,16 @@ import torch
6
7
from vllm.utils.torch_utils import is_quantized_kv_cache
8
from vllm.v1.attention.backend import AttentionType
9
-from vllm.v1.attention.backends.fa_utils import is_flash_attn_varlen_func_available
9
+from vllm.v1.attention.backends.fa_utils import (
10
+ get_flash_attn_version,
11
+ is_flash_attn_varlen_func_available,
12
+)
13
from vllm.v1.attention.ops.triton_reshape_and_cache_flash import (
14
triton_reshape_and_cache_flash_diffkv,
15
)
16
17
if is_flash_attn_varlen_func_available():
18
from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func
16
-from vllm.logger import init_logger
19
from vllm.v1.attention.backends.utils import get_kv_cache_layout
20
21
from .flash_attn import (
@@ -23,8 +25,6 @@ from .flash_attn import (
25
cascade_attention,
26
)
27
26
-logger = init_logger(__name__)
27
-
28
29
class FlashAttentionDiffKVBackend(FlashAttentionBackend):
30
# Default to 128 for this backend
@@ -86,6 +86,20 @@ class FlashAttentionDiffKVBackend(FlashAttentionBackend):
86
87
88
class FlashAttentionDiffKVImpl(FlashAttentionImpl):
89
+ vllm_flash_attn_version: int | None
90
+
91
+ def __init__(self, *args, **kwargs) -> None:
92
+ super().__init__(*args, **kwargs)
93
+ # Re-derive the FA version with diff-kv context so that
94
+ # get_flash_attn_version can apply the FA3 -> FA4 upgrade rule
95
+ # for sinks + hdim != hdim_v.
96
+ self.vllm_flash_attn_version = get_flash_attn_version(
97
+ requires_alibi=self.alibi_slopes is not None,
98
+ head_size=self.head_size,
99
+ head_size_v=FlashAttentionDiffKVBackend.head_size_v,
100
+ has_sinks=self.sinks is not None,
101
+ )
102
+
103
def do_kv_cache_update(
104
self,
105
layer: torch.nn.Module,
vllm/v1/kv_cache_interface.py
+14
@@ -356,6 +356,20 @@ class ChunkedLocalAttentionSpec(AttentionSpec):
356
@dataclass(frozen=True, kw_only=True)
357
class SlidingWindowSpec(AttentionSpec):
358
sliding_window: int
359
+ head_size_v: int = None # type: ignore[assignment]
360
+
361
+ def __post_init__(self):
362
+ if self.head_size_v is None:
363
+ object.__setattr__(self, "head_size_v", self.head_size)
364
+
365
+ @property
366
+ def real_page_size_bytes(self) -> int:
367
+ return (
368
+ self.block_size
369
+ * self.num_kv_heads
370
+ * (self.head_size + self.head_size_v)
371
+ * get_dtype_size(self.dtype)
372
+ )
373
374
def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
375
assert vllm_config.parallel_config.decode_context_parallel_size == 1, (
vllm/vllm_flash_attn/flash_attn_interface.py
+1
@@ -387,6 +387,7 @@ def flash_attn_varlen_func(
387
num_splits=num_splits,
388
return_lse=return_softmax_lse,
389
out=out,
390
+ learnable_sink=s_aux,
391
)
392
else:
393
raise ValueError(f"Unsupported FA version: {fa_version}")