[ROCm] Remove redundant AITER fused_qk_rmsnorm probe (avoids config-time HIP init) (#47992)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Jul 22, 2026 at 05:58 UTC
387189c42997b27e2c04b5d97ef8190ffa2bf909
3 files changed
+16
-77
vllm/_aiter_ops.py
+11
-68
@@ -409,37 +409,6 @@ def _rocm_aiter_fused_topk_fake(
409
410
# Cache whether aiter supports FP8 MLA parameters
411
_AITER_MLA_SUPPORTS_FP8: bool | None = None
412
-_AITER_HAS_FUSED_QK_RMSNORM: bool | None = None
413
-
414
-
415
-def check_aiter_fused_qk_rmsnorm() -> bool:
416
- """Check if aiter provides fused_qk_rmsnorm.
417
-
418
- Supports both the new private name ``_fused_qk_rmsnorm``
419
- (AITER >= PR #2958) and the old public name ``fused_qk_rmsnorm``
420
- (AITER >= PR #2442).
421
-
422
- TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
423
- AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
424
- """
425
- global _AITER_HAS_FUSED_QK_RMSNORM
426
- if _AITER_HAS_FUSED_QK_RMSNORM is None:
427
- try:
428
- from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
429
- _fused_qk_rmsnorm,
430
- )
431
-
432
- _AITER_HAS_FUSED_QK_RMSNORM = True
433
- except (ImportError, ModuleNotFoundError, AttributeError):
434
- try:
435
- from aiter.ops.fused_qk_norm_rope_cache_quant import ( # noqa: F401
436
- fused_qk_rmsnorm,
437
- )
438
-
439
- _AITER_HAS_FUSED_QK_RMSNORM = True
440
- except (ImportError, ModuleNotFoundError, AttributeError):
441
- _AITER_HAS_FUSED_QK_RMSNORM = False
442
- return _AITER_HAS_FUSED_QK_RMSNORM
412
413
414
def _check_aiter_mla_fp8_support() -> bool:
@@ -1267,43 +1236,17 @@ def _fused_mla_dual_rms_norm_impl(
1236
x1_epsilon: float,
1237
x2_epsilon: float,
1238
) -> tuple[torch.Tensor, torch.Tensor]:
1270
- try:
1271
- import aiter.ops.fused_qk_norm_rope_cache_quant as aiter_ops
1272
- except (ImportError, ModuleNotFoundError, AttributeError) as exc:
1273
- raise ImportError(
1274
- "fused_qk_rmsnorm requires AITer >= PR #2442. "
1275
- "Please upgrade aiter or disable the "
1276
- "fuse_mla_dual_rms_norm pass."
1277
- ) from exc
1278
-
1279
- if hasattr(aiter_ops, "_fused_qk_rmsnorm"):
1280
- return aiter_ops._fused_qk_rmsnorm(
1281
- q_out=None,
1282
- q=x1,
1283
- q_weight=x1_weight,
1284
- q_eps=x1_epsilon,
1285
- k_out=None,
1286
- k=x2,
1287
- k_weight=x2_weight,
1288
- k_eps=x2_epsilon,
1289
- )
1290
-
1291
- # TODO(rbrugaro-amd): remove the legacy fused_qk_rmsnorm path once
1292
- # AITER stabilizes the API (https://github.com/ROCm/aiter/issues/3207).
1293
- if hasattr(aiter_ops, "fused_qk_rmsnorm"):
1294
- return aiter_ops.fused_qk_rmsnorm(
1295
- q=x1,
1296
- q_weight=x1_weight,
1297
- q_eps=x1_epsilon,
1298
- k=x2,
1299
- k_weight=x2_weight,
1300
- k_eps=x2_epsilon,
1301
- )
1302
-
1303
- raise ImportError(
1304
- "fused_qk_rmsnorm requires AITer >= PR #2442. "
1305
- "Please upgrade aiter or disable the "
1306
- "fuse_mla_dual_rms_norm pass."
1239
+ from aiter.ops.fused_qk_norm_rope_cache_quant import _fused_qk_rmsnorm
1240
+
1241
+ return _fused_qk_rmsnorm(
1242
+ q_out=None,
1243
+ q=x1,
1244
+ q_weight=x1_weight,
1245
+ q_eps=x1_epsilon,
1246
+ k_out=None,
1247
+ k=x2,
1248
+ k_weight=x2_weight,
1249
+ k_eps=x2_epsilon,
1250
)
1251
1252
vllm/compilation/passes/pass_manager.py
+2
-6
@@ -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 check_aiter_fused_qk_rmsnorm, rocm_aiter_ops
10
+from vllm._aiter_ops import 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
@@ -177,11 +177,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
177
self.passes += [ScatterSplitReplacementPass(config)]
178
self.passes += [QkNormRopeKvCacheFusionPass(config)]
179
180
- if (
181
- self.pass_config.fuse_mla_dual_rms_norm
182
- and rocm_aiter_ops.is_enabled()
183
- and check_aiter_fused_qk_rmsnorm()
184
- ):
180
+ if self.pass_config.fuse_mla_dual_rms_norm and rocm_aiter_ops.is_enabled():
181
self.passes += [MLADualRMSNormFusionPass(config)]
182
183
if self.pass_config.fuse_rope_kvcache:
vllm/config/vllm.py
+3
-3
@@ -187,10 +187,10 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
187
188
189
def enable_mla_dual_rms_norm_fusion(cfg: "VllmConfig") -> bool:
190
- """Enable MLA dual RMS norm fusion when AITer has fused_qk_rmsnorm."""
191
- from vllm._aiter_ops import check_aiter_fused_qk_rmsnorm, rocm_aiter_ops
190
+ """Enable MLA dual RMS norm fusion on ROCm with AITER."""
191
+ from vllm._aiter_ops import rocm_aiter_ops
192
193
- return rocm_aiter_ops.is_enabled() and check_aiter_fused_qk_rmsnorm()
193
+ return rocm_aiter_ops.is_enabled()
194
195
196
def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: