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 9, 23:00 UTC
[Attention][MLA] Make FLASHINFER_MLA the default MLA backend on Blackwell, and TRTLLM the default prefill (#32615)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Matthew Bonanni committed
Jan 19, 2026 at 18:41 UTC
1a1fc3bbc062e3ad289447e6a13ce57971618b43
6 files changed
+19
-18
.buildkite/test-amd.yaml
+1
-1
@@ -455,7 +455,7 @@ steps:
455
- vllm/v1/attention
456
- tests/v1/attention
457
commands:
458
- - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
458
+ - pytest -v -s v1/attention
459
460
- label: V1 Test others (CPU) # 5 mins
461
mirror_hardwares: [amdexperimental, amdproduction, amdtentative]
.buildkite/test-pipeline.yaml
+1
-1
@@ -399,7 +399,7 @@ steps:
399
- vllm/v1/attention
400
- tests/v1/attention
401
commands:
402
- - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
402
+ - pytest -v -s v1/attention
403
404
- label: V1 Test others (CPU) # 5 mins
405
source_file_dependencies:
.buildkite/test_areas/attention.yaml
+1
-1
@@ -18,4 +18,4 @@ steps:
18
- vllm/v1/attention
19
- tests/v1/attention
20
commands:
21
- - VLLM_DISABLE_FLASHINFER_PREFILL=1 pytest -v -s v1/attention # TODO: FI prefill is bugged and causes incorrectness, fix this
21
+ - pytest -v -s v1/attention
vllm/config/attention.py
+1
-1
@@ -35,7 +35,7 @@ class AttentionConfig:
35
use_cudnn_prefill: bool = False
36
"""Whether to use cudnn prefill."""
37
38
- use_trtllm_ragged_deepseek_prefill: bool = False
38
+ use_trtllm_ragged_deepseek_prefill: bool = True
39
"""Whether to use TRTLLM ragged deepseek prefill."""
40
41
use_trtllm_attention: bool | None = None
vllm/model_executor/layers/attention/mla_attention.py
+11
-10
@@ -450,7 +450,6 @@ def use_flashinfer_prefill() -> bool:
450
not vllm_config.attention_config.disable_flashinfer_prefill
451
and flashinfer_available
452
and not vllm_config.attention_config.use_cudnn_prefill
453
- and not vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill
453
and current_platform.is_device_capability_family(100)
454
)
455
@@ -1324,25 +1323,27 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1323
def __init__(self, *args, **kwargs) -> None:
1324
super().__init__(*args, **kwargs)
1325
1327
- if use_flashinfer_prefill():
1328
- logger.debug_once("Using FlashInfer prefill for MLA")
1329
- self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi
1330
- self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi
1331
- self._pad_v = False
1332
- elif use_trtllm_ragged_deepseek_prefill():
1333
- logger.debug_once("Using TRT-LLM ragged DeepSeek prefill for MLA")
1326
+ if use_trtllm_ragged_deepseek_prefill():
1327
+ logger.info_once(
1328
+ "Using TRT-LLM ragged DeepSeek prefill for MLA", scope="local"
1329
+ )
1330
self._run_prefill_context_chunk = (
1331
self._run_prefill_context_chunk_trtllm_ragged
1332
)
1333
self._run_prefill_new_tokens = self._run_prefill_new_tokens_trtllm_ragged
1334
self._pad_v = False
1335
+ elif use_flashinfer_prefill():
1336
+ logger.info_once("Using FlashInfer prefill for MLA", scope="local")
1337
+ self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi
1338
+ self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi
1339
+ self._pad_v = False
1340
elif use_cudnn_prefill():
1340
- logger.debug_once("Using CUDNN prefill for MLA")
1341
+ logger.info_once("Using CUDNN prefill for MLA", scope="local")
1342
self._run_prefill_context_chunk = self._run_prefill_context_chunk_cudnn
1343
self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn
1344
self._pad_v = False
1345
else: # Use FlashAttention
1345
- logger.debug_once("Using FlashAttention prefill for MLA")
1346
+ logger.info_once("Using FlashAttention prefill for MLA", scope="local")
1347
self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa
1348
self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa
1349
vllm/platforms/cuda.py
+4
-4
@@ -50,8 +50,8 @@ def _get_backend_priorities(
50
if use_mla:
51
if device_capability.major == 10:
52
return [
53
- AttentionBackendEnum.CUTLASS_MLA,
53
AttentionBackendEnum.FLASHINFER_MLA,
54
+ AttentionBackendEnum.CUTLASS_MLA,
55
AttentionBackendEnum.FLASH_ATTN_MLA,
56
AttentionBackendEnum.FLASHMLA,
57
AttentionBackendEnum.TRITON_MLA,
@@ -183,12 +183,12 @@ class CudaPlatformBase(Platform):
183
if vllm_config.attention_config.backend is None:
184
# Default case
185
if cls.is_device_capability_family(100) and not use_sparse:
186
- # Blackwell => Force CutlassMLA (unless sparse, i.e. DSv3.2).
187
- use_cutlass_mla = True
186
+ # Blackwell => Force FlashInfer MLA (unless sparse, i.e. DSv3.2).
187
+ use_flashinfer_mla = True
188
# Set the backend in AttentionConfig so it's used during
189
# backend selection
190
vllm_config.attention_config.backend = (
191
- AttentionBackendEnum.CUTLASS_MLA
191
+ AttentionBackendEnum.FLASHINFER_MLA
192
)
193
else:
194
# Not Blackwell