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 3, 13:00 UTC
[ROCm][Bugfix] Fix fused_mla_dual_rms_norm for AITER API rename _fused_qk_rmsnorm (#42606)
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Rita Brugarolas committed
May 15, 2026 at 13:50 UTC
bd9dbe60601c986b50260f299fe279d057d7d89f
2 files changed
+56
-16
vllm/_aiter_ops.py
+50
-14
@@ -414,17 +414,32 @@ _AITER_HAS_FUSED_QK_RMSNORM: bool | None = None
414
415
416
def check_aiter_fused_qk_rmsnorm() -> bool:
417
- """Check if aiter provides fused_qk_rmsnorm (requires AITer >= PR #2442)."""
417
+ """Check if aiter provides fused_qk_rmsnorm.
418
+
419
+ Supports both the new private name ``_fused_qk_rmsnorm``
420
+ (AITER >= PR #2958) and the old public name ``fused_qk_rmsnorm``
421
+ (AITER >= PR #2442).
422
+
423
+ TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
424
+ AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
425
+ """
426
global _AITER_HAS_FUSED_QK_RMSNORM
427
if _AITER_HAS_FUSED_QK_RMSNORM is None:
428
try:
429
from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
422
- fused_qk_rmsnorm,
430
+ _fused_qk_rmsnorm,
431
)
432
433
_AITER_HAS_FUSED_QK_RMSNORM = True
434
except (ImportError, ModuleNotFoundError, AttributeError):
427
- _AITER_HAS_FUSED_QK_RMSNORM = False
435
+ try:
436
+ from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
437
+ fused_qk_rmsnorm,
438
+ )
439
+
440
+ _AITER_HAS_FUSED_QK_RMSNORM = True
441
+ except (ImportError, ModuleNotFoundError, AttributeError):
442
+ _AITER_HAS_FUSED_QK_RMSNORM = False
443
return _AITER_HAS_FUSED_QK_RMSNORM
444
445
@@ -1066,21 +1081,42 @@ def _fused_mla_dual_rms_norm_impl(
1081
x2_epsilon: float,
1082
) -> tuple[torch.Tensor, torch.Tensor]:
1083
try:
1069
- from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm
1070
- except (ImportError, ModuleNotFoundError) as exc:
1084
+ import aiter.ops.fused_qk_norm_rope_cache_quant as aiter_ops
1085
+ except (ImportError, ModuleNotFoundError, AttributeError) as exc:
1086
raise ImportError(
1072
- "fused_qk_rmsnorm requires a newer AITer version "
1073
- "(>= PR #2442). Please upgrade aiter or disable the "
1087
+ "fused_qk_rmsnorm requires AITer >= PR #2442. "
1088
+ "Please upgrade aiter or disable the "
1089
"fuse_mla_dual_rms_norm pass."
1090
) from exc
1091
1077
- return fused_qk_rmsnorm(
1078
- q=x1,
1079
- q_weight=x1_weight,
1080
- q_eps=x1_epsilon,
1081
- k=x2,
1082
- k_weight=x2_weight,
1083
- k_eps=x2_epsilon,
1092
+ if hasattr(aiter_ops, "_fused_qk_rmsnorm"):
1093
+ return aiter_ops._fused_qk_rmsnorm(
1094
+ q_out=None,
1095
+ q=x1,
1096
+ q_weight=x1_weight,
1097
+ q_eps=x1_epsilon,
1098
+ k_out=None,
1099
+ k=x2,
1100
+ k_weight=x2_weight,
1101
+ k_eps=x2_epsilon,
1102
+ )
1103
+
1104
+ # TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
1105
+ # AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
1106
+ if hasattr(aiter_ops, "fused_qk_rmsnorm"):
1107
+ return aiter_ops.fused_qk_rmsnorm(
1108
+ q=x1,
1109
+ q_weight=x1_weight,
1110
+ q_eps=x1_epsilon,
1111
+ k=x2,
1112
+ k_weight=x2_weight,
1113
+ k_eps=x2_epsilon,
1114
+ )
1115
+
1116
+ raise ImportError(
1117
+ "fused_qk_rmsnorm requires AITer >= PR #2442. "
1118
+ "Please upgrade aiter or disable the "
1119
+ "fuse_mla_dual_rms_norm pass."
1120
)
1121
1122
vllm/compilation/passes/pass_manager.py
+6
-2
@@ -7,7 +7,7 @@ from typing import Any, ParamSpec, TypeVar
7
from torch import fx as fx
8
9
from vllm import envs
10
-from vllm._aiter_ops import rocm_aiter_ops
10
+from vllm._aiter_ops import check_aiter_fused_qk_rmsnorm, rocm_aiter_ops
11
from vllm.compilation.passes.utility.post_cleanup import PostCleanupPass
12
from vllm.config import VllmConfig, set_current_vllm_config
13
from vllm.logger import init_logger
@@ -169,7 +169,11 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
169
if rocm_aiter_ops.is_enabled():
170
self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
171
172
- if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
172
+ if (
173
+ self.pass_config.fuse_mla_dual_rms_norm
174
+ and rocm_aiter_ops.is_enabled()
175
+ and check_aiter_fused_qk_rmsnorm()
176
+ ):
177
self.passes += [MLADualRMSNormFusionPass(config)]
178
179
if self.pass_config.fuse_rope_kvcache: