@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 01923eec70
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 7, 16:30 UTC

[ROCm][Quantization] GPT OSS Upstream MoE wmxfp4_afp8 with static scales (#30357)

Signed-off-by: Aleksandr Malyshev <maleksan@amd.com> Co-authored-by: Aleksandr Malyshev <maleksan@amd.com>

Aleksandr Malyshev committed Feb 26, 2026 at 14:50 UTC 01923eec7092fd5b718cb9b45eb6df152abe9296
3 files changed +315 -37
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+134 -3
@@ -6,6 +6,7 @@ import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm import _custom_ops as ops
9 +from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.logger import init_logger
11 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
12 from vllm.model_executor.layers.fused_moe.config import (
@@ -178,7 +179,40 @@ def triton_kernel_moe_forward(
179 apply_router_weight_on_input: bool = False,
180 global_num_experts: int = -1,
181 expert_map: torch.Tensor | None = None,
182 + unpadded_N_w1=None,
183 + unpadded_K_w1=None,
184 + unpadded_N_w2=None,
185 + unpadded_K_w2=None,
186 ) -> torch.Tensor:
187 + if (
188 + quant_config is not None
189 + and quant_config.use_mxfp4_w4a8
190 + and rocm_aiter_ops.is_enabled()
191 + ):
192 + from aiter.ops.triton.moe_routing.routing import routing as aiter_routing
193 +
194 + routing_data, gather_idx, scatter_idx = aiter_routing(
195 + gating_output, topk, sm_first=not renormalize
196 + )
197 + return triton_kernel_fused_mxfp4_w4a8_experts(
198 + None,
199 + hidden_states,
200 + w1,
201 + w2,
202 + routing_data,
203 + gather_idx,
204 + scatter_idx,
205 + activation=activation.value,
206 + quant_config=quant_config,
207 + apply_router_weight_on_input=apply_router_weight_on_input,
208 + global_num_experts=global_num_experts,
209 + expert_map=expert_map,
210 + unpadded_N_w1=unpadded_N_w1,
211 + unpadded_K_w1=unpadded_K_w1,
212 + unpadded_N_w2=unpadded_N_w2,
213 + unpadded_K_w2=unpadded_K_w2,
214 + )
215 +
216 if expert_map is not None:
217 # With expert parallelism, legacy_routing produces routing data
218 # using global expert IDs which don't correspond to local weight
@@ -210,6 +244,9 @@ def triton_kernel_moe_forward(
244 effective_global_num_experts = global_num_experts
245
246 output = torch.empty_like(hidden_states)
247 + effective_quant_config = (
248 + quant_config if quant_config is not None else FUSED_MOE_UNQUANTIZED_CONFIG
249 + )
250
251 return triton_kernel_fused_experts(
252 output,
@@ -221,7 +258,7 @@ def triton_kernel_moe_forward(
258 scatter_idx,
259 topk=topk,
260 activation=activation,
224 - quant_config=quant_config,
261 + quant_config=effective_quant_config,
262 apply_router_weight_on_input=apply_router_weight_on_input,
263 global_num_experts=effective_global_num_experts,
264 expert_map=effective_expert_map,
@@ -252,8 +289,7 @@ def triton_kernel_fused_experts(
289 assert activation == MoEActivation.SWIGLUOAI, (
290 "Only SWIGLUOAI activation is supported"
291 )
255 - if quant_config is None:
256 - quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
292 + assert quant_config is not None
293
294 # type check, uint8 means mxfp4
295 assert hidden_states.dtype == torch.bfloat16
@@ -330,6 +366,98 @@ def triton_kernel_fused_experts(
366 return output_tensor
367
368
369 +# This is a triton implementation of the fused_experts function
370 +def triton_kernel_fused_mxfp4_w4a8_experts(
371 + output_tensor: torch.Tensor,
372 + hidden_states: torch.Tensor,
373 + w1, # Tensor or triton_kernels.Tensor
374 + w2, # Tensor or triton_kernels.Tensor
375 + routing_data, # RoutingData
376 + gather_indx, # GatherIndx
377 + scatter_indx, # ScatterIndx
378 + activation: str = "silu",
379 + quant_config: FusedMoEQuantConfig | None = None,
380 + swiglu_alpha: float = 1.702,
381 + swiglu_limit: float = 7.0,
382 + apply_router_weight_on_input: bool = False,
383 + global_num_experts: int = -1,
384 + expert_map: torch.Tensor | None = None,
385 + a1q_scale: torch.Tensor | None = None,
386 + unpadded_N_w1=None,
387 + unpadded_K_w1=None,
388 + unpadded_N_w2=None,
389 + unpadded_K_w2=None,
390 +) -> torch.Tensor:
391 + assert quant_config is not None
392 + # type check, uint8 means mxfp4
393 + assert hidden_states.dtype == torch.bfloat16
394 + assert quant_config.w1_bias is None or quant_config.w1_bias.dtype == torch.float32
395 + assert quant_config.w2_bias is None or quant_config.w2_bias.dtype == torch.float32
396 +
397 + # Shape check, only check non-mxfp4
398 + assert hidden_states.shape[-1] == w1.shape[-2]
399 + assert w2.shape[-1] == w1.shape[1]
400 +
401 + E, _, N = w1.shape
402 +
403 + if global_num_experts == -1:
404 + global_num_experts = E
405 +
406 + gammas = routing_data.gate_scal if routing_data else None
407 +
408 + from aiter.ops.triton.moe_op_gemm_a8w4 import moe_gemm_a8w4
409 + from aiter.ops.triton.quant_moe import downcast_to_static_fp8
410 +
411 + assert quant_config.w1_precision is not None, (
412 + "w1_precision in quant config can't be None"
413 + )
414 + assert quant_config.w2_precision is not None, (
415 + "w2_precision in quant config can't be None"
416 + )
417 +
418 + hidden_states = downcast_to_static_fp8(
419 + hidden_states, quant_config.w1_precision.flex_ctx.lhs_data.scale
420 + )
421 +
422 + intermediate_cache1 = moe_gemm_a8w4(
423 + hidden_states,
424 + w1.storage.data,
425 + None,
426 + quant_config.w1_precision.weight_scale.storage.data,
427 + quant_config.w1_precision.flex_ctx.lhs_data.scale,
428 + quant_config.w2_precision.flex_ctx.lhs_data.scale,
429 + quant_config.w1_bias,
430 + routing_data,
431 + gather_indx=gather_indx,
432 + gammas=gammas if apply_router_weight_on_input else None,
433 + swizzle_mx_scale="CDNA4_SCALE",
434 + out_dtype=torch.float8_e4m3fn,
435 + apply_swiglu=True,
436 + alpha=swiglu_alpha,
437 + limit=swiglu_limit,
438 + unpadded_N=unpadded_N_w1,
439 + unpadded_K=unpadded_K_w1,
440 + )
441 +
442 + intermediate_cache3 = moe_gemm_a8w4(
443 + intermediate_cache1,
444 + w2.storage.data,
445 + None,
446 + quant_config.w2_precision.weight_scale.storage.data,
447 + quant_config.w2_precision.flex_ctx.lhs_data.scale,
448 + None,
449 + quant_config.w2_bias,
450 + routing_data,
451 + scatter_indx=scatter_indx,
452 + gammas=None if apply_router_weight_on_input else gammas,
453 + swizzle_mx_scale="CDNA4_SCALE",
454 + unpadded_N=unpadded_N_w2,
455 + unpadded_K=unpadded_K_w2,
456 + )
457 +
458 + return intermediate_cache3
459 +
460 +
461 def make_routing_data(
462 topk_ids: torch.Tensor,
463 topk_weights: torch.Tensor,
@@ -520,6 +648,9 @@ class OAITritonExperts(BaseOAITritonExperts):
648 expert_tokens_meta: mk.ExpertTokensMetadata | None,
649 apply_router_weight_on_input: bool,
650 ):
651 + if self.quant_config is None:
652 + self.quant_config: FusedMoEQuantConfig = FUSED_MOE_UNQUANTIZED_CONFIG
653 +
654 if expert_map is not None:
655 topk_ids = expert_map[topk_ids]
656
vllm/model_executor/layers/fused_moe/layer.py
+8 -5
@@ -525,16 +525,18 @@ class FusedMoE(CustomOp):
525
526 # Round up hidden size before creating moe_config.
527 # This way moe_config is created with the correct hidden_size from the start.
528 + unpadded_hidden_size = hidden_size
529 + self.model_type = (
530 + self.vllm_config.model_config.hf_config.model_type
531 + if self.vllm_config.model_config is not None
532 + else None
533 + )
534 hidden_size = maybe_roundup_hidden_size(
535 hidden_size=hidden_size,
536 act_dtype=moe_in_dtype,
537 moe_parallel_config=self.moe_parallel_config,
538 is_lora_enabled=vllm_config.lora_config is not None,
533 - model_type=(
534 - self.vllm_config.model_config.hf_config.model_type
535 - if self.vllm_config.model_config is not None
536 - else None
537 - ),
539 + model_type=self.model_type,
540 is_mxfp4_quant=(
541 quant_config is not None and quant_config.is_mxfp4_quant(prefix, self)
542 ),
@@ -610,6 +612,7 @@ class FusedMoE(CustomOp):
612 moe_quant_params = {
613 "num_experts": self.local_num_experts,
614 "hidden_size": hidden_size,
615 + "unpadded_hidden_size": unpadded_hidden_size,
616 "intermediate_size_per_partition": self.intermediate_size_per_partition,
617 "params_dtype": params_dtype,
618 "weight_loader": self.weight_loader,
vllm/model_executor/layers/quantization/quark/quark_moe.py
+173 -29
@@ -5,8 +5,8 @@ from typing import Any
5
6 import torch
7
8 -import vllm.envs as envs
8 from vllm import _custom_ops as ops
9 +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
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.mxfp4 import (
32 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
33 prepare_fp8_moe_layer_for_marlin,
34 )
35 +from vllm.model_executor.layers.quantization.utils.mxfp4_utils import _swizzle_mxfp4
36 from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
37 OCP_MX_BLOCK_SIZE,
38 OCP_MX_Scheme,
@@ -49,7 +50,11 @@ from vllm.utils.math_utils import round_up
50
51 logger = init_logger(__name__)
52
52 -__all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkOCP_MX_MoEMethod"]
53 +__all__ = [
54 + "QuarkMoEMethod",
55 + "QuarkOCP_MX_MoEMethod",
56 + "QuarkOCP_MX_MoEMethod_OSS",
57 +]
58
59
60 class QuarkMoEMethod(FusedMoEMethodBase):
@@ -71,14 +76,30 @@ class QuarkMoEMethod(FusedMoEMethodBase):
76 "output_tensors and bias "
77 "quantized are not supported"
78 )
79 +
80 weight_config = layer_quant_config.get("weight")
81 input_config = layer_quant_config.get("input_tensors")
82 +
83 if quant_config._is_fp8_w4a8(weight_config, input_config):
84 return QuarkW4A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
85 elif quant_config._is_fp8_w8a8(weight_config, input_config):
86 return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
87 elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config):
81 - return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config)
88 + emulate = not current_platform.supports_mx() or not (
89 + rocm_aiter_ops.is_fused_moe_enabled()
90 + )
91 + if (
92 + input_config.get("dtype") == "fp8_e4m3"
93 + and not input_config.get("is_dynamic")
94 + and not emulate
95 + ):
96 + return QuarkOCP_MX_MoEMethod_OSS(
97 + weight_config, input_config, module.moe_config
98 + )
99 + else:
100 + return QuarkOCP_MX_MoEMethod(
101 + weight_config, input_config, module.moe_config
102 + )
103 else:
104 raise RuntimeError("Unsupported FusedMoe scheme")
105
@@ -706,13 +727,11 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
727 get_current_vllm_config().model_config.hf_config, "model_type", None
728 )
729
709 - self._emulate = (
730 + self.emulate = (
731 not current_platform.supports_mx()
732 or not self.ocp_mx_scheme.startswith("w_mxfp4")
733 ) and (self.mxfp4_backend is None or not self.use_rocm_aiter_moe)
734
714 - self.emulate = True if self.model_type == "gpt_oss" else self._emulate
715 -
735 if self.emulate:
736 logger.warning_once(
737 f"The current mode (supports_mx={current_platform.supports_mx()}, "
@@ -753,6 +772,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
772 )
773
774 params_dtype = torch.uint8
775 + self.intermediate_size_per_partition = intermediate_size_per_partition
776 if self.model_type == "gpt_oss":
777 if current_platform.is_rocm():
778 intermediate_size_per_partition_after_pad = round_up(
@@ -765,6 +785,10 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
785 else:
786 intermediate_size_per_partition_after_pad = intermediate_size_per_partition
787
788 + self.unpadded_hidden_size = extra_weight_attrs.get(
789 + "unpadded_hidden_size", hidden_size
790 + )
791 +
792 # WEIGHTS
793 w13_weight = torch.nn.Parameter(
794 torch.empty(
@@ -991,30 +1015,20 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1015 shared_experts_input: torch.Tensor | None,
1016 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1017 if not self.emulate:
994 - if (
995 - self.model_type == "gpt_oss"
996 - and self.mxfp4_backend == Mxfp4Backend.TRITON
997 - ):
998 - raise NotImplementedError(
999 - "Triton kernel implemented fused MoE for GPT_OSS model "
1000 - "in Quark(MoE) format is not integrated or provided yet."
1001 - )
1002 -
1003 - else:
1004 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1005 - rocm_aiter_fused_experts,
1006 - )
1018 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1019 + rocm_aiter_fused_experts,
1020 + )
1021
1008 - return rocm_aiter_fused_experts(
1009 - x,
1010 - layer.w13_weight,
1011 - layer.w2_weight,
1012 - topk_weights=topk_weights,
1013 - topk_ids=topk_ids,
1014 - activation=layer.activation,
1015 - quant_config=self.moe_quant_config,
1016 - expert_map=layer.expert_map,
1017 - )
1022 + return rocm_aiter_fused_experts(
1023 + x,
1024 + layer.w13_weight,
1025 + layer.w2_weight,
1026 + topk_weights=topk_weights,
1027 + topk_ids=topk_ids,
1028 + activation=layer.activation,
1029 + quant_config=self.moe_quant_config,
1030 + expert_map=layer.expert_map,
1031 + )
1032 else:
1033 from vllm.model_executor.layers.fused_moe import fused_experts
1034
@@ -1031,3 +1045,133 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1045 expert_map=layer.expert_map,
1046 quant_config=self.moe_quant_config,
1047 )
1048 +
1049 +
1050 +class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
1051 + def __init__(
1052 + self,
1053 + weight_config: dict[str, Any],
1054 + input_config: dict[str, Any],
1055 + moe: FusedMoEConfig,
1056 + ):
1057 + super().__init__(weight_config, input_config, moe)
1058 +
1059 + def process_weights_after_loading(self, layer):
1060 + from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
1061 +
1062 + w13_bias = layer.w13_bias.to(torch.float32)
1063 + w2_bias = layer.w2_bias.to(torch.float32)
1064 +
1065 + layer.w13_bias = torch.nn.Parameter(w13_bias, requires_grad=False)
1066 + layer.w2_bias = torch.nn.Parameter(w2_bias, requires_grad=False)
1067 +
1068 + # FIXME warp need to be adjusted based on batch size
1069 + # only apply to batched mode
1070 + if self.moe.use_ep:
1071 + num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
1072 + else:
1073 + num_warps = 8
1074 +
1075 + w13_weight, w13_flex, w13_scale = _swizzle_mxfp4(
1076 + layer.w13_weight, layer.w13_weight_scale, num_warps
1077 + )
1078 + w2_weight, w2_flex, w2_scale = _swizzle_mxfp4(
1079 + layer.w2_weight, layer.w2_weight_scale, num_warps
1080 + )
1081 +
1082 + self.w13_weight_triton_tensor = w13_weight
1083 + self.w2_weight_triton_tensor = w2_weight
1084 +
1085 + # need to delete the original weights to save memory on single GPU
1086 + del layer.w13_weight
1087 + del layer.w2_weight
1088 + layer.w13_weight = None
1089 + layer.w2_weight = None
1090 + torch.cuda.empty_cache()
1091 +
1092 + if self.static_input_scales:
1093 + if layer.w13_input_scale is None or layer.w2_input_scale is None:
1094 + raise ValueError(
1095 + "QuantConfig has static quantization, but found "
1096 + "activation scales are None."
1097 + )
1098 + if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
1099 + layer.w2_input_scale
1100 + ):
1101 + logger.warning_once(
1102 + "Found input_scales that are not equal for "
1103 + "fp8 MoE layer. Using the maximum across experts "
1104 + "for each layer."
1105 + )
1106 +
1107 + layer.w13_input_scale = torch.nn.Parameter(
1108 + layer.w13_input_scale.max().to(torch.float32), requires_grad=False
1109 + )
1110 + layer.w2_input_scale = torch.nn.Parameter(
1111 + layer.w2_input_scale.max().to(torch.float32), requires_grad=False
1112 + )
1113 +
1114 + from triton_kernels.numerics import InFlexData
1115 +
1116 + lhs_data13 = InFlexData(scale=layer.w13_input_scale)
1117 + lhs_data2 = InFlexData(scale=layer.w2_input_scale)
1118 +
1119 + self.w13_precision_config = PrecisionConfig(
1120 + weight_scale=w13_scale,
1121 + flex_ctx=FlexCtx(rhs_data=w13_flex, lhs_data=lhs_data13),
1122 + )
1123 +
1124 + self.w2_precision_config = PrecisionConfig(
1125 + weight_scale=w2_scale,
1126 + flex_ctx=FlexCtx(rhs_data=w2_flex, lhs_data=lhs_data2),
1127 + )
1128 +
1129 + def get_fused_moe_quant_config(
1130 + self, layer: torch.nn.Module
1131 + ) -> FusedMoEQuantConfig | None:
1132 + return mxfp4_w4a8_moe_quant_config(
1133 + w1_scale=self.w13_precision_config,
1134 + w2_scale=self.w2_precision_config,
1135 + a1_scale=layer.w13_input_scale,
1136 + a2_scale=layer.w2_input_scale,
1137 + w1_bias=layer.w13_bias,
1138 + w2_bias=layer.w2_bias,
1139 + block_shape=None,
1140 + )
1141 +
1142 + @property
1143 + def is_monolithic(self) -> bool:
1144 + return True
1145 +
1146 + def apply_monolithic(
1147 + self,
1148 + layer: torch.nn.Module,
1149 + x: torch.Tensor,
1150 + router_logits: torch.Tensor,
1151 + expert_map: torch.Tensor | None = None,
1152 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1153 + if layer.enable_eplb:
1154 + raise NotImplementedError(
1155 + "EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet."
1156 + )
1157 +
1158 + from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
1159 + triton_kernel_moe_forward,
1160 + )
1161 +
1162 + return triton_kernel_moe_forward(
1163 + hidden_states=x,
1164 + w1=self.w13_weight_triton_tensor,
1165 + w2=self.w2_weight_triton_tensor,
1166 + gating_output=router_logits,
1167 + topk=layer.top_k,
1168 + renormalize=layer.renormalize,
1169 + global_num_experts=layer.global_num_experts,
1170 + expert_map=expert_map,
1171 + quant_config=self.moe_quant_config,
1172 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
1173 + unpadded_N_w1=self.intermediate_size_per_partition * 2,
1174 + unpadded_K_w1=self.unpadded_hidden_size,
1175 + unpadded_N_w2=self.unpadded_hidden_size,
1176 + unpadded_K_w2=self.intermediate_size_per_partition,
1177 + )