@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1f5ec2889c

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

[mla] Support fused FP8/NVFP4 output quantization in MLA attention (#35792) (#36205)

Signed-off-by: Carl You <4531192+carlyou@users.noreply.github.com> Signed-off-by: Carl Y <4531192+carlyou@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Carl Y committed Apr 2, 2026 at 18:16 UTC 1f5ec2889c4176593c2fedae3b33b6244c996d29

12 files changed +928 -17

.buildkite/test_areas/compile.yaml

@@ -72,6 +72,7 @@ steps:

72 - vllm/v1/attention/backends/flashinfer.py

73 - vllm/compilation/ # TODO(luka) limit to vllm/compilation/passes

74 - tests/compile/passes/test_fusion_attn.py

75 + - tests/compile/passes/test_mla_attn_quant_fusion.py

76 - tests/compile/passes/test_silu_mul_quant_fusion.py

77 - tests/compile/passes/distributed/test_fusion_all_reduce.py

78 - tests/compile/fullgraph/test_full_graph.py

@@ -79,6 +80,7 @@ steps:

80 # b200 runners are limited, so we limit the tests to the minimum set only supported on Blackwell

81 - nvidia-smi

82 - pytest -v -s tests/compile/passes/test_fusion_attn.py -k FLASHINFER

83 + - pytest -v -s tests/compile/passes/test_mla_attn_quant_fusion.py

84 - pytest -v -s tests/compile/passes/test_silu_mul_quant_fusion.py

85 # this runner has 2 GPUs available even though num_devices=2 is not set

86 - pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py

docs/design/fusions.md

+18 -2

@@ -22,6 +22,7 @@ or just on the low or high end.

22 | ------------------------------------------------------------------------------ | ---------------------------- | ---------------------------------------------- | ------------------------------ | ------------------ | --------- | ------------ |

23 | [AllReduce + RMSNorm](#allreduce--rmsnorm-fuse_allreduce_rms) | `fuse_allreduce_rms` | All-reduce → RMSNorm (+residual_add) (→ quant) | O2 (Hopper/Blackwell + TP > 1) | 5-20% | No | Low |

26 | [RoPE + KV-Cache Update](#rope--kv-cache-update-fuse_rope_kvcache) | `fuse_rope_kvcache` | Rotary embedding → KV cache write | O2 (ROCm/AITER only) | 2-4% | No | Low |

27 | [QK Norm + RoPE](#qk-norm--rope-enable_qk_norm_rope_fusion) | `enable_qk_norm_rope_fusion` | Q/K RMSNorm → rotary embedding | Off by default | 2-3% | No | Low |

@@ -40,6 +41,7 @@ The table below lists the quantization schemes supported by each fusion on each

41 | ---------------------------- | ---------------------------------------- | ---------------------------------------- | ---------------------------------------- | ------------- | ---------------------------------------- |

42 | `fuse_allreduce_rms` | FP16/BF16, FP8 static, NVFP4 | FP16/BF16, FP8 static | — | — | — |

45 | `fuse_rope_kvcache` | — | — | — | — | FP16/BF16 |

@@ -129,7 +131,8 @@ on SM90/SM100) and configurable via `PassConfig.fi_allreduce_fusion_max_size_mb`

131 explicitly. It requires the full model graph to be visible (Inductor partition or `splitting_ops=[]`).

132

133 **What it fuses.** Fuses the attention output quantization directly after the attention computation,

132 -eliminating a full-precision memory round-trip of the attention output. Patterns covered:

134 +eliminating a full-precision memory round-trip of the attention output. This fusion supports both

135 +standard `Attention` and `MLAAttention` (used by DeepSeek-V2/V3/R1 models). Patterns covered:

136

137 `Attention → FP8 static quant`:

138

@@ -142,11 +145,24 @@ eliminating a full-precision memory round-trip of the attention output. Patterns

145

146 - `FLASHINFER`: CUDA sm100+ with FlashInfer installed

147

148 +`MLAAttention → FP8 static quant` / `MLAAttention → NVFP4 dynamic quant`:

149 +

150 +The MLA fusion operates at the graph level on the `unified_mla_attention_with_output` op and works

151 +with all MLA decode and prefill backend combinations. Unlike standard `Attention` backends (where

152 +the kernel writes FP8 output directly), no MLA prefill or decode backend currently supports direct

153 +FP8/FP4 output. The fusion writes to an intermediate buffer and quantizes in a separate step, so

154 +there is no memory round-trip elimination yet.

155 +

156 +!!! info

157 + The MLA attention fusion is not expected to yield a measurable speedup yet.

158 + This will improve once MLA prefill/decode kernels support direct FP8/FP4 output.

159 +

160 Other attention backends do not support fused output quantization yet.

161

162 **Code locations.**

163

149 -- Pass: [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py)

164 +- Pass (Attention): [`vllm/compilation/passes/fusion/attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/attn_quant_fusion.py)

165 +- Pass (MLAAttention): [`vllm/compilation/passes/fusion/mla_attn_quant_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/mla_attn_quant_fusion.py)

166 - Attention backends: [`vllm/v1/attention/backends/`](https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/)

167

168 ### RoPE + KV-Cache Update (`fuse_rope_kvcache`)

tests/compile/fusions_e2e/conftest.py

+9 -3

@@ -84,10 +84,14 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):

84 rocm_aiter_ops.refresh_env_variables()

86 # Filter here to reduce code duplication

87 + backend_name = attn_backend.backend.name.lower()

88 requires_mla = "deepseek" in model_name.lower()

88 - is_mla = "mla" in attn_backend.backend.name.lower()

89 + is_mla = "mla" in backend_name

90 + # DeepSeek V3.2 uses sparse MLA

91 + requires_sparse = "v3.2" in model_name.lower()

92 + is_sparse = "sparse" in backend_name

90 - if requires_mla != is_mla:

94 + if requires_mla != is_mla or requires_sparse != is_sparse:

95 pytest.skip(

96 f"Incompatible model '{model_name}' and "

97 f"attention backend '{attn_backend.backend.name}'"

@@ -231,7 +235,9 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):

235 )

236

237 elif match_name == "attn_quant_fusion":

234 - actual_match = match_table.get(match_name, 0)

238 + actual_match = match_table.get(

239 + "attn_quant_fusion", 0

240 + ) + match_table.get("mla_attn_quant_fusion", 0)

241 assert actual_match == expected_matches * n_expected, (

242 f"Could not find {expected_matches * n_expected} "

243 f"{match_name} (found {actual_match})."

tests/compile/fusions_e2e/models.py

+32 -1

@@ -58,6 +58,15 @@ TRITON_MLA_ATTN = pytest.param(

58 id="TRITON_MLA",

59 )

61 +FLASHMLA_SPARSE_ATTN = pytest.param(

62 + AttentionBackendCase(backend=AttentionBackendEnum.FLASHMLA_SPARSE),

63 + id="FLASHMLA_SPARSE",

64 + marks=pytest.mark.skipif(

65 + not is_blackwell(),

66 + reason="FlashMLA Sparse requires Blackwell",

67 + ),

68 +)

69 +

70 # Models

71 llama3_8b = ModelFusionInfo(

72 model_name="meta-llama/Llama-3.1-8B-Instruct",

@@ -141,6 +150,18 @@ qwen3_a3b_fp8 = ModelFusionInfo(

150 ),

151 )

152

153 +deepseek_coder_v2_lite_fp8 = ModelFusionInfo(

154 + model_name="RedHatAI/DeepSeek-Coder-V2-Lite-Instruct-FP8",

155 + matches=lambda n_layers: Matches(

156 + # first_k_dense_replace=1; MoE hides most rms+quant sites

157 + rms_quant_fusion=1,

158 + act_quant_fusion=min(1, n_layers), # dense layers only

159 + # MLA attn + static FP8 quant

160 + attn_quant_fusion=n_layers,

161 + ar_rms_fusion=n_layers * 2 + 1,

162 + ),

163 +)

164 +

165 deepseek_v3_fp8 = ModelFusionInfo(

166 model_name="deepseek-ai/DeepSeek-V3",

167 matches=lambda n_layers: Matches(

@@ -152,7 +173,7 @@ deepseek_v3_fp8 = ModelFusionInfo(

173 rms_quant_fusion=n_layers * 2 + min(3, n_layers), # add for 3 dense layers

174 # silu+block quant

175 act_quant_fusion=min(3, n_layers), # dense layers only

155 - # MLA attn + quant not supported yet:

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 ar_rms_fusion=n_layers * 2 + 1,

@@ -162,6 +183,16 @@ deepseek_v3_fp8 = ModelFusionInfo(

183 ),

184 )

185

186 +deepseek_v32_fp4 = ModelFusionInfo(

187 + model_name="nvidia/DeepSeek-V3.2-NVFP4",

188 + matches=lambda n_layers: Matches(

189 + rms_quant_fusion=0,

190 + act_quant_fusion=0,

191 + attn_quant_fusion=n_layers,

192 + ar_rms_fusion=n_layers * 2 + 1,

193 + ),

194 +)

195 +

196 gpt_oss_20b = ModelFusionInfo(

197 model_name="openai/gpt-oss-20b",

198 matches=lambda n_layers: Matches(

tests/compile/fusions_e2e/test_tp1_quant.py

+9 -2

@@ -18,11 +18,14 @@ from .common import (

18 from .models import (

19 FLASHINFER_ATTN,

20 FLASHINFER_MLA_ATTN,

21 + FLASHMLA_SPARSE_ATTN,

22 ROCM_AITER_UNIFIED_ATTN,

23 ROCM_ATTN,

24 TRITON_ATTN,

25 TRITON_MLA_ATTN,

26 + deepseek_coder_v2_lite_fp8,

27 deepseek_v3_fp8,

28 + deepseek_v32_fp4,

29 llama3_8b_fp4,

30 llama3_8b_fp8,

31 llama4_scout_fp4,

@@ -37,6 +40,7 @@ from .models import (

40 (*llama3_8b_fp8, False),

41 (*qwen3_a3b_fp8, False),

42 (*qwen3_a3b_fp8, True),

43 + (*deepseek_coder_v2_lite_fp8, False),

44 (*deepseek_v3_fp8, False),

45 (*deepseek_v3_fp8, True),

46 pytest.param(

@@ -144,9 +148,12 @@ def test_tp1_fp8_fusions(

148

149 @pytest.mark.parametrize(

150 "model_name, matches_fn, model_kwargs, hf_overrides",

147 - [llama3_8b_fp4, llama4_scout_fp4],

151 + [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],

152 +)

153 +@pytest.mark.parametrize(

154 + "attn_backend",

155 + [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],

156 )

149 -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])

157 @pytest.mark.parametrize("n_layers", [6])

158 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))

159 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)

tests/compile/fusions_e2e/test_tp2_ar_rms.py

+15 -3

@@ -18,8 +18,11 @@ from .common import (

18 from .models import (

19 FLASHINFER_ATTN,

20 FLASHINFER_MLA_ATTN,

21 + FLASHMLA_SPARSE_ATTN,

22 TRITON_ATTN,

23 + deepseek_coder_v2_lite_fp8,

24 deepseek_v3_fp8,

25 + deepseek_v32_fp4,

26 gpt_oss_20b,

27 llama3_8b,

28 llama3_8b_fp4,

@@ -37,7 +40,13 @@ pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only tes

40 @pytest.mark.parametrize(

41 "model_name, matches_fn, model_kwargs, hf_overrides",

42 # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported

40 - [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8, deepseek_v3_fp8],

43 + [

44 + llama3_8b_fp8,

45 + llama4_scout_fp8,

46 + qwen3_a3b_fp8,

47 + deepseek_coder_v2_lite_fp8,

48 + deepseek_v3_fp8,

49 + ],

50 )

51 @pytest.mark.parametrize(

52 "attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]

@@ -104,9 +113,12 @@ def test_tp2_ar_rms_fp8_fusions(

113 @multi_gpu_test(num_gpus=2)

114 @pytest.mark.parametrize(

115 "model_name, matches_fn, model_kwargs, hf_overrides",

107 - [llama3_8b_fp4, llama4_scout_fp4],

116 + [llama3_8b_fp4, llama4_scout_fp4, deepseek_v32_fp4],

117 +)

118 +@pytest.mark.parametrize(

119 + "attn_backend",

120 + [FLASHINFER_ATTN, FLASHMLA_SPARSE_ATTN],

121 )

109 -@pytest.mark.parametrize("attn_backend", [FLASHINFER_ATTN])

122 @pytest.mark.parametrize("n_layers", [4])

123 @pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm"))

124 @pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)

tests/compile/passes/test_mla_attn_quant_fusion.py new

+508

@@ -0,0 +1,508 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import copy

4 +

5 +import pytest

6 +import torch._dynamo

7 +

8 +from tests.compile.backend import LazyInitPass, TestBackend

9 +from tests.utils import TestFP8Layer, flat_product

10 +from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata

11 +from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant

12 +from vllm.compilation.passes.fusion.matcher_utils import QUANT_OPS

13 +from vllm.compilation.passes.fusion.mla_attn_quant_fusion import (

14 + MLA_ATTN_OP,

15 + MLAAttnQuantFusionPass,

16 +)

17 +from vllm.compilation.passes.fx_utils import find_op_nodes

18 +from vllm.compilation.passes.utility.noop_elimination import NoOpEliminationPass

19 +from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass

20 +from vllm.config import (

21 + AttentionConfig,

22 + CacheConfig,

23 + CompilationConfig,

24 + CompilationMode,

25 + ModelConfig,

26 + PassConfig,

27 + SchedulerConfig,

28 + VllmConfig,

29 + set_current_vllm_config,

30 +)

31 +from vllm.forward_context import get_forward_context, set_forward_context

32 +from vllm.model_executor.layers.attention import MLAAttention

33 +from vllm.model_executor.layers.linear import ColumnParallelLinear

34 +from vllm.model_executor.layers.quantization.fp8 import Fp8Config

35 +from vllm.model_executor.layers.quantization.modelopt import ModelOptNvFp4Config

36 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

37 + QuantKey,

38 + kFp8StaticTensorSym,

39 + kNvfp4Dynamic,

40 +)

41 +from vllm.platforms import current_platform

42 +from vllm.v1.attention.backend import AttentionMetadata

43 +from vllm.v1.attention.backends.registry import AttentionBackendEnum

44 +from vllm.v1.kv_cache_interface import MLAAttentionSpec

45 +

46 +FP8_DTYPE = current_platform.fp8_dtype()

47 +FP4_DTYPE = torch.uint8

48 +

49 +

50 +class MLAAttentionQuantPatternModel(torch.nn.Module):

51 + """Base model for MLA AttentionQuantPattern fusion."""

52 +

53 + def __init__(

54 + self,

55 + num_heads: int,

56 + qk_nope_head_dim: int,

57 + qk_rope_head_dim: int,

58 + v_head_dim: int,

59 + kv_lora_rank: int,

60 + kv_cache_dtype: torch.dtype,

61 + device: torch.device,

62 + vllm_config: VllmConfig,

63 + **kwargs,

64 + ):

65 + super().__init__()

66 + self.num_heads = num_heads

67 + self.qk_nope_head_dim = qk_nope_head_dim

68 + self.qk_rope_head_dim = qk_rope_head_dim

69 + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

70 + self.v_head_dim = v_head_dim

71 + self.kv_lora_rank = kv_lora_rank

72 + self.output_dim = num_heads * v_head_dim

73 + self.head_size = kv_lora_rank + qk_rope_head_dim

74 + self.kv_cache_dtype = kv_cache_dtype

75 + self.device = device

76 + self.vllm_config = vllm_config

77 +

78 + # Create kv_b_proj (ColumnParallelLinear) on device.

79 + # Reuse weights from prior model instance when available, because

80 + # ColumnParallelLinear may get NaN from recycled CUDA memory after

81 + # torch.compile runs in the same process.

82 + kv_b_proj = ColumnParallelLinear(

83 + input_size=kv_lora_rank,

84 + output_size=num_heads * (qk_nope_head_dim + v_head_dim),

85 + bias=False,

86 + prefix="model.layers.0.self_attn.kv_b_proj",

87 + ).to(device)

88 + kv_b_proj_weight = kwargs.get("kv_b_proj_weight")

89 + if kv_b_proj_weight is not None:

90 + kv_b_proj.weight.data.copy_(kv_b_proj_weight)

91 + elif kv_b_proj.weight.data.isnan().any():

92 + # Sanitize NaN from recycled CUDA memory

93 + kv_b_proj.weight.data.normal_()

94 +

95 + # Create MLAAttention

96 + self.mla_attn = MLAAttention(

97 + num_heads=num_heads,

98 + scale=1.0 / (self.qk_head_dim**0.5),

99 + qk_nope_head_dim=qk_nope_head_dim,

100 + qk_rope_head_dim=qk_rope_head_dim,

101 + v_head_dim=v_head_dim,

102 + q_lora_rank=None,

103 + kv_lora_rank=kv_lora_rank,

104 + kv_b_proj=kv_b_proj,

105 + cache_config=vllm_config.cache_config,

106 + quant_config=self.quant_config,

107 + prefix="model.layers.0.self_attn.attn",

108 + )

109 + self.mla_attn._k_scale = self.mla_attn._k_scale.to(device)

110 + self.mla_attn._v_scale = self.mla_attn._v_scale.to(device)

111 +

112 + # Initialize W_UK_T and W_UV from kv_b_proj weights

113 + self.mla_attn.process_weights_after_loading(torch.get_default_dtype())

114 + self.kv_b_proj_weight = kv_b_proj.weight.data.clone()

115 +

116 + self.block_size = 16

117 +

118 + # Initialize MLA MetadataBuilder

119 + self.builder = self.mla_attn.attn_backend.get_builder_cls()(

120 + kv_cache_spec=MLAAttentionSpec(

121 + block_size=self.block_size,

122 + num_kv_heads=1,

123 + head_size=self.head_size,

124 + dtype=self.kv_cache_dtype,

125 + ),

126 + layer_names=[self.mla_attn.layer_name],

127 + vllm_config=self.vllm_config,

128 + device=self.device,

129 + )

130 +

131 + def build_attn_metadata(self, batch_size: int) -> AttentionMetadata:

132 + """Initialize MLA attention metadata.

133 +

134 + NOTE: Uses decode-only batch (query_len=1 per request). The prefill

135 + (forward_mha) path is not separately tested here because it requires

136 + FlashAttention availability and different input tensor shapes. The

137 + quant logic in forward_impl is identical for both paths — it quantizes

138 + the full output[:num_actual_toks] buffer after both forward_mha and

139 + forward_mqa have written their results.

140 + """

141 +

142 + batch_spec = BatchSpec(seq_lens=[1] * batch_size, query_lens=[1] * batch_size)

143 + common_attn_metadata = create_common_attn_metadata(

144 + batch_spec, self.block_size, self.device, arange_block_indices=True

145 + )

146 +

147 + max_blocks = (max(batch_spec.seq_lens) + self.block_size - 1) // self.block_size

148 + num_blocks = batch_size * max_blocks

149 +

150 + # MLA KV cache is 3D: (num_blocks, block_size, head_size)

151 + attn_backend = self.mla_attn.attn_backend

152 + kv_cache_shape = attn_backend.get_kv_cache_shape(

153 + num_blocks, self.block_size, 1, self.head_size

154 + )

155 + try:

156 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()

157 + except (AttributeError, NotImplementedError):

158 + kv_cache_stride_order = tuple(range(len(kv_cache_shape)))

159 +

160 + ordered_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)

161 + inv_order = [

162 + kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))

163 + ]

164 +

165 + raw_tensor = torch.zeros(

166 + ordered_shape, dtype=self.kv_cache_dtype, device=self.device

167 + )

168 + kv_cache = raw_tensor.permute(*inv_order)

169 +

170 + self.mla_attn.kv_cache = kv_cache

171 +

172 + self.attn_metadata = self.builder.build(

173 + common_prefix_len=0, common_attn_metadata=common_attn_metadata

174 + )

175 +

176 + return self.attn_metadata

177 +

178 +

179 +class TestMLAAttentionFp8StaticQuantPatternModel(MLAAttentionQuantPatternModel):

180 + """Test model for MLA Attention + FP8 static quant fusion."""

181 +

182 + quant_key = kFp8StaticTensorSym

183 + quant_config = Fp8Config()

184 +

185 + def __init__(self, *args, **kwargs):

186 + super().__init__(*args, **kwargs)

187 +

188 + self.fp8_linear = TestFP8Layer(

189 + weight_shape=(self.output_dim, self.output_dim),

190 + activation_quant_key=self.quant_key,

191 + weight_quant_key=self.quant_key,

192 + device=self.device,

193 + )

194 +

195 + w = kwargs.get("w")

196 + if w is not None:

197 + self.fp8_linear.weight = w["weight"]

198 + self.fp8_linear.weight_scale = w["wscale"]

199 + self.fp8_linear.input_scale = w["scale"]

200 +

201 + self.w = {

202 + "weight": self.fp8_linear.weight,

203 + "wscale": self.fp8_linear.weight_scale,

204 + "scale": self.fp8_linear.input_scale,

205 + }

206 +

207 + def forward(

208 + self,

209 + q: torch.Tensor,

210 + kv_c_normed: torch.Tensor,

211 + k_pe: torch.Tensor,

212 + ):

213 + """Forward pass that creates the MLA attention + FP8 quant pattern."""

214 + attn_output = self.mla_attn(

215 + q,

216 + kv_c_normed,

217 + k_pe,

218 + output_shape=(q.shape[0], self.output_dim),

219 + )

220 + return self.fp8_linear(attn_output)

221 +

222 +

223 +class TestMLAAttentionNvfp4QuantPatternModel(MLAAttentionQuantPatternModel):

224 + """Test model for MLA Attention + NVFP4 quant fusion."""

225 +

226 + quant_key = kNvfp4Dynamic

227 + quant_config = ModelOptNvFp4Config(

228 + is_checkpoint_nvfp4_serialized=False,

229 + kv_cache_quant_algo=None,

230 + exclude_modules=[],

231 + )

232 +

233 + def __init__(self, *args, **kwargs):

234 + super().__init__(*args, **kwargs)

235 +

236 + self.w = kwargs.get(

237 + "w",

238 + {

239 + "weight": torch.randint(

240 + 256,

241 + (self.output_dim, self.output_dim // 2),

242 + dtype=FP4_DTYPE,

243 + device=self.device,

244 + ),

245 + "wscale_swizzled": torch.randn(

246 + self.output_dim, self.output_dim // 16

247 + ).to(dtype=FP8_DTYPE, device=self.device),

248 + "wscale": torch.tensor([500], dtype=torch.float32, device=self.device),

249 + "scale": torch.tensor([0.002], dtype=torch.float32, device=self.device),

250 + },

251 + )

252 +

253 + def forward(

254 + self,

255 + q: torch.Tensor,

256 + kv_c_normed: torch.Tensor,

257 + k_pe: torch.Tensor,

258 + ):

259 + """Forward pass that creates the MLA attention + NVFP4 quant pattern."""

260 + attn_output = self.mla_attn(

261 + q,

262 + kv_c_normed,

263 + k_pe,

264 + output_shape=(q.shape[0], self.output_dim),

265 + )

266 + quant_output, output_block_scale = scaled_fp4_quant(

267 + attn_output, 1 / self.w["scale"]

268 + )

269 + return cutlass_scaled_fp4_mm(

270 + a=quant_output,

271 + b=self.w["weight"],

272 + block_scale_a=output_block_scale,

273 + block_scale_b=self.w["wscale_swizzled"],

274 + alpha=self.w["scale"] * self.w["wscale"],

275 + out_dtype=attn_output.dtype,

276 + )

277 +

278 +

279 +def is_nvfp4_supported():

280 + return current_platform.has_device_capability(100)

281 +

282 +

283 +# MLA test configuration

284 +MLA_DIMS: list[tuple[int, int, int, int, int]] = []

285 +PATTERN_TEST_MODELS_MLA_FP8: list[tuple[str, type]] = []

286 +PATTERN_TEST_MODELS_MLA_FP4: list[tuple[str, type]] = []

287 +BACKENDS_MLA_FP8: list[AttentionBackendEnum] = []

288 +BACKENDS_MLA_FP4: list[AttentionBackendEnum] = []

289 +

290 +if current_platform.is_cuda():

291 + # (num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank)

292 + MLA_DIMS = [(16, 128, 64, 128, 512)]

293 + PATTERN_TEST_MODELS_MLA_FP8 = [

294 + (

295 + "deepseek-ai/DeepSeek-V2-Lite",

296 + TestMLAAttentionFp8StaticQuantPatternModel,

297 + )

298 + ]

299 + PATTERN_TEST_MODELS_MLA_FP4 = [

300 + (

301 + "deepseek-ai/DeepSeek-V2-Lite",

302 + TestMLAAttentionNvfp4QuantPatternModel,

303 + )

304 + ]

305 + BACKENDS_MLA_FP8 = [AttentionBackendEnum.TRITON_MLA]

306 + BACKENDS_MLA_FP4 = [AttentionBackendEnum.TRITON_MLA]

307 +

308 +

309 +@pytest.mark.parametrize(

310 + "num_heads, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, kv_lora_rank",

311 + MLA_DIMS,

312 +)

313 +@pytest.mark.parametrize("batch_size", [7, 256] if current_platform.is_cuda() else [8])

314 +@pytest.mark.parametrize("dtype", [torch.bfloat16])

315 +@pytest.mark.parametrize(

316 + "backend, model_name, model_class, custom_ops",

317 + list(

318 + flat_product(

319 + BACKENDS_MLA_FP8,

320 + PATTERN_TEST_MODELS_MLA_FP8,

321 + ["+quant_fp8", "-quant_fp8"],

322 + )

323 + )

324 + + list(flat_product(BACKENDS_MLA_FP4, PATTERN_TEST_MODELS_MLA_FP4, [""])),

325 +)

326 +@pytest.mark.skipif(

327 + not current_platform.is_cuda_alike(), reason="Only test ROCm or CUDA"

328 +)

329 +@pytest.mark.skipif(not current_platform.supports_fp8(), reason="Need FP8")

330 +def test_mla_attention_quant_pattern(

331 + num_heads: int,

332 + qk_nope_head_dim: int,

333 + qk_rope_head_dim: int,

334 + v_head_dim: int,

335 + kv_lora_rank: int,

336 + batch_size: int,

337 + dtype: torch.dtype,

338 + custom_ops: str,

339 + model_name: str,

340 + model_class: type[MLAAttentionQuantPatternModel],

341 + backend: AttentionBackendEnum,

342 + dist_init,

343 + monkeypatch,

344 + use_fresh_inductor_cache,

345 +):

346 + """Test MLA AttentionQuantPattern fusion pass"""

347 + if (

348 + model_class is TestMLAAttentionNvfp4QuantPatternModel

349 + and not is_nvfp4_supported()

350 + ):

351 + pytest.skip("NVFP4 is not supported on this GPU (requires SM 100+).")

352 +

353 + monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")

354 +

355 + custom_ops_list = custom_ops.split(",") if custom_ops else []

356 +

357 + device = torch.device("cuda:0")

358 + torch.set_default_dtype(dtype)

359 + torch.manual_seed(42)

360 +

361 + model_config = ModelConfig(

362 + model=model_name,

363 + max_model_len=2048,

364 + dtype=dtype,

365 + )

366 + vllm_config = VllmConfig(

367 + model_config=model_config,

368 + scheduler_config=SchedulerConfig(

369 + max_num_seqs=1024,

370 + max_model_len=model_config.max_model_len,

371 + is_encoder_decoder=model_config.is_encoder_decoder,

372 + ),

373 + compilation_config=CompilationConfig(

374 + mode=CompilationMode.VLLM_COMPILE,

375 + custom_ops=custom_ops_list,

376 + ),

377 + cache_config=CacheConfig(cache_dtype="auto"),

378 + attention_config=AttentionConfig(backend=backend),

379 + )

380 +

381 + # MLA inputs: q(B, N, qk_head_dim), kv_c_normed(B, L), k_pe(B, 1, R)

382 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

383 + q = torch.randn(batch_size, num_heads, qk_head_dim, dtype=dtype, device=device)

384 + kv_c_normed = torch.randn(batch_size, kv_lora_rank, dtype=dtype, device=device)

385 + k_pe = torch.randn(batch_size, 1, qk_rope_head_dim, dtype=dtype, device=device)

386 +

387 + # Mark first dimension as dynamic

388 + torch._dynamo.mark_dynamic(q, 0)

389 + torch._dynamo.mark_dynamic(kv_c_normed, 0)

390 + torch._dynamo.mark_dynamic(k_pe, 0)

391 +

392 + # Run model without fusion

393 + vllm_config_unfused = copy.deepcopy(vllm_config)

394 + with (

395 + set_current_vllm_config(vllm_config_unfused),

396 + set_forward_context(attn_metadata=None, vllm_config=vllm_config_unfused),

397 + ):

398 + model_unfused = model_class(

399 + num_heads=num_heads,

400 + qk_nope_head_dim=qk_nope_head_dim,

401 + qk_rope_head_dim=qk_rope_head_dim,

402 + v_head_dim=v_head_dim,

403 + kv_lora_rank=kv_lora_rank,

404 + kv_cache_dtype=dtype,

405 + device=device,

406 + vllm_config=vllm_config_unfused,

407 + )

408 + model_unfused = model_unfused.to(device)

409 + # HACK: See #131044

410 + result_unfused_0 = model_unfused(q, kv_c_normed, k_pe) # noqa: F841

411 +

412 + forward_ctx = get_forward_context()

413 + forward_ctx.attn_metadata = model_unfused.build_attn_metadata(batch_size)

414 +

415 + compiled_unfused = torch.compile(model_unfused, fullgraph=True)

416 + result_unfused = compiled_unfused(q, kv_c_normed, k_pe)

417 +

418 + # Run model with attn fusion enabled

419 + vllm_config.compilation_config.pass_config = PassConfig(

420 + fuse_attn_quant=True, eliminate_noops=True

421 + )

422 + with (

423 + set_current_vllm_config(vllm_config),

424 + set_forward_context(attn_metadata=None, vllm_config=vllm_config),

425 + ):

426 + model_fused = model_class(

427 + num_heads=num_heads,

428 + qk_nope_head_dim=qk_nope_head_dim,

429 + qk_rope_head_dim=qk_rope_head_dim,

430 + v_head_dim=v_head_dim,

431 + kv_lora_rank=kv_lora_rank,

432 + kv_cache_dtype=dtype,

433 + device=device,

434 + vllm_config=vllm_config,

435 + w=model_unfused.w,

436 + kv_b_proj_weight=model_unfused.kv_b_proj_weight,

437 + )

438 + model_fused = model_fused.to(device)

439 +

440 + forward_ctx = get_forward_context()

441 + forward_ctx.attn_metadata = model_fused.build_attn_metadata(batch_size)

442 +

443 + # Create test backend with fusion passes

444 + noop_pass = NoOpEliminationPass(vllm_config)

445 + attn_pass = LazyInitPass(MLAAttnQuantFusionPass, vllm_config)

446 + cleanup_pass = PostCleanupPass(vllm_config)

447 +

448 + test_backend = TestBackend(noop_pass, attn_pass, cleanup_pass)

449 + # HACK: See https://github.com/vllm-project/vllm/issues/31044

450 + result_fused_0 = model_fused(q, kv_c_normed, k_pe) # noqa: F841

451 +

452 + compiled_fused = torch.compile(

453 + model_fused, backend=test_backend, fullgraph=True

454 + )

455 +

456 + result_fused = compiled_fused(q, kv_c_normed, k_pe)

457 +

458 + # Check attn fusion support

459 + quant_key: QuantKey = model_class.quant_key

460 + attn_fusion_supported = [

461 + layer.impl.fused_output_quant_supported(quant_key)

462 + for key, layer in vllm_config.compilation_config.static_forward_context.items()

463 + if isinstance(layer, MLAAttention)

464 + ]

465 + assert sum(attn_fusion_supported) == len(attn_fusion_supported), (

466 + "All MLA layers should support attention fusion"

467 + )

468 +

469 + # Check quantization ops in the graph

470 + quant_op = (

471 + torch.ops.aten.reciprocal

472 + if "-quant_fp8" in custom_ops_list

473 + else QUANT_OPS[quant_key]

474 + )

475 + test_backend.check_before_ops([quant_op], fully_replaced=quant_key is kNvfp4Dynamic)

476 +

477 + assert attn_pass.pass_.matched_count == sum(attn_fusion_supported)

478 +

479 + # Check MLA attention ops in the graph

480 + attn_nodes_pre = list(find_op_nodes(MLA_ATTN_OP, test_backend.graph_pre_pass))

481 + attn_nodes_post = list(find_op_nodes(MLA_ATTN_OP, test_backend.graph_post_pass))

482 +

483 + assert len(attn_nodes_pre) > 0, "Should have MLA attention nodes before fusion"

484 + assert len(attn_nodes_pre) == len(attn_nodes_post), (

485 + "Should have same number of MLA attention nodes before and after fusion"

486 + )

487 + assert attn_nodes_pre[0].kwargs.get("output_scale") is None, (

488 + "MLA attention should not have output_scale before fusion"

489 + )

490 + assert attn_nodes_post[0].kwargs.get("output_scale") is not None, (

491 + "MLA attention should have output_scale after fusion"

492 + )

493 +

494 + assert attn_nodes_pre[0].kwargs.get("output_block_scale") is None, (

495 + "MLA attention should not have output_block_scale before fusion"

496 + )

497 +

498 + if quant_key.dtype == FP8_DTYPE:

499 + assert attn_nodes_post[0].kwargs.get("output_block_scale") is None, (

500 + "MLA attention should not have output_block_scale after FP8 fusion"

501 + )

502 + elif quant_key.dtype == FP4_DTYPE:

503 + assert attn_nodes_post[0].kwargs.get("output_block_scale") is not None, (

504 + "MLA attention should have output_block_scale after FP4 fusion"

505 + )

506 +

507 + # Check numerical correctness

508 + torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)

vllm/compilation/passes/fusion/mla_attn_quant_fusion.py new

+262

@@ -0,0 +1,262 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +from collections.abc import Callable

5 +

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

10 +from vllm.config import VllmConfig, get_layers_from_vllm_config

11 +from vllm.logger import init_logger

12 +from vllm.model_executor.layers.attention.mla_attention import MLAAttention

13 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

14 + kFp8StaticTensorSym,

15 + kNvfp4Dynamic,

16 +)

17 +from vllm.platforms import current_platform

18 +

19 +from ..vllm_inductor_pass import VllmFusionPatternMatcherPass, VllmPatternReplacement

20 +from .matcher_utils import MatcherQuantFP8

21 +from .rms_quant_fusion import QUANT_OPS

22 +

23 +logger = init_logger(__name__)

24 +

25 +FP8_DTYPE = current_platform.fp8_dtype()

26 +FP4_DTYPE = torch.uint8

27 +

28 +MLA_ATTN_OP = torch.ops.vllm.unified_mla_attention_with_output.default

29 +

30 +

31 +class MLAAttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):

32 + """

33 + Fusion for MLA Attention+Fp8StaticQuant.

34 +

35 + Matches the pattern: MLA attention -> static FP8 quant, and replaces

36 + it with MLA attention(output_scale=scale, output=fp8_buffer).

37 + """

38 +

39 + def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None:

40 + self._layer_name = layer.layer_name

41 + self._num_heads = layer.num_heads

42 + self._v_head_dim = layer.v_head_dim

43 + self._kv_lora_rank = layer.kv_lora_rank

44 + self._qk_rope_head_dim = layer.qk_rope_head_dim

45 + self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim

46 + self._output_dim = layer.num_heads * layer.v_head_dim

47 + self._dtype = dtype

48 + self._quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)

49 +

50 + @property

51 + def pattern(self) -> Callable[..., torch.Tensor]:

52 + def _pattern(

53 + q: torch.Tensor,

54 + kv_c_normed: torch.Tensor,

55 + k_pe: torch.Tensor,

56 + output_attn: torch.Tensor,

57 + scale: torch.Tensor,

58 + kv_cache_dummy_dep: torch.Tensor,

59 + ) -> torch.Tensor:

60 + at1 = auto_functionalized(

61 + MLA_ATTN_OP,

62 + q=q,

63 + kv_c_normed=kv_c_normed,

64 + k_pe=k_pe,

65 + output=output_attn,

66 + layer_name=self._layer_name,

67 + output_scale=None,

68 + output_block_scale=None,

69 + kv_cache_dummy_dep=kv_cache_dummy_dep,

70 + )

71 + # MLA output is already 2D (T, N*V), no reshape needed

72 + return self._quant_matcher(at1[1], scale)[0]

73 +

74 + return _pattern

75 +

76 + @property

77 + def replacement(self) -> Callable[..., torch.Tensor]:

78 + def _replacement(

79 + q: torch.Tensor,

80 + kv_c_normed: torch.Tensor,

81 + k_pe: torch.Tensor,

82 + output_attn: torch.Tensor,

83 + scale: torch.Tensor,

84 + kv_cache_dummy_dep: torch.Tensor,

85 + ) -> torch.Tensor:

86 + # MLA output in quant_dtype

87 + output_attn = torch.empty(

88 + [q.shape[0], self._output_dim],

89 + dtype=FP8_DTYPE,

90 + device=q.device,

91 + )

92 + at1 = auto_functionalized(

93 + MLA_ATTN_OP,

94 + q=q,

95 + kv_c_normed=kv_c_normed,

96 + k_pe=k_pe,

97 + output=output_attn,

98 + layer_name=self._layer_name,

99 + output_scale=scale,

100 + output_block_scale=None,

101 + kv_cache_dummy_dep=kv_cache_dummy_dep,

102 + )

103 + return at1[1]

104 +

105 + return _replacement

106 +

107 + def get_inputs(self) -> list[torch.Tensor]:

108 + return [

109 + self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),

110 + self.empty(5, self._kv_lora_rank, dtype=self._dtype),

111 + self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),

112 + self.empty(5, self._output_dim, dtype=self._dtype),

113 + self.empty_fp32(1, 1),

114 + self.empty(0, dtype=self._dtype),

115 + ]

