@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 379689d533

Fork 0

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 7, 02:45 UTC

[Perf] Support FP8 KV cache for Flashinfer MLA Sparse (#35891)

Wei Zhao committed Mar 7, 2026 at 16:51 UTC 379689d533642cfc1d3ab2cf4dc02f09a8318a5f

8 files changed +89 -17

docs/design/attention_backends.md

+1 -1

@@ -206,7 +206,7 @@ configuration.

206 |---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|

207 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |

208 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |

209 -| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |

209 +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |

210 | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |

211 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |

212 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |

tests/v1/attention/test_mla_backends.py

+18 -2

@@ -327,6 +327,12 @@ class MockSparseMLAAttentionLayer:

327 self._k_scale_float = 1.0

328 self._v_scale_float = 1.0

329

330 + self._decode_concat_quant_fp8_op = _DecodeConcatQuantFP8(

331 + static=True,

332 + group_shape=GroupShape.PER_TENSOR,

333 + compile_native=True,

334 + )

335 +

336 def forward_impl(

337 self,

338 q: torch.Tensor,

@@ -338,6 +344,7 @@ class MockSparseMLAAttentionLayer:

344 ) -> torch.Tensor:

345 """Forward for sparse MLA - uses forward_mqa for all tokens."""

346 kv_cache_dtype = getattr(self.impl, "kv_cache_dtype", "auto")

347 + fp8_attention = kv_cache_dtype.startswith("fp8")

348

349 # Write to KV cache

350 if kv_cache.numel() > 0:

@@ -350,6 +357,9 @@ class MockSparseMLAAttentionLayer:

357 scale=self._k_scale,

358 )

359

360 + if fp8_attention and kv_cache_dtype != "fp8_ds_mla":

361 + kv_cache = kv_cache.view(current_platform.fp8_dtype())

362 +

363 num_tokens = q.shape[0]

364

365 # Sparse MLA uses forward_mqa for all tokens

@@ -367,8 +377,14 @@ class MockSparseMLAAttentionLayer:

377 # Convert from (N, B, L) to (B, N, L)

378 mqa_ql_nope = mqa_ql_nope.transpose(0, 1)

379

370 - # Pass as tuple to forward_mqa

371 - mqa_q = (mqa_ql_nope, mqa_q_pe)

380 + if fp8_attention and self.impl.supports_quant_query_input:

381 + assert mqa_ql_nope.shape[0] == mqa_q_pe.shape[0]

382 + assert mqa_ql_nope.shape[1] == mqa_q_pe.shape[1]

383 + mqa_q = self._decode_concat_quant_fp8_op(

384 + mqa_ql_nope, mqa_q_pe, self._q_scale

385 + )

386 + else:

387 + mqa_q = (mqa_ql_nope, mqa_q_pe)

388

389 attn_out, _ = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)

390

tests/v1/attention/test_sparse_mla_backends.py

+11 -1

