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] Hotfix: guard MLA dual RMS norm fusion against older AITer versions (#40386)
Signed-off-by: Rita Brugarolas Brufau <rita.brugarolasbrufau@amd.com>
Rita Brugarolas committed
Apr 20, 2026 at 14:20 UTC
21b086d0aaea612bdf3fef0c313513eac8350083
2 files changed
+27
-4
vllm/_aiter_ops.py
+24
-1
@@ -376,6 +376,22 @@ def _rocm_aiter_fused_topk_fake(
376
377
# Cache whether aiter supports FP8 MLA parameters
378
_AITER_MLA_SUPPORTS_FP8: bool | None = None
379
+_AITER_HAS_FUSED_QK_RMSNORM: bool | None = None
380
+
381
+
382
+def check_aiter_fused_qk_rmsnorm() -> bool:
383
+ """Check if aiter provides fused_qk_rmsnorm (requires AITer >= PR #2442)."""
384
+ global _AITER_HAS_FUSED_QK_RMSNORM
385
+ if _AITER_HAS_FUSED_QK_RMSNORM is None:
386
+ try:
387
+ from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
388
+ fused_qk_rmsnorm,
389
+ )
390
+
391
+ _AITER_HAS_FUSED_QK_RMSNORM = True
392
+ except (ImportError, ModuleNotFoundError, AttributeError):
393
+ _AITER_HAS_FUSED_QK_RMSNORM = False
394
+ return _AITER_HAS_FUSED_QK_RMSNORM
395
396
397
def _check_aiter_mla_fp8_support() -> bool:
@@ -970,7 +986,14 @@ def _fused_mla_dual_rms_norm_impl(
986
x1_epsilon: float,
987
x2_epsilon: float,
988
) -> tuple[torch.Tensor, torch.Tensor]:
973
- from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm
989
+ try:
990
+ from aiter.ops.fused_qk_norm_rope_cache_quant import fused_qk_rmsnorm
991
+ except (ImportError, ModuleNotFoundError) as exc:
992
+ raise ImportError(
993
+ "fused_qk_rmsnorm requires a newer AITer version "
994
+ "(>= PR #2442). Please upgrade aiter or disable the "
995
+ "fuse_mla_dual_rms_norm pass."
996
+ ) from exc
997
998
return fused_qk_rmsnorm(
999
q=x1,
vllm/config/vllm.py
+3
-3
@@ -166,10 +166,10 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
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
169
+ """Enable MLA dual RMS norm fusion when AITer has fused_qk_rmsnorm."""
170
+ from vllm._aiter_ops import check_aiter_fused_qk_rmsnorm, rocm_aiter_ops
171
172
- return rocm_aiter_ops.is_enabled()
172
+ return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm()
173
174
175
OPTIMIZATION_LEVEL_00 = {