@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7aea73d83d

[ROCm][Quark][6/N] Use MXFP4 linear kernel abstraction for `aiter` backend (#49348)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>

fxmarty-amd committed Jul 28, 2026 at 04:35 UTC 7aea73d83d6064449ff8147899de161e6eb68a20
6 files changed +370 -166
tests/evals/gsm8k/configs/Qwen3-1.7B-MXFP4.yaml new
+5
@@ -0,0 +1,5 @@
1 +model_name: "amd-quark/Qwen3-1.7B-MXFP4"
2 +accuracy_threshold: 0.27
3 +num_questions: 1319
4 +num_fewshot: 5
5 +server_args: "--enforce-eager --max-model-len 4096"
tests/kernels/quantization/test_mxfp4_kernel_selection.py new
+130
@@ -0,0 +1,130 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for MXFP4 linear kernel selection logic (CPU-only)
4 +
5 +Run `pytest tests/kernels/quantization/test_mxfp4_kernel_selection.py`.
6 +"""
7 +
8 +from unittest.mock import patch
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.kernels.linear import (
14 + AiterMxfp4LinearKernel,
15 + MxFp4LinearKernel,
16 + MxFp4LinearLayerConfig,
17 + init_mxfp4_linear_kernel,
18 + register_linear_kernel,
19 +)
20 +from vllm.platforms import PlatformEnum
21 +
22 +pytestmark = pytest.mark.cpu_test
23 +
24 +
25 +def test_can_implement_is_abstract():
26 + """Test that can_implement()/is_supported() are properly defined."""
27 + assert hasattr(MxFp4LinearKernel, "can_implement")
28 + assert hasattr(MxFp4LinearKernel, "is_supported")
29 +
30 +
31 +def test_aiter_kernel_is_supported_requires_native_mx_support():
32 + """AiterMxfp4LinearKernel must not be selected on platforms without
33 + native MX compute, even if AITER itself is importable."""
34 + with patch(
35 + "vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform.supports_mx",
36 + return_value=False,
37 + ):
38 + is_supported, reason = AiterMxfp4LinearKernel.is_supported()
39 + assert not is_supported
40 + assert reason
41 +
42 +
43 +class OOTMxFp4LinearKernel(MxFp4LinearKernel):
44 + @classmethod
45 + def is_supported(
46 + cls, compute_capability: int | None = None
47 + ) -> tuple[bool, str | None]:
48 + return True, None
49 +
50 + @classmethod
51 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
52 + return True, None
53 +
54 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
55 + pass
56 +
57 + def apply_weights(
58 + self,
59 + layer: torch.nn.Module,
60 + x: torch.Tensor,
61 + bias: torch.Tensor | None = None,
62 + ) -> torch.Tensor:
63 + pass
64 +
65 +
66 +@patch("vllm.model_executor.kernels.linear.current_platform")
67 +def test_init_mxfp4_linear_kernel_dispatches_to_registered_kernel(platform_mock):
68 + """init_mxfp4_linear_kernel should select a registered kernel that
69 + reports itself as supported, and construct it with a fresh config."""
70 + platform_mock._enum = PlatformEnum.OOT
71 + register_linear_kernel(OOTMxFp4LinearKernel, PlatformEnum.OOT, "mxfp4")
72 +
73 + kernel = init_mxfp4_linear_kernel()
74 +
75 + assert isinstance(kernel, OOTMxFp4LinearKernel)
76 + assert kernel.config == MxFp4LinearLayerConfig()
77 +
78 +
79 +class UnsupportedMxFp4LinearKernel(MxFp4LinearKernel):
80 + @classmethod
81 + def is_supported(
82 + cls, compute_capability: int | None = None
83 + ) -> tuple[bool, str | None]:
84 + return False, "never supported"
85 +
86 + @classmethod
87 + def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
88 + return True, None
89 +
90 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
91 + pass
92 +
93 + def apply_weights(
94 + self,
95 + layer: torch.nn.Module,
96 + x: torch.Tensor,
97 + bias: torch.Tensor | None = None,
98 + ) -> torch.Tensor:
99 + pass
100 +
101 +
102 +@patch("vllm.model_executor.kernels.linear.current_platform")
103 +def test_init_mxfp4_linear_kernel_raises_when_no_kernel_matches(platform_mock):
104 + platform_mock._enum = PlatformEnum.UNSPECIFIED
105 + register_linear_kernel(
106 + UnsupportedMxFp4LinearKernel, PlatformEnum.UNSPECIFIED, "mxfp4"
107 + )
108 +
109 + with pytest.raises(ValueError, match="Failed to find a kernel"):
110 + init_mxfp4_linear_kernel()
111 +
112 +
113 +@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.is_aiter_found_and_supported")
114 +@patch("vllm.model_executor.kernels.linear.mxfp4.aiter.current_platform")
115 +@patch("vllm.model_executor.kernels.linear.current_platform")
116 +def test_init_mxfp4_linear_kernel_raises_on_rocm_without_aiter(
117 + linear_platform_mock, aiter_platform_mock, is_aiter_found_and_supported_mock
118 +):
119 + """On ROCm, the only registered MXFP4 linear kernel is AITER-based.
120 + If AITER is not found/supported, no kernel should be selected."""
121 + linear_platform_mock._enum = PlatformEnum.ROCM
122 + aiter_platform_mock.supports_mx.return_value = True
123 + is_aiter_found_and_supported_mock.return_value = False
124 +
125 + with pytest.raises(
126 + ValueError,
127 + match="(?s)Failed to find a kernel.*"
128 + "AITER not found or not supported on the current platform",
129 + ):
130 + init_mxfp4_linear_kernel()
tests/quantization/test_quark.py
+42
@@ -17,6 +17,7 @@ import pytest
17 import torch
18 from packaging import version
19
20 +from vllm._aiter_ops import is_aiter_found_and_supported
21 from vllm.model_executor.layers.quantization.quark.quark import ( # noqa: E501
22 QuarkLinearMethod,
23 QuarkW8A8Fp8,
@@ -26,6 +27,9 @@ from vllm.model_executor.layers.quantization.quark.quark_moe import ( # noqa: E
27 QuarkW4A8Fp8MoEMethod,
28 QuarkW8A8Int8MoEMethod,
29 )
30 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
31 + quant_dequant_mxfp4,
32 +)
33 from vllm.model_executor.layers.quantization.utils.quant_utils import (
34 is_layer_skipped,
35 )
@@ -52,6 +56,8 @@ QUARK_MXFP4_AVAILABLE = find_spec("quark") is not None and version.parse(
56 importlib.metadata.version("amd-quark")
57 ) >= version.parse(QUARK_MXFP4_MIN_VERSION)
58
59 +AITER_AVAILABLE = is_aiter_found_and_supported()
60 +
61 DEVICE_TYPE = current_platform.device_type
62
63 if QUARK_MXFP4_AVAILABLE:
@@ -487,6 +493,42 @@ def test_mxfp4_dequant_kernel_match_quark(
493 assert torch.equal(out_hip, out_torch)
494
495
496 +@pytest.mark.skipif(
497 + not QUARK_MXFP4_AVAILABLE,
498 + reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
499 +)
500 +@pytest.mark.skipif(
501 + not AITER_AVAILABLE,
502 + reason="AITER is not found or not supported on the current platform",
503 +)
504 +@pytest.mark.parametrize("float_dtype", [torch.bfloat16, torch.float16])
505 +@pytest.mark.parametrize("scalings", [[2.3, 0.03, 7.3, 0.1, 0.004, 17.3, 1e4, 1e-4]])
506 +def test_mxfp4_dynamic_quant_match_quark(
507 + float_dtype: torch.dtype, scalings: list[float]
508 +):
509 + """`AiterMxfp4LinearKernel` quantizes weights dynamically through AITER's
510 + `dynamic_mxfp4_quant`, while the emulation path quantizes/dequantizes
511 + through Quark's `qdq_mxfp4`. Check that both agree on the same input.
512 + """
513 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
514 +
515 + torch.manual_seed(0)
516 +
517 + hidden_size = 32 * 64
518 + inp = (torch.rand(48, hidden_size, dtype=float_dtype, device=DEVICE_TYPE) - 0.5) * 2
519 + for i in range(hidden_size // 32):
520 + inp[:, i * 32 : (i + 1) * 32] = (
521 + inp[:, i * 32 : (i + 1) * 32] * scalings[i % len(scalings)]
522 + )
523 +
524 + x_q, x_s = dynamic_mxfp4_quant(inp)
525 + out_dynamic_quant = dq_mxfp4_torch(x_q, x_s, float_dtype)
526 +
527 + out_quark_qdq = quant_dequant_mxfp4(inp)
528 +
529 + assert torch.equal(out_dynamic_quant, out_quark_qdq)
530 +
531 +
532 # Unit tests for ``is_layer_skipped`` fused-name handling.
533
534 FUSED_MAPPING = {
vllm/model_executor/kernels/linear/__init__.py
+8
@@ -74,6 +74,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import (
74 MxFp4LinearKernel,
75 MxFp4LinearLayerConfig,
76 )
77 +from vllm.model_executor.kernels.linear.mxfp4.aiter import (
78 + AiterMxfp4LinearKernel,
79 +)
80 from vllm.model_executor.kernels.linear.mxfp4.flashinfer import (
81 FlashInferMxFp4LinearKernel,
82 )
@@ -274,6 +277,7 @@ _LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
277 AiterFp8BlockScaledMMKernel,
278 AiterPerTokenFp8ScaledMMLinearKernel,
279 AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
280 + AiterMxfp4LinearKernel,
281 },
282 "machete": {
283 MacheteLinearKernel,
@@ -469,6 +473,9 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = {
473 MarlinMxFp4LinearKernel,
474 HummingMxFp4LinearKernel,
475 ],
476 + PlatformEnum.ROCM: [
477 + AiterMxfp4LinearKernel,
478 + ],
479 PlatformEnum.XPU: [
480 XPUMxFp4LinearKernel,
481 ],
@@ -1079,6 +1086,7 @@ __all__ = [
1086 "init_mxfp4_linear_kernel",
1087 "MxFp4LinearKernel",
1088 "MxFp4LinearLayerConfig",
1089 + "AiterMxfp4LinearKernel",
1090 "FlashInferMxFp4LinearKernel",
1091 "MarlinMxFp4LinearKernel",
1092 "FlashInferCutedslMxfp8LinearKernel",
vllm/model_executor/kernels/linear/mxfp4/aiter.py new
+177
@@ -0,0 +1,177 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +from torch.nn.parameter import Parameter
6 +
7 +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
8 +from vllm.platforms import current_platform
9 +
10 +from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
11 +
12 +# NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP
13 +# which can force the engine core to spawn instead of fork.
14 +# is_aiter_found_and_supported() checks platform + arch + library availability via
15 +# find_spec/amdsmi, so it stays HIP-free.
16 +# Actual aiter imports are deferred to the functions/methods that need them,
17 +# where HIP initialization is expected.
18 +if is_aiter_found_and_supported():
19 + from vllm.utils.torch_utils import direct_register_custom_op
20 +
21 + def gemm_with_dynamic_quant(
22 + x: torch.Tensor,
23 + weight: torch.Tensor,
24 + weight_scale: torch.Tensor,
25 + rocm_use_aiter_fp4_asm_gemm: bool = False,
26 + out_dtype: torch.dtype | None = torch.bfloat16,
27 + x_scales: torch.Tensor | None = None,
28 + ) -> torch.Tensor:
29 + from aiter.ops.triton.gemm_afp4wfp4 import (
30 + gemm_afp4wfp4,
31 + gemm_afp4wfp4_preshuffled_weight_scales,
32 + )
33 + from aiter.ops.triton.quant import dynamic_mxfp4_quant
34 +
35 + if rocm_use_aiter_fp4_asm_gemm:
36 + from aiter import gemm_a4w4, per_1x32_f4_quant_hip
37 +
38 + M = x.shape[0]
39 + N = weight.shape[0]
40 + K = weight.shape[1]
41 + if rocm_use_aiter_fp4_asm_gemm:
42 + if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K):
43 + if x_scales is None:
44 + # use hip quant kernel for performance
45 + if M >= 32:
46 + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
47 + else:
48 + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False)
49 + else:
50 + x_q = x
51 + x_s = x_scales
52 +
53 + if M >= 32:
54 + x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1)
55 + else:
56 + x_s = x_s[:M, ...].view(torch.uint8)
57 +
58 + y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
59 + gemm_afp4wfp4_preshuffled_weight_scales(
60 + x_q.view(torch.uint8),
61 + weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
62 + x_s,
63 + weight_scale.view(torch.uint8).view(
64 + weight_scale.shape[0] // 32, -1
65 + ),
66 + out_dtype,
67 + y,
68 + )
69 + else:
70 + if x_scales is None:
71 + # use hip quant kernel for performance
72 + x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
73 + else:
74 + x_q = x
75 + x_s = x_scales
76 +
77 + y = gemm_a4w4(
78 + x_q,
79 + weight.view(x_q.dtype),
80 + x_s,
81 + weight_scale.view(x_s.dtype),
82 + dtype=out_dtype,
83 + bpreshuffle=True,
84 + )
85 + return y[:M]
86 + else:
87 + if x_scales is None:
88 + x_q, x_s = dynamic_mxfp4_quant(x)
89 + else:
90 + x_q = x
91 + x_s = x_scales
92 + y = torch.empty(
93 + x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
94 + )
95 +
96 + gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
97 + return y
98 +
99 + def gemm_with_dynamic_quant_fake(
100 + x: torch.Tensor,
101 + weight: torch.Tensor,
102 + weight_scale: torch.Tensor,
103 + x_scales: torch.Tensor = None,
104 + rocm_use_aiter_fp4_asm_gemm: bool = False,
105 + out_dtype: torch.dtype | None = torch.bfloat16,
106 + ) -> torch.Tensor:
107 + return torch.empty(
108 + (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device
109 + )
110 +
111 + direct_register_custom_op(
112 + op_name="gemm_with_dynamic_quant",
113 + op_func=gemm_with_dynamic_quant,
114 + mutates_args=[],
115 + fake_impl=gemm_with_dynamic_quant_fake,
116 + dispatch_key=current_platform.dispatch_key,
117 + )
118 +
119 +
120 +class AiterMxfp4LinearKernel(MxFp4LinearKernel):
121 + """AITER-based native MXFP4 GEMM kernel for ROCm."""
122 +
123 + def __init__(self, config: MxFp4LinearLayerConfig) -> None:
124 + super().__init__(config)
125 + self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()
126 + self.out_dtype = torch.get_default_dtype()
127 +
128 + @classmethod
129 + def is_supported(
130 + cls, compute_capability: int | None = None
131 + ) -> tuple[bool, str | None]:
132 + if not current_platform.supports_mx():
133 + return False, "current platform does not support native MXFP4 computation"
134 + if is_aiter_found_and_supported():
135 + return True, None
136 + return False, "AITER not found or not supported on the current platform"
137 +
138 + @classmethod
139 + def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
140 + return True, None
141 +
142 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
143 + if self.use_asm_gemm:
144 + from aiter.ops.shuffle import shuffle_weight
145 +
146 + weight_scale = layer.weight_scale.data
147 + sm, sn = weight_scale.shape
148 + weight_scale = weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1)
149 + weight_scale = weight_scale.permute(0, 3, 5, 2, 4, 1, 6).contiguous()
150 + weight_scale = weight_scale.view(sm, sn)
151 + layer.weight_scale = Parameter(weight_scale, requires_grad=False)
152 +
153 + layer.weight = Parameter(
154 + shuffle_weight(layer.weight.data, layout=(16, 16)),
155 + requires_grad=False,
156 + )
157 + else:
158 + layer.weight_scale = Parameter(
159 + layer.weight_scale.data.T.contiguous(), requires_grad=False
160 + )
161 +
162 + def apply_weights(
163 + self,
164 + layer: torch.nn.Module,
165 + x: torch.Tensor,
166 + bias: torch.Tensor | None = None,
167 + ) -> torch.Tensor:
168 + y = torch.ops.vllm.gemm_with_dynamic_quant(
169 + x,
170 + layer.weight,
171 + layer.weight_scale,
172 + self.use_asm_gemm,
173 + self.out_dtype,
174 + )
175 + if bias is not None:
176 + y = y + bias
177 + return y
vllm/model_executor/layers/quantization/quark/schemes/quark_ocp_mx.py
+8 -166
@@ -9,8 +9,8 @@ from typing import Any
9 import torch
10 import torch.nn.functional as F
11
12 -from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
12 from vllm.logger import init_logger
13 +from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
14 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
15 dequant_mxfp4,
16 quant_dequant_mxfp4,
@@ -36,120 +36,6 @@ from .quark_scheme import QuarkScheme
36 logger = init_logger(__name__)
37
38
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 -
48 - def gemm_with_dynamic_quant(
49 - x: torch.Tensor,
50 - weight: torch.Tensor,
51 - weight_scale: torch.Tensor,
52 - rocm_use_aiter_fp4_asm_gemm: bool = False,
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]
68 - if rocm_use_aiter_fp4_asm_gemm:
69 - if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K):
70 - if x_scales is None:
71 - # use hip quant kernel for performance
72 - if M >= 32:
73 - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
74 - else:
75 - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False)
76 - else:
77 - x_q = x
78 - x_s = x_scales
79 -
80 - if M >= 32:
81 - x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1)
82 - else:
83 - x_s = x_s[:M, ...].view(torch.uint8)
84 -
85 - y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
86 - gemm_afp4wfp4_preshuffled_weight_scales(
87 - x_q.view(torch.uint8),
88 - weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
89 - x_s,
90 - weight_scale.view(torch.uint8).view(
91 - weight_scale.shape[0] // 32, -1
92 - ),
93 - out_dtype,
94 - y,
95 - )
96 - else:
97 - if x_scales is None:
98 - # use hip quant kernel for performance
99 - x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
100 - else:
101 - x_q = x
102 - x_s = x_scales
103 -
104 - y = gemm_a4w4(
105 - x_q,
106 - weight.view(x_q.dtype),
107 - x_s,
108 - weight_scale.view(x_s.dtype),
109 - dtype=out_dtype,
110 - bpreshuffle=True,
111 - )
112 - return y[:M]
113 - else:
114 - if x_scales is None:
115 - x_q, x_s = dynamic_mxfp4_quant(x)
116 - else:
117 - x_q = x
118 - x_s = x_scales
119 - y = torch.empty(
120 - x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
121 - )
122 -
123 - gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
124 - return y
125 -
126 - def gemm_with_dynamic_quant_fake(
127 - x: torch.Tensor,
128 - weight: torch.Tensor,
129 - weight_scale: torch.Tensor,
130 - x_scales: torch.Tensor = None,
131 - rocm_use_aiter_fp4_asm_gemm: bool = False,
132 - out_dtype: torch.dtype | None = torch.bfloat16,
133 - ) -> torch.Tensor:
134 - return torch.empty(
135 - (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device
136 - )
137 -
138 - direct_register_custom_op(
139 - op_name="gemm_with_dynamic_quant",
140 - op_func=gemm_with_dynamic_quant,
141 - mutates_args=[],
142 - fake_impl=gemm_with_dynamic_quant_fake,
143 - dispatch_key=current_platform.dispatch_key,
144 - )
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 -
39 class QuarkOCP_MX(QuarkScheme):
40 def __init__(
41 self,
@@ -157,8 +43,6 @@ class QuarkOCP_MX(QuarkScheme):
43 input_quant_spec: dict[str, Any] | None,
44 dynamic_mxfp4_quant: bool = False,
45 ):
160 - self.out_dtype = torch.get_default_dtype()
161 - self.qscheme = "per_group"
46 self.weight_quant_spec = weight_quant_spec
47 self.input_quant_spec = input_quant_spec
48 self.dynamic_mxfp4_quant = dynamic_mxfp4_quant
@@ -211,17 +95,10 @@ class QuarkOCP_MX(QuarkScheme):
95 self.input_dtype != "mxfp4" or self.weight_dtype != "mxfp4"
96 )
97
214 - self.rocm_use_aiter_fp4_asm_gemm = (
215 - rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()
216 - )
217 -
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 "
223 - "https://github.com/ROCm/aiter for installation details."
224 - )
98 + # TODO: Move emulation code path as a kernel, and always
99 + # use init_mxfp4_linear_kernel.
100 + if not self.emulate:
101 + self.ocp_mx_linear = init_mxfp4_linear_kernel()
102
103 if not current_platform.supports_mx():
104 logger.warning_once(
@@ -268,7 +145,7 @@ class QuarkOCP_MX(QuarkScheme):
145 from aiter.ops.triton.quant import dynamic_mxfp4_quant
146
147 w_q, w_s = dynamic_mxfp4_quant(layer.weight)
271 - layer.weight_scale = torch.nn.Parameter(w_s.T.contiguous(), requires_grad=False)
148 + layer.weight_scale = torch.nn.Parameter(w_s, requires_grad=False)
149 layer.weight = torch.nn.Parameter(w_q, requires_grad=False)
150
151 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
@@ -284,31 +161,7 @@ class QuarkOCP_MX(QuarkScheme):
161 else:
162 if self.dynamic_mxfp4_quant:
163 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
293 - weight_scale_shuffle = weight_scale_shuffle.view(
294 - sm // 32, 2, 16, sn // 8, 2, 4, 1
295 - )
296 - weight_scale_shuffle = weight_scale_shuffle.permute(
297 - 0, 3, 5, 2, 4, 1, 6
298 - ).contiguous()
299 - weight_scale_shuffle = weight_scale_shuffle.view(sm, sn)
300 - layer.weight_scale = torch.nn.Parameter(
301 - weight_scale_shuffle, requires_grad=False
302 - )
303 -
304 - # shuffle weight
305 - weight_shuffle = layer.weight.data
306 - weight_shuffle = shuffle_weight(weight_shuffle, layout=(16, 16))
307 - layer.weight = torch.nn.Parameter(weight_shuffle, requires_grad=False)
308 - else:
309 - layer.weight_scale = torch.nn.Parameter(
310 - layer.weight_scale.data.T.contiguous(), requires_grad=False
311 - )
164 + self.ocp_mx_linear.process_weights_after_loading(layer)
165
166 def create_weights(
167 self,
@@ -375,15 +228,4 @@ class QuarkOCP_MX(QuarkScheme):
228 dq_w = self.dequant_func(layer.weight, layer.weight_scale, x.dtype)
229 qdq_x = self.quant_dequant_func(x)
230 return F.linear(qdq_x, dq_w, bias)
378 - y = torch.ops.vllm.gemm_with_dynamic_quant(
379 - x,
380 - layer.weight,
381 - layer.weight_scale,
382 - self.rocm_use_aiter_fp4_asm_gemm,
383 - self.out_dtype,
384 - )
385 - # gemm_with_dynamic_quant has no bias argument; add it here so the
386 - # native path matches F.linear (e.g. qkv_proj with qkv_bias=True).
387 - if bias is not None:
388 - y = y + bias
389 - return y
231 + return self.ocp_mx_linear.apply_weights(layer, x, bias)