116 +

117 +

118 +class MLAAttnNvfp4QuantPattern(

119 + VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]

120 +):

121 + """

122 + Fusion for MLA Attention+Nvfp4Quant.

123 +

124 + Matches the pattern: MLA attention -> NVFP4 quant, and replaces

125 + it with MLA attention(output_scale=scale, output_block_scale=block_scale,

126 + output=fp4_buffer).

127 + """

128 +

129 + def __init__(self, layer: MLAAttention, dtype: torch.dtype) -> None:

130 + self._layer_name = layer.layer_name

131 + self._num_heads = layer.num_heads

132 + self._v_head_dim = layer.v_head_dim

133 + self._kv_lora_rank = layer.kv_lora_rank

134 + self._qk_rope_head_dim = layer.qk_rope_head_dim

135 + self._qk_head_dim = layer.qk_nope_head_dim + layer.qk_rope_head_dim

136 + self._output_dim = layer.num_heads * layer.v_head_dim

137 + self._dtype = dtype

138 + self._QUANT_OP = QUANT_OPS[kNvfp4Dynamic]

139 +

140 + @property

141 + def pattern(

142 + self,

143 + ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:

144 + def _pattern(

145 + q: torch.Tensor,

146 + kv_c_normed: torch.Tensor,

147 + k_pe: torch.Tensor,

148 + output_attn: torch.Tensor,

149 + input_scale: torch.Tensor,

150 + kv_cache_dummy_dep: torch.Tensor,

151 + ) -> tuple[torch.Tensor, torch.Tensor]:

152 + at1 = auto_functionalized(

153 + MLA_ATTN_OP,

154 + q=q,

155 + kv_c_normed=kv_c_normed,

156 + k_pe=k_pe,

157 + output=output_attn,

158 + layer_name=self._layer_name,

159 + output_scale=None,

160 + output_block_scale=None,

161 + kv_cache_dummy_dep=kv_cache_dummy_dep,

162 + )

163 + # Replicate what scaled_fp4_quant() does: allocate output

164 + # tensors inline then call the .out variant.

165 + output_quant, output_scale = create_fp4_output_tensors(

166 + at1[1].shape[0], at1[1].shape[1], at1[1].device, True

167 + )

168 + at2 = auto_functionalized(

169 + self._QUANT_OP,

170 + input=at1[1],

171 + input_scale=input_scale,

172 + is_sf_swizzled_layout=True,

173 + output=output_quant,

174 + output_scale=output_scale,

175 + )

176 + output_scale_view = torch.ops.aten.view.dtype(at2[2], FP8_DTYPE)

177 + return at2[1], output_scale_view

178 +

179 + return _pattern

180 +

181 + @property

182 + def replacement(

183 + self,

184 + ) -> Callable[..., tuple[torch.Tensor, torch.Tensor]]:

185 + def _replacement(

186 + q: torch.Tensor,

187 + kv_c_normed: torch.Tensor,

188 + k_pe: torch.Tensor,

189 + output_attn: torch.Tensor,

190 + input_scale: torch.Tensor,

191 + kv_cache_dummy_dep: torch.Tensor,

192 + ) -> tuple[torch.Tensor, torch.Tensor]:

193 + # MLA output in quant_dtype (FP4 packed as uint8)

194 + output_attn = torch.empty(

195 + [q.shape[0], self._output_dim // 2],

196 + dtype=FP4_DTYPE,

197 + device=q.device,

198 + )

199 + # attention output block scale

200 + output_scale = create_fp4_output_tensors(

201 + q.shape[0], self._output_dim, q.device, True

202 + )[1]

203 + output_scale_view = torch.ops.aten.view.dtype(output_scale, FP8_DTYPE)

204 + at2 = auto_functionalized(

205 + MLA_ATTN_OP,

206 + q=q,

207 + kv_c_normed=kv_c_normed,

208 + k_pe=k_pe,

209 + output=output_attn,

210 + layer_name=self._layer_name,

211 + output_scale=input_scale,

212 + output_block_scale=output_scale_view,

213 + kv_cache_dummy_dep=kv_cache_dummy_dep,

214 + )

215 + return at2[1], at2[2]

216 +

217 + return _replacement

218 +

219 + def get_inputs(self) -> list[torch.Tensor]:

220 + return [

221 + self.empty(5, self._num_heads, self._qk_head_dim, dtype=self._dtype),

222 + self.empty(5, self._kv_lora_rank, dtype=self._dtype),

223 + self.empty(5, 1, self._qk_rope_head_dim, dtype=self._dtype),

224 + self.empty(5, self._output_dim, dtype=self._dtype),

225 + self.empty_fp32(1, 1),

226 + self.empty(0, dtype=self._dtype),

227 + ]

228 +

229 +

230 +class MLAAttnQuantFusionPass(VllmFusionPatternMatcherPass):

231 + """

232 + This pass fuses post-attention quantization onto MLA attention if supported.

233 +

234 + It uses the pattern matcher and matches each MLA layer manually, as strings

235 + cannot be wildcarded. This also lets us check support on attention layers

236 + upon registration instead of during pattern matching.

237 + """

238 +

239 + def __init__(self, config: VllmConfig) -> None:

240 + super().__init__(config, "mla_attn_quant_fusion")

241 +

242 + dtype = config.model_config.dtype

243 + layers = list(get_layers_from_vllm_config(config, MLAAttention).values())

244 +

245 + if len(layers) == 0:

246 + logger.warning(

247 + "MLA attention + quant fusion is enabled, but no MLA "

248 + "attention layers were found in "

249 + "CompilationConfig.static_forward_context "

250 + "so no fusion patterns were registered."

251 + )

252 +

253 + for layer in layers:

254 + if layer.impl.fused_output_quant_supported(kFp8StaticTensorSym):

255 + self.register(MLAAttnFp8StaticQuantPattern(layer, dtype))

256 +

257 + if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):

258 + for layer in layers:

259 + if layer.impl.fused_output_quant_supported(kNvfp4Dynamic):

260 + self.register(MLAAttnNvfp4QuantPattern(layer, dtype))

261 +

262 + self.dump_patterns(config, self.pm_pass)

vllm/compilation/passes/pass_manager.py

@@ -27,6 +27,7 @@ if rocm_aiter_ops.is_enabled():

27 if current_platform.is_cuda_alike():

28 from .fusion.act_quant_fusion import ActivationQuantFusionPass

29 from .fusion.attn_quant_fusion import AttnQuantFusionPass

30 + from .fusion.mla_attn_quant_fusion import MLAAttnQuantFusionPass

31 from .fusion.qk_norm_rope_fusion import QKNormRoPEFusionPass

32 from .fusion.rms_quant_fusion import RMSNormQuantFusionPass

33 from .fusion.rope_kvcache_fusion import RopeKVCacheFusionPass

@@ -157,6 +158,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

158

159 if self.pass_config.fuse_attn_quant:

160 self.passes += [AttnQuantFusionPass(config)]

161 + self.passes += [MLAAttnQuantFusionPass(config)]

162

163 if self.pass_config.enable_qk_norm_rope_fusion:

164 self.passes += [SplitCoalescingPass(config)]

vllm/config/compilation.py

+1 -1

@@ -121,7 +121,7 @@ class PassConfig:

121 fuse_act_quant: bool = None # type: ignore[assignment]

122 """Fuse the custom SiluMul + quant ops."""

123 fuse_attn_quant: bool = None # type: ignore[assignment]

124 - """Fuse the custom attention + quant ops."""

124 + """Fuse the custom Attention and MLAAttention + quant ops."""

125 eliminate_noops: bool = Field(default=True)

126 """Eliminate no-op ops."""

127 enable_sp: bool = None # type: ignore[assignment]

vllm/model_executor/layers/attention/mla_attention.py

+49 -5

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

449 group_shape=GroupShape.PER_TENSOR,

450 compile_native=True,

451 )

