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 4, 20:00 UTC
[compile] mla + group fp8 fusion (#38877)
Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Carl Y committed
Apr 21, 2026 at 20:16 UTC
4506319a286272e0ec4c4938691490ce6b33e832
9 files changed
+511
-59
docs/design/fusions.md
+2
-2
@@ -44,7 +44,7 @@ The table below lists the quantization schemes supported by each fusion on each
44
| `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |
45
| `fuse_minimax_qk_norm`\* | FP16/BF16 | FP16/BF16 | FP16/BF16 | FP16/BF16 | — |
46
| `fuse_attn_quant`\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static\* |
47
-| `fuse_attn_quant` (MLA)\* | FP8 static\*, NVFP4\* | FP8 static\* | FP8 static\* | — | FP8 static(untested)\* |
47
+| `fuse_attn_quant` (MLA)\* | FP8 static\*, FP8 per-group\*, NVFP4\* | FP8 static\*, FP8 per-group\* | FP8 static\*, FP8 per-group\* | — | FP8 static\* (untested) |
48
| `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |
49
| `enable_qk_norm_rope_fusion` | FP16/BF16 | FP16/BF16 | FP16/BF16† | FP16/BF16† | — |
50
| `enable_sp` | FP16/BF16, FP8 static† | FP16/BF16, FP8 static | FP16/BF16† | FP16/BF16† | — |
@@ -152,7 +152,7 @@ standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patt
152
153
- `FLASHINFER`: CUDA sm100+ with FlashInfer installed
154
155
-`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:
155
+`MLAAttention → FP8 static, FP8 per-group, NVFP4 dynamic quant`
156
157
The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works
158
with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where
tests/compile/fusions_e2e/conftest.py
+16
@@ -116,6 +116,22 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
116
model_kwargs["attention_config"] = {"backend": attn_backend.backend.name}
117
model_kwargs["tensor_parallel_size"] = tp_size
118
119
+ # Sparse MLA models (DSv3.2) hit an over-strict inductor assertion in
120
+ # decompose_auto_functionalized when +rotary_embedding is forced into
121
+ # the compile graph. Disable qk_norm+rope fusion (which auto-enables
122
+ # +rotary_embedding) for this combo to avoid the known torch bug.
123
+ # TODO: remove once upstream torch fix lands.
124
+ if requires_sparse:
125
+ if "pass_config" in compilation_config:
126
+ compilation_config["pass_config"].enable_qk_norm_rope_fusion = False
127
+ matches_check = [m for m in matches_check if m != "norm_rope_fusion"]
128
+ # DSv3.2 sparse indexer uses persistent_topk with k=config.index_topk
129
+ # (2048 for the default config). max_model_len must be >= index_topk
130
+ # or the topk kernel raises "k out of range" at runtime.
131
+ model_kwargs["max_model_len"] = max(
132
+ model_kwargs.get("max_model_len", 0), 2048
133
+ )
134
+
135
# Always compile the full graph instead of piecewise
136
if not compilation_config["use_inductor_graph_partition"]:
137
compilation_config["splitting_ops"] = []
tests/compile/fusions_e2e/models.py
+19
-5
@@ -59,7 +59,10 @@ TRITON_MLA_ATTN = pytest.param(
59
)
60
61
FLASHMLA_SPARSE_ATTN = pytest.param(
62
- AttentionBackendCase(backend=AttentionBackendEnum.FLASHMLA_SPARSE),
62
+ AttentionBackendCase(
63
+ backend=AttentionBackendEnum.FLASHMLA_SPARSE,
64
+ model_kwargs=dict(kv_cache_dtype="fp8_ds_mla"),
65
+ ),
66
id="FLASHMLA_SPARSE",
67
marks=pytest.mark.skipif(
68
not is_blackwell(),
@@ -173,9 +176,8 @@ deepseek_v3_fp8 = ModelFusionInfo(
176
rms_quant_fusion=n_layers * 2 + min(3, n_layers), # add for 3 dense layers
177
# silu+block quant
178
act_quant_fusion=min(3, n_layers), # dense layers only
176
- # MLA attn + per-group FP8 quant not supported yet:
177
- # https://github.com/vllm-project/vllm/issues/35792
178
- attn_quant_fusion=0,
179
+ # MLA attn + per-group FP8 quant
180
+ attn_quant_fusion=n_layers,
181
ar_rms_fusion=n_layers * 2 + 1,
182
# TODO
183
# sequence_parallel= n_layers * 2 + 1,
@@ -183,11 +185,23 @@ deepseek_v3_fp8 = ModelFusionInfo(
185
),
186
)
187
188
+deepseek_r1_fp4 = ModelFusionInfo(
189
+ model_name="nvidia/DeepSeek-R1-0528-NVFP4-v2",
190
+ matches=lambda n_layers: Matches(
191
+ rms_quant_fusion=0,
192
+ act_quant_fusion=min(3, n_layers),
193
+ attn_quant_fusion=n_layers,
194
+ ar_rms_fusion=n_layers * 2 + 1,
195
+ ),
196
+)
197
+
198
deepseek_v32_fp4 = ModelFusionInfo(
199
model_name="nvidia/DeepSeek-V3.2-NVFP4",
200
matches=lambda n_layers: Matches(
201
rms_quant_fusion=0,
190
- act_quant_fusion=0,
202
+ # silu+quant on dense layers only; MoE hides the act+quant site
203
+ act_quant_fusion=min(3, n_layers),
204
+ # MLA attn + NVFP4 output quant fuses on sparse MLA output path
205
attn_quant_fusion=n_layers,
206
ar_rms_fusion=n_layers * 2 + 1,
207
),
tests/compile/fusions_e2e/test_tp1_quant.py
+3
-2
@@ -24,6 +24,7 @@ from .models import (
24
TRITON_ATTN,
25
TRITON_MLA_ATTN,
26
deepseek_coder_v2_lite_fp8,
27
+ deepseek_r1_fp4,
28
deepseek_v3_fp8,
29
deepseek_v32_fp4,
30
llama3_8b_fp4,
@@ -148,11 +149,11 @@ def test_tp1_fp8_fusions(
149
150
@pytest.mark.parametrize(
151
"model_name, matches_fn, model_kwargs, hf_overrides",
151
- [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],
152
+ [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4],
153
)
154
@pytest.mark.parametrize(
155
"attn_backend",
155
- [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],
156
+ [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN],
157
)
158
@pytest.mark.parametrize("n_layers", [6])
159
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+3
-2
@@ -21,6 +21,7 @@ from .models import (
21
FLASHMLA_SPARSE_ATTN,
22
TRITON_ATTN,
23
deepseek_coder_v2_lite_fp8,
24
+ deepseek_r1_fp4,
25
deepseek_v3_fp8,
26
deepseek_v32_fp4,
27
gpt_oss_20b,
@@ -113,11 +114,11 @@ def test_tp2_ar_rms_fp8_fusions(
114
@multi_gpu_test(num_gpus=2)
115
@pytest.mark.parametrize(
116
"model_name, matches_fn, model_kwargs, hf_overrides",
116
- [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],
117
+ [llama3_8b_fp4, llama4_scout_fp4, deepseek_r1_fp4, deepseek_v32_fp4],
118
)
119
@pytest.mark.parametrize(
120
"attn_backend",
120
- [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],
121
+ [FLASHINFER_ATTN, FLASHINFER_MLA_ATTN, FLASHMLA_SPARSE_ATTN],
122
)
123
@pytest.mark.parametrize("n_layers", [4])
124
@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))
tests/compile/passes/test_mla_attn_quant_fusion.py
+99
-18
@@ -29,12 +29,18 @@ from vllm.config import (
29
set_current_vllm_config,
30
)
31
from vllm.forward_context import get_forward_context, set_forward_context
32
+from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
33
+ CutlassFp8BlockScaledMMKernel,
34
+)
35
from vllm.model_executor.layers.attention import MLAAttention
36
from vllm.model_executor.layers.linear import ColumnParallelLinear
37
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
38
from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config
39
from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
+ GroupShape,
41
QuantKey,
42
+ create_fp8_quant_key,
43
+ kFp8Dynamic128Sym,
44
kFp8StaticTensorSym,
45
kNvfp4Dynamic,
46
)
@@ -279,6 +285,67 @@ class TestMLAAttentionNvfp4QuantPatternModel(MLAAttentionQuantPatternModel):
285
)
286
287
288
+class TestMLAAttentionFp8GroupQuantPatternModel(MLAAttentionQuantPatternModel):
289
+ """Test model for MLA Attention + per-group FP8 (block quant) fusion."""
290
+
291
+ quant_key = kFp8Dynamic128Sym
292
+ quant_config = Fp8Config(
293
+ is_checkpoint_fp8_serialized=True,
294
+ weight_block_size=[128, 128],
295
+ )
296
+
297
+ def __init__(self, *args, **kwargs):
298
+ super().__init__(*args, **kwargs)
299
+
300
+ weight_quant_key = create_fp8_quant_key(
301
+ static=True, group_shape=GroupShape(128, 128)
302
+ )
303
+ device = kwargs.get("device", torch.device("cuda:0"))
304
+
305
+ # Subclass to set weight_block_size before process_weights_after_loading
306
+ class _BlockFP8Layer(TestFP8Layer):
307
+ def __init__(self, *a, **kw):
308
+ self.weight_block_size = [128, 128]
309
+ super().__init__(*a, **kw)
310
+
311
+ # Force CutlassFp8BlockScaledMMKernel to ensure the graph uses
312
+ # per_token_group_fp8_quant (not the deepgemm packed variant).
313
+ self.block_fp8_linear = _BlockFP8Layer(
314
+ weight_shape=(self.output_dim, self.output_dim),
315
+ activation_quant_key=self.quant_key,
316
+ weight_quant_key=weight_quant_key,
317
+ input_dtype=self.dtype,
318
+ device=device,
319
+ force_kernel=CutlassFp8BlockScaledMMKernel,
320
+ )
321
+
322
+ w = kwargs.get("w")
323
+ if w is not None:
324
+ self.block_fp8_linear.weight = w["weight"]
325
+ # Block-wise uses weight_scale_inv, not weight_scale
326
+ self.block_fp8_linear.weight_scale_inv = w["wscale"]
327
+
328
+ self.w = {
329
+ "weight": self.block_fp8_linear.weight,
330
+ "wscale": self.block_fp8_linear.weight_scale_inv,
331
+ }
332
+
333
+ def forward(
334
+ self,
335
+ q: torch.Tensor,
336
+ kv_c_normed: torch.Tensor,
337
+ k_pe: torch.Tensor,
338
+ ):
339
+ """Forward pass: MLA attention -> block FP8 linear (group quant)."""
340
+ attn_output = self.mla_attn(
341
+ q,
342
+ kv_c_normed,
343
+ k_pe,
344
+ output_shape=(q.shape[0], self.output_dim),
345
+ )
346
+ return self.block_fp8_linear(attn_output)
347
+
348
+
349
def is_nvfp4_supported():
350
return current_platform.has_device_capability(100)
351
@@ -286,6 +353,7 @@ def is_nvfp4_supported():
353
# MLA test configuration
354
MLA_DIMS: list[tuple[int, int, int, int, int]] = []
355
PATTERN_TEST_MODELS_MLA_FP8: list[tuple[str, type]] = []
356
+PATTERN_TEST_MODELS_MLA_GROUP_FP8: list[tuple[str, type]] = []
357
PATTERN_TEST_MODELS_MLA_FP4: list[tuple[str, type]] = []
358
BACKENDS_MLA_FP8: list[AttentionBackendEnum] = []
359
BACKENDS_MLA_FP4: list[AttentionBackendEnum] = []
@@ -299,6 +367,12 @@ if current_platform.is_cuda():
367
TestMLAAttentionFp8StaticQuantPatternModel,
368
)
369
]
370
+ PATTERN_TEST_MODELS_MLA_GROUP_FP8 = [
371
+ (
372
+ "deepseek-ai/DeepSeek-V3",
373
+ TestMLAAttentionFp8GroupQuantPatternModel,
374
+ )
375
+ ]
376
PATTERN_TEST_MODELS_MLA_FP4 = [
377
(
378
"deepseek-ai/DeepSeek-V2-Lite",
@@ -324,6 +398,13 @@ if current_platform.is_cuda():
398
["+quant_fp8", "-quant_fp8"],
399
)
400
)
401
+ + list(
402
+ flat_product(
403
+ BACKENDS_MLA_FP8,
404
+ PATTERN_TEST_MODELS_MLA_GROUP_FP8,
405
+ ["+quant_fp8"],
406
+ )
407
+ )
408
+ list(flat_product(BACKENDS_MLA_FP4, PATTERN_TEST_MODELS_MLA_FP4, [""])),
409
)
410
@pytest.mark.skipif(
@@ -470,12 +551,13 @@ def test_mla_attention_quant_pattern(
551
)
552
553
# Check quantization ops in the graph
554
+ is_per_group = quant_key.scale.group_shape.is_per_group()
555
quant_op = (
556
torch.ops.aten.reciprocal
557
if "-quant_fp8" in custom_ops_list
558
else QUANT_OPS[quant_key]
559
)
478
- test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)
560
+ test_backend.check_before_ops([quant_op], fully_replaced=is_per_group)
561
562
assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)
563
@@ -487,25 +569,24 @@ def test_mla_attention_quant_pattern(
569
assert len(attn_nodes_pre) == len(attn_nodes_post), (
570
"Should have same number of MLA attention nodes before and after fusion"
571
)
490
- assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (
491
- "MLA attention should not have output_scale before fusion"
492
- )
493
- assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (
494
- "MLA attention should have output_scale after fusion"
495
- )
572
497
- assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (
498
- "MLA attention should not have output_block_scale before fusion"
499
- )
573
+ # Before fusion: neither scale should be set
574
+ assert attn_nodes_pre[0].kwargs.get("output_scale") is None
575
+ assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None
576
501
- if quant_key.dtype == FP8_DTYPE:
502
- assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (
503
- "MLA attention should not have output_block_scale after FP8 fusion"
504
- )
505
- elif quant_key.dtype == FP4_DTYPE:
506
- assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (
507
- "MLA attention should have output_block_scale after FP4 fusion"
508
- )
577
+ # After fusion: derive expected scale presence from quant_key properties.
578
+ # - output_scale: present for static quant or non-FP8 (NVFP4 carries input_scale)
579
+ # - output_block_scale: present when quant uses per-group/block scaling
580
+ has_output_scale = attn_nodes_post[0].kwargs.get("output_scale") is not None
581
+ has_block_scale = attn_nodes_post[0].kwargs.get("output_block_scale") is not None
582
+
583
+ expects_output_scale = quant_key.scale.static or quant_key.dtype != FP8_DTYPE
584
+ assert has_output_scale == expects_output_scale, (
585
+ f"output_scale: expected present={expects_output_scale}, got {has_output_scale}"
586
+ )
587
+ assert has_block_scale == is_per_group, (
588
+ f"output_block_scale: expected present={is_per_group}, got {has_block_scale}"
589
+ )
590
591
# Check numerical correctness
592
torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)
vllm/compilation/passes/fusion/mla_attn_quant_fusion.py
+259
-18
@@ -6,15 +6,18 @@ from collections.abc import Callable
6
import torch
7
from torch._higher_order_ops.auto_functionalize import auto_functionalized
8
9
-from vllm._custom_ops import create_fp4_output_tensors
9
from vllm.config import VllmConfig, get_layers_from_vllm_config
10
from vllm.logger import init_logger
11
from vllm.model_executor.layers.attention.mla_attention import MLAAttention
12
from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ QuantKey,
14
+ kFp8Dynamic64Sym,
15
+ kFp8Dynamic128Sym,
16
kFp8StaticTensorSym,
17
kNvfp4Dynamic,
18
)
19
from vllm.platforms import current_platform
20
+from vllm.utils.math_utils import round_up
21
from vllm.utils.torch_utils import _USE_LAYERNAME, _encode_layer_name
22
23
from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement
@@ -203,6 +206,8 @@ class MLAAttnNvfp4QuantPattern(
206
kv_c_normed,
207
k_pe,
208
output_attn,
209
+ output_quant,
210
+ output_scale,
211
input_scale,
212
kv_cache_dummy_dep,
213
layer_name,
@@ -218,9 +223,6 @@ class MLAAttnNvfp4QuantPattern(
223
output_block_scale=None,
224
kv_cache_dummy_dep=kv_cache_dummy_dep,
225
)
221
- output_quant, output_scale = create_fp4_output_tensors(
222
- at1[1].shape[0], at1[1].shape[1], at1[1].device, True
223
- )
226
at2 = auto_functionalized(
227
self._QUANT_OP,
228
input=at1[1],
@@ -235,7 +237,14 @@ class MLAAttnNvfp4QuantPattern(
237
return _pattern_with_ln
238
239
def _pattern(
238
- q, kv_c_normed, k_pe, output_attn, input_scale, kv_cache_dummy_dep
240
+ q,
241
+ kv_c_normed,
242
+ k_pe,
243
+ output_attn,
244
+ output_quant,
245
+ output_scale,
246
+ input_scale,
247
+ kv_cache_dummy_dep,
248
):
249
at1 = auto_functionalized(
250
MLA_ATTN_OP,
@@ -248,11 +257,6 @@ class MLAAttnNvfp4QuantPattern(
257
output_block_scale=None,
258
kv_cache_dummy_dep=kv_cache_dummy_dep,
259
)
251
- # Replicate what scaled_fp4_quant() does: allocate output
252
- # tensors inline then call the .out variant.
253
- output_quant, output_scale = create_fp4_output_tensors(
254
- at1[1].shape[0], at1[1].shape[1], at1[1].device, True
255
- )
260
at2 = auto_functionalized(
261
self._QUANT_OP,
262
input=at1[1],
@@ -279,6 +283,8 @@ class MLAAttnNvfp4QuantPattern(
283
kv_c_normed,
284
k_pe,
285
output_attn,
286
+ _output_quant,
287
+ output_scale,
288
input_scale,
289
kv_cache_dummy_dep,
290
layer_name,
@@ -289,9 +295,6 @@ class MLAAttnNvfp4QuantPattern(
295
dtype=FP4_DTYPE,
296
device=q.device,
297
)
292
- output_scale = create_fp4_output_tensors(
293
- q.shape[0], self._output_dim, q.device, True
294
- )[1]
298
output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
299
at2 = auto_functionalized(
300
MLA_ATTN_OP,
@@ -309,7 +312,14 @@ class MLAAttnNvfp4QuantPattern(
312
return _replacement_with_ln
313
314
def _replacement(
312
- q, kv_c_normed, k_pe, output_attn, input_scale, kv_cache_dummy_dep
315
+ q,
316
+ kv_c_normed,
317
+ k_pe,
318
+ output_attn,
319
+ _output_quant,
320
+ output_scale,
321
+ input_scale,
322
+ kv_cache_dummy_dep,
323
):
324
# MLA output in quant_dtype (FP4 packed as uint8)
325
output_attn = torch.empty(
@@ -317,10 +327,6 @@ class MLAAttnNvfp4QuantPattern(
327
dtype=FP4_DTYPE,
328
device=q.device,
329
)
320
- # attention output block scale
321
- output_scale = create_fp4_output_tensors(
322
- q.shape[0], self._output_dim, q.device, True
323
- )[1]
330
output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)
331
at2 = auto_functionalized(
332
MLA_ATTN_OP,
@@ -343,6 +349,8 @@ class MLAAttnNvfp4QuantPattern(
349
self.empty(5, self._kv_lora_rank, dtype=self._dtype),
350
self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),
351
self.empty(5, self._output_dim, dtype=self._dtype),
352
+ self.empty(5, self._output_dim // 2, dtype=FP4_DTYPE),
353
+ self.empty_i32(128, round_up(self._output_dim // 16, 4)),
354
self.empty_fp32(1, 1),
355
self.empty(0, dtype=self._dtype),
356
]
@@ -351,6 +359,218 @@ class MLAAttnNvfp4QuantPattern(
359
return inputs
360
361
362
+class MLAAttnFp8GroupQuantPattern(
363
+ VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
364
+):
365
+ """
366
+ Fusion for MLA Attention+Fp8GroupQuant (per-group dynamic FP8).
367
+
368
+ Matches the pattern: MLA attention -> per_token_group_fp8_quant, and
369
+ replaces it with MLA attention(output_block_scale=group_scale_buffer).
370
+ Used by models with block FP8 quantization (e.g. DeepSeek V3).
371
+ """
372
+
373
+ def __init__(
374
+ self,
375
+ layer: MLAAttention,
376
+ dtype: torch.dtype,
377
+ quant_key: QuantKey,
378
+ has_col_major_scales: bool,
379
+ is_e8m0: bool,
380
+ is_tma_aligned: bool,
381
+ ) -> None:
382
+ self._layer_name = layer.layer_name
383
+ self._num_heads = layer.num_heads
384
+ self._v_head_dim = layer.v_head_dim
385
+ self._kv_lora_rank = layer.kv_lora_rank
386
+ self._qk_rope_head_dim = layer.qk_rope_head_dim
387
+ self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim
388
+ self._output_dim = layer.num_heads * layer.v_head_dim
389
+ self._dtype = dtype
390
+ self._layer = layer
391
+ self._group_size = quant_key.scale.group_shape[1]
392
+ self._has_col_major_scales = has_col_major_scales
393
+ self._is_e8m0 = is_e8m0
394
+ self._is_tma_aligned = is_tma_aligned
395
+
396
+ self._quant_matcher = MatcherQuantFP8(
397
+ quant_key,
398
+ has_col_major_scales=has_col_major_scales,
399
+ is_e8m0=is_e8m0,
400
+ is_tma_aligned=is_tma_aligned,
401
+ )
402
+
403
+ @property
404
+ def pattern(
405
+ self,
406
+ ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
407
+ _ln = _encode_layer_name(self._layer_name)
408
+
409
+ if _USE_LAYERNAME:
410
+
411
+ def _pattern_with_ln( # type: ignore[misc]
412
+ q,
413
+ kv_c_normed,
414
+ k_pe,
415
+ output_attn,
416
+ kv_cache_dummy_dep,
417
+ scale,
418
+ layer_name,
419
+ ):
420
+ at1 = auto_functionalized(
421
+ MLA_ATTN_OP,
422
+ q=q,
423
+ kv_c_normed=kv_c_normed,
424
+ k_pe=k_pe,
425
+ output=output_attn,
426
+ layer_name=layer_name,
427
+ output_scale=None,
428
+ output_block_scale=None,
429
+ kv_cache_dummy_dep=kv_cache_dummy_dep,
430
+ )
431
+ attn_out = at1[1]
432
+ result = torch.empty(
433
+ attn_out.shape, device=attn_out.device, dtype=FP8_DTYPE
434
+ )
435
+ finfo = torch.finfo(FP8_DTYPE)
436
+ _, result, scale = auto_functionalized(
437
+ self._quant_matcher.QUANT_OP,
438
+ input=attn_out,
439
+ output_q=result,
440
+ output_s=scale,
441
+ group_size=self._group_size,
442
+ eps=1e-10,
443
+ fp8_min=finfo.min,
444
+ fp8_max=finfo.max,
445
+ scale_ue8m0=self._is_e8m0,
446
+ dummy_is_scale_transposed=self._has_col_major_scales,
447
+ dummy_is_tma_aligned=self._is_tma_aligned,
448
+ )
449
+ return result, scale
450
+
451
+ return _pattern_with_ln
452
+
453
+ def _pattern(
454
+ q: torch.Tensor,
455
+ kv_c_normed: torch.Tensor,
456
+ k_pe: torch.Tensor,
457
+ output_attn: torch.Tensor,
458
+ kv_cache_dummy_dep: torch.Tensor,
459
+ scale: torch.Tensor,
460
+ ) -> tuple[torch.Tensor, torch.Tensor]:
461
+ at1 = auto_functionalized(
462
+ MLA_ATTN_OP,
463
+ q=q,
464
+ kv_c_normed=kv_c_normed,
465
+ k_pe=k_pe,
466
+ output=output_attn,
467
+ layer_name=_ln,
468
+ output_scale=None,
469
+ output_block_scale=None,
470
+ kv_cache_dummy_dep=kv_cache_dummy_dep,
471
+ )
472
+ attn_out = at1[1]
473
+ result = torch.empty(
474
+ attn_out.shape, device=attn_out.device, dtype=FP8_DTYPE
475
+ )
476
+ finfo = torch.finfo(FP8_DTYPE)
477
+ _, result, scale = auto_functionalized(
478
+ self._quant_matcher.QUANT_OP,
479
+ input=attn_out,
480
+ output_q=result,
481
+ output_s=scale,
482
+ group_size=self._group_size,
483
+ eps=1e-10,
484
+ fp8_min=finfo.min,
485
+ fp8_max=finfo.max,
486
+ scale_ue8m0=self._is_e8m0,
487
+ dummy_is_scale_transposed=self._has_col_major_scales,
488
+ dummy_is_tma_aligned=self._is_tma_aligned,
489
+ )
490
+ return result, scale
491
+
492
+ return _pattern
493
+
494
+ @property
495
+ def replacement(
496
+ self,
497
+ ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:
498
+ _ln = _encode_layer_name(self._layer_name)
499
+
500
+ if _USE_LAYERNAME:
501
+
502
+ def _replacement_with_ln( # type: ignore[misc]
503
+ q,
504
+ kv_c_normed,
505
+ k_pe,
506
+ output_attn,
507
+ kv_cache_dummy_dep,
508
+ scale,
509
+ layer_name,
510
+ ):
511
+ output_attn = torch.empty(
512
+ [q.shape[0], self._output_dim],
513
+ dtype=FP8_DTYPE,
514
+ device=q.device,
515
+ )
516
+ at1 = auto_functionalized(
517
+ MLA_ATTN_OP,
518
+ q=q,
519
+ kv_c_normed=kv_c_normed,
520
+ k_pe=k_pe,
521
+ output=output_attn,
522
+ layer_name=layer_name,
523
+ output_scale=None,
524
+ output_block_scale=scale,
525
+ kv_cache_dummy_dep=kv_cache_dummy_dep,
526
+ quant_group_size=self._group_size,
527
+ quant_scale_ue8m0=self._is_e8m0,
528
+ quant_col_major=self._has_col_major_scales,
529
+ quant_tma_aligned=self._is_tma_aligned,
530
+ )
531
+ return at1[1], at1[2]
532
+
533
+ return _replacement_with_ln
534
+
535
+ def _replacement(q, kv_c_normed, k_pe, output_attn, kv_cache_dummy_dep, scale):
536
+ output_attn = torch.empty(
537
+ [q.shape[0], self._output_dim],
538
+ dtype=FP8_DTYPE,
539
+ device=q.device,
540
+ )
541
+ at1 = auto_functionalized(
542
+ MLA_ATTN_OP,
543
+ q=q,
544
+ kv_c_normed=kv_c_normed,
545
+ k_pe=k_pe,
546
+ output=output_attn,
547
+ layer_name=_ln,
548
+ output_scale=None,
549
+ output_block_scale=scale,
550
+ kv_cache_dummy_dep=kv_cache_dummy_dep,
551
+ quant_group_size=self._group_size,
552
+ quant_scale_ue8m0=self._is_e8m0,
553
+ quant_col_major=self._has_col_major_scales,
554
+ quant_tma_aligned=self._is_tma_aligned,
555
+ )
556
+ return at1[1], at1[2]
557
+
558
+ return _replacement
559
+
560
+ def get_inputs(self) -> list[torch.Tensor]:
561
+ inputs: list = [
562
+ self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),
563
+ self.empty(5, self._kv_lora_rank, dtype=self._dtype),
564
+ self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),
565
+ self.empty(5, self._output_dim, dtype=self._dtype),
566
+ self.empty(0, dtype=self._dtype),
567
+ self._quant_matcher.empty_f32(1, 1),
568
+ ]
569
+ if _USE_LAYERNAME:
570
+ inputs.append(_encode_layer_name(self._layer_name))
571
+ return inputs
572
+
573
+
574
class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):
575
"""
576
This pass fuses post-attention quantization onto MLA attention if supported.
@@ -389,4 +609,25 @@ class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):
609
if _USE_LAYERNAME:
610
break
611
612
+ # Per-group FP8 (block quant) — register all flag combinations.
613
+ if current_platform.is_cuda():
614
+ for quant_key in [kFp8Dynamic128Sym, kFp8Dynamic64Sym]:
615
+ for col_major in [True, False]:
616
+ for is_e8m0 in [True, False]:
617
+ for tma_aligned in [False, True]:
618
+ for layer in layers:
619
+ if layer.impl.fused_output_quant_supported(quant_key):
620
+ self.register(
621
+ MLAAttnFp8GroupQuantPattern(
622
+ layer,
623
+ dtype,
624
+ quant_key,
625
+ col_major,
626
+ is_e8m0,
627
+ tma_aligned,
628
+ )
629
+ )
630
+ if _USE_LAYERNAME:
631
+ break
632
+
633
self.dump_patterns(config, self.pm_pass)
vllm/model_executor/layers/attention/mla_attention.py
+96
-10
@@ -234,7 +234,12 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
234
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
235
from vllm.model_executor.layers.quantization.utils.quant_utils import (
236
GroupShape,
237
+ QuantKey,
238
get_and_maybe_dequant_weights,
239
+ kFp8Dynamic64Sym,
240
+ kFp8Dynamic128Sym,
241
+ kFp8StaticTensorSym,
242
+ kNvfp4Dynamic,
243
)
244
from vllm.platforms import current_platform
245
from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory
@@ -276,6 +281,44 @@ from vllm.v1.kv_cache_interface import (
281
282
logger = init_logger(__name__)
283
284
+_FP8_DTYPE = current_platform.fp8_dtype()
285
+
286
+
287
+def _detect_output_quant_key(
288
+ output: torch.Tensor,
289
+ output_scale: torch.Tensor | None,
290
+ output_block_scale: torch.Tensor | None,
291
+ output_dim: int,
292
+) -> QuantKey | None:
293
+ """Detect the output quantization key from fusion pass parameters.
294
+
295
+ Returns the appropriate QuantKey, or None if no quantization is needed.
296
+ Detection is based on output dtype and which scale tensors are present.
297
+ """
298
+ if output_scale is None and output_block_scale is None:
299
+ return None
300
+ if output_block_scale is not None:
301
+ if output.dtype == _FP8_DTYPE:
302
+ # Per-group FP8 uses block scales only, not a separate output_scale
303
+ assert output_scale is None
304
+ # Infer group size from scale shape
305
+ num_groups = output_block_scale.shape[-1]
306
+ group_size = output_dim // num_groups
307
+ if group_size == 128:
308
+ return kFp8Dynamic128Sym
309
+ elif group_size == 64:
310
+ return kFp8Dynamic64Sym
311
+ else:
312
+ raise ValueError(
313
+ f"Unsupported group FP8 group_size={group_size} "
314
+ f"(output_dim={output_dim}, num_groups={num_groups}). "
315
+ f"Only group_size 128 and 64 are supported."
316
+ )
317
+ # output_scale None implies MXFP4, not supported
318
+ assert output_scale is not None
319
+ return kNvfp4Dynamic
320
+ return kFp8StaticTensorSym
321
+
322
323
class MLAAttention(nn.Module, AttentionLayerBase):
324
"""Multi-Head Latent Attention layer.
@@ -549,9 +592,17 @@ class MLAAttention(nn.Module, AttentionLayerBase):
592
output: torch.Tensor,
593
output_scale: torch.Tensor | None = None,
594
output_block_scale: torch.Tensor | None = None,
595
+ quant_group_size: int | None = None,
596
+ quant_scale_ue8m0: bool | None = None,
597
+ quant_col_major: bool | None = None,
598
+ quant_tma_aligned: bool | None = None,
599
) -> torch.Tensor:
553
- use_quant = output_scale is not None or output_block_scale is not None
554
- if use_quant:
600
+ assert output is not None, "Output tensor must be provided."
601
+
602
+ quant_key = _detect_output_quant_key(
603
+ output, output_scale, output_block_scale, self.num_heads * self.v_head_dim
604
+ )
605
+ if quant_key is not None:
606
# The fusion pass has allocated output with quantized dtype
607
# (FP8 or uint8 for FP4). We can't write into it directly,
608
# so we swap in a temp buffer for computation, then quantize
@@ -582,7 +633,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
633
# The zero fill is required when used with DP + EP
634
# to ensure all ranks within a DP group compute the
635
# same expert outputs.
585
- if use_quant:
636
+ if quant_key is not None:
637
return quant_output.fill_(0)
638
return output.fill_(0)
639
@@ -724,18 +775,41 @@ class MLAAttention(nn.Module, AttentionLayerBase):
775
# v_up projection
776
self._v_up_proj(attn_out, out=mqa_output_slice)
777
727
- if use_quant:
778
+ if quant_key is not None:
779
# Quantize the BF16 computation result into the quantized output
780
actual = output[:num_actual_toks]
730
- if output_block_scale is not None:
781
+ if quant_key == kNvfp4Dynamic:
782
# NVFP4: two FP4 values packed into one uint8
783
+ assert output_block_scale is not None
784
fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale)
785
quant_output[:num_actual_toks].copy_(fp4_data)
734
- output_block_scale.copy_(fp4_scales)
735
- else:
786
+ output_block_scale[: fp4_scales.shape[0]].copy_(fp4_scales)
787
+ elif quant_key in (kFp8Dynamic128Sym, kFp8Dynamic64Sym):
788
+ # Per-group FP8
789
+ assert output_block_scale is not None
790
+ assert quant_group_size is not None, (
791
+ "Group FP8 output quant requested but "
792
+ "quant_group_size not passed through custom op"
793
+ )
794
+ finfo = torch.finfo(_FP8_DTYPE)
795
+ torch.ops._C.per_token_group_fp8_quant(
796
+ actual,
797
+ quant_output[:num_actual_toks],
798
+ output_block_scale[:num_actual_toks],
799
+ quant_group_size,
800
+ 1e-10, # eps
801
+ finfo.min,
802
+ finfo.max,
803
+ quant_scale_ue8m0,
804
+ quant_col_major,
805
+ quant_tma_aligned,
806
+ )
807
+ elif quant_key == kFp8StaticTensorSym:
808
# Static FP8 quantization
809
fp8_data, _ = self._quant_fp8_op(actual, output_scale)
810
quant_output[:num_actual_toks].copy_(fp8_data)
811
+ else:
812
+ raise ValueError(f"Unsupported quant_key: {quant_key}")
813
return quant_output
814
815
return output_padded
@@ -980,6 +1054,10 @@ def unified_mla_attention_with_output(
1054
output_scale: torch.Tensor | None = None,
1055
output_block_scale: torch.Tensor | None = None,
1056
kv_cache_dummy_dep: torch.Tensor | None = None,
1057
+ quant_group_size: int | None = None,
1058
+ quant_scale_ue8m0: bool | None = None,
1059
+ quant_col_major: bool | None = None,
1060
+ quant_tma_aligned: bool | None = None,
1061
) -> None:
1062
# kv_cache_dummy_dep is not used but accepting it creates a data dependency
1063
# that ensures torch.compile preserves ordering between KV cache update and
@@ -996,6 +1074,10 @@ def unified_mla_attention_with_output(
1074
output=output,
1075
output_scale=output_scale,
1076
output_block_scale=output_block_scale,
1077
+ quant_group_size=quant_group_size,
1078
+ quant_scale_ue8m0=quant_scale_ue8m0,
1079
+ quant_col_major=quant_col_major,
1080
+ quant_tma_aligned=quant_tma_aligned,
1081
)
1082
1083
@@ -1008,6 +1090,10 @@ def unified_mla_attention_with_output_fake(
1090
output_scale: torch.Tensor | None = None,
1091
output_block_scale: torch.Tensor | None = None,
1092
kv_cache_dummy_dep: torch.Tensor | None = None,
1093
+ quant_group_size: int | None = None,
1094
+ quant_scale_ue8m0: bool | None = None,
1095
+ quant_col_major: bool | None = None,
1096
+ quant_tma_aligned: bool | None = None,
1097
) -> None:
1098
return
1099
@@ -2078,13 +2164,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2164
"""
2165
2166
def fused_output_quant_supported(self, quant_key):
2081
- from vllm.model_executor.layers.quantization.utils.quant_utils import (
2167
+ return quant_key in (
2168
kFp8StaticTensorSym,
2169
kNvfp4Dynamic,
2170
+ kFp8Dynamic128Sym,
2171
+ kFp8Dynamic64Sym,
2172
)
2173
2086
- return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)
2087
-
2174
def __init__(
2175
self,
2176
num_heads: int,
vllm/v1/attention/backend.py
+14
-2
@@ -11,6 +11,8 @@ import torch
11
from typing_extensions import deprecated
12
13
from vllm.model_executor.layers.quantization.utils.quant_utils import (
14
+ kFp8Dynamic64Sym,
15
+ kFp8Dynamic128Sym,
16
kFp8StaticTensorSym,
17
kNvfp4Dynamic,
18
)
@@ -880,7 +882,12 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
882
Since MLA quantization is done manually in forward_impl (common code),
883
all MLA backends support it by default.
884
"""
883
- return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)
885
+ return quant_key in (
886
+ kFp8StaticTensorSym,
887
+ kNvfp4Dynamic,
888
+ kFp8Dynamic128Sym,
889
+ kFp8Dynamic64Sym,
890
+ )
891
892
def do_kv_cache_update(
893
self,
@@ -918,7 +925,12 @@ class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
925
Since MLA quantization is done manually in forward_impl (common code),
926
all MLA backends support it by default.
927
"""
921
- return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)
928
+ return quant_key in (
929
+ kFp8StaticTensorSym,
930
+ kNvfp4Dynamic,
931
+ kFp8Dynamic128Sym,
932
+ kFp8Dynamic64Sym,
933
+ )
934
935
@abstractmethod
936
def __init__(