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

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 4, 22:30 UTC

[ROCm] Add MLA dual RMS norm fusion (Q, KV) pass for DeepSeek/Kimi-K2 (#39242)

Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>

Rita Brugarolas committed Apr 20, 2026 at 07:56 UTC fb5635d3f90635ce9d1acbc975baab6e911a262b

8 files changed +361 -1

docs/design/fusions.md

+40

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

36 ## Support Matrix

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

54 | `fuse_act_padding` | — | — | — | — | FP16/BF16 |

55 +| `fuse_mla_dual_rms_norm` | — | — | — | — | BF16 |

57 \* `fuse_attn_quant` support depends on the attention backend in use; not all backends support

58 fused quantization output. See the [`fuse_attn_quant` section](#attention--quantization-fuse_attn_quant)

@@ -381,6 +383,44 @@ when the hidden size is 2880 and AITER Triton GEMMs *not* enabled.

383

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

385

386 +### MLA Dual RMSNorm (`fuse_mla_dual_rms_norm`)

387 +

388 +!!! info

389 + ROCm/AITER-only. Targeted at DeepSeek-V3 / Kimi-K2 MLA attention.

390 +

391 +!!! note

392 + When the native implementation of `rms_norm` is used (the default on CUDA and

393 + ROCm for now), Inductor's built-in fusion already handles merging these norms

394 + automatically. This explicit pass targets the case where AITER's custom

395 + `rms_norm` op is active, which Inductor cannot fuse on its own.

396 +

397 +**What it fuses.** Fuses the paired `q_a_layernorm` and `kv_a_layernorm` RMS norm

398 +operations in MLA attention into a single `fused_qk_rmsnorm` HIP kernel call via AITER,

399 +reducing kernel launch overhead from 2 launches to 1 per MLA layer.

400 +

401 +```text

402 +# Unfused:

403 +q_c, kv_lora = split(projected, [q_dim, kv_dim])

404 +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])

405 +q_c = rms_norm(q_c, q_weight, eps)

406 +kv_c = rms_norm(kv_c, kv_weight, eps)

407 +

408 +# Fused:

409 +q_c, kv_lora = split(projected, [q_dim, kv_dim])

410 +kv_c, k_pe = split(kv_lora, [kv_c_dim, k_pe_dim])

411 +q_normed, kv_normed = fused_mla_dual_rms_norm(

412 + q_c, q_weight, kv_c, kv_weight, eps1, eps2)

413 +```

414 +

415 +Requires: AMD ROCm with AITER enabled. Enabled by default at optimization level O1 and above

416 +when AITER is available.

417 +

418 +**Code locations.**

419 +

420 +- Pass: [`vllm/compilation/passes/fusion/rocm_aiter_fusion.py`](https://github.com/vllm-project/vllm/blob/main/vllm/compilation/passes/fusion/rocm_aiter_fusion.py) (`MLADualRMSNormFusionPass`)

421 +- Custom op: [`vllm/_aiter_ops.py`](https://github.com/vllm-project/vllm/blob/main/vllm/_aiter_ops.py) (`fused_mla_dual_rms_norm`)

422 +- AITER kernel: [`fused_qk_rmsnorm`](https://github.com/ROCm/aiter/pull/2442)

423 +

424 ## See Also

425

426 - [Optimization Levels](optimization_levels.md) — high-level presets that set

docs/design/optimization_levels.md

@@ -56,6 +56,7 @@ Fusions:

56 - `-cc.pass_config.fuse_norm_quant=True`*

57 - `-cc.pass_config.fuse_act_quant=True`*

58 - `-cc.pass_config.fuse_act_padding=True`†

59 +- `-cc.pass_config.fuse_mla_dual_rms_norm=True`†

61 \* These fusions are only enabled when either op is using a custom kernel, otherwise Inductor fusion is better.</br>

62 † These fusions are ROCm-only and require AITER.

tests/compile/passes/test_fuse_mla_dual_rms_norm.py new

+148

@@ -0,0 +1,148 @@

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

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

3 +"""

4 +Unit test for the MLADualRMSNormFusionPass.

5 +

6 +The pass fuses paired q/kv RMS norms in MLA attention into a single

7 +fused_mla_dual_rms_norm op backed by AITER's fused_qk_rmsnorm kernel.

8 +"""

9 +

10 +import pytest

11 +import torch

12 +

13 +import vllm.config

14 +from tests.compile.backend import TestBackend

15 +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops

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

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

18 +from vllm.config import (

19 + CompilationConfig,

20 + CompilationMode,

21 + ModelConfig,

22 + PassConfig,

23 + VllmConfig,

24 +)

25 +from vllm.model_executor.layers.layernorm import RMSNorm

26 +

27 +# MLA attention geometry for DeepSeek-V3 / Kimi-K2

28 +Q_DIM = 1536

29 +KV_C_DIM = 512

30 +K_PE_DIM = 64

31 +EPS = 1e-6

32 +

33 +

34 +class MLADualRMSNormTestModel(torch.nn.Module):

35 + """

36 + Minimal model reproducing the MLA dual RMS norm pattern:

37 + linear -> split([q_dim, kv_dim])

38 + +-- q_c (getitem 0) -> rms_norm(q_w, eps) -> linear

39 + +-- kv_lora (getitem 1) -> split([kv_c_dim, k_pe_dim])

40 + +-- kv_c (getitem 0) -> rms_norm(kv_w, eps)

41 + +-- k_pe

42 + """

43 +

44 + def __init__(

45 + self,

46 + hidden_size: int,

47 + q_dim: int = Q_DIM,

48 + kv_c_dim: int = KV_C_DIM,

49 + k_pe_dim: int = K_PE_DIM,

50 + eps: float = EPS,

51 + ):

52 + super().__init__()

53 + self.q_dim = q_dim

54 + self.kv_dim = kv_c_dim + k_pe_dim

55 + self.kv_c_dim = kv_c_dim

56 + self.k_pe_dim = k_pe_dim

57 +

58 + self.proj = torch.nn.Linear(hidden_size, q_dim + self.kv_dim, bias=False)

59 + self.q_norm = RMSNorm(q_dim, eps=eps)

60 + self.kv_norm = RMSNorm(kv_c_dim, eps=eps)

61 + self.q_b_proj = torch.nn.Linear(q_dim, hidden_size, bias=False)

62 +

63 + def forward(

64 + self, x: torch.Tensor

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

66 + # Avoid graph input being a direct arg to a matched pattern node

67 + x = torch.relu(x)

68 +

69 + projected = self.proj(x)

70 +

71 + q_c, kv_lora = projected.split([self.q_dim, self.kv_dim], dim=-1)

72 + kv_c, k_pe = kv_lora.split([self.kv_c_dim, self.k_pe_dim], dim=-1)

73 +

74 + q_normed = self.q_norm(q_c)

75 + kv_normed = self.kv_norm(kv_c)

76 +

77 + q_out = self.q_b_proj(q_normed)

78 + return q_out, kv_normed, k_pe

79 +

80 + def ops_in_model_before(self):

81 + return [torch.ops.vllm_ir.rms_norm.default]

82 +

83 + def ops_in_model_after(self):

84 + return [torch.ops.vllm.fused_mla_dual_rms_norm.default]

85 +

86 +

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

88 +@pytest.mark.parametrize("hidden_size", [7168])

89 +@pytest.mark.skipif(

90 + not is_aiter_found_and_supported(),

91 + reason="Only test on ROCm with AITER installed and supported",

92 +)

93 +def test_fuse_mla_dual_rms_norm(

94 + dtype: torch.dtype,

95 + hidden_size: int,

96 + monkeypatch: pytest.MonkeyPatch,

97 +):

98 + torch._dynamo.reset()

99 +

100 + vllm_config = VllmConfig(

101 + model_config=ModelConfig(dtype=dtype),

102 + compilation_config=CompilationConfig(

103 + mode=CompilationMode.VLLM_COMPILE,

104 + custom_ops=["+rms_norm"],

105 + pass_config=PassConfig(

106 + fuse_mla_dual_rms_norm=True,

107 + eliminate_noops=True,

108 + ),

109 + ),

110 + )

111 +

112 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:

113 + from vllm.compilation.passes.fusion.rocm_aiter_fusion import (

114 + MLADualRMSNormFusionPass,

115 + )

116 +

117 + torch.set_default_device("cuda")

118 + torch.set_default_dtype(dtype)

119 + torch.manual_seed(42)

120 +

121 + m.setenv("VLLM_ROCM_USE_AITER", "1")

122 + rocm_aiter_ops.refresh_env_variables()

123 +

124 + fusion_pass = MLADualRMSNormFusionPass(vllm_config)

125 + passes = [

126 + NoOpEliminationPass(vllm_config),

127 + fusion_pass,

128 + PostCleanupPass(vllm_config),

129 + ]

130 + backend = TestBackend(*passes)

131 + model = MLADualRMSNormTestModel(hidden_size)

132 +

133 + x = torch.randn(1, hidden_size)

134 + torch._dynamo.mark_dynamic(x, 0)

135 +

136 + outputs_unfused = model(x)

137 +

138 + model_fused = torch.compile(model, backend=backend)

139 + outputs_fused = model_fused(x)

140 +

141 + torch.testing.assert_close(outputs_unfused, outputs_fused, atol=1e-2, rtol=1e-2)

142 +

143 + assert fusion_pass.matched_count == 1, (

144 + f"Expected 1 fused pair, got {fusion_pass.matched_count}"

145 + )

146 +

147 + backend.check_before_ops(model.ops_in_model_before())

148 + backend.check_after_ops(model.ops_in_model_after())

vllm/_aiter_ops.py

+42

@@ -962,6 +962,37 @@ def _rocm_aiter_triton_add_rmsnorm_pad_fake(

962 return out, residual_out

963

964

965 +def _fused_mla_dual_rms_norm_impl(

966 + x1: torch.Tensor,

967 + x1_weight: torch.Tensor,

968 + x2: torch.Tensor,

969 + x2_weight: torch.Tensor,

970 + x1_epsilon: float,

971 + x2_epsilon: float,

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

973 + from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm

974 +

975 + return fused_qk_rmsnorm(

976 + q=x1,

977 + q_weight=x1_weight,

978 + q_eps=x1_epsilon,

979 + k=x2,

980 + k_weight=x2_weight,

981 + k_eps=x2_epsilon,

982 + )

983 +

984 +

985 +def _fused_mla_dual_rms_norm_fake(

986 + x1: torch.Tensor,

987 + x1_weight: torch.Tensor,

988 + x2: torch.Tensor,

989 + x2_weight: torch.Tensor,

990 + x1_epsilon: float,

991 + x2_epsilon: float,

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

993 + return (torch.empty_like(x1), torch.empty_like(x2))

994 +

995 +

996 def _rocm_aiter_gemm_a8wfp4_impl(

997 x: torch.Tensor,

998 w: torch.Tensor,

@@ -1491,6 +1522,13 @@ class rocm_aiter_ops:

1522 fake_impl=_triton_rotary_embedding_fake,

1523 )

1524

1525 + direct_register_custom_op(

1526 + op_name="fused_mla_dual_rms_norm",

1527 + op_func=_fused_mla_dual_rms_norm_impl,

1528 + mutates_args=[],

1529 + fake_impl=_fused_mla_dual_rms_norm_fake,

1530 + )

1531 +

1532 _OPS_REGISTERED = True

1533

1534 @staticmethod

@@ -1537,6 +1575,10 @@ class rocm_aiter_ops:

1575 def get_triton_rotary_embedding_op() -> OpOverload:

1576 return torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default

1577

1578 + @staticmethod

1579 + def get_fused_mla_dual_rms_norm_op() -> OpOverload:

1580 + return torch.ops.vllm.fused_mla_dual_rms_norm.default

1581 +

1582 @staticmethod

1583 def rms_norm(

1584 x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float

vllm/compilation/passes/fusion/rocm_aiter_fusion.py

+106 -1

@@ -1,5 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 +from collections.abc import Callable

4 from typing import Any

6 import torch

@@ -7,6 +8,7 @@ import torch._inductor.pattern_matcher as pm

8 from torch import fx

9 from torch._inductor.pattern_matcher import PatternMatcherPass

11 +import vllm.ir.ops

12 import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401

13 from vllm._aiter_ops import rocm_aiter_ops

14 from vllm.config import VllmConfig

@@ -20,7 +22,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

22 from vllm.platforms import current_platform

24 from ..inductor_pass import enable_fake_mode

23 -from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

25 +from ..vllm_inductor_pass import (

26 + VllmFusionPatternMatcherPass,

27 + VllmInductorPass,

28 + VllmPatternMatcherPass,

29 + VllmPatternReplacement,

30 +)

31 from .act_quant_fusion import ActivationQuantPattern

32 from .matcher_utils import (

33 MatcherFusedAddRMSNorm,

@@ -512,3 +519,101 @@ class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass):

519

520 def uuid(self) -> str:

521 return VllmInductorPass.hash_source(self, AddAiterRMSNormPadPattern)

522 +

523 +

524 +class MLADualRMSNormPattern(

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

526 +):

527 + """

528 + Fuse paired q_a_layernorm + kv_a_layernorm in MLA attention into

529 + AITER's ``fused_qk_rmsnorm`` HIP kernel.

530 +

531 + Target FX-graph pattern (unfused, ``vllm_ir`` stage)::

532 +

533 + gemm -> split_with_sizes([q_dim, kv_dim])

534 + +-- q_c -> vllm_ir.rms_norm(q_c, q_w, eps)

535 + +-- kv_lora -> split_with_sizes([kv_c_dim, k_pe_dim])

536 + +-- kv_c -> vllm_ir.rms_norm(kv_c, kv_w, eps)

537 + +-- k_pe

538 +

539 + The pattern covers the connected subgraph rooted at the first

540 + ``split_with_sizes`` (which produces ``q_c`` and ``kv_lora``),

541 + through the two ``rms_norm`` calls, and the ``k_pe`` passthrough.

542 + """

543 +

544 + def __init__(self, epsilon: float) -> None:

545 + self._epsilon = epsilon

546 +

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

548 + q_dim, kv_c_dim, k_pe_dim = 8, 4, 2

549 + return [

550 + self.empty_bf16(5, q_dim + kv_c_dim + k_pe_dim),

551 + self.empty_bf16(q_dim),

552 + self.empty_bf16(kv_c_dim),

553 + ]

554 +

555 + @property

556 + def pattern(

557 + self,

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

559 + eps = self._epsilon

560 +

561 + def _pattern(

562 + projected: torch.Tensor,

563 + q_weight: torch.Tensor,

564 + kv_weight: torch.Tensor,

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

566 + q_dim = q_weight.shape[0]

567 + kv_dim = projected.shape[-1] - q_dim

568 + kv_c_dim = kv_weight.shape[0]

569 + k_pe_dim = kv_dim - kv_c_dim

570 + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1)

571 + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1)

572 + q_normed = vllm.ir.ops.rms_norm(q_c, q_weight, eps)

573 + kv_normed = vllm.ir.ops.rms_norm(kv_c, kv_weight, eps)

574 + return q_normed, kv_normed, k_pe

575 +

576 + return _pattern

577 +

578 + @property

579 + def replacement(

580 + self,

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

582 + eps = self._epsilon

583 +

584 + def _replacement(

585 + projected: torch.Tensor,

586 + q_weight: torch.Tensor,

587 + kv_weight: torch.Tensor,

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

589 + q_dim = q_weight.shape[0]

590 + kv_dim = projected.shape[-1] - q_dim

591 + kv_c_dim = kv_weight.shape[0]

592 + k_pe_dim = kv_dim - kv_c_dim

593 + q_c, kv_lora = projected.split([q_dim, kv_dim], dim=-1)

594 + kv_c, k_pe = kv_lora.split([kv_c_dim, k_pe_dim], dim=-1)

595 + q_normed, kv_normed = torch.ops.vllm.fused_mla_dual_rms_norm(

596 + q_c,

597 + q_weight,

598 + kv_c,

599 + kv_weight,

600 + eps,

601 + eps,

602 + )

603 + return q_normed, kv_normed, k_pe

604 +

605 + return _replacement

606 +

607 +

608 +class MLADualRMSNormFusionPass(VllmFusionPatternMatcherPass):

609 + """

610 + Post-grad PatternMatcher pass that fuses paired q / kv RMS norms in

611 + MLA attention into ``fused_mla_dual_rms_norm`` backed by aiter's

612 + ``fused_qk_rmsnorm`` HIP kernel.

613 + """

614 +

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

616 + super().__init__(config, "mla_dual_rms_norm_fusion_pass")

617 +

618 + for epsilon in [1e-5, 1e-6]:

619 + self.register(MLADualRMSNormPattern(epsilon))

vllm/compilation/passes/pass_manager.py

@@ -19,6 +19,7 @@ from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

20 if rocm_aiter_ops.is_enabled():

21 from .fusion.rocm_aiter_fusion import (

22 + MLADualRMSNormFusionPass,

23 RocmAiterRMSNormQuantFusionPass,

24 RocmAiterSiluMulFp8GroupQuantFusionPass,

25 RocmAiterTritonAddRMSNormPadFusionPass,

@@ -155,6 +156,9 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

156 if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():

157 self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]

158

159 + if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():

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

161 +

162 if self.pass_config.fuse_rope_kvcache:

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

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

vllm/config/compilation.py

@@ -142,6 +142,8 @@ class PassConfig:

142 # ROCm/AITER specific fusions

143 fuse_act_padding: bool = None # type: ignore[assignment]

144 """Fuse the custom RMSNorm + padding ops."""

145 + fuse_mla_dual_rms_norm: bool = None # type: ignore[assignment]

146 + """Fuse paired q/kv RMS norms in MLA attention."""

147 fuse_rope_kvcache: bool = None # type: ignore[assignment]

148 """Fuse the QK rope + KV cache ops."""

149

@@ -224,6 +226,7 @@ class PassConfig:

226 "fuse_gemm_comms",

227 "fuse_allreduce_rms",

228 "fuse_act_padding",

229 + "fuse_mla_dual_rms_norm",

230 "fuse_rope_kvcache",

231 mode="wrap",

232 )

@@ -270,6 +273,12 @@ class PassConfig:

273 "The fusion will be disabled."

274 )

275 self.fuse_act_padding = False

276 + if self.fuse_mla_dual_rms_norm and not current_platform.is_rocm():

277 + logger.warning_once(

278 + "MLA dual RMS norm fusion requires ROCm/AITER. "

279 + "The fusion will be disabled."

280 + )

281 + self.fuse_mla_dual_rms_norm = False

282 if self.fuse_rope_kvcache and not current_platform.is_rocm():

283 logger.warning_once(

284 "KV cache fusion currently only enabled on ROCm. "

vllm/config/vllm.py

+11

@@ -165,6 +165,13 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:

165 )

166

167

168 +def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool:

169 + """Enable MLA dual RMS norm fusion when AITer is available."""

170 + from vllm._aiter_ops import rocm_aiter_ops

171 +

172 + return rocm_aiter_ops.is_enabled()

173 +

174 +

175 OPTIMIZATION_LEVEL_00 = {

176 "compilation_config": {

177 "pass_config": {

@@ -175,6 +182,7 @@ OPTIMIZATION_LEVEL_00 = {

182 "enable_sp": False,

183 "fuse_gemm_comms": False,

184 "fuse_act_padding": False,

185 + "fuse_mla_dual_rms_norm": False,

186 "fuse_rope_kvcache": False,

187 },

188 "cudagraph_mode": CUDAGraphMode.NONE,

@@ -194,6 +202,7 @@ OPTIMIZATION_LEVEL_01 = {

202 "enable_sp": False,

203 "fuse_gemm_comms": False,

204 "fuse_act_padding": enable_norm_pad_fusion,

205 + "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,

206 "fuse_rope_kvcache": False,

207 },

208 "cudagraph_mode": CUDAGraphMode.PIECEWISE,

@@ -213,6 +222,7 @@ OPTIMIZATION_LEVEL_02 = {

222 "enable_sp": IS_DENSE,

223 "fuse_gemm_comms": IS_DENSE,

224 "fuse_act_padding": enable_norm_pad_fusion,

225 + "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,

226 "fuse_rope_kvcache": enable_rope_kvcache_fusion,

227 },

228 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,

@@ -232,6 +242,7 @@ OPTIMIZATION_LEVEL_03 = {

242 "enable_sp": IS_DENSE,

243 "fuse_gemm_comms": IS_DENSE,

244 "fuse_act_padding": enable_norm_pad_fusion,

245 + "fuse_mla_dual_rms_norm": enable_mla_dual_rms_norm_fusion,

246 "fuse_rope_kvcache": enable_rope_kvcache_fusion,

247 },

248 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,