452 + self._quant_fp8_op = QuantFP8(

453 + static=True,

454 + group_shape=GroupShape.PER_TENSOR,

455 + compile_native=True,

456 + )

457

458 @property

459 def chunked_prefill_workspace_size(self) -> int:

@@ -545,9 +550,19 @@ class MLAAttention(nn.Module, AttentionLayerBase):

550 ) -> torch.Tensor:

551 assert output is not None, "Output tensor must be provided."

552

548 - if output_scale is not None or output_block_scale is not None:

549 - raise NotImplementedError(

550 - "fused output quantization is not yet supported for MLA"

553 + use_quant = output_scale is not None or output_block_scale is not None

554 + if use_quant:

555 + # The fusion pass has allocated output with quantized dtype

556 + # (FP8 or uint8 for FP4). We can't write into it directly,

557 + # so we swap in a temp buffer for computation, then quantize

558 + # into the real output at the end.

559 + # NOTE(carlyou): this is temporary until kernels support fp8 output

560 + quant_output = output

561 + output = torch.empty(

562 + output.shape[0],

563 + self.num_heads * self.v_head_dim,

564 + dtype=q.dtype,

565 + device=output.device,

566 )

567

568 if attn_metadata is None:

@@ -567,6 +582,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):

582 # The zero fill is required when used with DP + EP

