[ROCm][Quantization][4/N] refactor quark_moe fp8 w/ oracle (#43721)
Signed-off-by: Bowen Bao <bowenbao@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Bowen Bao committed
Jun 22, 2026 at 15:58 UTC
c97e8f99d69d1bdf02110f01fe3f5e00dcae71fa
5 files changed
+91
-104
tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-FP8.yaml
new
+6
@@ -0,0 +1,6 @@
1
+model_name: "amd/Qwen3-30B-A3B-Thinking-2507-FP8"
2
+accuracy_threshold: 0.81
3
+num_questions: 1319
4
+num_fewshot: 5
5
+max_tokens: 1024
6
+server_args: "--max-model-len 4096 --gpu-memory-utilization 0.85"
tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml
new
+6
@@ -0,0 +1,6 @@
1
+model_name: "amd/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8"
2
+accuracy_threshold: 0.81
3
+num_questions: 1319
4
+num_fewshot: 5
5
+max_tokens: 1024
6
+server_args: "--max-model-len 4096 --gpu-memory-utilization 0.85"
tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt
+2
@@ -3,3 +3,5 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml
3
Qwen1.5-MoE-W4A16-CT.yaml
4
DeepSeek-V2-Lite-Instruct-FP8.yaml
5
Qwen3-Next-FP8-EP2_MI355.yaml
6
+Qwen3-30B-A3B-Thinking-2507-FP8.yaml
7
+Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml
tests/evals/gsm8k/test_gsm8k_correctness.py
+1
@@ -48,6 +48,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict:
48
results = evaluate_gsm8k(
49
num_questions=eval_config["num_questions"],
50
num_shots=eval_config["num_fewshot"],
51
+ max_tokens=eval_config.get("max_tokens", 256),
52
host=host,
53
port=port,
54
request_timeout_seconds=request_timeout_seconds,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+76
-104
@@ -7,7 +7,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
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import get_current_vllm_config
12
from vllm.logger import init_logger
@@ -15,7 +14,6 @@ from vllm.model_executor.layers.fused_moe import (
14
FusedMoEConfig,
15
FusedMoEMethodBase,
16
FusedMoeWeightScaleSupported,
18
- MoEActivation,
17
RoutedExperts,
18
SharedExperts,
19
)
@@ -28,7 +26,13 @@ from vllm.model_executor.layers.fused_moe.config import (
26
mxfp4_w4a16_moe_quant_config,
27
ocp_mx_moe_quant_config,
28
)
31
-from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe
29
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
30
+ Fp8MoeBackend,
31
+ convert_to_fp8_moe_kernel_format,
32
+ make_fp8_moe_kernel,
33
+ make_fp8_moe_quant_config,
34
+ select_fp8_moe_backend,
35
+)
36
from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import (
37
TRITON_BACKENDS,
38
Mxfp4MoeBackend,
@@ -45,15 +49,15 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
49
make_nvfp4_moe_quant_config,
50
select_nvfp4_moe_backend,
51
)
48
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
49
- prepare_fp8_moe_layer_for_marlin,
50
-)
52
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
53
OCP_MX_BLOCK_SIZE,
54
OCP_MX_Scheme,
55
)
56
from vllm.model_executor.layers.quantization.utils.quant_utils import (
57
GroupShape,
58
+ kFp8DynamicTensorSym,
59
+ kFp8DynamicTokenSym,
60
+ kFp8StaticChannelSym,
61
kFp8StaticTensorSym,
62
kMxfp4Dynamic,
63
kNvfp4Dynamic,
@@ -66,7 +70,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
70
)
71
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
72
from vllm.platforms import current_platform
69
-from vllm.scalar_type import scalar_types
73
74
logger = init_logger(__name__)
75
@@ -163,17 +166,22 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
166
"channelwise, dynamic per token quantization."
167
)
168
166
- # For GPUs that lack FP8 hardware support, we can leverage the Marlin
167
- # kernel for fast weight-only FP8 quantization
168
- self.use_marlin = (
169
- not current_platform.has_device_capability(89)
170
- or envs.VLLM_TEST_FORCE_FP8_MARLIN
171
- )
172
- # Disable marlin for rocm
173
- if current_platform.is_rocm():
174
- self.use_marlin = False
169
+ # Determine quant keys for oracle backend selection
170
+ if per_channel:
171
+ weight_key = kFp8StaticChannelSym
172
+ activation_key = kFp8DynamicTokenSym
173
+ elif self.static_input_scales:
174
+ weight_key = kFp8StaticTensorSym
175
+ activation_key = kFp8StaticTensorSym
176
+ else:
177
+ weight_key = kFp8StaticTensorSym
178
+ activation_key = kFp8DynamicTensorSym
179
176
- self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
180
+ self.fp8_backend, self.experts_cls = select_fp8_moe_backend(
181
+ config=moe,
182
+ weight_key=weight_key,
183
+ activation_key=activation_key,
184
+ )
185
186
self.model_type = getattr(
187
get_current_vllm_config().model_config.hf_config, "model_type", None
@@ -407,50 +415,51 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
415
layer.w2_weight_scale = torch.nn.Parameter(
416
w2_weight_scale, requires_grad=False
417
)
410
- # Property to determine if AITER is used
411
- if self.rocm_aiter_moe_enabled:
412
- # reshaping weights is required for aiter moe kernel.
413
- shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
414
- layer.w13_weight.data, layer.w2_weight.data
415
- )
418
+ self._setup_kernel(layer)
419
417
- layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
418
- layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
419
-
420
- elif self.use_marlin:
421
- w13_weight, w2_weight, w13_weight_scale, w2_weight_scale = (
422
- prepare_fp8_moe_layer_for_marlin(
423
- layer,
424
- layer.w13_weight,
425
- layer.w2_weight,
426
- layer.w13_weight_scale,
427
- layer.w2_weight_scale,
428
- )
429
- )
430
- # TODO(rob): once we apply refactor to Quark, switch to using
431
- # replace_parameter for compatibility with reloading in RL.
432
- layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
433
- layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
434
- layer.w13_weight_scale = torch.nn.Parameter(
435
- w13_weight_scale, requires_grad=False
436
- )
437
- layer.w2_weight_scale = torch.nn.Parameter(
438
- w2_weight_scale, requires_grad=False
439
- )
420
+ def _setup_kernel(self, layer: RoutedExperts) -> None:
421
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
422
+ fp8_backend=self.fp8_backend,
423
+ layer=layer,
424
+ w13=layer.w13_weight,
425
+ w2=layer.w2_weight,
426
+ w13_scale=layer.w13_weight_scale,
427
+ w2_scale=layer.w2_weight_scale,
428
+ w13_input_scale=layer.w13_input_scale,
429
+ w2_input_scale=layer.w2_input_scale,
430
+ )
431
+ replace_parameter(layer, "w13_weight", w13)
432
+ replace_parameter(layer, "w2_weight", w2)
433
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
434
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
435
441
- def get_fused_moe_quant_config(
442
- self, layer: RoutedExperts
443
- ) -> FusedMoEQuantConfig | None:
444
- return fp8_w8a8_moe_quant_config(
436
+ if self.fp8_backend == Fp8MoeBackend.AITER:
437
+ layer.w13_weight.is_shuffled = True
438
+ layer.w2_weight.is_shuffled = True
439
+
440
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
441
+ assert self.moe_quant_config is not None
442
+ assert self.experts_cls is not None
443
+ self.moe_kernel = make_fp8_moe_kernel(
444
+ moe_quant_config=self.moe_quant_config,
445
+ moe_config=self.moe,
446
+ fp8_backend=self.fp8_backend,
447
+ experts_cls=self.experts_cls,
448
+ routing_tables=layer._expert_routing_tables(),
449
+ )
450
+
451
+ def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
452
+ return make_fp8_moe_quant_config(
453
+ fp8_backend=self.fp8_backend,
454
w1_scale=layer.w13_weight_scale,
455
w2_scale=layer.w2_weight_scale,
456
a1_scale=layer.w13_input_scale,
457
a2_scale=layer.w2_input_scale,
449
- w1_bias=layer.w13_bias,
450
- w2_bias=layer.w2_bias,
458
+ w1_bias=getattr(layer, "w13_bias", None),
459
+ w2_bias=getattr(layer, "w2_bias", None),
460
per_act_token_quant=self.input_qscheme == "per_channel",
461
per_out_ch_quant=self.weight_qscheme == "per_channel",
453
- gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
462
+ swiglu_limit=getattr(layer, "swiglu_limit", None),
463
)
464
465
def apply(
@@ -462,57 +471,20 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
471
shared_experts: SharedExperts | None,
472
shared_experts_input: torch.Tensor | None,
473
) -> torch.Tensor:
465
- if self.rocm_aiter_moe_enabled:
466
- from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import (
467
- rocm_aiter_fused_experts,
468
- )
469
-
470
- return rocm_aiter_fused_experts(
471
- hidden_states=x,
472
- w1=layer.w13_weight,
473
- w2=layer.w2_weight,
474
- topk_weights=topk_weights,
475
- topk_ids=topk_ids,
476
- activation=layer.activation,
477
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
478
- quant_config=self.moe_quant_config,
479
- moe_config=layer.moe_config,
480
- expert_map=layer.expert_map,
481
- )
482
- elif self.use_marlin:
483
- assert layer.activation == MoEActivation.SILU, (
484
- f"{layer.activation} not supported for Marlin MoE."
485
- )
486
- return fused_marlin_moe(
487
- x,
488
- layer.w13_weight,
489
- layer.w2_weight,
490
- None,
491
- None,
492
- layer.w13_weight_scale,
493
- layer.w2_weight_scale,
494
- topk_weights,
495
- topk_ids,
496
- quant_type_id=scalar_types.float8_e4m3fn.id,
497
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
498
- global_num_experts=layer.global_num_experts,
499
- expert_map=layer.expert_map,
500
- )
501
- else:
502
- from vllm.model_executor.layers.fused_moe import fused_experts
503
-
504
- return fused_experts(
505
- hidden_states=x,
506
- w1=layer.w13_weight,
507
- w2=layer.w2_weight,
508
- topk_weights=topk_weights,
509
- topk_ids=topk_ids,
510
- activation=layer.activation,
511
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
512
- global_num_experts=layer.global_num_experts,
513
- expert_map=layer.expert_map,
514
- quant_config=self.moe_quant_config,
515
- )
474
+ assert self.moe_kernel is not None
475
+ return self.moe_kernel.apply(
476
+ hidden_states=x,
477
+ w1=layer.w13_weight,
478
+ w2=layer.w2_weight,
479
+ topk_weights=topk_weights,
480
+ topk_ids=topk_ids,
481
+ activation=layer.activation,
482
+ global_num_experts=layer.global_num_experts,
483
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
484
+ expert_map=layer.expert_map,
485
+ shared_experts=shared_experts,
486
+ shared_experts_input=shared_experts_input,
487
+ )
488
489
490
class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):