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, 13:15 UTC
[Perf] Set Flashinfer sparse MLA as default backend for FP8 kv cache (#37252)
Signed-off-by: wzhao18 <wzhao18.sz@gmail.com>
Wei Zhao committed
Mar 17, 2026 at 16:09 UTC
b36adfa349cfab0e79f3d736d5e5413bd3ee19f5
3 files changed
+61
-20
docs/design/attention_backends.md
+4
-2
@@ -127,8 +127,8 @@ Priority is **1 = highest** (tried first).
127
| 3 | `FLASH_ATTN_MLA` |
128
| 4 | `FLASHMLA` |
129
| 5 | `TRITON_MLA` |
130
-| 6 | `FLASHMLA_SPARSE` |
131
-| 7 | `FLASHINFER_MLA_SPARSE` |
130
+| 6 | `FLASHINFER_MLA_SPARSE`**\*** |
131
+| 7 | `FLASHMLA_SPARSE` |
132
133
**Ampere/Hopper (SM 8.x-9.x):**
134
@@ -140,6 +140,8 @@ Priority is **1 = highest** (tried first).
140
| 4 | `TRITON_MLA` |
141
| 5 | `FLASHMLA_SPARSE` |
142
143
+> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
144
+>
145
> **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
146
147
## Legend
tools/pre_commit/generate_attention_backend_docs.py
+26
-3
@@ -1262,14 +1262,23 @@ When no backend is specified (the default):
1262
"""
1263
1264
1265
-def _priority_table(title: str, backends: list[str]) -> list[str]:
1265
+def _priority_table(
1266
+ title: str,
1267
+ backends: list[str],
1268
+ annotations: dict[str, str] | None = None,
1269
+) -> list[str]:
1270
"""Generate a priority table for a list of backends."""
1271
+
1272
+ def _fmt(b: str) -> str:
1273
+ suffix = annotations.get(b, "") if annotations else ""
1274
+ return f"`{b}`{suffix}"
1275
+
1276
return [
1277
f"**{title}:**",
1278
"",
1279
"| Priority | Backend |",
1280
"| -------- | ------- |",
1272
- *[f"| {i} | `{b}` |" for i, b in enumerate(backends, 1)],
1281
+ *[f"| {i} | {_fmt(b)} |" for i, b in enumerate(backends, 1)],
1282
"",
1283
]
1284
@@ -1298,11 +1307,25 @@ def generate_priority_section(priorities: dict[str, list[str]]) -> str:
1307
1308
lines.extend(["### MLA Attention (DeepSeek-style)", ""])
1309
1310
+ mla_sm100_annotations = {
1311
+ "FLASHINFER_MLA_SPARSE": "**\\***",
1312
+ }
1313
if "mla_sm100" in priorities:
1302
- lines.extend(_priority_table(sm100, priorities["mla_sm100"]))
1314
+ lines.extend(
1315
+ _priority_table(sm100, priorities["mla_sm100"], mla_sm100_annotations)
1316
+ )
1317
if "mla_default" in priorities:
1318
lines.extend(_priority_table(ampere, priorities["mla_default"]))
1319
1320
+ if "mla_sm100" in priorities:
1321
+ lines.append(
1322
+ "> **\\*** For sparse MLA, FP8 KV cache always prefers "
1323
+ "`FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` "
1324
+ "is preferred for low query-head counts (<= 16), while "
1325
+ "`FLASHMLA_SPARSE` is preferred otherwise."
1326
+ )
1327
+ lines.append(">")
1328
+
1329
lines.append(
1330
"> **Note:** ROCm and CPU platforms have their own selection logic. "
1331
"See the platform-specific documentation for details."
vllm/platforms/cuda.py
+31
-15
@@ -4,6 +4,8 @@
4
pynvml. However, it should not initialize cuda context.
5
"""
6
7
+from __future__ import annotations
8
+
9
import os
10
from collections.abc import Callable
11
from datetime import timedelta
@@ -49,21 +51,34 @@ def _get_backend_priorities(
51
use_mla: bool,
52
device_capability: DeviceCapability,
53
num_heads: int | None = None,
54
+ kv_cache_dtype: CacheDType | None = None,
55
) -> list[AttentionBackendEnum]:
56
"""Get backend priorities with lazy import to avoid circular dependency."""
57
if use_mla:
58
if device_capability.major == 10:
56
- # Prefer FlashInfer at low head counts (FlashMLA uses padding)
57
- if num_heads is not None and num_heads <= 16:
59
+ # Sparse MLA backend priorities
60
+ # See https://github.com/vllm-project/vllm/issues/35807 for
61
+ # benchmark results
62
+ if kv_cache_dtype is not None and kv_cache_dtype.startswith("fp8"):
63
+ # Prefer FlashInfer for fp8 kv cache
64
sparse_backends = [
65
AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
66
AttentionBackendEnum.FLASHMLA_SPARSE,
67
]
68
else:
63
- sparse_backends = [
64
- AttentionBackendEnum.FLASHMLA_SPARSE,
65
- AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
66
- ]
69
+ # BF16 KV Cache
70
+ # Prefer FlashInfer at low head counts (FlashMLA uses padding)
71
+ if num_heads is not None and num_heads <= 16:
72
+ sparse_backends = [
73
+ AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
74
+ AttentionBackendEnum.FLASHMLA_SPARSE,
75
+ ]
76
+ else:
77
+ sparse_backends = [
78
+ AttentionBackendEnum.FLASHMLA_SPARSE,
79
+ AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
80
+ ]
81
+
82
return [
83
AttentionBackendEnum.FLASHINFER_MLA,
84
AttentionBackendEnum.CUTLASS_MLA,
@@ -165,7 +180,7 @@ class CudaPlatformBase(Platform):
180
pass
181
182
@classmethod
168
- def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
183
+ def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
184
parallel_config = vllm_config.parallel_config
185
model_config = vllm_config.model_config
186
@@ -198,11 +213,11 @@ class CudaPlatformBase(Platform):
213
def get_valid_backends(
214
cls,
215
device_capability: DeviceCapability,
201
- attn_selector_config: "AttentionSelectorConfig",
216
+ attn_selector_config: AttentionSelectorConfig,
217
num_heads: int | None = None,
218
) -> tuple[
204
- list[tuple["AttentionBackendEnum", int]],
205
- dict["AttentionBackendEnum", tuple[int, list[str]]],
219
+ list[tuple[AttentionBackendEnum, int]],
220
+ dict[AttentionBackendEnum, tuple[int, list[str]]],
221
]:
222
valid_backends_priorities = []
223
invalid_reasons: dict[AttentionBackendEnum, tuple[int, list[str]]] = {}
@@ -211,6 +226,7 @@ class CudaPlatformBase(Platform):
226
attn_selector_config.use_mla,
227
device_capability,
228
num_heads,
229
+ attn_selector_config.kv_cache_dtype,
230
)
231
for priority, backend in enumerate(backend_priorities):
232
try:
@@ -231,8 +247,8 @@ class CudaPlatformBase(Platform):
247
@classmethod
248
def get_attn_backend_cls(
249
cls,
234
- selected_backend: "AttentionBackendEnum | None",
235
- attn_selector_config: "AttentionSelectorConfig",
250
+ selected_backend: AttentionBackendEnum | None,
251
+ attn_selector_config: AttentionSelectorConfig,
252
num_heads: int | None = None,
253
) -> str:
254
device_capability = cls.get_device_capability()
@@ -324,7 +340,7 @@ class CudaPlatformBase(Platform):
340
return selected_backend.get_path()
341
342
@classmethod
327
- def get_supported_vit_attn_backends(cls) -> list["AttentionBackendEnum"]:
343
+ def get_supported_vit_attn_backends(cls) -> list[AttentionBackendEnum]:
344
if cls.has_device_capability(80):
345
return [
346
AttentionBackendEnum.FLASH_ATTN,
@@ -345,8 +361,8 @@ class CudaPlatformBase(Platform):
361
cls,
362
head_size: int,
363
dtype: torch.dtype,
348
- backend: "AttentionBackendEnum | None" = None,
349
- ) -> "AttentionBackendEnum":
364
+ backend: AttentionBackendEnum | None = None,
365
+ ) -> AttentionBackendEnum:
366
if backend is not None:
367
assert backend in cls.get_supported_vit_attn_backends(), (
368
f"Backend {backend} is not supported for vit attention. "