From c97e8f99d69d1bdf02110f01fe3f5e00dcae71fa Mon Sep 17 00:00:00 2001 From: Bowen Bao Date: Mon, 22 Jun 2026 15:58:03 -0700 Subject: [PATCH] [ROCm][Quantization][4/N] refactor quark_moe fp8 w/ oracle (#43721) Signed-off-by: Bowen Bao Co-authored-by: Andreas Karatzas --- .../Qwen3-30B-A3B-Thinking-2507-FP8.yaml | 6 + .../Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml | 6 + .../configs/models-mi3xx-fp8-and-mixed.txt | 2 + tests/evals/gsm8k/test_gsm8k_correctness.py | 1 + .../layers/quantization/quark/quark_moe.py | 178 ++++++++---------- 5 files changed, 90 insertions(+), 103 deletions(-) create mode 100644 tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-FP8.yaml create mode 100644 tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml diff --git a/tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-FP8.yaml b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-FP8.yaml new file mode 100644 index 00000000000..7ec5b825c2b --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-FP8.yaml @@ -0,0 +1,6 @@ +model_name: "amd/Qwen3-30B-A3B-Thinking-2507-FP8" +accuracy_threshold: 0.81 +num_questions: 1319 +num_fewshot: 5 +max_tokens: 1024 +server_args: "--max-model-len 4096 --gpu-memory-utilization 0.85" diff --git a/tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml new file mode 100644 index 00000000000..6095cef535c --- /dev/null +++ b/tests/evals/gsm8k/configs/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml @@ -0,0 +1,6 @@ +model_name: "amd/Qwen3-30B-A3B-Thinking-2507-PTPC-FP8" +accuracy_threshold: 0.81 +num_questions: 1319 +num_fewshot: 5 +max_tokens: 1024 +server_args: "--max-model-len 4096 --gpu-memory-utilization 0.85" diff --git a/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt b/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt index f1122008f59..bcd00044bc0 100644 --- a/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt +++ b/tests/evals/gsm8k/configs/models-mi3xx-fp8-and-mixed.txt @@ -3,3 +3,5 @@ Qwen2.5-VL-3B-Instruct-FP8-dynamic.yaml Qwen1.5-MoE-W4A16-CT.yaml DeepSeek-V2-Lite-Instruct-FP8.yaml Qwen3-Next-FP8-EP2_MI355.yaml +Qwen3-30B-A3B-Thinking-2507-FP8.yaml +Qwen3-30B-A3B-Thinking-2507-PTPC-FP8.yaml diff --git a/tests/evals/gsm8k/test_gsm8k_correctness.py b/tests/evals/gsm8k/test_gsm8k_correctness.py index cd90d71669a..e841048b4a8 100644 --- a/tests/evals/gsm8k/test_gsm8k_correctness.py +++ b/tests/evals/gsm8k/test_gsm8k_correctness.py @@ -48,6 +48,7 @@ def run_gsm8k_eval(eval_config: dict, server_url: str) -> dict: results = evaluate_gsm8k( num_questions=eval_config["num_questions"], num_shots=eval_config["num_fewshot"], + max_tokens=eval_config.get("max_tokens", 256), host=host, port=port, request_timeout_seconds=request_timeout_seconds, diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 703fc815015..5af7a519900 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -7,7 +7,6 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import _custom_ops as ops -from vllm import envs from vllm._aiter_ops import rocm_aiter_ops from vllm.config import get_current_vllm_config from vllm.logger import init_logger @@ -15,7 +14,6 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported, - MoEActivation, RoutedExperts, SharedExperts, ) @@ -28,7 +26,13 @@ from vllm.model_executor.layers.fused_moe.config import ( mxfp4_w4a16_moe_quant_config, ocp_mx_moe_quant_config, ) -from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe +from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + convert_to_fp8_moe_kernel_format, + make_fp8_moe_kernel, + make_fp8_moe_quant_config, + select_fp8_moe_backend, +) from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( TRITON_BACKENDS, Mxfp4MoeBackend, @@ -45,15 +49,15 @@ from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( make_nvfp4_moe_quant_config, select_nvfp4_moe_backend, ) -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import ( - prepare_fp8_moe_layer_for_marlin, -) from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import ( OCP_MX_BLOCK_SIZE, OCP_MX_Scheme, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( GroupShape, + kFp8DynamicTensorSym, + kFp8DynamicTokenSym, + kFp8StaticChannelSym, kFp8StaticTensorSym, kMxfp4Dynamic, kNvfp4Dynamic, @@ -66,7 +70,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform -from vllm.scalar_type import scalar_types logger = init_logger(__name__) @@ -163,17 +166,22 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): "channelwise, dynamic per token quantization." ) - # For GPUs that lack FP8 hardware support, we can leverage the Marlin - # kernel for fast weight-only FP8 quantization - self.use_marlin = ( - not current_platform.has_device_capability(89) - or envs.VLLM_TEST_FORCE_FP8_MARLIN - ) - # Disable marlin for rocm - if current_platform.is_rocm(): - self.use_marlin = False + # Determine quant keys for oracle backend selection + if per_channel: + weight_key = kFp8StaticChannelSym + activation_key = kFp8DynamicTokenSym + elif self.static_input_scales: + weight_key = kFp8StaticTensorSym + activation_key = kFp8StaticTensorSym + else: + weight_key = kFp8StaticTensorSym + activation_key = kFp8DynamicTensorSym - self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + self.fp8_backend, self.experts_cls = select_fp8_moe_backend( + config=moe, + weight_key=weight_key, + activation_key=activation_key, + ) self.model_type = getattr( get_current_vllm_config().model_config.hf_config, "model_type", None @@ -407,50 +415,51 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): layer.w2_weight_scale = torch.nn.Parameter( w2_weight_scale, requires_grad=False ) - # Property to determine if AITER is used - if self.rocm_aiter_moe_enabled: - # reshaping weights is required for aiter moe kernel. - shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights( - layer.w13_weight.data, layer.w2_weight.data - ) + self._setup_kernel(layer) - layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False) + def _setup_kernel(self, layer: RoutedExperts) -> None: + w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format( + fp8_backend=self.fp8_backend, + layer=layer, + w13=layer.w13_weight, + w2=layer.w2_weight, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_input_scale=layer.w13_input_scale, + w2_input_scale=layer.w2_input_scale, + ) + replace_parameter(layer, "w13_weight", w13) + replace_parameter(layer, "w2_weight", w2) + replace_parameter(layer, "w13_weight_scale", w13_scale) + replace_parameter(layer, "w2_weight_scale", w2_scale) - elif self.use_marlin: - w13_weight, w2_weight, w13_weight_scale, w2_weight_scale = ( - prepare_fp8_moe_layer_for_marlin( - layer, - layer.w13_weight, - layer.w2_weight, - layer.w13_weight_scale, - layer.w2_weight_scale, - ) - ) - # TODO(rob): once we apply refactor to Quark, switch to using - # replace_parameter for compatibility with reloading in RL. - layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False) - layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False) - layer.w13_weight_scale = torch.nn.Parameter( - w13_weight_scale, requires_grad=False - ) - layer.w2_weight_scale = torch.nn.Parameter( - w2_weight_scale, requires_grad=False - ) + if self.fp8_backend == Fp8MoeBackend.AITER: + layer.w13_weight.is_shuffled = True + layer.w2_weight.is_shuffled = True - def get_fused_moe_quant_config( - self, layer: RoutedExperts - ) -> FusedMoEQuantConfig | None: - return fp8_w8a8_moe_quant_config( + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + assert self.experts_cls is not None + self.moe_kernel = make_fp8_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + fp8_backend=self.fp8_backend, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + ) + + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig: + return make_fp8_moe_quant_config( + fp8_backend=self.fp8_backend, w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, a1_scale=layer.w13_input_scale, a2_scale=layer.w2_input_scale, - w1_bias=layer.w13_bias, - w2_bias=layer.w2_bias, + w1_bias=getattr(layer, "w13_bias", None), + w2_bias=getattr(layer, "w2_bias", None), per_act_token_quant=self.input_qscheme == "per_channel", per_out_ch_quant=self.weight_qscheme == "per_channel", - gemm1_clamp_limit=getattr(layer, "swiglu_limit", None), + swiglu_limit=getattr(layer, "swiglu_limit", None), ) def apply( @@ -462,57 +471,20 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - if self.rocm_aiter_moe_enabled: - from vllm.model_executor.layers.fused_moe.experts.rocm_aiter_moe import ( - rocm_aiter_fused_experts, - ) - - return rocm_aiter_fused_experts( - hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=layer.activation, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - quant_config=self.moe_quant_config, - moe_config=layer.moe_config, - expert_map=layer.expert_map, - ) - elif self.use_marlin: - assert layer.activation == MoEActivation.SILU, ( - f"{layer.activation} not supported for Marlin MoE." - ) - return fused_marlin_moe( - x, - layer.w13_weight, - layer.w2_weight, - None, - None, - layer.w13_weight_scale, - layer.w2_weight_scale, - topk_weights, - topk_ids, - quant_type_id=scalar_types.float8_e4m3fn.id, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - ) - else: - from vllm.model_executor.layers.fused_moe import fused_experts - - return fused_experts( - hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=layer.activation, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - quant_config=self.moe_quant_config, - ) + assert self.moe_kernel is not None + return self.moe_kernel.apply( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, + ) class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):