583 # to ensure all ranks within a DP group compute the

584 # same expert outputs.

585 + if use_quant:

586 + return quant_output.fill_(0)

587 return output.fill_(0)

588

589 if self.impl.dcp_world_size == -1:

@@ -706,6 +723,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):

723

724 # v_up projection

725 self._v_up_proj(attn_out, out=mqa_output_slice)

726 +

727 + if use_quant:

728 + # Quantize the BF16 computation result into the quantized output

729 + actual = output[:num_actual_toks]

730 + if output_block_scale is not None:

731 + # NVFP4: two FP4 values packed into one uint8

732 + fp4_data, fp4_scales = ops.scaled_fp4_quant(actual, output_scale)

733 + quant_output[:num_actual_toks].copy_(fp4_data)

734 + output_block_scale.copy_(fp4_scales)

735 + else:

736 + # Static FP8 quantization

737 + fp8_data, _ = self._quant_fp8_op(actual, output_scale)

738 + quant_output[:num_actual_toks].copy_(fp8_data)

739 + return quant_output

740 +

741 return output_padded

742

743 def process_weights_after_loading(self, act_dtype: torch.dtype):

@@ -2069,6 +2101,14 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2101 understand this class

2102 """

2103

2104 + def fused_output_quant_supported(self, quant_key):

2105 + from vllm.model_executor.layers.quantization.utils.quant_utils import (

2106 + kFp8StaticTensorSym,

2107 + kNvfp4Dynamic,

2108 + )

2109 +

2110 + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)

2111 +

2112 def __init__(

2113 self,

2114 num_heads: int,

@@ -2513,8 +2553,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2553 if hasattr(self.kv_b_proj, "weight")

2554 else self.kv_b_proj.params_dtype

2555 )

2516 - if use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype():

2517 - kv_c_normed = kv_c_normed.to(_kv_b_proj_w_dtype)

2556 + # For NVFP4, weights are packed uint8 — keep input in model dtype

2557 + # since the NVFP4 linear layer quantizes internally.

2558 + if (

2559 + use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype()

2560 + ) and _kv_b_proj_w_dtype != torch.uint8:

2561 + kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype)

2562

2563 k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)

2564 kv_nope = self.kv_b_proj(kv_c_normed)[0].view(

vllm/v1/attention/backend.py

+21

@@ -10,6 +10,11 @@ import numpy as np

10 import torch

11 from typing_extensions import deprecated

13 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

14 + kFp8StaticTensorSym,

15 + kNvfp4Dynamic,

16 +)

17 +

18 if TYPE_CHECKING:

19 from vllm.config import VllmConfig

20 from vllm.config.cache import CacheDType

@@ -873,6 +878,14 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):

878 """MQA-style decode forward pass."""

879 raise NotImplementedError

880

881 + def fused_output_quant_supported(self, quant_key: "QuantKey"):

882 + """

883 + Does this attention implementation support fused output quantization.

884 + Since MLA quantization is done manually in forward_impl (common code),

885 + all MLA backends support it by default.

886 + """

887 + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)

888 +

889 def do_kv_cache_update(

890 self,

891 kv_c_normed: torch.Tensor,

@@ -903,6 +916,14 @@ class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):

916 They do not support prefill (MHA-style) attention.

917 """

918

919 + def fused_output_quant_supported(self, quant_key: "QuantKey"):

920 + """

921 + Does this attention implementation support fused output quantization.

922 + Since MLA quantization is done manually in forward_impl (common code),

923 + all MLA backends support it by default.

924 + """

925 + return quant_key in (kFp8StaticTensorSym, kNvfp4Dynamic)

926 +

927 @abstractmethod

928 def __init__(

929 self,