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 6, 19:15 UTC
[Attention] Use FA4 for MLA prefill (#34732)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Mar 12, 2026 at 12:10 UTC
f444c05c3267ed26f1fd52822d60479b81b2b829
9 files changed
+413
-78
benchmarks/attention_benchmarks/benchmark.py
+103
-29
@@ -59,7 +59,9 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
59
"""Run MLA benchmark with appropriate backend."""
60
from mla_runner import run_mla_benchmark as run_mla
61
62
- return run_mla(config.backend, config, **kwargs)
62
+ return run_mla(
63
+ config.backend, config, prefill_backend=config.prefill_backend, **kwargs
64
+ )
65
66
67
def run_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
@@ -440,14 +442,21 @@ def main():
442
# Backend selection
443
parser.add_argument(
444
"--backends",
445
+ "--decode-backends",
446
nargs="+",
444
- help="Backends to benchmark (flash, triton, flashinfer, cutlass_mla, "
447
+ help="Decode backends to benchmark (flash, triton, flashinfer, cutlass_mla, "
448
"flashinfer_mla, flashattn_mla, flashmla)",
449
)
450
parser.add_argument(
451
"--backend",
452
help="Single backend (alternative to --backends)",
453
)
454
+ parser.add_argument(
455
+ "--prefill-backends",
456
+ nargs="+",
457
+ help="Prefill backends to compare (fa2, fa3, fa4). "
458
+ "Uses the first decode backend for impl construction.",
459
+ )
460
461
# Batch specifications
462
parser.add_argument(
@@ -502,7 +511,7 @@ def main():
511
512
# Override args with YAML values, but CLI args take precedence
513
# Check if CLI provided backends (they would be non-None and not default)
505
- cli_backends_provided = args.backends is not None or args.backend is not None
514
+ cli_backends_provided = args.backend is not None or args.backends is not None
515
516
# Backend(s) - only use YAML if CLI didn't specify
517
if not cli_backends_provided:
@@ -512,6 +521,12 @@ def main():
521
elif "backends" in yaml_config:
522
args.backends = yaml_config["backends"]
523
args.backend = None
524
+ elif "decode_backends" in yaml_config:
525
+ args.backends = yaml_config["decode_backends"]
526
+ args.backend = None
527
+
528
+ # Prefill backends (e.g., ["fa3", "fa4"])
529
+ args.prefill_backends = yaml_config.get("prefill_backends", None)
530
531
# Check for special modes
532
if "mode" in yaml_config:
@@ -613,7 +628,10 @@ def main():
628
629
# Determine backends
630
backends = args.backends or ([args.backend] if args.backend else ["flash"])
631
+ prefill_backends = getattr(args, "prefill_backends", None)
632
console.print(f"Backends: {', '.join(backends)}")
633
+ if prefill_backends:
634
+ console.print(f"Prefill backends: {', '.join(prefill_backends)}")
635
console.print(f"Batch specs: {', '.join(args.batch_specs)}")
636
console.print()
637
@@ -850,37 +868,93 @@ def main():
868
869
else:
870
# Normal mode: compare backends
853
- total = len(backends) * len(args.batch_specs)
871
+ decode_results = []
872
+ prefill_results = []
873
855
- with tqdm(total=total, desc="Benchmarking") as pbar:
856
- for spec in args.batch_specs:
857
- for backend in backends:
858
- config = BenchmarkConfig(
859
- backend=backend,
860
- batch_spec=spec,
861
- num_layers=args.num_layers,
862
- head_dim=args.head_dim,
863
- num_q_heads=args.num_q_heads,
864
- num_kv_heads=args.num_kv_heads,
865
- block_size=args.block_size,
866
- device=args.device,
867
- repeats=args.repeats,
868
- warmup_iters=args.warmup_iters,
869
- profile_memory=args.profile_memory,
870
- )
874
+ # Run decode backend comparison
875
+ if not prefill_backends:
876
+ # No prefill backends specified: compare decode backends as before
877
+ total = len(backends) * len(args.batch_specs)
878
872
- result = run_benchmark(config)
873
- all_results.append(result)
879
+ with tqdm(total=total, desc="Benchmarking") as pbar:
880
+ for spec in args.batch_specs:
881
+ for backend in backends:
882
+ config = BenchmarkConfig(
883
+ backend=backend,
884
+ batch_spec=spec,
885
+ num_layers=args.num_layers,
886
+ head_dim=args.head_dim,
887
+ num_q_heads=args.num_q_heads,
888
+ num_kv_heads=args.num_kv_heads,
889
+ block_size=args.block_size,
890
+ device=args.device,
891
+ repeats=args.repeats,
892
+ warmup_iters=args.warmup_iters,
893
+ profile_memory=args.profile_memory,
894
+ )
895
875
- if not result.success:
876
- console.print(f"[red]Error {backend} {spec}: {result.error}[/]")
896
+ result = run_benchmark(config)
897
+ decode_results.append(result)
898
878
- pbar.update(1)
899
+ if not result.success:
900
+ console.print(
901
+ f"[red]Error {backend} {spec}: {result.error}[/]"
902
+ )
903
880
- # Display results
881
- console.print("\n[bold green]Results:[/]")
882
- formatter = ResultsFormatter(console)
883
- formatter.print_table(all_results, backends)
904
+ pbar.update(1)
905
+
906
+ console.print("\n[bold green]Results:[/]")
907
+ formatter = ResultsFormatter(console)
908
+ formatter.print_table(decode_results, backends)
909
+
910
+ # Run prefill backend comparison
911
+ if prefill_backends:
912
+ # Use first decode backend for impl construction
913
+ decode_backend = backends[0]
914
+ total = len(prefill_backends) * len(args.batch_specs)
915
+
916
+ console.print(
917
+ f"[yellow]Prefill comparison mode: "
918
+ f"using {decode_backend} for decode impl[/]"
919
+ )
920
+
921
+ with tqdm(total=total, desc="Prefill benchmarking") as pbar:
922
+ for spec in args.batch_specs:
923
+ for pb in prefill_backends:
924
+ config = BenchmarkConfig(
925
+ backend=decode_backend,
926
+ batch_spec=spec,
927
+ num_layers=args.num_layers,
928
+ head_dim=args.head_dim,
929
+ num_q_heads=args.num_q_heads,
930
+ num_kv_heads=args.num_kv_heads,
931
+ block_size=args.block_size,
932
+ device=args.device,
933
+ repeats=args.repeats,
934
+ warmup_iters=args.warmup_iters,
935
+ profile_memory=args.profile_memory,
936
+ prefill_backend=pb,
937
+ )
938
+
939
+ result = run_benchmark(config)
940
+
941
+ # Label result with prefill backend name for display
942
+ labeled_config = replace(result.config, backend=pb)
943
+ result = replace(result, config=labeled_config)
944
+ prefill_results.append(result)
945
+
946
+ if not result.success:
947
+ console.print(f"[red]Error {pb} {spec}: {result.error}[/]")
948
+
949
+ pbar.update(1)
950
+
951
+ console.print("\n[bold green]Prefill Backend Results:[/]")
952
+ formatter = ResultsFormatter(console)
953
+ formatter.print_table(
954
+ prefill_results, prefill_backends, compare_to_fastest=True
955
+ )
956
+
957
+ all_results = decode_results + prefill_results
958
959
# Save results
960
if all_results:
benchmarks/attention_benchmarks/common.py
+2
@@ -77,6 +77,7 @@ class MockKVBProj:
77
self.qk_nope_head_dim = qk_nope_head_dim
78
self.v_head_dim = v_head_dim
79
self.out_dim = qk_nope_head_dim + v_head_dim
80
+ self.weight = torch.empty(0, dtype=torch.bfloat16)
81
82
def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:
83
"""
@@ -213,6 +214,7 @@ class BenchmarkConfig:
214
use_cuda_graphs: bool = False
215
216
# MLA-specific
217
+ prefill_backend: str | None = None
218
kv_lora_rank: int | None = None
219
qk_nope_head_dim: int | None = None
220
qk_rope_head_dim: int | None = None
benchmarks/attention_benchmarks/configs/mla_prefill.yaml
+89
-25
@@ -1,4 +1,19 @@
1
-# MLA prefill-only benchmark configuration for sparse backends
1
+# MLA prefill backend comparison
2
+#
3
+# Compares all available MLA prefill backends:
4
+# FA backends: fa2, fa3, fa4 (FlashAttention versions)
5
+# Non-FA: flashinfer, cudnn, trtllm (Blackwell-only, require flashinfer)
6
+#
7
+# Uses cutlass_mla as the decode backend for impl construction
8
+# (only the prefill path is exercised).
9
+#
10
+# Backends that aren't available on the current platform will report errors
11
+# in the results table (e.g., fa3 on Blackwell, cudnn without artifactory).
12
+#
13
+# Usage:
14
+# python benchmark.py --config configs/mla_prefill.yaml
15
+
16
+description: "MLA prefill backend comparison"
17
18
model:
19
name: "deepseek-v3"
@@ -12,20 +27,25 @@ model:
27
v_head_dim: 128
28
block_size: 128
29
15
-# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
16
-# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
17
-model_parameter_sweep:
18
- param_name: "num_q_heads"
19
- values: [128, 64, 32, 16]
20
- label_format: "{backend}_{value}h"
30
+# model:
31
+# name: "deepseek-v2-lite"
32
+# num_layers: 27
33
+# num_q_heads: 16
34
+# num_kv_heads: 1
35
+# head_dim: 576
36
+# kv_lora_rank: 512
37
+# qk_nope_head_dim: 128
38
+# qk_rope_head_dim: 64
39
+# v_head_dim: 128
40
+# block_size: 128
41
42
batch_specs:
43
# Pure prefill
24
- - "1q512"
25
- - "1q1k"
26
- - "1q2k"
27
- - "1q4k"
28
- - "1q8k"
44
+ - "q512"
45
+ - "q1k"
46
+ - "q2k"
47
+ - "q4k"
48
+ - "q8k"
49
50
# Batched pure prefill
51
- "2q512"
@@ -44,19 +64,63 @@ batch_specs:
64
- "8q4k"
65
- "8q8k"
66
47
- # Extend
48
- - "1q512s4k"
49
- - "1q512s8k"
50
- - "1q1ks8k"
51
- - "1q2ks8k"
52
- - "1q2ks16k"
53
- - "1q4ks16k"
67
+ # Chunked prefill / extend
68
+ # Short context
69
+ - "q128s1k"
70
+ - "q256s2k"
71
+ - "q512s4k"
72
+ - "q1ks4k"
73
+ - "q2ks8k"
74
+ - "2q128s1k"
75
+ - "2q256s2k"
76
+ - "2q512s4k"
77
+ - "2q1ks4k"
78
+ - "2q2ks8k"
79
+ - "4q128s1k"
80
+ - "4q256s2k"
81
+ - "4q512s4k"
82
+ - "4q1ks4k"
83
+ - "4q2ks8k"
84
+ - "8q128s1k"
85
+ - "8q256s2k"
86
+ - "8q512s4k"
87
+ - "8q1ks4k"
88
+
89
+ # Medium context
90
+ - "q128s16k"
91
+ - "q512s16k"
92
+ - "q1ks16k"
93
+ - "q2ks16k"
94
+ - "2q128s16k"
95
+ - "2q512s16k"
96
+ - "2q1ks16k"
97
+ - "2q2ks16k"
98
+ - "4q128s16k"
99
+ - "4q512s16k"
100
+ - "4q1ks16k"
101
+ - "4q2ks16k"
102
+
103
+ # Long context
104
+ - "q128s64k"
105
+ - "q512s64k"
106
+ - "q1ks64k"
107
+ - "q2ks64k"
108
+ - "2q128s64k"
109
+ - "2q512s64k"
110
+ - "2q1ks64k"
111
+ - "2q2ks64k"
112
+
113
+decode_backends:
114
+ - CUTLASS_MLA
115
55
-backends:
56
- - FLASHMLA_SPARSE
57
- - FLASHINFER_MLA_SPARSE
116
+prefill_backends:
117
+ - fa2
118
+ - fa3
119
+ - fa4
120
+ - flashinfer
121
+ - cudnn
122
+ - trtllm
123
124
device: "cuda:0"
60
-repeats: 10
61
-warmup_iters: 3
62
-profile_memory: true
125
+repeats: 20
126
+warmup_iters: 5
benchmarks/attention_benchmarks/configs/mla_sparse_prefill.yaml
new
+62
@@ -0,0 +1,62 @@
1
+# MLA prefill-only benchmark configuration for sparse backends
2
+
3
+model:
4
+ name: "deepseek-v3"
5
+ num_layers: 60
6
+ num_q_heads: 128
7
+ num_kv_heads: 1
8
+ head_dim: 576
9
+ kv_lora_rank: 512
10
+ qk_nope_head_dim: 128
11
+ qk_rope_head_dim: 64
12
+ v_head_dim: 128
13
+ block_size: 128
14
+
15
+# Model parameter sweep: simulate tensor parallelism by varying num_q_heads
16
+# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads
17
+model_parameter_sweep:
18
+ param_name: "num_q_heads"
19
+ values: [128, 64, 32, 16]
20
+ label_format: "{backend}_{value}h"
21
+
22
+batch_specs:
23
+ # Pure prefill
24
+ - "1q512"
25
+ - "1q1k"
26
+ - "1q2k"
27
+ - "1q4k"
28
+ - "1q8k"
29
+
30
+ # Batched pure prefill
31
+ - "2q512"
32
+ - "2q1k"
33
+ - "2q2k"
34
+ - "2q4k"
35
+ - "2q8k"
36
+ - "4q512"
37
+ - "4q1k"
38
+ - "4q2k"
39
+ - "4q4k"
40
+ - "4q8k"
41
+ - "8q512"
42
+ - "8q1k"
43
+ - "8q2k"
44
+ - "8q4k"
45
+ - "8q8k"
46
+
47
+ # Extend
48
+ - "1q512s4k"
49
+ - "1q512s8k"
50
+ - "1q1ks8k"
51
+ - "1q2ks8k"
52
+ - "1q2ks16k"
53
+ - "1q4ks16k"
54
+
55
+backends:
56
+ - FLASHMLA_SPARSE
57
+ - FLASHINFER_MLA_SPARSE
58
+
59
+device: "cuda:0"
60
+repeats: 10
61
+warmup_iters: 3
62
+profile_memory: true
benchmarks/attention_benchmarks/mla_runner.py
+143
-14
@@ -62,6 +62,7 @@ def create_minimal_vllm_config(
62
max_num_seqs: int = 256,
63
mla_dims: dict | None = None,
64
index_topk: int | None = None,
65
+ prefill_backend: str | None = None,
66
) -> VllmConfig:
67
"""
68
Create minimal VllmConfig for MLA benchmarks.
@@ -75,6 +76,9 @@ def create_minimal_vllm_config(
76
setup_mla_dims(model_name)
77
index_topk: Optional topk value for sparse MLA backends. If provided,
78
the config will include index_topk for sparse attention.
79
+ prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
80
+ "cudnn", "trtllm"). Configures the attention config to
81
+ force the specified prefill backend.
82
83
Returns:
84
VllmConfig for benchmarking
@@ -163,7 +167,7 @@ def create_minimal_vllm_config(
167
168
compilation_config = CompilationConfig()
169
166
- return VllmConfig(
170
+ vllm_config = VllmConfig(
171
model_config=model_config,
172
cache_config=cache_config,
173
parallel_config=parallel_config,
@@ -171,9 +175,84 @@ def create_minimal_vllm_config(
175
compilation_config=compilation_config,
176
)
177
178
+ if prefill_backend is not None:
179
+ prefill_cfg = get_prefill_backend_config(prefill_backend)
180
+ if prefill_cfg["flash_attn_version"] is not None:
181
+ vllm_config.attention_config.flash_attn_version = prefill_cfg[
182
+ "flash_attn_version"
183
+ ]
184
+ vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[
185
+ "disable_flashinfer_prefill"
186
+ ]
187
+ vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[
188
+ "use_cudnn_prefill"
189
+ ]
190
+ vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = prefill_cfg[
191
+ "use_trtllm_ragged_deepseek_prefill"
192
+ ]
193
+
194
+ return vllm_config
195
+
196
197
# ============================================================================
176
-# Backend Configuration
198
+# Prefill Backend Configuration
199
+# ============================================================================
200
+
201
+# Maps prefill backend names to attention config overrides.
202
+# FA backends set flash_attn_version and disable non-FA paths.
203
+# Non-FA backends enable their specific path and disable others.
204
+_PREFILL_BACKEND_CONFIG: dict[str, dict] = {
205
+ "fa2": {
206
+ "flash_attn_version": 2,
207
+ "disable_flashinfer_prefill": True,
208
+ "use_cudnn_prefill": False,
209
+ "use_trtllm_ragged_deepseek_prefill": False,
210
+ },
211
+ "fa3": {
212
+ "flash_attn_version": 3,
213
+ "disable_flashinfer_prefill": True,
214
+ "use_cudnn_prefill": False,
215
+ "use_trtllm_ragged_deepseek_prefill": False,
216
+ },
217
+ "fa4": {
218
+ "flash_attn_version": 4,
219
+ "disable_flashinfer_prefill": True,
220
+ "use_cudnn_prefill": False,
221
+ "use_trtllm_ragged_deepseek_prefill": False,
222
+ },
223
+ "flashinfer": {
224
+ "flash_attn_version": None,
225
+ "disable_flashinfer_prefill": False,
226
+ "use_cudnn_prefill": False,
227
+ "use_trtllm_ragged_deepseek_prefill": False,
228
+ },
229
+ "cudnn": {
230
+ "flash_attn_version": None,
231
+ "disable_flashinfer_prefill": True,
232
+ "use_cudnn_prefill": True,
233
+ "use_trtllm_ragged_deepseek_prefill": False,
234
+ },
235
+ "trtllm": {
236
+ "flash_attn_version": None,
237
+ "disable_flashinfer_prefill": True,
238
+ "use_cudnn_prefill": False,
239
+ "use_trtllm_ragged_deepseek_prefill": True,
240
+ },
241
+}
242
+
243
+
244
+def get_prefill_backend_config(prefill_backend: str) -> dict:
245
+ """Get attention config overrides for a prefill backend."""
246
+ if prefill_backend not in _PREFILL_BACKEND_CONFIG:
247
+ raise ValueError(
248
+ f"Unknown prefill backend: {prefill_backend!r}. "
249
+ f"Available: {list(_PREFILL_BACKEND_CONFIG.keys())}"
250
+ )
251
+ return _PREFILL_BACKEND_CONFIG[prefill_backend]
252
+
253
+
254
+# ============================================================================
255
+# Decode Backend Configuration
256
# ============================================================================
257
258
@@ -203,6 +282,7 @@ def _get_backend_config(backend: str) -> dict:
282
Returns:
283
Dict with backend configuration
284
"""
285
+ from vllm.v1.attention.backend import MultipleOf
286
from vllm.v1.attention.backends.registry import AttentionBackendEnum
287
288
try:
@@ -219,8 +299,8 @@ def _get_backend_config(backend: str) -> dict:
299
block_sizes = backend_class.get_supported_kernel_block_sizes()
300
# Use first supported block size (backends typically support one for MLA)
301
block_size = block_sizes[0] if block_sizes else None
222
- if hasattr(block_size, "value"):
223
- # Handle MultipleOf enum
302
+ if isinstance(block_size, MultipleOf):
303
+ # No fixed block size; fall back to config value
304
block_size = None
305
306
# Check if sparse via class method if available
@@ -676,16 +756,11 @@ def _run_single_benchmark(
756
if is_sparse and indexer is not None:
757
indexer.fill_random_indices(total_q, max_kv_len)
758
679
- # Determine which forward method to use
680
- if is_sparse:
681
- # Sparse backends use forward_mqa
759
+ # Determine which forward method to use based on metadata
760
+ if metadata.decode is not None:
761
forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)
683
- elif metadata.decode is not None:
684
- forward_fn = lambda: impl._forward_decode(
685
- decode_inputs, kv_cache, metadata, layer
686
- )
762
elif metadata.prefill is not None:
688
- forward_fn = lambda: impl._forward_prefill(
763
+ forward_fn = lambda: impl.forward_mha(
764
prefill_inputs["q"],
765
prefill_inputs["k_c_normed"],
766
prefill_inputs["k_pe"],
@@ -732,6 +807,7 @@ def _run_mla_benchmark_batched(
807
backend: str,
808
configs_with_params: list[tuple], # [(config, threshold, num_splits), ...]
809
index_topk: int = 2048,
810
+ prefill_backend: str | None = None,
811
) -> list[BenchmarkResult]:
812
"""
813
Unified batched MLA benchmark runner for all backends.
@@ -743,11 +819,13 @@ def _run_mla_benchmark_batched(
819
to avoid setup/teardown overhead.
820
821
Args:
746
- backend: Backend name
822
+ backend: Backend name (decode backend used for impl construction)
823
configs_with_params: List of (config, threshold, num_splits) tuples
824
- threshold: reorder_batch_threshold (FlashAttn/FlashMLA only)
825
- num_splits: num_kv_splits (CUTLASS only)
826
index_topk: Topk value for sparse MLA backends (default 2048)
827
+ prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
828
+ When set, forces the specified FlashAttention version for prefill.
829
830
Returns:
831
List of BenchmarkResult objects
@@ -780,11 +858,25 @@ def _run_mla_benchmark_batched(
858
block_size=block_size,
859
mla_dims=mla_dims, # Use custom dims from config or default
860
index_topk=index_topk if is_sparse else None,
861
+ prefill_backend=prefill_backend,
862
)
863
864
results = []
865
866
with set_current_vllm_config(vllm_config):
867
+ # Clear cached prefill backend detection functions so they re-evaluate
868
+ # with the current VllmConfig. These are @functools.cache decorated and
869
+ # would otherwise return stale results from a previous backend's config.
870
+ from vllm.model_executor.layers.attention.mla_attention import (
871
+ use_cudnn_prefill,
872
+ use_flashinfer_prefill,
873
+ use_trtllm_ragged_deepseek_prefill,
874
+ )
875
+
876
+ use_flashinfer_prefill.cache_clear()
877
+ use_cudnn_prefill.cache_clear()
878
+ use_trtllm_ragged_deepseek_prefill.cache_clear()
879
+
880
# Create backend impl, layer, builder, and indexer (reused across benchmarks)
881
impl, layer, builder_instance, indexer = _create_backend_impl(
882
backend_cfg,
@@ -794,6 +886,38 @@ def _run_mla_benchmark_batched(
886
index_topk=index_topk if is_sparse else None,
887
)
888
889
+ # Verify the actual prefill backend matches what was requested
890
+ if prefill_backend is not None:
891
+ prefill_cfg = get_prefill_backend_config(prefill_backend)
892
+ fa_version = prefill_cfg["flash_attn_version"]
893
+
894
+ if fa_version is not None:
895
+ # FA backend: verify the impl's FA version
896
+ actual_fa_version = getattr(impl, "vllm_flash_attn_version", None)
897
+ if actual_fa_version != fa_version:
898
+ raise RuntimeError(
899
+ f"Prefill backend '{prefill_backend}' requested FA "
900
+ f"version {fa_version}, but the impl is using FA "
901
+ f"version {actual_fa_version}. Check "
902
+ f"vllm/v1/attention/backends/fa_utils.py."
903
+ )
904
+ else:
905
+ # Non-FA backend: verify the builder picked the right path
906
+ expected_flags = {
907
+ "flashinfer": "_use_fi_prefill",
908
+ "cudnn": "_use_cudnn_prefill",
909
+ "trtllm": "_use_trtllm_ragged_prefill",
910
+ }
911
+ flag_name = expected_flags.get(prefill_backend)
912
+ if flag_name and not getattr(builder_instance, flag_name, False):
913
+ raise RuntimeError(
914
+ f"Prefill backend '{prefill_backend}' was requested "
915
+ f"but the metadata builder did not enable it. This "
916
+ f"usually means a dependency is missing (e.g., "
917
+ f"flashinfer not installed) or the platform doesn't "
918
+ f"support it."
919
+ )
920
+
921
# Run each benchmark with the shared impl
922
for config, threshold, num_splits in configs_with_params:
923
# Set threshold for this benchmark (FlashAttn/FlashMLA only)
@@ -844,6 +968,7 @@ def run_mla_benchmark(
968
reorder_batch_threshold: int | None = None,
969
num_kv_splits: int | None = None,
970
index_topk: int = 2048,
971
+ prefill_backend: str | None = None,
972
) -> BenchmarkResult | list[BenchmarkResult]:
973
"""
974
Unified MLA benchmark runner for all backends.
@@ -861,6 +986,8 @@ def run_mla_benchmark(
986
(single config mode only)
987
num_kv_splits: Number of KV splits for CUTLASS (single config mode only)
988
index_topk: Topk value for sparse MLA backends (default 2048)
989
+ prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
990
+ When set, forces the specified FlashAttention version for prefill.
991
992
Returns:
993
BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)
@@ -884,7 +1011,9 @@ def run_mla_benchmark(
1011
return_single = True
1012
1013
# Use unified batched execution
887
- results = _run_mla_benchmark_batched(backend, configs_with_params, index_topk)
1014
+ results = _run_mla_benchmark_batched(
1015
+ backend, configs_with_params, index_topk, prefill_backend=prefill_backend
1016
+ )
1017
1018
# Return single result or list based on input
1019
return results[0] if return_single else results
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -39,7 +39,7 @@ else()
39
FetchContent_Declare(
40
vllm-flash-attn
41
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
42
- GIT_TAG 140c00c0241bb60cc6e44e7c1be9998d4b20d8d2
42
+ GIT_TAG 1488682bb545f7d020e958a33116b1419d1cfc83
43
GIT_PROGRESS TRUE
44
# Don't share the vllm-flash-attn build between build types
45
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
vllm/config/attention.py
+2
-2
@@ -30,14 +30,14 @@ class AttentionConfig:
30
use_cudnn_prefill: bool = False
31
"""Whether to use cudnn prefill."""
32
33
- use_trtllm_ragged_deepseek_prefill: bool = True
33
+ use_trtllm_ragged_deepseek_prefill: bool = False
34
"""Whether to use TRTLLM ragged deepseek prefill."""
35
36
use_trtllm_attention: bool | None = None
37
"""If set to True/False, use or don't use the TRTLLM attention backend
38
in flashinfer. If None, auto-detect the attention backend in flashinfer."""
39
40
- disable_flashinfer_prefill: bool = False
40
+ disable_flashinfer_prefill: bool = True
41
"""Whether to disable flashinfer prefill."""
42
43
disable_flashinfer_q_quantization: bool = False
vllm/model_executor/layers/attention/mla_attention.py
+8
-7
@@ -1282,8 +1282,6 @@ def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool:
1282
1283
@functools.cache
1284
def use_flashinfer_prefill() -> bool:
1285
- # For blackwell default to flashinfer prefill if it's available since
1286
- # it is faster than FA2.
1285
from vllm.config import get_current_vllm_config
1286
1287
vllm_config = get_current_vllm_config()
@@ -2154,13 +2152,16 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2152
2153
# For MLA the v head dim is smaller than qk head dim so we pad out
2154
# v with 0s to match the qk head dim for attention backends that do
2157
- # not support different headdims
2158
- # We don't need to pad V if we are on a hopper system with FA3
2155
+ # not support different headdims.
2156
+ # FA3 on Hopper (SM90) and FA4 natively handle diff headdims.
2157
device_capability = current_platform.get_device_capability()
2158
self._pad_v = self.vllm_flash_attn_version is None or not (
2161
- self.vllm_flash_attn_version == 3
2162
- and device_capability is not None
2163
- and device_capability[0] == 9
2159
+ (
2160
+ self.vllm_flash_attn_version == 3
2161
+ and device_capability is not None
2162
+ and device_capability[0] == 9
2163
+ )
2164
+ or self.vllm_flash_attn_version == 4
2165
)
2166
2167
self.dcp_world_size: int = -1
vllm/v1/attention/backends/fa_utils.py
+3
@@ -125,11 +125,14 @@ def get_flash_attn_version(
125
# FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
126
# supported head dimensions.
127
# See: https://github.com/Dao-AILab/flash-attention/issues/1959
128
+ # Exception: hdim 192 is supported for MLA's diff-headdim case
129
+ # (qk=192, v=128), added upstream in commits 1a15733e/1b36ab19.
130
if (
131
fa_version == 4
132
and device_capability.major >= 10
133
and head_size is not None
134
and head_size > 128
135
+ and head_size != 192
136
):
137
logger.warning_once(
138
"FA4 on Blackwell does not support head_size=%d due to TMEM "