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 5, 14:45 UTC
[ROCm][Quantization][1/N] Refactor quark_moe w_mxfp4 w/ oracle (#38774)
Signed-off-by: Bowen Bao <bowenbao@amd.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>
Bowen Bao committed
Apr 2, 2026 at 20:29 UTC
103f0de565f21123822283f3e8accc0fd58a22e2
6 files changed
+170
-15
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml
new
+8
@@ -0,0 +1,8 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
4
+metric_threshold: 0.568
5
+reasoning_effort: low
6
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend aiter"
7
+env:
8
+ VLLM_ROCM_USE_AITER: "1"
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml
new
+6
@@ -0,0 +1,6 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+model_name: amd/gpt-oss-20b-w-mxfp4-a-bf16
4
+metric_threshold: 0.568
5
+reasoning_effort: low
6
+server_args: "--attention-backend ROCM_AITER_UNIFIED_ATTN --moe-backend triton"
\ No newline at end of file
tests/evals/gpt_oss/configs/gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml
renamed
tests/evals/gpt_oss/configs/models-gfx950.txt
+3
-1
@@ -1,4 +1,6 @@
1
# GFX950 model configurations for GPQA evaluation
2
# Tests different environment variable combinations
3
gpt-oss-20b-rocm-baseline.yaml
4
-gpt-oss-20b-rocm-mxfp4-fp8.yaml
\ No newline at end of file
4
+gpt-oss-20b-rocm-quark-mxfp4-bf16-aiter.yaml
5
+gpt-oss-20b-rocm-quark-mxfp4-bf16-triton.yaml
6
+gpt-oss-20b-rocm-quark-mxfp4-fp8-triton.yaml
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+7
-7
@@ -54,8 +54,8 @@ class Mxfp4MoeBackend(Enum):
54
# Marlin
55
BATCHED_MARLIN = "BATCHED_MARLIN"
56
MARLIN = "MARLIN"
57
- # ROCm AITER (CK)
58
- CK = "CK"
57
+ # ROCm AITER
58
+ AITER = "AITER"
59
# Triton
60
TRITON = "TRITON"
61
TRITON_UNFUSED = "TRITON_UNFUSED"
@@ -130,7 +130,7 @@ def backend_to_kernel_cls(
130
131
return [BatchedMarlinExperts]
132
133
- elif backend == Mxfp4MoeBackend.CK:
133
+ elif backend == Mxfp4MoeBackend.AITER:
134
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
135
AiterExperts,
136
)
@@ -155,7 +155,7 @@ def map_mxfp4_backend(runner_backend: str) -> Mxfp4MoeBackend:
155
"flashinfer_cutlass_afp8": Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_MXFP8,
156
"triton": Mxfp4MoeBackend.TRITON,
157
"marlin": Mxfp4MoeBackend.MARLIN,
158
- "ck": Mxfp4MoeBackend.CK,
158
+ "aiter": Mxfp4MoeBackend.AITER,
159
"xpu": Mxfp4MoeBackend.XPU,
160
}
161
if backend := mapping.get(runner_backend):
@@ -173,7 +173,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
173
"""
174
_AVAILABLE_BACKENDS = [
175
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
176
- Mxfp4MoeBackend.CK,
176
+ Mxfp4MoeBackend.AITER,
177
Mxfp4MoeBackend.TRITON,
178
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
179
Mxfp4MoeBackend.TRITON_UNFUSED,
@@ -656,7 +656,7 @@ def convert_to_mxfp4_moe_kernel_format(
656
w2_bias,
657
)
658
659
- elif mxfp4_backend == Mxfp4MoeBackend.CK:
659
+ elif mxfp4_backend == Mxfp4MoeBackend.AITER:
660
from vllm._aiter_ops import rocm_aiter_ops
661
662
if w13_bias is not None:
@@ -794,7 +794,7 @@ def make_mxfp4_moe_quant_config(
794
Mxfp4MoeBackend.TRITON_UNFUSED,
795
Mxfp4MoeBackend.FLASHINFER_TRTLLM_MXFP4_BF16,
796
Mxfp4MoeBackend.FLASHINFER_CUTLASS_MXFP4_BF16,
797
- Mxfp4MoeBackend.CK,
797
+ Mxfp4MoeBackend.AITER,
798
):
799
return mxfp4_w4a16_moe_quant_config(
800
w1_bias=w1_bias,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+146
-7
@@ -5,6 +5,7 @@ from typing import Any
5
6
import torch
7
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm import _custom_ops as ops
10
from vllm import envs
11
from vllm._aiter_ops import rocm_aiter_ops
@@ -27,7 +28,11 @@ from vllm.model_executor.layers.fused_moe.config import (
28
)
29
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
30
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
31
+ TRITON_BACKENDS,
32
Mxfp4MoeBackend,
33
+ convert_to_mxfp4_moe_kernel_format,
34
+ make_mxfp4_moe_kernel,
35
+ make_mxfp4_moe_quant_config,
36
mxfp4_round_up_hidden_size_and_intermediate_size,
37
select_mxfp4_moe_backend,
38
)
@@ -47,7 +52,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
52
normalize_e4m3fn_to_e4m3fnuz,
53
per_tensor_dequantize,
54
)
50
-from vllm.model_executor.utils import set_weight_attrs
55
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
56
from vllm.platforms import current_platform
57
from vllm.scalar_type import scalar_types
58
@@ -699,9 +704,16 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
704
f"Please check that the combination is supported in OCP_MX_Scheme."
705
)
706
702
- self.mxfp4_backend: Mxfp4MoeBackend | None = None
707
+ self.mxfp4_backend: Mxfp4MoeBackend = Mxfp4MoeBackend.NONE
708
+ self.experts_cls: type[mk.FusedMoEExperts] | None = None
709
+ self.moe_kernel: mk.FusedMoEKernel | None = None
710
+
711
+ # Used for triton kernel precision configs
712
+ self.w13_precision_config = None
713
+ self.w2_precision_config = None
714
+
715
if self.ocp_mx_scheme == "w_mxfp4":
704
- self.mxfp4_backend, _ = select_mxfp4_moe_backend(moe)
716
+ self.mxfp4_backend, self.experts_cls = select_mxfp4_moe_backend(moe)
717
elif self.ocp_mx_scheme.startswith("w_mxfp4"):
718
# TODO(bowenbao): refactor and introduce backends for other OCP MX schemes.
719
self.mxfp4_backend = Mxfp4MoeBackend.NONE
@@ -738,9 +750,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
750
not current_platform.supports_mx()
751
or not self.ocp_mx_scheme.startswith("w_mxfp4")
752
) and (
741
- self.mxfp4_backend is None
742
- or self.mxfp4_backend is Mxfp4MoeBackend.NONE
743
- or not self.use_rocm_aiter_moe
753
+ self.mxfp4_backend is Mxfp4MoeBackend.NONE or not self.use_rocm_aiter_moe
754
)
755
756
if self.emulate:
@@ -944,11 +954,23 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
954
w2_input_scale, requires_grad=False
955
)
956
947
- # secondly, process mxfp weights
957
+ # For w_mxfp4, use oracle functions
958
+ if (
959
+ self.ocp_mx_scheme == "w_mxfp4"
960
+ and self.mxfp4_backend != Mxfp4MoeBackend.NONE
961
+ ):
962
+ self._setup_kernel_via_oracle(layer)
963
+ return
964
+
965
+ # TODO(bowenbao): gradually migrate to oracles.
966
+ # secondly, process mxfp weights for other schemes
967
if self.emulate:
968
+ # Build quant config for emulation path
969
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
970
torch.accelerator.empty_cache()
971
return
972
973
+ # Existing AITER path for w_mxfp4_a_mxfp4 and other schemes
974
from aiter.utility.fp4_utils import e8m0_shuffle
975
976
# Pre-shuffle weight scales
@@ -980,11 +1002,87 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1002
layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
1003
layer.w13_weight.is_shuffled = True
1004
layer.w2_weight.is_shuffled = True
1005
+
1006
+ # Build quant config for AITER path
1007
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1008
torch.accelerator.empty_cache()
1009
1010
+ def _setup_kernel_via_oracle(self, layer: FusedMoE):
1011
+ """Setup kernel using oracle functions for w_mxfp4 scheme."""
1012
+ w13 = layer.w13_weight
1013
+ w2 = layer.w2_weight
1014
+ w13_scale = layer.w13_weight_scale
1015
+ w2_scale = layer.w2_weight_scale
1016
+ w13_bias = getattr(layer, "w13_bias", None)
1017
+ w2_bias = getattr(layer, "w2_bias", None)
1018
+
1019
+ # Convert weights to kernel format
1020
+ w13, w2, w13_scale, w2_scale, w13_bias, w2_bias = (
1021
+ convert_to_mxfp4_moe_kernel_format(
1022
+ mxfp4_backend=self.mxfp4_backend,
1023
+ layer=layer,
1024
+ w13_weight=w13,
1025
+ w2_weight=w2,
1026
+ w13_weight_scale=w13_scale,
1027
+ w2_weight_scale=w2_scale,
1028
+ w13_bias=w13_bias,
1029
+ w2_bias=w2_bias,
1030
+ )
1031
+ )
1032
+
1033
+ # For TRITON backends, weights are wrapped tensors from triton_kernels
1034
+ # that don't support .detach(). Manually assign parameters.
1035
+ if self.mxfp4_backend not in TRITON_BACKENDS:
1036
+ replace_parameter(layer, "w13_weight", w13)
1037
+ replace_parameter(layer, "w2_weight", w2)
1038
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
1039
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
1040
+ else:
1041
+ layer.w13_weight = w13
1042
+ layer.w2_weight = w2
1043
+ self.w13_precision_config = w13_scale
1044
+ self.w2_precision_config = w2_scale
1045
+
1046
+ if w13_bias is not None and w2_bias is not None:
1047
+ replace_parameter(layer, "w13_bias", w13_bias)
1048
+ replace_parameter(layer, "w2_bias", w2_bias)
1049
+
1050
+ # Build quant config and kernel
1051
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1052
+ if self.moe_quant_config is not None and self.experts_cls is not None:
1053
+ self.moe_kernel = make_mxfp4_moe_kernel(
1054
+ moe_quant_config=self.moe_quant_config,
1055
+ moe_config=self.moe,
1056
+ mxfp4_backend=self.mxfp4_backend,
1057
+ experts_cls=self.experts_cls,
1058
+ routing_tables=layer._maybe_init_expert_routing_tables(),
1059
+ shared_experts=layer.shared_experts,
1060
+ )
1061
+
1062
def get_fused_moe_quant_config(
1063
self, layer: torch.nn.Module
1064
) -> FusedMoEQuantConfig | None:
1065
+ # For w_mxfp4 with oracle backend, use oracle function
1066
+ if (
1067
+ self.ocp_mx_scheme == "w_mxfp4"
1068
+ and self.mxfp4_backend != Mxfp4MoeBackend.NONE
1069
+ ):
1070
+ w1_scale = layer.w13_weight_scale
1071
+ w2_scale = layer.w2_weight_scale
1072
+ if self.mxfp4_backend in TRITON_BACKENDS:
1073
+ w1_scale = self.w13_precision_config
1074
+ w2_scale = self.w2_precision_config
1075
+ return make_mxfp4_moe_quant_config(
1076
+ mxfp4_backend=self.mxfp4_backend,
1077
+ w1_scale=w1_scale,
1078
+ w2_scale=w2_scale,
1079
+ w1_bias=getattr(layer, "w13_bias", None),
1080
+ w2_bias=getattr(layer, "w2_bias", None),
1081
+ )
1082
+
1083
+ # Existing code for other schemes
1084
+ # TODO(bowenbao): kept for emulation fallback, to be refactored into
1085
+ # dedicated emulation backend.
1086
if self.ocp_mx_scheme == "w_mxfp4":
1087
return mxfp4_w4a16_moe_quant_config(
1088
w1_scale=layer.w13_weight_scale,
@@ -1020,6 +1118,12 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1118
block_shape=None,
1119
)
1120
1121
+ @property
1122
+ def is_monolithic(self) -> bool:
1123
+ if self.moe_kernel is not None:
1124
+ return self.moe_kernel.is_monolithic
1125
+ return False
1126
+
1127
def apply(
1128
self,
1129
layer: FusedMoE,
@@ -1028,6 +1132,22 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1132
topk_ids: torch.Tensor,
1133
shared_experts_input: torch.Tensor | None,
1134
) -> torch.Tensor:
1135
+ # For w_mxfp4 with oracle kernel
1136
+ if self.moe_kernel is not None:
1137
+ return self.moe_kernel.apply(
1138
+ hidden_states=x,
1139
+ w1=layer.w13_weight,
1140
+ w2=layer.w2_weight,
1141
+ topk_weights=topk_weights,
1142
+ topk_ids=topk_ids,
1143
+ activation=layer.activation,
1144
+ global_num_experts=layer.global_num_experts,
1145
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1146
+ expert_map=layer.expert_map,
1147
+ shared_experts_input=shared_experts_input,
1148
+ )
1149
+
1150
+ # Existing code for emulation/AITER paths
1151
if not self.emulate:
1152
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1153
rocm_aiter_fused_experts,
@@ -1061,6 +1181,25 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1181
quant_config=self.moe_quant_config,
1182
)
1183
1184
+ def apply_monolithic(
1185
+ self,
1186
+ layer: FusedMoE,
1187
+ x: torch.Tensor,
1188
+ router_logits: torch.Tensor,
1189
+ ) -> torch.Tensor:
1190
+ assert self.is_monolithic
1191
+ assert self.moe_kernel is not None
1192
+ return self.moe_kernel.apply_monolithic(
1193
+ hidden_states=x,
1194
+ w1=layer.w13_weight,
1195
+ w2=layer.w2_weight,
1196
+ router_logits=router_logits,
1197
+ activation=layer.activation,
1198
+ global_num_experts=layer.global_num_experts,
1199
+ expert_map=layer.expert_map,
1200
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1201
+ )
1202
+
1203
1204
class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
1205
def __init__(