[ROCm][CI] Avoid HIP init at config time via lazy aiter import in Quark OCP-MX (#48015)
Signed-off-by: Dino Music <Dino.Music@amd.com>
music-dino committed
Jul 16, 2026 at 18:47 UTC
7d56fe2adc93d20a01be30eaa275e156ecdd4468
1 file changed
+29
-21
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+29
-21
@@ -9,7 +9,7 @@ from typing import Any
9
import torch
10
import torch.nn.functional as F
11
12
-from vllm._aiter_ops import rocm_aiter_ops
12
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
15
dequant_mxfp4,
@@ -36,19 +36,15 @@ from .quark_scheme import QuarkScheme
36
logger = init_logger(__name__)
37
38
39
-try:
40
- from aiter.ops.shuffle import shuffle_weight
41
- from aiter.ops.triton.gemm_afp4wfp4 import (
42
- gemm_afp4wfp4,
43
- gemm_afp4wfp4_preshuffled_weight_scales,
44
- )
45
- from aiter.ops.triton.quant import dynamic_mxfp4_quant
46
-
39
+# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP
40
+# which can force the engine core to spawn instead of fork.
41
+# is_aiter_found_and_supported() checks platform + arch + library availability via
42
+# find_spec/amdsmi, so it stays HIP-free.
43
+# Actual aiter imports are deferred to the functions/methods that need them,
44
+# where HIP initialization is expected.
45
+if is_aiter_found_and_supported():
46
from vllm.utils.torch_utils import direct_register_custom_op
47
49
- if rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled():
50
- from aiter import gemm_a4w4, per_1x32_f4_quant_hip
51
-
48
def gemm_with_dynamic_quant(
49
x: torch.Tensor,
50
weight: torch.Tensor,
@@ -57,6 +53,15 @@ try:
53
out_dtype: torch.dtype | None = torch.bfloat16,
54
x_scales: torch.Tensor | None = None,
55
) -> torch.Tensor:
56
+ from aiter.ops.triton.gemm_afp4wfp4 import (
57
+ gemm_afp4wfp4,
58
+ gemm_afp4wfp4_preshuffled_weight_scales,
59
+ )
60
+ from aiter.ops.triton.quant import dynamic_mxfp4_quant
61
+
62
+ if rocm_use_aiter_fp4_asm_gemm:
63
+ from aiter import gemm_a4w4, per_1x32_f4_quant_hip
64
+
65
M = x.shape[0]
66
N = weight.shape[0]
67
K = weight.shape[1]
@@ -137,13 +142,12 @@ try:
142
fake_impl=gemm_with_dynamic_quant_fake,
143
dispatch_key=current_platform.dispatch_key,
144
)
140
-except (ImportError, AttributeError, RuntimeError):
141
- if current_platform.is_rocm():
142
- logger.warning(
143
- "AITER is not found or QuarkOCP_MX is not supported on the current "
144
- "platform. QuarkOCP_MX quantization will not be available."
145
- )
146
- dynamic_mxfp4_quant = gemm_afp4wfp4 = None
145
+elif current_platform.is_rocm():
146
+ logger.warning(
147
+ "AITER is not found or not supported on the current platform, "
148
+ "QuarkOCP_MX will fall back to emulation."
149
+ "Native MXFP4/MXFP6 acceleration will not be available."
150
+ )
151
152
153
class QuarkOCP_MX(QuarkScheme):
@@ -211,8 +215,8 @@ class QuarkOCP_MX(QuarkScheme):
215
rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()
216
)
217
214
- if not self.emulate and (dynamic_mxfp4_quant is None or gemm_afp4wfp4 is None):
215
- # Currently need these kernels if not emulating
218
+ if not self.emulate and not is_aiter_found_and_supported():
219
+ # Currently need AITER kernels if not emulating
220
raise NotImplementedError(
221
f"{self.__class__.__name__} requires AITER to be installed "
222
"for non-emulation mode! Please refer to "
@@ -261,6 +265,8 @@ class QuarkOCP_MX(QuarkScheme):
265
def process_dynamic_mxfp4_weights_after_loading(
266
self, layer: torch.nn.Module
267
) -> None:
268
+ from aiter.ops.triton.quant import dynamic_mxfp4_quant
269
+
270
w_q, w_s = dynamic_mxfp4_quant(layer.weight)
271
layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False)
272
layer.weight = torch.nn.Parameter(w_q, requires_grad=False)
@@ -279,6 +285,8 @@ class QuarkOCP_MX(QuarkScheme):
285
if self.dynamic_mxfp4_quant:
286
self.process_dynamic_mxfp4_weights_after_loading(layer)
287
elif self.rocm_use_aiter_fp4_asm_gemm:
288
+ from aiter.ops.shuffle import shuffle_weight
289
+
290
# shuffle weight scale
291
weight_scale_shuffle = layer.weight_scale.data
292
sm, sn = weight_scale_shuffle.shape