@@ -191,6 +191,16 @@ def test_sparse_backend_decode_correctness(

191 if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:

192 pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")

193

194 + if (

195 + backend_cls == FlashMLASparseBackend

196 + and kv_cache_dtype.startswith("fp8")

197 + and kv_cache_dtype != "fp8_ds_mla"

198 + ):

199 + pytest.skip(

200 + "FlashMLA Sparse Attention backend fp8 only supports "

201 + "fp8_ds_mla kv-cache dtype"

202 + )

203 +

204 supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()

205 if block_size not in supported_block_sizes:

206 pytest.skip(

@@ -419,7 +429,7 @@ def test_sparse_backend_decode_correctness(

429 num_blocks=vllm_config.cache_config.num_gpu_blocks,

430 common_attn_metadata=common_attn_metadata,

431 randomize_blocks=False,

422 - kv_cache_dtype=kv_cache_dtype if use_fp8_ds_mla_quantization else "auto",

432 + kv_cache_dtype=kv_cache_dtype,

433 scale=kv_cache_scale,

434 )

435

tools/pre_commit/generate_attention_backend_docs.py

+15 -1

@@ -49,6 +49,11 @@ MLA_ATTENTION_FILE = (

49 # Backends to skip during doc generation

50 SKIP_BACKENDS = {"CUSTOM", "TORCH_SDPA"}

52 +BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = {

53 + # fp8 is an alias for fp8_ds_mla for FlashMLA Sparse

54 + "FLASHMLA_SPARSE": {"fp8"},

55 +}

56 +

58 def is_relevant_file(filepath: str) -> bool:

59 """Check if a file matches any of the relevant patterns."""

@@ -546,10 +551,19 @@ def analyze_backend(backend_name: str, class_path: str) -> dict[str, Any] | None

551 tree, impl_class_name, "can_return_lse_for_decode", False, file_path

552 )

553

554 + kv_cache_dtypes = parse_kv_cache_dtypes(class_node)

555 + if backend_name in BACKEND_KV_DTYPE_EXCLUDES:

556 + excluded = BACKEND_KV_DTYPE_EXCLUDES[backend_name]

557 + kv_cache_dtypes = ", ".join(

558 + d

559 + for d in (d.strip() for d in kv_cache_dtypes.split(","))

560 + if d not in excluded

561 + )

562 +

563 return {

564 "name": backend_name,

565 "dtypes": parse_supported_dtypes(class_node),

552 - "kv_cache_dtypes": parse_kv_cache_dtypes(class_node),

566 + "kv_cache_dtypes": kv_cache_dtypes,

567 "block_sizes": parse_block_sizes(class_node),

568 "head_sizes": parse_head_sizes(class_node),

569 "attn_types": parse_attention_types(class_node),

vllm/model_executor/layers/attention/mla_attention.py

+30 -5

@@ -331,11 +331,6 @@ class MLAAttention(nn.Module, AttentionLayerBase):

331 calculate_kv_scales = False

332 self.quant_config = quant_config

333

334 - # Initialize KV cache quantization attributes

335 - self.kv_cache_dtype = kv_cache_dtype

336 - self.calculate_kv_scales = calculate_kv_scales

337 - _init_kv_cache_quant(self, quant_config, prefix)

338 -

334 dtype = torch.get_default_dtype()

335 self.attn_backend = get_attn_backend(

336 self.head_size,

@@ -347,6 +342,36 @@ class MLAAttention(nn.Module, AttentionLayerBase):

342 num_heads=self.num_heads,

343 )

344

345 + # FlashMLA Sparse Attention fp8 backend uses "fp8_ds_mla" kv-cache format

346 + # Automatically convert fp8 kv-cache format to "fp8_ds_mla"

347 + if (

348 + self.attn_backend.get_name() == "FLASHMLA_SPARSE"

349 + and kv_cache_dtype.startswith("fp8")

350 + and kv_cache_dtype != "fp8_ds_mla"

351 + ):

352 + assert cache_config is not None

353 + cache_config.cache_dtype = "fp8_ds_mla"

354 + kv_cache_dtype = "fp8_ds_mla"

355 + logger.info_once(

356 + "Using DeepSeek's fp8_ds_mla KV cache format. To use standard "

357 + "fp8 kv-cache format, please set `--attention-backend "

358 + "FLASHINFER_MLA_SPARSE`"

359 + )

360 +

361 + if (

362 + self.attn_backend.get_name() == "FLASHINFER_MLA_SPARSE"

363 + and kv_cache_dtype.startswith("fp8")

364 + ):

365 + logger.info_once(

366 + "Using standard fp8 KV cache format. To use DeepSeek's fp8_ds_mla "

367 + "KV cache format, please set `--attention-backend FLASHMLA_SPARSE`"

368 + )

369 +

370 + # Initialize KV cache quantization attributes

371 + self.kv_cache_dtype = kv_cache_dtype

372 + self.calculate_kv_scales = calculate_kv_scales

373 + _init_kv_cache_quant(self, quant_config, prefix)

374 +

375 if (

376 cache_config is not None

377 and cache_config.enable_prefix_caching

vllm/model_executor/models/config.py

-7

@@ -31,20 +31,13 @@ class VerifyAndUpdateConfig:

31 class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):

32 @classmethod

33 def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:

34 - """

35 - Updated fp8 cache to custom "fp8_ds_mla" format for DeepSeekV32

36 - """

34 hf_config = vllm_config.model_config.hf_config

36 # Mirror the check in vllm/model_executor/models/deepseek_v2.py

37 is_v32 = hasattr(hf_config, "index_topk")

38 assert is_v32

43 - # For DeepSeekV3.2, a custom fp8 format is used when fp8 kv-cache is enabled.

40 cache_config = vllm_config.cache_config

45 - if cache_config.cache_dtype.startswith("fp8"):

46 - cache_config.cache_dtype = "fp8_ds_mla"

47 - logger.info("Using custom fp8 kv-cache format for DeepSeekV3.2")

41 if cache_config.cache_dtype == "bfloat16":

42 cache_config.cache_dtype = "auto"

43 logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -63,6 +63,8 @@ class FlashInferMLASparseBackend(AttentionBackend):

63 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

64 "auto",

65 "bfloat16",

66 + "fp8",

67 + "fp8_e4m3",

68 ]

70 @staticmethod

@@ -304,6 +306,11 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata

306 self.bmm1_scale: float | None = None

307 self.bmm2_scale: float | None = None

308

309 + # fp8 query quantization is required when using fp8 kv_cache,

310 + # as the TRTLLM-GEN sparse MLA kernel requires matching dtypes

311 + # for query and kv_cache (mixed bf16+fp8 is not supported).

312 + self.supports_quant_query_input = True

313 +

314 def forward_mqa(

315 self,

316 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -83,6 +83,7 @@ class FlashMLASparseBackend(AttentionBackend):

83 "auto",

84 "bfloat16",

85 "fp8_ds_mla",

86 + "fp8", # alias for fp8_ds_mla

87 ]

89 @staticmethod

@@ -567,6 +568,12 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

568 )

569 self.fp8_decode_padded_heads = self._compute_fp8_decode_padded_heads(num_heads)

570

571 + if kv_cache_dtype.startswith("fp8"):

572 + assert kv_cache_dtype == "fp8_ds_mla", (

573 + "FlashMLA Sparse Attention backend fp8 only supports "

574 + "fp8_ds_mla kv-cache dtype"

575 + )

576 +

577 if kv_cache_dtype == "fp8_ds_mla":

578 # Reserve workspace during initialization

579 vllm_config = get_current_vllm_config()