diff --git a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py index 70a7e81cc4..ea63d1bcf1 100644 --- a/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py +++ b/vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py @@ -9,7 +9,7 @@ from typing import Any import torch import torch.nn.functional as F -from vllm._aiter_ops import rocm_aiter_ops +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.mxfp4_utils import ( dequant_mxfp4, @@ -36,19 +36,15 @@ from .quark_scheme import QuarkScheme logger = init_logger(__name__) -try: - from aiter.ops.shuffle import shuffle_weight - from aiter.ops.triton.gemm_afp4wfp4 import ( - gemm_afp4wfp4, - gemm_afp4wfp4_preshuffled_weight_scales, - ) - from aiter.ops.triton.quant import dynamic_mxfp4_quant - +# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP +# which can force the engine core to spawn instead of fork. +# is_aiter_found_and_supported() checks platform + arch + library availability via +# find_spec/amdsmi, so it stays HIP-free. +# Actual aiter imports are deferred to the functions/methods that need them, +# where HIP initialization is expected. +if is_aiter_found_and_supported(): from vllm.utils.torch_utils import direct_register_custom_op - if rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled(): - from aiter import gemm_a4w4, per_1x32_f4_quant_hip - def gemm_with_dynamic_quant( x: torch.Tensor, weight: torch.Tensor, @@ -57,6 +53,15 @@ try: out_dtype: torch.dtype | None = torch.bfloat16, x_scales: torch.Tensor | None = None, ) -> torch.Tensor: + from aiter.ops.triton.gemm_afp4wfp4 import ( + gemm_afp4wfp4, + gemm_afp4wfp4_preshuffled_weight_scales, + ) + from aiter.ops.triton.quant import dynamic_mxfp4_quant + + if rocm_use_aiter_fp4_asm_gemm: + from aiter import gemm_a4w4, per_1x32_f4_quant_hip + M = x.shape[0] N = weight.shape[0] K = weight.shape[1] @@ -137,13 +142,12 @@ try: fake_impl=gemm_with_dynamic_quant_fake, dispatch_key=current_platform.dispatch_key, ) -except (ImportError, AttributeError, RuntimeError): - if current_platform.is_rocm(): - logger.warning( - "AITER is not found or QuarkOCP_MX is not supported on the current " - "platform. QuarkOCP_MX quantization will not be available." - ) - dynamic_mxfp4_quant = gemm_afp4wfp4 = None +elif current_platform.is_rocm(): + logger.warning( + "AITER is not found or not supported on the current platform, " + "QuarkOCP_MX will fall back to emulation." + "Native MXFP4/MXFP6 acceleration will not be available." + ) class QuarkOCP_MX(QuarkScheme): @@ -211,8 +215,8 @@ class QuarkOCP_MX(QuarkScheme): rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled() ) - if not self.emulate and (dynamic_mxfp4_quant is None or gemm_afp4wfp4 is None): - # Currently need these kernels if not emulating + if not self.emulate and not is_aiter_found_and_supported(): + # Currently need AITER kernels if not emulating raise NotImplementedError( f"{self.__class__.__name__} requires AITER to be installed " "for non-emulation mode! Please refer to " @@ -261,6 +265,8 @@ class QuarkOCP_MX(QuarkScheme): def process_dynamic_mxfp4_weights_after_loading( self, layer: torch.nn.Module ) -> None: + from aiter.ops.triton.quant import dynamic_mxfp4_quant + w_q, w_s = dynamic_mxfp4_quant(layer.weight) layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False) layer.weight = torch.nn.Parameter(w_q, requires_grad=False) @@ -279,6 +285,8 @@ class QuarkOCP_MX(QuarkScheme): if self.dynamic_mxfp4_quant: self.process_dynamic_mxfp4_weights_after_loading(layer) elif self.rocm_use_aiter_fp4_asm_gemm: + from aiter.ops.shuffle import shuffle_weight + # shuffle weight scale weight_scale_shuffle = layer.weight_scale.data sm, sn = weight_scale_shuffle.shape