@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a2c8fc6657

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 3, 10:30 UTC

[ROCm][Quantization][3/N] Refactor quark_moe w4a4 w/ oracle (#41436)

Signed-off-by: Bowen Bao <bowenbao@amd.com>

Bowen Bao committed May 18, 2026 at 10:46 UTC a2c8fc66573664395f491a94da1882fdf92e034b

8 files changed +224 -158

tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml new

+12

@@ -0,0 +1,12 @@

1 +model_name: "amd/Qwen3.5-35B-A3B-MXFP4"

2 +accuracy_threshold: 0.89

3 +tolerance: 0.03

4 +num_questions: 1319

5 +num_fewshot: 5

6 +server_args: >-

7 + --max-model-len 4096

8 + --tensor-parallel-size 2

9 + --gpu-memory-utilization 0.35

10 + --moe-backend aiter

11 +env:

12 + VLLM_ROCM_USE_AITER: "1"

tests/evals/gsm8k/configs/Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml renamed

+3 -1

@@ -1,8 +1,10 @@

1 model_name: "amd/Qwen3.5-35B-A3B-MXFP4"

2 -accuracy_threshold: 0.82

2 +accuracy_threshold: 0.89

3 tolerance: 0.03

4 num_questions: 1319

5 num_fewshot: 5

6 server_args: >-

7 --max-model-len 4096

8 --tensor-parallel-size 2

9 + --moe-backend emulation

10 + --gpu-memory-utilization 0.35

tests/evals/gsm8k/configs/models-mi3xx.txt

+2 -1

@@ -3,4 +3,5 @@ DeepSeek-R1-DP_MI325.yaml

3 DeepSeek-V3.2-TP_MI325.yaml

4 DeepSeek-V3.2-DP_MI325.yaml

5 Qwen3-30B-A3B-NVFP4.yaml

6 -Qwen3.5-35B-A3B-MXFP4-TP2.yaml

\ No newline at end of file

6 +Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml

7 +Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml

\ No newline at end of file

tests/evals/gsm8k/configs/models-qwen35-mi355.txt

+2 -1

@@ -1,2 +1,3 @@

1 Qwen3.5-35B-A3B-DEP2.yaml

2 -Qwen3.5-35B-A3B-MXFP4-TP2.yaml

2 +Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml

3 +Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml

tests/quantization/test_gfx950_moe.py

+86 -2

@@ -1,6 +1,90 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""

4 +Tests for MXFP4 MoE oracle backend selection on mi355x (GFX950).

6 +These tests run on real hardware — no mocks. Skipped on non-GFX950 platforms.

7 +"""

5 -def test_mi355_moe():

6 - print("TODO: add tests for Mi355 MoE quantization")

9 +import pytest

10 +import torch

11 +

12 +from vllm.model_executor.layers.fused_moe.config import (

13 + FusedMoEConfig,

14 + FusedMoEParallelConfig,

15 + RoutingMethodType,

16 +)

17 +from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (

18 + Mxfp4MoeBackend,

19 + select_mxfp4_moe_backend,

20 +)

21 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

22 + kMxfp4Dynamic,

23 +)

24 +from vllm.platforms import current_platform

25 +

26 +ROCM_AVAILABLE = current_platform.is_rocm()

27 +ROCM_GFX950 = False

28 +ROCM_AITER_AVAILABLE = False

29 +

30 +if ROCM_AVAILABLE:

31 + from vllm._aiter_ops import rocm_aiter_ops

32 + from vllm.platforms.rocm import on_gfx950

33 +

34 + ROCM_GFX950 = on_gfx950()

35 + ROCM_AITER_AVAILABLE = rocm_aiter_ops.is_fused_moe_enabled()

36 +

37 +

38 +def _make_w4a4_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:

39 + from vllm.model_executor.layers.fused_moe.activation import MoEActivation

40 +

41 + return FusedMoEConfig(

42 + num_experts=8,

43 + experts_per_token=2,

44 + hidden_dim=256,

45 + intermediate_size_per_partition=256,

46 + num_local_experts=8,

47 + num_logical_experts=8,

48 + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),

49 + activation=MoEActivation.SILU,

50 + in_dtype=torch.bfloat16,

51 + device="cuda",

52 + routing_method=RoutingMethodType.Renormalize,

53 + moe_backend=moe_backend,

54 + )

55 +

56 +

57 +@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")

58 +@pytest.mark.skipif(not ROCM_AITER_AVAILABLE, reason="Requires AITER enabled")

59 +def test_w4a4_dispatches_to_aiter():

60 + """With AITER enabled + GFX950, W4A4 selects AITER_MXFP4_MXFP4."""

61 + config = _make_w4a4_moe_config()

62 + backend, experts_cls = select_mxfp4_moe_backend(

63 + config, activation_key=kMxfp4Dynamic

64 + )

65 + assert backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4

66 + assert experts_cls is not None

67 +

68 +

69 +@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")

70 +@pytest.mark.skipif(

71 + ROCM_AITER_AVAILABLE,

72 + reason="Test requires AITER disabled (unset VLLM_ROCM_USE_AITER)",

73 +)

74 +def test_w4a4_raises_without_aiter_and_no_moe_backend():

75 + """Without AITER and no --moe-backend, raises NotImplementedError

76 + with hint to use --moe-backend emulation."""

77 + config = _make_w4a4_moe_config()

78 + with pytest.raises(NotImplementedError, match="--moe-backend emulation"):

79 + select_mxfp4_moe_backend(config, activation_key=kMxfp4Dynamic)

80 +

81 +

82 +@pytest.mark.skipif(not ROCM_GFX950, reason="Requires GFX950 (mi355x)")

83 +def test_w4a4_dispatches_to_emulation_with_moe_backend():

84 + """With --moe-backend emulation, W4A4 selects EMULATION."""

85 + config = _make_w4a4_moe_config(moe_backend="emulation")

86 + backend, experts_cls = select_mxfp4_moe_backend(

87 + config, activation_key=kMxfp4Dynamic

88 + )

89 + assert backend == Mxfp4MoeBackend.EMULATION

90 + assert experts_cls is not None

vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py

+17

@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

26 kFp8Static128BlockSym,

27 kFp8StaticChannelSym,

28 kFp8StaticTensorSym,

29 + kMxfp4Dynamic,

30 kMxfp4Static,

31 )

@@ -377,6 +378,21 @@ class AiterExperts(mk.FusedMoEExpertsModular):

378 def activation_format() -> mk.FusedMoEActivationFormat:

379 return mk.FusedMoEActivationFormat.Standard

380

381 + @staticmethod

382 + def is_supported_config(

383 + cls, moe_config, weight_key, activation_key, activation_format

384 + ):

385 + is_supported, reason = super().is_supported_config(

386 + cls, moe_config, weight_key, activation_key, activation_format

387 + )

388 + if not is_supported and not rocm_aiter_ops.is_fused_moe_enabled():

389 + reason = (

390 + f"{reason}. AITER MoE is not enabled — "

391 + "set VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1 "

392 + "to enable it"

393 + )

394 + return is_supported, reason

395 +

396 @staticmethod

397 def _supports_current_device() -> bool:

398 return rocm_aiter_ops.is_fused_moe_enabled()

@@ -397,6 +413,7 @@ class AiterExperts(mk.FusedMoEExpertsModular):

413 (kFp8StaticTensorSym, kFp8DynamicTensorSym),

414 (kFp8StaticChannelSym, kFp8DynamicTokenSym),

415 (kMxfp4Static, None),

416 + (kMxfp4Static, kMxfp4Dynamic),

417 ]

418 if (weight_key, activation_key) not in SUPPORTED_W_A:

419 return False

vllm/model_executor/layers/fused_moe/oracle/mxfp4.py

+76 -2

@@ -31,6 +31,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

31 QuantKey,

32 kFp8Dynamic128Sym,

33 kFp8StaticTensorSym,

34 + kMxfp4Dynamic,

35 kMxfp4Static,

36 kMxfp8Dynamic,

37 )

@@ -74,6 +75,7 @@ class Mxfp4MoeBackend(Enum):

75 # Keep the legacy name as an alias while the ROCm split backend rename settles.

76 AITER = "AITER_MXFP4_BF16"

77 AITER_MXFP4_FP8 = "AITER_MXFP4_FP8" # W4A8: triton kernel

78 + AITER_MXFP4_MXFP4 = "AITER_MXFP4_MXFP4" # W4A4: CK kernel

79 # Triton

80 TRITON = "TRITON"

81 TRITON_UNFUSED = "TRITON_UNFUSED"

@@ -91,6 +93,7 @@ class Mxfp4MoeBackend(Enum):

93 AITER_BACKENDS = (

94 Mxfp4MoeBackend.AITER_MXFP4_BF16,

95 Mxfp4MoeBackend.AITER_MXFP4_FP8,

96 + Mxfp4MoeBackend.AITER_MXFP4_MXFP4,

97 )

@@ -195,6 +198,13 @@ def backend_to_kernel_cls(

198

199 return [AiterW4A8ExpertsMonolithic]

200

201 + elif backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:

202 + from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (

203 + AiterExperts,

204 + )

205 +

206 + return [AiterExperts]

207 +

208 elif backend == Mxfp4MoeBackend.XPU:

209 from vllm.model_executor.layers.fused_moe.experts.xpu_moe import XPUExpertsMXFp4

210

@@ -241,8 +251,10 @@ def map_mxfp4_backend(runner_backend: MoEBackend) -> list[Mxfp4MoeBackend]:

251 "aiter": [

252 Mxfp4MoeBackend.AITER_MXFP4_BF16,

253 Mxfp4MoeBackend.AITER_MXFP4_FP8,

254 + Mxfp4MoeBackend.AITER_MXFP4_MXFP4,

255 ],

256 "aiter_mxfp4_fp8": [Mxfp4MoeBackend.AITER_MXFP4_FP8],

257 + "aiter_mxfp4_mxfp4": [Mxfp4MoeBackend.AITER_MXFP4_MXFP4],

258 "xpu": [Mxfp4MoeBackend.XPU],

259 "cpu": [Mxfp4MoeBackend.CPU],

260 "emulation": [Mxfp4MoeBackend.EMULATION],

@@ -263,6 +275,7 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:

275 Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_MXFP8,

276 Mxfp4MoeBackend.AITER_MXFP4_BF16,

277 Mxfp4MoeBackend.AITER_MXFP4_FP8,

278 + Mxfp4MoeBackend.AITER_MXFP4_MXFP4,

279 Mxfp4MoeBackend.TRITON,

280 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,

281 Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,

@@ -272,7 +285,6 @@ def _get_priority_backends_for_gpt_oss() -> list[Mxfp4MoeBackend]:

285 Mxfp4MoeBackend.MARLIN,

286 Mxfp4MoeBackend.BATCHED_MARLIN,

287 Mxfp4MoeBackend.XPU,

275 - Mxfp4MoeBackend.EMULATION,

288 ]

289 return _AVAILABLE_BACKENDS

290

@@ -308,6 +320,8 @@ def _backend_activation_key(backend: Mxfp4MoeBackend) -> QuantKey | None:

320 return kMxfp8Dynamic

321 if backend == Mxfp4MoeBackend.AITER_MXFP4_FP8:

322 return kFp8StaticTensorSym

323 + if backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:

324 + return kMxfp4Dynamic

325 return None # BF16 activation

326

327

@@ -379,6 +393,7 @@ def _filter_by_activation(

393 b

394 for b in backends

395 if _backend_activation_key(b) == requested_activation_key

396 + or b == Mxfp4MoeBackend.EMULATION

397 ]

398 bf16 = [b for b in backends if _backend_activation_key(b) is None]

399 return bf16 if bf16 else backends

@@ -556,7 +571,12 @@ def select_mxfp4_moe_backend(

571

572 if current_platform.is_cuda() or current_platform.is_rocm():

573 raise NotImplementedError(

559 - "No MXFP4 MoE backend supports the deployment configuration."

574 + "No MXFP4 MoE backend supports the deployment configuration. "

575 + f"weight_key=kMxfp4Static, activation_key={activation_key}. "

576 + "Native backends require specific hardware. "

577 + "Set `VLLM_LOGGING_LEVEL=DEBUG` to see detailed unsupported reasons. "

578 + "To use the emulation backend for research/debugging, pass "

579 + "--moe-backend emulation."

580 )

581

582 return Mxfp4MoeBackend.NONE, None

@@ -957,6 +977,49 @@ def convert_gpt_oss_weight_to_mxfp4_moe_kernel_format(

977 w2_bias,

978 )

979

980 + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:

981 + from vllm._aiter_ops import rocm_aiter_ops

982 +

983 + if w13_bias is not None:

984 + w13_bias = w13_bias.data.to(torch.float32)

985 + if w2_bias is not None:

986 + w2_bias = w2_bias.data.to(torch.float32)

987 +

988 + # e8m0_shuffle on weight scales (GFX950 swizzle layout)

989 + from aiter.utility.fp4_utils import e8m0_shuffle

990 +

991 + s0, s1, _ = w13_weight_scale.shape

992 + w13_weight_scale.data = e8m0_shuffle(w13_weight_scale.view(s0 * s1, -1)).view(

993 + s0, s1, -1

994 + )

995 +

996 + s0, s1, _ = w2_weight_scale.shape

997 + w2_weight_scale.data = e8m0_shuffle(w2_weight_scale.view(s0 * s1, -1)).view(

998 + s0, s1, -1

999 + )

1000 +

1001 + # View as native FP4 dtype

1002 + fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None)

1003 + if fp4_dtype is not None:

1004 + w13_weight.data = w13_weight.data.view(fp4_dtype)

1005 + w2_weight.data = w2_weight.data.view(fp4_dtype)

1006 +

1007 + # Shuffle weights for AITER CK kernel

1008 + shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1009 + w13_weight, w2_weight

1010 + )

1011 + shuffled_w13.is_shuffled = True

1012 + shuffled_w2.is_shuffled = True

1013 +

1014 + return (

1015 + shuffled_w13,

1016 + shuffled_w2,

1017 + w13_weight_scale,

1018 + w2_weight_scale,

1019 + w13_bias,

1020 + w2_bias,

1021 + )

1022 +

1023 elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_BF16:

1024 from vllm._aiter_ops import rocm_aiter_ops

1025

@@ -1541,6 +1604,17 @@ def make_mxfp4_moe_quant_config(

1604 block_shape=None,

1605 gemm1_clamp_limit=swiglu_limit,

1606 )

1607 + elif mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:

1608 + return ocp_mx_moe_quant_config(

1609 + quant_dtype="mxfp4",

1610 + w1_bias=w1_bias,

1611 + w2_bias=w2_bias,

1612 + w1_scale=w1_scale,

1613 + w2_scale=w2_scale,

1614 + gemm1_alpha=gemm1_alpha,

1615 + gemm1_beta=gemm1_beta,

1616 + gemm1_clamp_limit=swiglu_limit,

1617 + )

1618 elif mxfp4_backend in (

1619 Mxfp4MoeBackend.MARLIN,

1620 Mxfp4MoeBackend.BATCHED_MARLIN,

vllm/model_executor/layers/quantization/quark/quark_moe.py

+26 -151

@@ -55,6 +55,7 @@ from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (

55 from vllm.model_executor.layers.quantization.utils.quant_utils import (

56 GroupShape,

57 kFp8StaticTensorSym,

58 + kMxfp4Dynamic,

59 kNvfp4Dynamic,

60 kNvfp4Static,

61 )

@@ -1042,6 +1043,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1043 self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(

1044 moe, activation_key=kFp8StaticTensorSym

1045 )

1046 + elif self.ocp_mx_scheme == "w_mxfp4_a_mxfp4":

1047 + # W4A4: MXFP4 weights + MXFP4 activations

1048 + self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(

1049 + moe, activation_key=kMxfp4Dynamic

1050 + )

1051

1052 # Validation for unsupported schemes

1053 if any(

@@ -1061,45 +1067,19 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1067 "Please open an issue."

1068 )

1069

1064 - self.use_rocm_aiter_moe = rocm_aiter_ops.is_fused_moe_enabled()

1065 -

1070 self.model_type = getattr(

1071 get_current_vllm_config().model_config.hf_config, "model_type", None

1072 )

1073

1070 - # TODO: Remove once all OCP MX schemes use the kernel abstraction

1071 - _AITER_NATIVE_OCP_MX_SCHEMES = ("w_mxfp4", "w_mxfp4_a_mxfp4", "w_mxfp4_a_fp8")

1072 - self.emulate = (

1073 - not current_platform.supports_mx()

1074 - or self.ocp_mx_scheme not in _AITER_NATIVE_OCP_MX_SCHEMES

1075 - ) and (

1076 - self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe

1077 - )

1078 -

1079 - if self.emulate:

1080 - # We use the same code path between MXFP4/MXFP6 emulation.

1074 + # If no native backend available, use emulation.

1075 + if self.mxfp4_backend is Mxfp4MoeBackend.NONE:

1076 self.mxfp4_backend = Mxfp4MoeBackend.EMULATION

1077

1083 - # TODO: Remove `self.mxfp4_backend != Mxfp4MoeBackend.NONE` and make it so that

1084 - # all MXFP4 backends use the kernel abstraction.

1085 - if self.mxfp4_backend != Mxfp4MoeBackend.NONE:

1086 - self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]

1078 + self.experts_cls = backend_to_kernel_cls(self.mxfp4_backend)[0]

1079

1088 - # Log backend selection

1089 - if self.mxfp4_backend != Mxfp4MoeBackend.NONE:

1090 - logger.info_once(

1091 - f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}"

1092 - )

1093 - elif self.emulate:

1094 - logger.warning_once(

1095 - f"The current mode (supports_mx={current_platform.supports_mx()}, "

1096 - f"use_rocm_aiter_moe={self.use_rocm_aiter_moe}, "

1097 - f"ocp_mx_scheme={self.ocp_mx_scheme}) "

1098 - "does not support native MXFP4/MXFP6 "

1099 - "computation. Simulated weight dequantization and activation "

1100 - "QDQ (quantize and dequantize) will be used, with the linear "

1101 - "layers computed in high precision."

1102 - )

1080 + logger.info_once(

1081 + f"Using {self.mxfp4_backend.value} backend for {self.ocp_mx_scheme}"

1082 + )

1083

1084 def maybe_roundup_sizes(

1085 self,

@@ -1243,97 +1223,8 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1223 layer.w13_input_scale = None

1224 layer.w2_input_scale = None

1225

1246 - def process_weights_after_loading(self, layer: RoutedExperts):

1247 - # For MXFP4 schemes with native backend, use oracle

1248 - if self.mxfp4_backend != Mxfp4MoeBackend.NONE:

1249 - self._setup_kernel(layer)

1250 - return

1251 -

1252 - if self.static_input_scales and self.input_dtype == "fp8":

1253 - # firstly, process activations if fp8 static input

1254 - if layer.w13_input_scale is None or layer.w2_input_scale is None:

1255 - raise ValueError(

1256 - "QuantConfig has static quantization, but found "

1257 - "activation scales are None."

1258 - )

1259 - if not all_close_1d(layer.w13_input_scale) or not all_close_1d(

1260 - layer.w2_input_scale

1261 - ):

1262 - logger.warning_once(

1263 - "Found input_scales that are not equal for "

1264 - "fp8 MoE layer. Using the maximum across experts "

1265 - "for each layer. "

1266 - )

1267 - layer.w13_input_scale = torch.nn.Parameter(

1268 - layer.w13_input_scale.max(), requires_grad=False

1269 - )

1270 - layer.w2_input_scale = torch.nn.Parameter(

1271 - layer.w2_input_scale.max(), requires_grad=False

1272 - )

1273 -

1274 - if current_platform.is_fp8_fnuz():

1275 - # Normalize the weights and scales

1276 - _, _, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(

1277 - torch.empty_like(layer.w13_weight, dtype=torch.float8_e4m3fn),

1278 - torch.empty_like(

1279 - layer.w13_weight_scale, dtype=layer.w13_weight_scale.dtype

1280 - ),

1281 - layer.w13_input_scale,

1282 - )

1283 - _, _, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(

1284 - torch.empty_like(layer.w2_weight, dtype=torch.float8_e4m3fn),

1285 - torch.empty_like(

1286 - layer.w2_weight_scale, dtype=layer.w13_weight_scale.dtype

1287 - ),

1288 - layer.w2_input_scale,

1289 - )

1290 - # Reset the parameter

1291 - if w13_input_scale is not None:

1292 - layer.w13_input_scale = torch.nn.Parameter(

1293 - w13_input_scale, requires_grad=False

1294 - )

1295 - if w2_input_scale is not None:

1296 - layer.w2_input_scale = torch.nn.Parameter(

1297 - w2_input_scale, requires_grad=False

1298 - )

1299 -

1300 - # TODO(bowenbao): gradually migrate to oracles.

1301 - # Existing AITER path for w_mxfp4_a_mxfp4 and other schemes

1302 - from aiter.utility.fp4_utils import e8m0_shuffle

1303 -

1304 - # Pre-shuffle weight scales

1305 - s0, s1, _ = layer.w13_weight_scale.shape

1306 - w13_weight_scale = layer.w13_weight_scale.view(s0 * s1, -1)

1307 - w13_weight_scale = e8m0_shuffle(w13_weight_scale)

1308 - layer.w13_weight_scale.data = w13_weight_scale.view(s0, s1, -1)

1309 -

1310 - s0, s1, _ = layer.w2_weight_scale.shape

1311 - w2_weight_scale = layer.w2_weight_scale.view(s0 * s1, -1)

1312 - w2_weight_scale = e8m0_shuffle(w2_weight_scale)

1313 - layer.w2_weight_scale.data = w2_weight_scale.view(s0, s1, -1)

1314 -

1315 - if self.fp4_dtype is not None:

1316 - layer.w13_weight = torch.nn.Parameter(

1317 - layer.w13_weight.view(self.fp4_dtype),

1318 - requires_grad=layer.w13_weight.requires_grad,

1319 - )

1320 - layer.w2_weight = torch.nn.Parameter(

1321 - layer.w2_weight.view(self.fp4_dtype),

1322 - requires_grad=layer.w2_weight.requires_grad,

1323 - )

1324 - # Pre-shuffle weight

1325 - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(

1326 - layer.w13_weight.data, layer.w2_weight.data

1327 - )

1328 -

1329 - layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)

1330 - layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)

1331 - layer.w13_weight.is_shuffled = True

1332 - layer.w2_weight.is_shuffled = True

1333 -

1334 - # Build quant config for AITER path

1335 - self.moe_quant_config = self.get_fused_moe_quant_config(layer)

1336 - torch.accelerator.empty_cache()

1226 + def process_weights_after_loading(self, layer):

1227 + self._setup_kernel(layer)

1228

1229 def _setup_kernel(self, layer: RoutedExperts):

1230 """Setup kernel using oracle functions for MXFP4 schemes (W4A16, W4A8)."""

@@ -1375,6 +1266,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1266 replace_parameter(layer, "w13_bias", w13_bias)

1267 replace_parameter(layer, "w2_bias", w2_bias)

1268

1269 + if self.mxfp4_backend == Mxfp4MoeBackend.AITER_MXFP4_MXFP4:

1270 + layer.w13_weight.is_shuffled = True

1271 + layer.w2_weight.is_shuffled = True

1272 +

1273 torch.accelerator.empty_cache()

1274

1275 # Build quant config and kernel

@@ -1464,38 +1359,18 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1359 shared_experts: SharedExperts | None,

1360 shared_experts_input: torch.Tensor | None,

1361 ) -> torch.Tensor:

1467 - # For oracle-based kernels (W4A16, W4A8) or emulation kernel

1468 - if self.moe_kernel is not None:

1469 - return self.moe_kernel.apply(

1470 - hidden_states=x,

1471 - w1=layer.w13_weight,

1472 - w2=layer.w2_weight,

1473 - topk_weights=topk_weights,

1474 - topk_ids=topk_ids,

1475 - activation=layer.activation,

1476 - global_num_experts=layer.global_num_experts,

1477 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

1478 - expert_map=layer.expert_map,

1479 - shared_experts=shared_experts,

1480 - shared_experts_input=shared_experts_input,

1481 - )

1482 -

1483 - # AITER path

1484 - # TODO: Refactor this to use modular MOE kernel as well.

1485 - from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (

1486 - rocm_aiter_fused_experts,

1487 - )

1488 -

1489 - return rocm_aiter_fused_experts(

1490 - x,

1491 - layer.w13_weight,

1492 - layer.w2_weight,

1362 + assert self.moe_kernel is not None

1363 + return self.moe_kernel.apply(

1364 + hidden_states=x,

1365 + w1=layer.w13_weight,

1366 + w2=layer.w2_weight,

1367 topk_weights=topk_weights,

1368 topk_ids=topk_ids,

1369 activation=layer.activation,

1496 - quant_config=self.moe_quant_config,

1497 - moe_config=layer.moe_config,

1370 + global_num_experts=layer.global_num_experts,

1371 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

1372 expert_map=layer.expert_map,

1373 + shared_experts_input=shared_experts_input,

1374 )

1375

1376 def apply_monolithic(