Fix ROCm architecture import on non-ROCm platforms (#51357)
Signed-off-by: Wu, Xiaochang <xiaochang.wu@intel.com> Signed-off-by: Xiaochang Wu <xiaochang.wu@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Xiaochang Wu committed
Aug 7, 2026 at 20:32 UTC
d5aae2b4641c5091e604e235617e51e60a564710
2 files changed
+20
-9
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+7
-3
@@ -1286,7 +1286,11 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1286
1287
Supports DeepGEMM, FlashInfer, TRTLLM MXFP8, Triton and Marlin backends.
1288
"""
1289
- from vllm.platforms.rocm import on_gfx1250
1289
+ is_gfx1250 = False
1290
+ if current_platform.is_rocm():
1291
+ from vllm.platforms.rocm import on_gfx1250
1292
+
1293
+ is_gfx1250 = on_gfx1250()
1294
1295
if mxfp4_backend == Mxfp4MoeBackend.DEEPGEMM_MXFP4:
1296
w13_weight_scale, w2_weight_scale = _pack_deepgemm_mxfp4_scales(
@@ -1459,7 +1463,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1463
w2_bias,
1464
)
1465
1462
- elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not on_gfx1250():
1466
+ elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and not is_gfx1250:
1467
# Initially introduced for DeepSeekV4
1468
1469
if w13_bias is not None:
@@ -1520,7 +1524,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
1524
)
1525
1526
elif mxfp4_backend in TRITON_BACKENDS or (
1523
- mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
1527
+ mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
1528
):
1529
from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1530
vllm/model_executor/layers/quantization/mxfp4.py
+13
-6
@@ -39,6 +39,7 @@ from vllm.model_executor.layers.quantization.base_config import (
39
)
40
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
41
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
42
+from vllm.platforms import current_platform
43
44
logger = init_logger(__name__)
45
@@ -482,8 +483,6 @@ def _use_k3_situ_aiter(moe: FusedMoEConfig) -> bool:
483
K3 is weight-only MXFP4 (W4A16) with SiTU activation, which the generic
484
MXFP4 backend selector does not cover; route it to AITER on gfx950.
485
"""
485
- from vllm.platforms import current_platform
486
-
486
if not current_platform.is_rocm():
487
return False
488
from vllm._aiter_ops import rocm_aiter_ops
@@ -739,10 +738,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
738
739
# For TRITON backends, weights are wrapped tensors from triton_kernels
740
# that don't support .detach(). Manually assign parameters.
742
- from vllm.platforms.rocm import on_gfx1250
741
+ is_gfx1250 = False
742
+ if current_platform.is_rocm():
743
+ from vllm.platforms.rocm import on_gfx1250
744
+
745
+ is_gfx1250 = on_gfx1250()
746
747
uses_triton_weight_format = self.mxfp4_backend in TRITON_BACKENDS or (
745
- self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
748
+ self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
749
)
750
if not uses_triton_weight_format:
751
replace_parameter(layer, "w13_weight", w13)
@@ -830,10 +833,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
833
w2_bias = getattr(layer, "w2_bias", None)
834
swiglu_limit = getattr(layer, "swiglu_limit", None)
835
833
- from vllm.platforms.rocm import on_gfx1250
836
+ is_gfx1250 = False
837
+ if current_platform.is_rocm():
838
+ from vllm.platforms.rocm import on_gfx1250
839
+
840
+ is_gfx1250 = on_gfx1250()
841
842
if self.mxfp4_backend in TRITON_BACKENDS or (
836
- self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and on_gfx1250()
843
+ self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16 and is_gfx1250
844
):
845
# TRITON backends free w13/w2_weight_scale after swizzling; the
846
# swizzled scales live inside the precision configs instead.