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 14, 05:30 UTC
[Attention][FlashInfer] Enable FP8 FlashInfer (TRTLLM) MLA decode (#24705)
Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com>
Matthew Bonanni committed
Sep 12, 2025 at 17:45 UTC
7ba32aa60b7bc1a08e9340d041014cc153f7a441
8 files changed
+23
-10
tests/v1/attention/test_attention_backends.py
+1
@@ -178,6 +178,7 @@ class MockAttentionLayer:
178
self._k_scale = torch.tensor(1.0, device=device)
179
self._v_scale = torch.tensor(1.0, device=device)
180
# Add float versions for flashinfer
181
+ self._q_scale_float = 1.0
182
self._k_scale_float = 1.0
183
self._v_scale_float = 1.0
184
tests/v1/tpu/test_pallas.py
+2
@@ -33,10 +33,12 @@ def test_ragged_paged_attention():
33
)
34
35
class FakeAttentionLayer:
36
+ _q_scale_float: float
37
_k_scale_float: float
38
_v_scale_float: float
39
40
layer = FakeAttentionLayer()
41
+ layer._q_scale_float = 1.0
42
layer._k_scale_float = 1.0
43
layer._v_scale_float = 1.0
44
vllm/attention/backends/abstract.py
+1
@@ -240,6 +240,7 @@ class AttentionLayer(Protocol):
240
_q_scale: torch.Tensor
241
_k_scale: torch.Tensor
242
_v_scale: torch.Tensor
243
+ _q_scale_float: float
244
_k_scale_float: float
245
_v_scale_float: float
246
_prob_scale: torch.Tensor
vllm/envs.py
+2
@@ -476,6 +476,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
476
# - "FLASHINFER": use flashinfer
477
# - "FLASHMLA": use FlashMLA
478
# - "FLASH_ATTN_MLA": use FlashAttention for MLA
479
+ # - "FLASHINFER_MLA": use FlashInfer for MLA
480
+ # - "CUTLASS_MLA": use CUTLASS for MLA
481
"VLLM_ATTENTION_BACKEND":
482
lambda: os.getenv("VLLM_ATTENTION_BACKEND", None),
483
vllm/model_executor/layers/quantization/kv_cache.py
+2
@@ -88,6 +88,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):
88
"Setting it to k_scale. This only matters for "
89
"the flash-attn backend.")
90
layer._q_scale.copy_(k_scale)
91
+ layer._q_scale_float = k_scale
92
93
# These are used in the final Attention.forward()
94
layer._k_scale.copy_(k_scale)
@@ -124,6 +125,7 @@ class BaseKVCacheMethod(QuantizeMethodBase):
125
126
# These are used in the final Attention.forward()
127
layer._q_scale.copy_(q_scale)
128
+ layer._q_scale_float = q_scale
129
layer._prob_scale.copy_(prob_scale)
130
if layer.kv_cache_dtype == "fp8" and (q_scale == 1.0
131
or prob_scale == 1.0):
vllm/platforms/cuda.py
+4
-1
@@ -179,6 +179,7 @@ class CudaPlatformBase(Platform):
179
cache_config.block_size = 128
180
logger.info("Forcing kv cache block size to 128 for "
181
"CUTLASS_MLA backend.")
182
+
183
if use_flashinfer_mla and cache_config.block_size not in [32, 64]:
184
cache_config.block_size = 64
185
logger.info(
@@ -541,7 +542,9 @@ class CudaPlatformBase(Platform):
542
attention_backend = "FLASHMLA"
543
544
# Only FlashMLA and CUTLASS_MLA support fp8
544
- if attention_backend in ["FLASHMLA", "CUTLASS_MLA"]:
545
+ if attention_backend in [
546
+ "FLASHMLA", "CUTLASS_MLA", "FLASHINFER_MLA"
547
+ ]:
548
supported = True
549
else:
550
supported = (not fp8_attention)
vllm/v1/attention/backends/mla/common.py
-2
@@ -584,7 +584,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
584
window_left=self._global_hyperparameters.window_left,
585
logits_soft_cap=self._global_hyperparameters.logits_soft_cap,
586
q_data_type=self.model_config.dtype,
587
- kv_data_type=self.kv_cache_spec.dtype,
587
)
588
589
# Prepare context prefills
@@ -605,7 +604,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
604
logits_soft_cap=self._global_hyperparameters.
605
logits_soft_cap,
606
q_data_type=self.model_config.dtype,
608
- kv_data_type=self.kv_cache_spec.dtype,
607
)
608
609
prefill.prefill_main = self._fi_prefill_main
vllm/v1/attention/backends/mla/flashinfer_mla.py
+11
-7
@@ -6,8 +6,7 @@ from typing import Optional, Union
6
import torch
7
from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
8
9
-from vllm.attention.backends.abstract import (AttentionLayer, AttentionType,
10
- is_quantized_kv_cache)
9
+from vllm.attention.backends.abstract import AttentionLayer, AttentionType
10
from vllm.logger import init_logger
11
from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
12
MLACommonImpl,
@@ -69,11 +68,9 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
68
"are not implemented for "
69
"FlashInferMLAImpl")
70
72
- if is_quantized_kv_cache(self.kv_cache_dtype):
73
- raise NotImplementedError(
74
- "FlashInferMLA V1 with FP8 KV cache not yet supported")
75
-
71
self._workspace_buffer = g_fi_workspace
72
+ self.bmm1_scale: Optional[float] = None
73
+ self.bmm2_scale: Optional[float] = None
74
75
def _forward_decode(
76
self,
@@ -92,6 +89,12 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
89
# trtllm API requires extra dimension q_len_per_request for MTP
90
q = q.unsqueeze(1)
91
92
+ if self.bmm1_scale is None:
93
+ self.bmm1_scale = (layer._q_scale_float * layer._k_scale_float *
94
+ self.scale)
95
+ if self.bmm2_scale is None:
96
+ self.bmm2_scale = layer._v_scale_float
97
+
98
o = trtllm_batch_decode_with_kv_cache_mla(
99
query=q,
100
kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),
@@ -102,7 +105,8 @@ class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):
105
block_tables=attn_metadata.decode.block_table,
106
seq_lens=attn_metadata.decode.seq_lens,
107
max_seq_len=attn_metadata.max_seq_len,
105
- bmm1_scale=self.scale,
108
+ bmm1_scale=self.bmm1_scale,
109
+ bmm2_scale=self.bmm2_scale,
110
)
111
112
# TODO: Return LSE pending support from Flashinfer API: