@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b129136c7a
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 8, 14:13 UTC

[ROCm][Quantization] GPT_OSS in amd-quark format model loading and emulations (#29008)

Signed-off-by: xuebwang-amd <xuebwang@amd.com> Signed-off-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <robshaw@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

xuebwang-amd committed Feb 10, 2026 at 23:08 UTC b129136c7a7389133c923123a1ebd76c4401c94d
13 files changed +1094 -213
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+13 -7
@@ -22,7 +22,7 @@ from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
22 from triton_kernels.tensor_details import layout
23 from triton_kernels.testing import assert_close
24
25 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
25 +from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
26 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
27 triton_kernel_moe_forward,
28 )
@@ -298,12 +298,18 @@ def test_equiv(num_token, a_dtype, w_dtype, tp, workspace_init):
298 pc2,
299 ) = init_compute_data(M, K, N, E, a_dtype, w_dtype, num_warps=8)
300
301 - quant_config = FusedMoEQuantConfig.make(
302 - w1_bias=w1_bias_tri,
303 - w2_bias=w2_bias_tri,
304 - w1_scale=pc1,
305 - w2_scale=pc2,
306 - )
301 + if a_dtype == "bf16" and w_dtype == "mx4":
302 + quant_config = mxfp4_w4a16_moe_quant_config(
303 + w1_scale=pc1,
304 + w2_scale=pc2,
305 + w1_bias=w1_bias_tri,
306 + w2_bias=w2_bias_tri,
307 + )
308 + else:
309 + raise NotImplementedError(
310 + f"Quantization configuration for activation={a_dtype} and weight={w_dtype} "
311 + f"has not been implemented."
312 + )
313
314 out_triton_monolithic = triton_kernel_moe_forward(
315 hidden_states=x_tri,
tests/models/quantization/test_gpt_oss.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +End-to-end accuracy test for GPT-OSS model quantization.
5 +
6 +Config:
7 + Task: gsm8k_platinum
8 + Filter: flexible-extract
9 + n-shot: 5
10 + Metric: exact_match
11 +
12 +Run: pytest tests/models/quantization/test_gpt_oss.py
13 +"""
14 +
15 +import importlib
16 +import importlib.metadata
17 +from dataclasses import dataclass
18 +
19 +import huggingface_hub
20 +import lm_eval
21 +import pytest
22 +from packaging import version
23 +
24 +MODEL_ACCURACIES = {
25 + # Full quantization: attention linears and MoE linears
26 + "amd/gpt-oss-20b-WFP8-AFP8-KVFP8": 0.89,
27 + # MoE linears only quantization
28 + "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8": 0.89,
29 + # MoE linears only quantization
30 + # "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-MXFP4-KV-FP8": 0.90,
31 +}
32 +
33 +QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse(
34 + importlib.metadata.version("amd-quark")
35 +) >= version.parse("0.9.0")
36 +
37 +
38 +def has_huggingface_access(repo):
39 + try:
40 + huggingface_hub.list_repo_refs(repo)
41 + return True
42 + except huggingface_hub.errors.RepositoryNotFoundError:
43 + return False
44 +
45 +
46 +HF_HUB_AMD_ORG_ACCESS = all(
47 + [has_huggingface_access(model_name) for model_name in MODEL_ACCURACIES]
48 +)
49 +
50 +
51 +@dataclass
52 +class ModelCase:
53 + model_id: str
54 + tp: int
55 +
56 +
57 +@dataclass
58 +class EvaluationConfig:
59 + model_name: str
60 +
61 + def get_model_args(self, tp_size: int):
62 + return {
63 + "pretrained": self.model_name,
64 + "chat_template_args": {"reasoning_effort": "low"},
65 + "enable_thinking": True,
66 + "think_end_token": "200008",
67 + "tensor_parallel_size": tp_size,
68 + "dtype": "auto",
69 + "gpu_memory_utilization": 0.95,
70 + "trust_remote_code": False,
71 + "enable_prefix_caching": False,
72 + "enforce_eager": False,
73 + }
74 +
75 +
76 +@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
77 +@pytest.mark.skipif(
78 + not HF_HUB_AMD_ORG_ACCESS,
79 + reason="Read access to huggingface.co/amd is required for this test.",
80 +)
81 +@pytest.mark.parametrize("tp_size", [1, 2, 4, 8])
82 +@pytest.mark.parametrize("model_name, expected_accuracy", MODEL_ACCURACIES.items())
83 +def test_gpt_oss_attention_quantization(
84 + model_name: str, tp_size: int, expected_accuracy: float
85 +):
86 + model_args = EvaluationConfig(model_name).get_model_args(tp_size)
87 +
88 + extra_run_kwargs = {
89 + "gen_kwargs": {"max_gen_toks": 8000},
90 + "apply_chat_template": True,
91 + "fewshot_as_multiturn": True,
92 + "num_fewshot": 5,
93 + }
94 +
95 + lm_eval_out = lm_eval.simple_evaluate(
96 + model="vllm",
97 + model_args=model_args,
98 + tasks="gsm8k_platinum",
99 + batch_size="auto",
100 + **extra_run_kwargs,
101 + )
102 + measured_accuracy = float(
103 + lm_eval_out["results"]["gsm8k_platinum"]["exact_match,flexible-extract"]
104 + )
105 +
106 + rtol = 0.02
107 + assert (
108 + measured_accuracy - rtol < expected_accuracy
109 + and measured_accuracy + rtol > expected_accuracy
110 + ), f"Expected: {expected_accuracy} | Measured: {measured_accuracy}"
tests/models/quantization/test_gpt_oss_attn_quantization.py deleted
-80
@@ -1,80 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Test attention quantization of gpt-oss model.
4 -The qkv_proj and o_proj in self_attention can be either quantized or excluded.
5 -
6 -Run `pytest tests/models/quantization/test_gpt_oss_attn_quantization.py`.
7 -
8 -"""
9 -
10 -import importlib
11 -import importlib.metadata
12 -from dataclasses import dataclass
13 -
14 -import huggingface_hub
15 -import lm_eval
16 -import pytest
17 -from packaging import version
18 -
19 -MODEL_NAMES = ["amd/gpt-oss-20b-customized-attention-quantization"]
20 -
21 -QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse(
22 - importlib.metadata.version("amd-quark")
23 -) >= version.parse("0.8.99")
24 -
25 -
26 -def has_huggingface_access(repo):
27 - try:
28 - huggingface_hub.list_repo_refs(repo)
29 - return True
30 - except huggingface_hub.errors.RepositoryNotFoundError:
31 - return False
32 -
33 -
34 -HF_HUB_AMD_ORG_ACCESS = all(
35 - [has_huggingface_access(model_name) for model_name in MODEL_NAMES]
36 -)
37 -
38 -
39 -@dataclass
40 -class ModelCase:
41 - model_id: str
42 - tp: int
43 -
44 -
45 -@dataclass
46 -class EvaluationConfig:
47 - model_name: str
48 -
49 - def get_model_args(self) -> str:
50 - return (
51 - f"pretrained={self.model_name},"
52 - "tensor_parallel_size=4,dtype=auto,gpu_memory_utilization=0.9,trust_remote_code=False"
53 - )
54 -
55 -
56 -EXPECTED_ACCURACIES = {"arc_challenge": 0.20}
57 -
58 -
59 -@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
60 -@pytest.mark.skipif(
61 - not HF_HUB_AMD_ORG_ACCESS,
62 - reason="Read access to huggingface.co/amd is required for this test.",
63 -)
64 -@pytest.mark.parametrize("model_name", MODEL_NAMES)
65 -@pytest.mark.parametrize("task_name, expected_accuracy", EXPECTED_ACCURACIES.items())
66 -def test_gpt_oss_attention_quantization(
67 - model_name: str, task_name: str, expected_accuracy: float
68 -):
69 - measured_accuracy = lm_eval.simple_evaluate(
70 - model="vllm",
71 - model_args=EvaluationConfig(model_name).get_model_args(),
72 - tasks=task_name,
73 - batch_size="auto",
74 - )["results"][task_name]["acc,none"]
75 -
76 - rtol = 0.05
77 - assert (
78 - measured_accuracy - rtol < expected_accuracy
79 - and measured_accuracy + rtol > expected_accuracy
80 - ), f"Expected: {expected_accuracy} | Measured: {measured_accuracy}"
vllm/model_executor/layers/fused_moe/config.py
+36
@@ -386,6 +386,10 @@ class FusedMoEQuantConfig:
386 def use_nvfp4_w4a4(self) -> bool:
387 return self.quant_dtype == "nvfp4"
388
389 + @property
390 + def use_mxfp4_w4a8(self) -> bool:
391 + return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4"
392 +
393 def config_name(self, dtype: torch.dtype) -> str | None:
394 """
395 Return a string used to construct the filename that contains the
@@ -532,6 +536,8 @@ def fp8_w8a8_moe_quant_config(
536 w2_scale: torch.Tensor,
537 a1_scale: torch.Tensor | None = None,
538 a2_scale: torch.Tensor | None = None,
539 + w1_bias: torch.Tensor | None = None,
540 + w2_bias: torch.Tensor | None = None,
541 per_act_token_quant: bool = False,
542 per_out_ch_quant: bool = False,
543 block_shape: list[int] | None = None,
@@ -549,6 +555,8 @@ def fp8_w8a8_moe_quant_config(
555 g1_alphas=g1_alphas,
556 w2_scale=w2_scale,
557 g2_alphas=g2_alphas,
558 + w1_bias=w1_bias,
559 + w2_bias=w2_bias,
560 a1_scale=a1_scale,
561 a1_gscale=a1_gscale,
562 a2_scale=a2_scale,
@@ -564,6 +572,8 @@ def int8_w8a8_moe_quant_config(
572 w2_scale: torch.Tensor,
573 a1_scale: torch.Tensor | None,
574 a2_scale: torch.Tensor | None,
575 + w1_bias: torch.Tensor | None = None,
576 + w2_bias: torch.Tensor | None = None,
577 per_act_token_quant: bool = False,
578 ) -> FusedMoEQuantConfig:
579 """
@@ -575,6 +585,8 @@ def int8_w8a8_moe_quant_config(
585 w2_scale=w2_scale,
586 a1_scale=a1_scale,
587 a2_scale=a2_scale,
588 + w1_bias=w1_bias,
589 + w2_bias=w2_bias,
590 per_act_token_quant=per_act_token_quant,
591 per_out_ch_quant=False,
592 block_shape=None,
@@ -654,6 +666,26 @@ def mxfp4_mxfp8_moe_quant_config(
666 )
667
668
669 +def mxfp4_w4a8_moe_quant_config(
670 + w1_scale: Union[torch.Tensor, "PrecisionConfig"],
671 + w2_scale: Union[torch.Tensor, "PrecisionConfig"],
672 + a1_scale: torch.Tensor | None = None,
673 + a2_scale: torch.Tensor | None = None,
674 + w1_bias: torch.Tensor | None = None,
675 + w2_bias: torch.Tensor | None = None,
676 + block_shape: list[int] | None = None,
677 +) -> FusedMoEQuantConfig:
678 + """
679 + Construct a quant config for fp8 activations and mxfp4 weights.
680 + """
681 + return FusedMoEQuantConfig(
682 + _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None),
683 + _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None),
684 + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
685 + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
686 + )
687 +
688 +
689 def ocp_mx_moe_quant_config(
690 quant_dtype: str,
691 w1_scale: Union[torch.Tensor, "PrecisionConfig"],
@@ -691,6 +723,8 @@ def nvfp4_moe_quant_config(
723 a2_gscale: torch.Tensor,
724 w1_scale: torch.Tensor,
725 w2_scale: torch.Tensor,
726 + w1_bias: torch.Tensor | None = None,
727 + w2_bias: torch.Tensor | None = None,
728 ) -> FusedMoEQuantConfig:
729 """
730 Construct a quant config for mxfp4 activations and nvp4 weights.
@@ -699,6 +733,8 @@ def nvfp4_moe_quant_config(
733 "nvfp4",
734 w1_scale=w1_scale,
735 w2_scale=w2_scale,
736 + w1_bias=w1_bias,
737 + w2_bias=w2_bias,
738 a1_gscale=a1_gscale,
739 a2_gscale=a2_gscale,
740 g1_alphas=g1_alphas,
vllm/model_executor/layers/fused_moe/fused_moe.py
+12 -17
@@ -38,7 +38,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
38 )
39 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
40 from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
41 -from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme
41 from vllm.model_executor.layers.quantization.utils.quant_utils import (
42 QuantKey,
43 kFp8Dynamic128Sym,
@@ -1583,6 +1582,11 @@ def _get_config_quant_dtype(
1582 return "mxfp6_e3m2"
1583 elif ocp_mx_scheme in {"w_mxfp4_a_mxfp6_e2m3", "w_mxfp6_e2m3_a_mxfp6_e2m3"}:
1584 return "mxfp6_e2m3"
1585 + elif ocp_mx_scheme in {"w_mxfp4", "w_mxfp6_e3m2", "w_mxfp6_e2m3"}:
1586 + return torch.bfloat16
1587 + elif ocp_mx_scheme in {"w_mxfp4_a_fp8", "w_mxfp6_e3m2_a_fp8", "w_mxfp6_e2m3_a_fp8"}:
1588 + return torch.float8_e4m3fn
1589 +
1590 return None
1591
1592
@@ -1617,17 +1621,10 @@ def fused_experts_impl(
1621 if use_int4_w4a16:
1622 assert hidden_states.size(1) // 2 == w1.size(2), "Hidden size mismatch"
1623 elif ocp_mx_scheme is not None:
1620 - if ocp_mx_scheme in {
1621 - "w_mxfp4_a_mxfp4",
1622 - "w_mxfp4_a_mxfp6_e3m2",
1623 - "w_mxfp4_a_mxfp6_e2m3",
1624 - }:
1624 + if ocp_mx_scheme.startswith("w_mxfp4"):
1625 # 16bit activation and fp4x2 packed weight
1626 assert hidden_states.size(1) == w1.size(2) * 2, "hidden size mismatch"
1627 - elif ocp_mx_scheme in {
1628 - "w_mxfp6_e3m2_a_mxfp6_e3m2",
1629 - "w_mxfp6_e2m3_a_mxfp6_e2m3",
1630 - }:
1627 + elif ocp_mx_scheme.startswith("w_mxfp6"):
1628 assert hidden_states.size(1) == (w1.size(2) * 4) // 3, (
1629 "hidden size mismatch"
1630 )
@@ -1717,17 +1714,13 @@ def fused_experts_impl(
1714 # TODO: On platforms for which `current_platform.supports_mx()` is True
1715 # and for which we have a native OCP mx fused MOE kernel,
1716 # this dequantization step should not be done.
1720 - if ocp_mx_scheme in {
1721 - OCP_MX_Scheme.w_mxfp4_a_mxfp4,
1722 - OCP_MX_Scheme.w_mxfp4_a_mxfp6_e3m2,
1723 - OCP_MX_Scheme.w_mxfp4_a_mxfp6_e2m3,
1724 - }:
1717 + if ocp_mx_scheme.startswith("w_mxfp4"):
1718 # Weight has to be dequantized for mxfp4 emulation.
1719 w1 = dequant_mxfp4(w1, w1_scale, hidden_states.dtype)
1720 w1_scale = None
1721 w2 = dequant_mxfp4(w2, w2_scale, hidden_states.dtype)
1722 w2_scale = None
1730 - elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e3m2_a_mxfp6_e3m2:
1723 + elif ocp_mx_scheme.startswith("w_mxfp6_e3m2"):
1724 w1 = dequant_mxfp6(
1725 w1, w1_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
1726 )
@@ -1736,7 +1729,7 @@ def fused_experts_impl(
1729 w2, w2_scale, quant_dtype="fp6_e3m2", float_dtype=hidden_states.dtype
1730 )
1731 w2_scale = None
1739 - elif ocp_mx_scheme == OCP_MX_Scheme.w_mxfp6_e2m3_a_mxfp6_e2m3:
1732 + elif ocp_mx_scheme.startswith("w_mxfp6_e2m3"):
1733 w1 = dequant_mxfp6(
1734 w1, w1_scale, quant_dtype="fp6_e2m3", float_dtype=hidden_states.dtype
1735 )
@@ -1779,6 +1772,7 @@ def fused_experts_impl(
1772 quant_dtype=quant_dtype,
1773 per_act_token_quant=per_channel_quant,
1774 block_shape=block_shape,
1775 + ocp_mx_scheme=ocp_mx_scheme,
1776 )
1777
1778 # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k
@@ -1846,6 +1840,7 @@ def fused_experts_impl(
1840 quant_dtype=quant_dtype,
1841 per_act_token_quant=per_channel_quant,
1842 block_shape=block_shape,
1843 + ocp_mx_scheme=ocp_mx_scheme,
1844 )
1845
1846 if expert_map is not None:
vllm/model_executor/layers/fused_moe/layer.py
+24 -13
@@ -221,12 +221,14 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
221 )
222
223
224 +# TODO(rob): move this down to the kernel.
225 def maybe_roundup_hidden_size(
226 hidden_size: int,
227 act_dtype: torch.dtype,
227 - quant_config: QuantizationConfig | None,
228 moe_parallel_config: FusedMoEParallelConfig,
229 is_lora_enabled: bool,
230 + model_type: str | None,
231 + is_mxfp4_quant: bool,
232 ) -> int:
233 """
234 Given layer hidden size and MoE configurations, round up hidden_size
@@ -235,11 +237,12 @@ def maybe_roundup_hidden_size(
237 Args:
238 hidden_size: Layer hidden-size
239 act_dtype: Data type of the layer activations.
238 - quant_config: Fused MoE quantization configuration.
240 moe_parallel_config: Fused MoE parallelization strategy configuration.
241 is_lora_enabled: True if the engine is enabled with LoRA. This
242 is used in the case of mxfp4 quantization in selecting the
243 MxFP4Backend.
244 + model_type: for checking if gpt-oss
245 + is_mxfp4_quant: whether the layer is quantized with mxfp4
246
247 Return:
248 Rounded up hidden_size if rounding up is required based on the configs.
@@ -254,7 +257,7 @@ def maybe_roundup_hidden_size(
257 )
258
259 # we are padding globally so EP buffer allocation works
257 - if quant_config and quant_config.get_name() == "mxfp4":
260 + if model_type == "gpt_oss" and is_mxfp4_quant:
261 from vllm.model_executor.layers.quantization.mxfp4 import (
262 Mxfp4Backend,
263 get_mxfp4_backend,
@@ -398,15 +401,6 @@ class FusedMoE(CustomOp):
401 # Expert mapping used in self.load_weights
402 self.expert_mapping = expert_mapping
403
401 - # Round up hidden size if needed.
402 - hidden_size = maybe_roundup_hidden_size(
403 - hidden_size,
404 - moe_in_dtype,
405 - quant_config,
406 - self.moe_parallel_config,
407 - is_lora_enabled=self.vllm_config.lora_config is not None,
408 - )
409 -
404 # For smuggling this layer into the fused moe custom op
405 compilation_config = vllm_config.compilation_config
406 if prefix in compilation_config.static_forward_context:
@@ -508,7 +502,6 @@ class FusedMoE(CustomOp):
502 ), "Aiter Fused MoE kernel only supports expert_map with 0 and 1s."
503
504 assert intermediate_size % self.tp_size == 0
511 - self.hidden_size = hidden_size
505 self.intermediate_size_per_partition = intermediate_size // self.tp_size
506 self.reduce_results = reduce_results
507 self.renormalize = renormalize
@@ -548,6 +541,24 @@ class FusedMoE(CustomOp):
541 )
542 self.routing_method_type: RoutingMethodType = self.router.routing_method_type
543
544 + # Round up hidden size before creating moe_config.
545 + # This way moe_config is created with the correct hidden_size from the start.
546 + hidden_size = maybe_roundup_hidden_size(
547 + hidden_size=hidden_size,
548 + act_dtype=moe_in_dtype,
549 + moe_parallel_config=self.moe_parallel_config,
550 + is_lora_enabled=vllm_config.lora_config is not None,
551 + model_type=(
552 + self.vllm_config.model_config.hf_config.model_type
553 + if self.vllm_config.model_config is not None
554 + else None
555 + ),
556 + is_mxfp4_quant=(
557 + quant_config is not None and quant_config.is_mxfp4_quant(prefix, self)
558 + ),
559 + )
560 + self.hidden_size = hidden_size
561 +
562 self.moe_config: FusedMoEConfig = FusedMoEConfig(
563 num_experts=self.global_num_experts,
564 experts_per_token=top_k,
vllm/model_executor/layers/fused_moe/utils.py
+23
@@ -23,6 +23,9 @@ from vllm.model_executor.layers.quantization.utils.mxfp6_utils import (
23 from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
24 mxfp8_e4m3_quantize,
25 )
26 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27 + per_tensor_dequantize,
28 +)
29 from vllm.triton_utils import tl, triton
30 from vllm.utils.math_utils import cdiv
31 from vllm.utils.torch_utils import is_torch_equal_or_newer
@@ -241,7 +244,27 @@ def moe_kernel_quantize_input(
244 per_act_token_quant: bool,
245 block_shape: list[int] | None = None,
246 is_fp4_scale_swizzled: bool = True,
247 + ocp_mx_scheme: str | None = None,
248 ) -> tuple[torch.Tensor, torch.Tensor | None]:
249 + # Handle OCP MX scheme that requires QDQ (quantize-dequantize) for emulation
250 + if ocp_mx_scheme is not None:
251 + if ocp_mx_scheme in {"w_mxfp4", "w_mxfp4_a_mxfp4"}:
252 + pass # No QDQ needed for these schemes
253 + elif ocp_mx_scheme.endswith("a_fp8"):
254 + # Perform QDQ (quantize and dequantize) on activation for emulation
255 + # purpose, because there is no native kernel for weight in ocp_mx_scheme
256 + # and activation in FP8. The implementation is based on existing
257 + # non-emulation ops.
258 + qA, qA_scale = ops.scaled_fp8_quant(
259 + A, A_scale, use_per_token_if_dynamic=False
260 + )
261 + A = per_tensor_dequantize(qA, qA_scale).to(A.dtype)
262 + # After QDQ, we don't need further quantization
263 + return A, None
264 + # else: For other schemes (e.g., *_a_mxfp6_e3m2, *_a_mxfp6_e2m3),
265 + # weights are already dequantized, and we proceed with normal
266 + # activation quantization below.
267 +
268 if quant_dtype == torch.float8_e4m3fn:
269 return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
270 elif quant_dtype == torch.int8:
vllm/model_executor/layers/quantization/base_config.py
+16
@@ -168,3 +168,19 @@ class QuantizationConfig(ABC):
168 Interface to update values after config initialization.
169 """
170 pass
171 +
172 + def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
173 + """
174 + Determine if mxfp4 quantization will be used for this config.
175 +
176 + This allows hidden_size rounding to happen before moe_config creation
177 + without needing to instantiate quant_method first.
178 +
179 + Args:
180 + prefix: The layer prefix/name in the model
181 + layer: The layer module
182 +
183 + Returns:
184 + True if this config uses MXFP4 quantization, False otherwise
185 + """
186 + return False
vllm/model_executor/layers/quantization/mxfp4.py
+5
@@ -229,10 +229,15 @@ class Mxfp4Config(QuantizationConfig):
229 )
230 return None
231
232 + def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
233 + """MXFP4 config always uses MXFP4 quantization."""
234 + return True
235 +
236
237 class Mxfp4MoEMethod(FusedMoEMethodBase):
238 def __init__(self, moe: FusedMoEConfig):
239 super().__init__(moe)
240 + self.weight_dtype = "mxfp4"
241 self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
242
243 self.marlin_input_dtype = None
vllm/model_executor/layers/quantization/quark/quark.py
+47 -23
@@ -320,38 +320,45 @@ class QuarkConfig(QuantizationConfig):
320 # Only symmetric weight quantization supported.
321 return is_int8_dtype and is_tensor and is_weight_symmetric and is_static
322
323 - def _is_ocp_mx(
324 - self,
325 - weight_quant: dict[str, Any] | None,
326 - input_quant: dict[str, Any] | None,
323 + def _is_w_ocp_mx_a_x(
324 + self, weight_quant: dict[str, Any] | None, input_quant: dict[str, Any] | None
325 ) -> bool:
328 - # Confirm weights and input quantized.
329 - if weight_quant is None or input_quant is None:
326 + """
327 + This check returns True only if it is an OCP-MX weight quantization.
328 + The activation can be any data type (e.g., FP16/BF16, FP8, or OCP-MX format).
329 + The rationale for checking only the weight type is that
330 + the model loading concept and process primarily concerns the weights themselves.
331 + """
332 + # Confirm weights quantized.
333 + if weight_quant is None:
334 logger.debug(
331 - "Quark model is not in OCP MX format: "
332 - "weight_quant or input_quant not set"
335 + "Quark model's weight quantization is incompatible with OCP_MX format: "
336 + "weight_quant is not set."
337 )
338 return False
339
340 # Input and weight qscheme needs to be per group.
337 - if (
338 - weight_quant.get("qscheme") != "per_group"
339 - or input_quant.get("qscheme") != "per_group"
340 - ):
341 - logger.debug("Quark model is not in OCP MX format: not per_group")
341 + if weight_quant.get("qscheme") != "per_group":
342 + logger.debug(
343 + "Quark model's weight quantization is incompatible with OCP MX format: "
344 + "weight is not per_group."
345 + )
346 return False
347
348 # Input and weight group size needs to be 32.
345 - if weight_quant.get("group_size") != 32 or input_quant.get("group_size") != 32:
346 - logger.debug("Quark model is not in OCP MX format: not group_size=32")
349 + if weight_quant.get("group_size") != 32:
350 + logger.debug(
351 + "Quark model's weight quantization is incompatible with OCP MX format: "
352 + "group_size of weight is not 32."
353 + )
354 return False
355
356 # Activations and weight scales need to be in e8m0 format.
350 - if (
351 - weight_quant.get("scale_format") != "e8m0"
352 - or input_quant.get("scale_format") != "e8m0"
353 - ):
354 - logger.debug("Quark model is not in OCP MX format: not scale_format e8m0")
357 + if weight_quant.get("scale_format") != "e8m0":
358 + logger.debug(
359 + "Quark model's weight quantization is incompatible with OCP MX format: "
360 + "scale_format of weight is not e8m0."
361 + )
362 return False
363
364 # Input and weight dtypes need to be any of fp4,
@@ -360,14 +367,31 @@ class QuarkConfig(QuantizationConfig):
367 "fp4",
368 "fp6_e3m2",
369 "fp6_e2m3",
363 - } or input_quant.get("dtype") not in {"fp4", "fp6_e3m2", "fp6_e2m3"}:
370 + }:
371 logger.debug(
365 - "Quark model is not in OCP MX format: dtype not fp4, fp6_e3m2, fp6_e2m3"
372 + "Quark model's weight quantization is incompatible with OCP MX format: "
373 + "dtype is not in {fp4, fp6_e3m2, fp6_e2m3}."
374 )
375 return False
376
377 return True
378
379 + def is_mxfp4_quant(self, prefix: str, layer: torch.nn.Module) -> bool:
380 + """
381 + For Quark, determine if it's OCP MXFP4 by checking config directly.
382 + This allows hidden_size rounding to happen before moe_config creation.
383 + """
384 + layer_quant_config = self._find_matched_config(prefix, layer)
385 + weight_config = layer_quant_config.get("weight")
386 + input_config = layer_quant_config.get("input_tensors")
387 +
388 + return (
389 + self._is_w_ocp_mx_a_x(weight_config, input_config)
390 + and weight_config is not None
391 + and weight_config.get("dtype") == "fp4"
392 + and getattr(torch, "float4_e2m1fn_x2", None) is not None
393 + )
394 +
395 def _find_matched_config(
396 self, layer_name: str, module: torch.nn.Module
397 ) -> dict[str, Any]:
@@ -441,7 +465,7 @@ class QuarkConfig(QuantizationConfig):
465 is_static_input_scheme=True,
466 input_symmetric=input_config.get("symmetric"),
467 )
444 - elif self._is_ocp_mx(weight_config, input_config):
468 + elif self._is_w_ocp_mx_a_x(weight_config, input_config):
469 return QuarkOCP_MX(weight_config, input_config)
470
471 raise NotImplementedError(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+298 -54
@@ -8,6 +8,7 @@ import torch
8 import vllm.envs as envs
9 from vllm import _custom_ops as ops
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
13 from vllm.model_executor.layers.fused_moe import (
14 FusedMoE,
@@ -18,9 +19,15 @@ from vllm.model_executor.layers.fused_moe import (
19 from vllm.model_executor.layers.fused_moe.config import (
20 FusedMoEQuantConfig,
21 fp8_w8a8_moe_quant_config,
22 + mxfp4_w4a8_moe_quant_config,
23 + mxfp4_w4a16_moe_quant_config,
24 ocp_mx_moe_quant_config,
25 )
26 from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
27 +from vllm.model_executor.layers.quantization.mxfp4 import (
28 + Mxfp4Backend,
29 + get_mxfp4_backend,
30 +)
31 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
32 prepare_fp8_moe_layer_for_marlin,
33 )
@@ -37,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
44 from vllm.model_executor.utils import set_weight_attrs
45 from vllm.platforms import current_platform
46 from vllm.scalar_type import scalar_types
47 +from vllm.utils.math_utils import round_up
48
49 logger = init_logger(__name__)
50
@@ -46,6 +54,7 @@ __all__ = ["QuarkMoEMethod", "QuarkW8A8Fp8MoEMethod", "QuarkOCP_MX_MoEMethod"]
54 class QuarkMoEMethod(FusedMoEMethodBase):
55 def __init__(self, moe: FusedMoEConfig):
56 super().__init__(moe)
57 + self.has_bias = self.moe.has_bias
58
59 @staticmethod
60 def get_moe_method(
@@ -67,7 +76,7 @@ class QuarkMoEMethod(FusedMoEMethodBase):
76 return QuarkW4A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
77 elif quant_config._is_fp8_w8a8(weight_config, input_config):
78 return QuarkW8A8Fp8MoEMethod(weight_config, input_config, module.moe_config)
70 - elif quant_config._is_ocp_mx(weight_config, input_config):
79 + elif quant_config._is_w_ocp_mx_a_x(weight_config, input_config):
80 return QuarkOCP_MX_MoEMethod(weight_config, input_config, module.moe_config)
81 else:
82 raise RuntimeError("Unsupported FusedMoe scheme")
@@ -86,6 +95,10 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
95
96 self.weight_qscheme = self.weight_quant.get("qscheme")
97 self.input_qscheme = self.input_quant.get("qscheme")
98 + self.weight_dtype = self.weight_quant.get("dtype", "").replace(
99 + "fp8_e4m3", "fp8"
100 + )
101 + self.input_dtype = self.input_quant.get("dtype", "").replace("fp8_e4m3", "fp8")
102 per_tensor = (
103 self.weight_qscheme == "per_tensor" and self.input_qscheme == "per_tensor"
104 )
@@ -121,6 +134,10 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
134
135 self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
136
137 + self.model_type = getattr(
138 + get_current_vllm_config().model_config.hf_config, "model_type", None
139 + )
140 +
141 def create_weights(
142 self,
143 layer: torch.nn.Module,
@@ -166,9 +183,16 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
183 if self.weight_qscheme == "per_tensor":
184 # Allocate 2 scales for w1 and w3 respectively.
185 # They are combined to a single scale after weight loading.
169 - w13_weight_scale = torch.nn.Parameter(
170 - torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False
171 - )
186 + if self.model_type != "gpt_oss":
187 + w13_weight_scale = torch.nn.Parameter(
188 + torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False
189 + )
190 + else:
191 + # For gpt_oss, the w1(gate) & w3(up) are fused as one.
192 + # Therefore, only one weight scale for each expert.
193 + w13_weight_scale = torch.nn.Parameter(
194 + torch.ones(num_experts, 1, dtype=torch.float32), requires_grad=False
195 + )
196 layer.register_parameter("w13_weight_scale", w13_weight_scale)
197 w2_weight_scale = torch.nn.Parameter(
198 torch.ones(num_experts, dtype=torch.float32), requires_grad=False
@@ -220,6 +244,27 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
244 layer.w13_input_scale = None
245 layer.w2_input_scale = None
246
247 + if self.has_bias:
248 + w13_bias = torch.nn.Parameter(
249 + torch.zeros(
250 + num_experts,
251 + 2 * intermediate_size_per_partition,
252 + dtype=torch.float32,
253 + ),
254 + requires_grad=False,
255 + )
256 + layer.register_parameter("w13_bias", w13_bias)
257 + set_weight_attrs(w13_bias, extra_weight_attrs)
258 +
259 + w2_bias = torch.nn.Parameter(
260 + torch.zeros(num_experts, hidden_size, dtype=torch.float32),
261 + requires_grad=False,
262 + )
263 + layer.register_parameter("w2_bias", w2_bias)
264 + set_weight_attrs(w2_bias, extra_weight_attrs)
265 + else:
266 + layer.w13_bias, layer.w2_bias = None, None
267 +
268 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
269 # Fp8 moe kernels require a single activation scale.
270 # We take the max of all the scales in case they differ.
@@ -278,21 +323,40 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
323 assert layer.w13_weight_scale is not None
324 shard_size = layer.intermediate_size_per_partition
325 max_w13_scales = layer.w13_weight_scale.max(dim=1).values
281 - for expert_id in range(layer.local_num_experts):
282 - start = 0
283 - for shard_id in range(2):
326 +
327 + # For gpt_oss, w1 and w3 are fused into a single combined
328 + # gate_up_proj tensor with size 2*intermediate_size_per_partition
329 + # and only one scale per expert.
330 + # Process the entire weight tensor as one shard.
331 + if self.model_type == "gpt_oss":
332 + for expert_id in range(layer.local_num_experts):
333 + # Process all 2*intermediate_size_per_partition rows at once
334 dq_weight = per_tensor_dequantize(
285 - layer.w13_weight[expert_id][start : start + shard_size, :],
286 - layer.w13_weight_scale[expert_id][shard_id],
335 + layer.w13_weight[expert_id],
336 + layer.w13_weight_scale[expert_id][0],
337 )
288 - layer.w13_weight[expert_id][start : start + shard_size, :], _ = (
289 - ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
338 + layer.w13_weight[expert_id], _ = ops.scaled_fp8_quant(
339 + dq_weight, max_w13_scales[expert_id]
340 )
291 - start += shard_size
341 + else:
342 + # For non-gpt_oss, process w1 and w3 shards separately
343 + for expert_id in range(layer.local_num_experts):
344 + start = 0
345 + for shard_id in range(2):
346 + dq_weight = per_tensor_dequantize(
347 + layer.w13_weight[expert_id][start : start + shard_size, :],
348 + layer.w13_weight_scale[expert_id][shard_id],
349 + )
350 + (
351 + layer.w13_weight[expert_id][start : start + shard_size, :],
352 + _,
353 + ) = ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
354 + start += shard_size
355
356 layer.w13_weight_scale = torch.nn.Parameter(
357 max_w13_scales, requires_grad=False
358 )
359 +
360 # quark's scale is 1 dim.
361 elif self.weight_qscheme == "per_channel":
362 if self.act_quant_group_shape == GroupShape.PER_TOKEN:
@@ -343,6 +407,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
407 w2_scale=layer.w2_weight_scale,
408 a1_scale=layer.w13_input_scale,
409 a2_scale=layer.w2_input_scale,
410 + w1_bias=layer.w13_bias,
411 + w2_bias=layer.w2_bias,
412 per_act_token_quant=self.input_qscheme == "per_channel",
413 per_out_ch_quant=self.weight_qscheme == "per_channel",
414 )
@@ -563,7 +629,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
629 def __init__(
630 self,
631 weight_config: dict[str, Any],
566 - input_config: dict[str, Any],
632 + input_config: dict[str, Any] | None,
633 moe: FusedMoEConfig,
634 ):
635 super().__init__(moe)
@@ -571,35 +637,79 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
637 self.input_quant = input_config
638
639 weight_qscheme = self.weight_quant.get("qscheme")
574 - input_qscheme = self.input_quant.get("qscheme")
575 - if not (weight_qscheme == "per_group" and input_qscheme == "per_group"):
640 + if not weight_qscheme == "per_group":
641 raise ValueError(
642 "For MX(FP4) Fused MoE layers, only per-group scales "
578 - "for weights and activations are supported. Found "
579 - f"{weight_qscheme}, {input_qscheme}"
643 + f"for weights are supported. Found {weight_qscheme}."
644 ) # noqa E501
645
582 - self.static_input_scales = not self.input_quant.get("is_dynamic")
583 -
646 self.weight_dtype = self.weight_quant["dtype"].replace("fp", "mxfp")
585 - self.input_dtype = self.input_quant["dtype"].replace("fp", "mxfp")
647 + if self.input_quant is not None:
648 + input_quant = self.input_quant["dtype"]
649 + if input_quant in ["fp4", "fp6_e3m2", "fp6_e2m3"]:
650 + self.input_dtype = input_quant.replace("fp", "mxfp")
651 + elif input_quant == "fp8_e4m3":
652 + self.input_dtype = input_quant.replace("fp8_e4m3", "fp8")
653 + else:
654 + raise NotImplementedError(
655 + f"Current input dtype {input_quant} is not compatible \
656 + with OCP MX (weight) MoE quantization. Please open an issue"
657 + )
658 + else:
659 + self.input_dtype = None
660 +
661 self.fp4_dtype = getattr(torch, "float4_e2m1fn_x2", None)
662
663 self.ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
664 self.input_dtype, self.weight_dtype
665 )
666
592 - if self.static_input_scales:
667 + if self.ocp_mx_scheme is None:
668 + raise ValueError(
669 + f"Unsupported OCP MX dtype combination for MoE: "
670 + f"input_dtype={self.input_dtype}, weight_dtype={self.weight_dtype}. "
671 + f"Please check that the combination is supported in OCP_MX_Scheme."
672 + )
673 +
674 + self.mxfp4_backend: Mxfp4Backend | None = None
675 + if self.ocp_mx_scheme == "w_mxfp4":
676 + self.mxfp4_backend = get_mxfp4_backend(moe.is_lora_enabled)
677 +
678 + if self.input_quant is not None:
679 + self.static_input_scales = not self.input_quant.get("is_dynamic")
680 + else:
681 + self.static_input_scales = False
682 +
683 + if any(
684 + self.ocp_mx_scheme.endswith(a_scheme)
685 + for a_scheme in ["a_mxfp4", "a_mxfp6_e3m2", "a_mxfp6_e2m3"]
686 + ):
687 + if self.static_input_scales:
688 + raise NotImplementedError(
689 + "QuarkOCP_MX_MoEMethod with static input scales is currently "
690 + f"not implemented for OCP MX scheme {self.ocp_mx_scheme}. "
691 + "Please open an issue."
692 + )
693 + elif self.ocp_mx_scheme.endswith("a_fp8") and not self.static_input_scales:
694 raise NotImplementedError(
594 - "QuarkOCP_MX_MoEMethod with static input scales is currently "
595 - "not implemented. Please open an issue."
695 + "QuarkOCP_MX_MoEMethod with dynamic input scales is currently "
696 + f"not implemented for OCP MX scheme {self.ocp_mx_scheme}. "
697 + "Please open an issue."
698 )
699
700 self.use_rocm_aiter_moe = rocm_aiter_ops.is_fused_moe_enabled()
701
600 - self.emulate = not current_platform.supports_mx() or not (
601 - self.use_rocm_aiter_moe and self.ocp_mx_scheme == "w_mxfp4_a_mxfp4"
702 + self.model_type = getattr(
703 + get_current_vllm_config().model_config.hf_config, "model_type", None
704 )
705 +
706 + self._emulate = (
707 + not current_platform.supports_mx()
708 + or not self.ocp_mx_scheme.startswith("w_mxfp4")
709 + ) and (self.mxfp4_backend is None or not self.use_rocm_aiter_moe)
710 +
711 + self.emulate = True if self.model_type == "gpt_oss" else self._emulate
712 +
713 if self.emulate:
714 logger.warning_once(
715 f"The current mode (supports_mx={current_platform.supports_mx()}, "
@@ -640,12 +750,23 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
750 )
751
752 params_dtype = torch.uint8
753 + if self.model_type == "gpt_oss":
754 + if current_platform.is_rocm():
755 + intermediate_size_per_partition_after_pad = round_up(
756 + intermediate_size_per_partition, 256
757 + )
758 + else:
759 + intermediate_size_per_partition_after_pad = round_up(
760 + intermediate_size_per_partition, 64
761 + )
762 + else:
763 + intermediate_size_per_partition_after_pad = intermediate_size_per_partition
764
765 # WEIGHTS
766 w13_weight = torch.nn.Parameter(
767 torch.empty(
768 num_experts,
648 - 2 * intermediate_size_per_partition,
769 + 2 * intermediate_size_per_partition_after_pad,
770 self.get_packed_dim(hidden_size, self.weight_dtype),
771 dtype=params_dtype,
772 ),
@@ -659,7 +780,9 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
780 torch.empty(
781 num_experts,
782 hidden_size,
662 - self.get_packed_dim(intermediate_size_per_partition, self.weight_dtype),
783 + self.get_packed_dim(
784 + intermediate_size_per_partition_after_pad, self.weight_dtype
785 + ),
786 dtype=params_dtype,
787 ),
788 requires_grad=False,
@@ -672,7 +795,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
795 w13_weight_scale = torch.nn.Parameter(
796 torch.ones(
797 num_experts,
675 - 2 * intermediate_size_per_partition,
798 + 2 * intermediate_size_per_partition_after_pad,
799 hidden_size // OCP_MX_BLOCK_SIZE,
800 dtype=params_dtype,
801 ),
@@ -682,7 +805,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
805 torch.ones(
806 num_experts,
807 hidden_size,
685 - intermediate_size_per_partition // OCP_MX_BLOCK_SIZE,
808 + intermediate_size_per_partition_after_pad // OCP_MX_BLOCK_SIZE,
809 dtype=params_dtype,
810 ),
811 requires_grad=False,
@@ -693,8 +816,96 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
816 layer.register_parameter("w13_weight_scale", w13_weight_scale)
817 layer.register_parameter("w2_weight_scale", w2_weight_scale)
818
819 + if self.has_bias:
820 + w13_bias = torch.nn.Parameter(
821 + torch.zeros(
822 + num_experts,
823 + 2 * intermediate_size_per_partition_after_pad,
824 + dtype=torch.float32,
825 + ),
826 + requires_grad=False,
827 + )
828 + layer.register_parameter("w13_bias", w13_bias)
829 + set_weight_attrs(w13_bias, extra_weight_attrs)
830 +
831 + w2_bias = torch.nn.Parameter(
832 + torch.zeros(num_experts, hidden_size, dtype=torch.float32),
833 + requires_grad=False,
834 + )
835 + layer.register_parameter("w2_bias", w2_bias)
836 + set_weight_attrs(w2_bias, extra_weight_attrs)
837 + else:
838 + layer.w13_bias, layer.w2_bias = None, None
839 +
840 + # INPUT_SCALES
841 + if self.static_input_scales:
842 + w13_input_scale = torch.nn.Parameter(
843 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False
844 + )
845 + layer.register_parameter("w13_input_scale", w13_input_scale)
846 + set_weight_attrs(w13_input_scale, extra_weight_attrs)
847 +
848 + w2_input_scale = torch.nn.Parameter(
849 + torch.ones(num_experts, dtype=torch.float32), requires_grad=False
850 + )
851 + layer.register_parameter("w2_input_scale", w2_input_scale)
852 + set_weight_attrs(w2_input_scale, extra_weight_attrs)
853 + else:
854 + layer.w13_input_scale = None
855 + layer.w2_input_scale = None
856 +
857 def process_weights_after_loading(self, layer):
858 + if self.static_input_scales:
859 + # firstly, process activations if fp8 static input
860 + if layer.w13_input_scale is None or layer.w2_input_scale is None:
861 + raise ValueError(
862 + "QuantConfig has static quantization, but found "
863 + "activation scales are None."
864 + )
865 + if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
866 + layer.w2_input_scale
867 + ):
868 + logger.warning_once(
869 + "Found input_scales that are not equal for "
870 + "fp8 MoE layer. Using the maximum across experts "
871 + "for each layer. "
872 + )
873 + layer.w13_input_scale = torch.nn.Parameter(
874 + layer.w13_input_scale.max(), requires_grad=False
875 + )
876 + layer.w2_input_scale = torch.nn.Parameter(
877 + layer.w2_input_scale.max(), requires_grad=False
878 + )
879 +
880 + if current_platform.is_fp8_fnuz():
881 + # Normalize the weights and scales
882 + _, _, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
883 + torch.empty_like(layer.w13_weight, dtype=torch.float8_e4m3fnuz),
884 + torch.empty_like(
885 + layer.w13_weight_scale, dtype=layer.w13_weight_scale.dtype
886 + ),
887 + layer.w13_input_scale,
888 + )
889 + _, _, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
890 + torch.empty_like(layer.w2_weight, dtype=torch.float8_e4m3fnuz),
891 + torch.empty_like(
892 + layer.w2_weight_scale, dtype=layer.w13_weight_scale.dtype
893 + ),
894 + layer.w2_input_scale,
895 + )
896 + # Reset the parameter
897 + if w13_input_scale is not None:
898 + layer.w13_input_scale = torch.nn.Parameter(
899 + w13_input_scale, requires_grad=False
900 + )
901 + if w2_input_scale is not None:
902 + layer.w2_input_scale = torch.nn.Parameter(
903 + w2_input_scale, requires_grad=False
904 + )
905 +
906 + # secondly, process mxfp weights
907 if self.emulate:
908 + torch.cuda.empty_cache()
909 return
910
911 from aiter.utility.fp4_utils import e8m0_shuffle
@@ -725,15 +936,40 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
936 def get_fused_moe_quant_config(
937 self, layer: torch.nn.Module
938 ) -> FusedMoEQuantConfig | None:
728 - return ocp_mx_moe_quant_config(
729 - quant_dtype=self.input_dtype,
730 - weight_dtype=self.weight_dtype,
731 - w1_scale=layer.w13_weight_scale,
732 - w2_scale=layer.w2_weight_scale,
733 - a1_scale=None,
734 - a2_scale=None,
735 - block_shape=None,
736 - )
939 + if self.ocp_mx_scheme == "w_mxfp4":
940 + return mxfp4_w4a16_moe_quant_config(
941 + w1_scale=layer.w13_weight_scale,
942 + w2_scale=layer.w2_weight_scale,
943 + w1_bias=layer.w13_bias,
944 + w2_bias=layer.w2_bias,
945 + )
946 + elif self.ocp_mx_scheme == "w_mxfp4_a_fp8":
947 + return mxfp4_w4a8_moe_quant_config(
948 + w1_scale=layer.w13_weight_scale,
949 + w2_scale=layer.w2_weight_scale,
950 + a1_scale=layer.w13_input_scale,
951 + a2_scale=layer.w2_input_scale,
952 + w1_bias=layer.w13_bias,
953 + w2_bias=layer.w2_bias,
954 + block_shape=None,
955 + )
956 + elif self.ocp_mx_scheme in ["w_mxfp6_e3m2_a_fp8", "w_mxfp6_e2m3_a_fp8"]:
957 + raise NotImplementedError(
958 + "Currently there is no corresponding fused moe quant config configured "
959 + f"in vLLM for OCP MX scheme {self.ocp_mx_scheme}. Please open an issue."
960 + )
961 + else:
962 + return ocp_mx_moe_quant_config(
963 + quant_dtype=self.input_dtype,
964 + weight_dtype=self.weight_dtype,
965 + w1_scale=layer.w13_weight_scale,
966 + w2_scale=layer.w2_weight_scale,
967 + w1_bias=layer.w13_bias,
968 + w2_bias=layer.w2_bias,
969 + a1_scale=None,
970 + a2_scale=None,
971 + block_shape=None,
972 + )
973
974 def apply(
975 self,
@@ -743,24 +979,34 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
979 topk_ids: torch.Tensor,
980 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
981 if not self.emulate:
746 - from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
747 - rocm_aiter_fused_experts,
748 - )
982 + if (
983 + self.model_type == "gpt_oss"
984 + and self.mxfp4_backend == Mxfp4Backend.TRITON
985 + ):
986 + raise NotImplementedError(
987 + "Triton kernel implemented fused MoE for GPT_OSS model "
988 + "in Quark(MoE) format is not integrated or provided yet."
989 + )
990
750 - out = rocm_aiter_fused_experts(
751 - x,
752 - layer.w13_weight,
753 - layer.w2_weight,
754 - topk_weights=topk_weights,
755 - topk_ids=topk_ids,
756 - activation=layer.activation,
757 - quant_config=self.moe_quant_config,
758 - expert_map=layer.expert_map,
759 - )
991 + else:
992 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
993 + rocm_aiter_fused_experts,
994 + )
995 +
996 + return rocm_aiter_fused_experts(
997 + x,
998 + layer.w13_weight,
999 + layer.w2_weight,
1000 + topk_weights=topk_weights,
1001 + topk_ids=topk_ids,
1002 + activation=layer.activation,
1003 + quant_config=self.moe_quant_config,
1004 + expert_map=layer.expert_map,
1005 + )
1006 else:
1007 from vllm.model_executor.layers.fused_moe import fused_experts
1008
763 - out = fused_experts(
1009 + return fused_experts(
1010 x,
1011 layer.w13_weight,
1012 layer.w2_weight,
@@ -773,5 +1019,3 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):
1019 expert_map=layer.expert_map,
1020 quant_config=self.moe_quant_config,
1021 )
776 -
777 - return out
vllm/model_executor/layers/quantization/utils/ocp_mx_utils.py
+19 -1
@@ -20,26 +20,44 @@ SUPPORTED_OCP_MX_DTYPES = {"mxfp4", "mxfp6_e3m2", "mxfp6_e2m3"}
20
21
22 class OCP_MX_Scheme(str, Enum):
23 + w_mxfp4 = "w_mxfp4"
24 w_mxfp4_a_mxfp4 = "w_mxfp4_a_mxfp4"
25 w_mxfp4_a_mxfp6_e3m2 = "w_mxfp4_a_mxfp6_e3m2"
26 w_mxfp4_a_mxfp6_e2m3 = "w_mxfp4_a_mxfp6_e2m3"
27 + w_mxfp4_a_fp8 = "w_mxfp4_a_fp8"
28 + w_mxfp6_e3m2 = "w_mxfp6_e3m2"
29 w_mxfp6_e3m2_a_mxfp6_e3m2 = "w_mxfp6_e3m2_a_mxfp6_e3m2"
30 + w_mxfp6_e3m2_a_fp8 = "w_mxfp6_e3m2_a_fp8"
31 + w_mxfp6_e2m3 = "w_mxfp6_e2m3"
32 w_mxfp6_e2m3_a_mxfp6_e2m3 = "w_mxfp6_e2m3_a_mxfp6_e2m3"
33 + w_mxfp6_e2m3_a_fp8 = "w_mxfp6_e2m3_a_fp8"
34
35 @classmethod
36 def from_quant_dtype(cls, input_dtype: str | None, weight_dtype: str | None):
31 - if input_dtype not in OCP_MX_DTYPES or weight_dtype not in OCP_MX_DTYPES:
37 + if input_dtype not in OCP_MX_DTYPES and weight_dtype not in OCP_MX_DTYPES:
38 return None
39 + elif input_dtype is None and weight_dtype == "mxfp4":
40 + return cls.w_mxfp4
41 + elif input_dtype is None and weight_dtype == "mxfp6_e3m2":
42 + return cls.w_mxfp6_e3m2
43 + elif input_dtype is None and weight_dtype == "mxfp6_e2m3":
44 + return cls.w_mxfp6_e2m3
45 elif input_dtype == "mxfp4" and weight_dtype == "mxfp4":
46 return cls.w_mxfp4_a_mxfp4
47 elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp4":
48 return cls.w_mxfp4_a_mxfp6_e3m2
49 elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp4":
50 return cls.w_mxfp4_a_mxfp6_e2m3
51 + elif input_dtype == "fp8" and weight_dtype == "mxfp4":
52 + return cls.w_mxfp4_a_fp8
53 elif input_dtype == "mxfp6_e3m2" and weight_dtype == "mxfp6_e3m2":
54 return cls.w_mxfp6_e3m2_a_mxfp6_e3m2
55 + elif input_dtype == "fp8" and weight_dtype == "mxfp6_e3m2":
56 + return cls.w_mxfp6_e3m2_a_fp8
57 elif input_dtype == "mxfp6_e2m3" and weight_dtype == "mxfp6_e2m3":
58 return cls.w_mxfp6_e2m3_a_mxfp6_e2m3
59 + elif input_dtype == "fp8" and weight_dtype == "mxfp6_e2m3":
60 + return cls.w_mxfp6_e2m3_a_fp8
61 else:
62 logger.warning(
63 "input_dtype='%s' and"
vllm/model_executor/models/gpt_oss.py
+491 -18
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from collections.abc import Iterable
3 +import typing
4 +from collections.abc import Callable, Iterable
5
6 import torch
7 import torch.distributed as dist
@@ -25,13 +26,17 @@ from vllm.model_executor.layers.layernorm import RMSNorm
26 from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
27 from vllm.model_executor.layers.logits_processor import LogitsProcessor
28 from vllm.model_executor.layers.quantization import QuantizationConfig
29 +from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_BLOCK_SIZE
30 from vllm.model_executor.layers.rotary_embedding import get_rope
31 from vllm.model_executor.layers.utils import rocm_unquantized_gemm
32 from vllm.model_executor.layers.vocab_parallel_embedding import (
33 ParallelLMHead,
34 VocabParallelEmbedding,
35 )
34 -from vllm.model_executor.model_loader.weight_utils import default_weight_loader
36 +from vllm.model_executor.model_loader.weight_utils import (
37 + default_weight_loader,
38 + maybe_remap_kv_scale_name,
39 +)
40 from vllm.model_executor.models.utils import sequence_parallel_chunk
41 from vllm.platforms import current_platform
42 from vllm.sequence import IntermediateTensors
@@ -98,6 +103,7 @@ class OAIAttention(nn.Module):
103 head_size=self.head_dim,
104 total_num_heads=self.num_attention_heads,
105 total_num_kv_heads=self.num_key_value_heads,
106 + bias=True,
107 quant_config=quant_config,
108 prefix=f"{prefix}.qkv_proj",
109 )
@@ -105,6 +111,7 @@ class OAIAttention(nn.Module):
111 self.o_proj = RowParallelLinear(
112 input_size=self.num_attention_heads * self.head_dim,
113 output_size=self.hidden_size,
114 + bias=True,
115 quant_config=quant_config,
116 prefix=f"{prefix}.o_proj",
117 )
@@ -306,6 +313,19 @@ class GptOssModel(nn.Module):
313 return x, aux_hidden_states
314 return x
315
316 + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
317 + # Params for weights, weight scales, activation scales
318 + # (param_name, weight_name, expert_id, shard_id)
319 + # NOTE: this is only used for quark.
320 + return FusedMoE.make_expert_params_mapping(
321 + self,
322 + ckpt_gate_proj_name="w1",
323 + ckpt_down_proj_name="w2",
324 + ckpt_up_proj_name="w3",
325 + num_experts=self.config.num_local_experts,
326 + num_redundant_experts=0,
327 + )
328 +
329 def _load_weights_mxfp4(
330 self,
331 ep_rank_end: int,
@@ -318,7 +338,6 @@ class GptOssModel(nn.Module):
338 params_dict = dict(self.named_parameters())
339 loaded_params: set[str] = set()
340
321 - mxfp4_block = 32
341 use_ep = self.parallel_config.enable_expert_parallel
342 num_experts = self.config.num_local_experts
343
@@ -333,9 +352,11 @@ class GptOssModel(nn.Module):
352 )
353
354 intermediate_size = self.config.intermediate_size
336 - intermediate_size_block = intermediate_size // mxfp4_block
355 + intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE
356 per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size)
338 - per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block
357 + per_rank_intermediate_size = (
358 + per_rank_intermediate_size_block * OCP_MX_BLOCK_SIZE
359 + )
360
361 # Calculate common slicing bounds for current rank
362 tp_rank_start = tp_rank * per_rank_intermediate_size
@@ -370,7 +391,9 @@ class GptOssModel(nn.Module):
391 narrow_weight = weight[ep_rank_start:ep_rank_end, ...]
392 else:
393 narrow_weight = weight[
373 - ..., tp_rank_start // mxfp4_block : tp_rank_end // mxfp4_block
394 + ...,
395 + tp_rank_start // OCP_MX_BLOCK_SIZE : tp_rank_end
396 + // OCP_MX_BLOCK_SIZE,
397 ]
398
399 param = params_dict[name]
@@ -495,6 +518,449 @@ class GptOssModel(nn.Module):
518 loaded_params.add(name)
519 return loaded_params
520
521 + def _load_weights_quark(
522 + self,
523 + ep_rank_end: int,
524 + ep_rank_start: int,
525 + heads_per_rank: int,
526 + head_start: int,
527 + weights: Iterable[tuple[str, torch.Tensor]],
528 + stacked_params_mapping: list[tuple[str, ...]],
529 + ) -> set[str]:
530 + params_dict = dict(self.named_parameters())
531 + loaded_params: set[str] = set()
532 +
533 + use_ep = self.parallel_config.enable_expert_parallel
534 + num_experts = self.config.num_local_experts
535 +
536 + if use_ep:
537 + tp_rank = get_tensor_model_parallel_rank()
538 + tp_size = get_tensor_model_parallel_world_size()
539 + else:
540 + tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
541 + tp_size=get_tensor_model_parallel_world_size(),
542 + dp_size=get_dp_group().world_size,
543 + dp_rank=get_dp_group().rank_in_group,
544 + pcp_size=get_pcp_group().world_size,
545 + pcp_rank=get_pcp_group().rank_in_group,
546 + )
547 +
548 + def _get_moe_weight_dtype(layer_id: int = 0) -> str | None:
549 + """Helper function to get MoE quantization weight dtype.
550 +
551 + Args:
552 + layer_id: Layer index to check (default 0, as all layers should
553 + have the same quantization method)
554 +
555 + Returns:
556 + Weight dtype string (e.g., "mxfp4", "fp8") or None if not available
557 + """
558 + if hasattr(self.layers[layer_id].mlp.experts.quant_method, "weight_dtype"):
559 + return self.layers[layer_id].mlp.experts.quant_method.weight_dtype
560 + return None
561 +
562 + intermediate_size = self.config.intermediate_size
563 +
564 + moe_weight_dtype = _get_moe_weight_dtype(layer_id=0)
565 +
566 + if moe_weight_dtype == "mxfp4":
567 + # MXFP4 requires OCP_MX_BLOCK_SIZE alignment
568 + intermediate_size_block = intermediate_size // OCP_MX_BLOCK_SIZE
569 + per_rank_intermediate_size_block = cdiv(intermediate_size_block, tp_size)
570 + per_rank_intermediate_size = (
571 + per_rank_intermediate_size_block * OCP_MX_BLOCK_SIZE
572 + )
573 + else:
574 + # FP8 and other formats don't need alignment
575 + per_rank_intermediate_size = cdiv(intermediate_size, tp_size)
576 +
577 + tp_rank_start = tp_rank * per_rank_intermediate_size
578 + tp_rank_end = min((tp_rank + 1) * per_rank_intermediate_size, intermediate_size)
579 + expert_params_mapping = self.get_expert_mapping()
580 + for name, loaded_weight in weights:
581 + if is_pp_missing_parameter(name, self):
582 + continue
583 +
584 + layer_id, expert_id, fused_name = None, None, None
585 + moe_quant_method = None
586 + if "experts" in name:
587 + parts = name.split(".")
588 + ids = [s for s in parts if s.isdigit()]
589 +
590 + # for amd-quark format that each expert is seperated
591 + # need to extract the parameter name with experts fused.
592 + # example model: amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8
593 + if len(ids) == 2:
594 + layer_id, expert_id = int(ids[0]), int(ids[-1])
595 + parts.pop(len(parts) - 1 - parts[::-1].index(str(expert_id)))
596 + fused_name = ".".join(parts)
597 +
598 + # for openai mxfp4 format that all experts are combined
599 + # no need to extract the parameter name with experts fused.
600 + # models: openai/gpt-oss-20b, openai/gpt-oss-120b
601 + elif len(ids) == 1:
602 + layer_id, expert_id = int(ids[0]), None
603 + fused_name = name
604 +
605 + else:
606 + raise NameError(
607 + f"Layer {name} contains more than 2 numeric indices. This is "
608 + "an unexpected condition. Please open an issue if encountered."
609 + )
610 +
611 + moe_quant_method = _get_moe_weight_dtype(layer_id=layer_id)
612 +
613 + def kv_cache_scale_loader(
614 + quant_config: QuantizationConfig,
615 + name: str,
616 + params_dict: dict[str, typing.Any],
617 + weight: torch.Tensor,
618 + default_weight_loader: Callable[..., None],
619 + loaded_params: set[str],
620 + ) -> tuple[bool, set[str]]:
621 + """
622 + Load KV cache output scales.
623 + Returns:
624 + Tuple of (bool, set):
625 + - bool: True if KV-cache scale was loaded into loaded_params
626 + - set: Updated set of loaded_params if True else the original set
627 + """
628 + # load explicit cached KV output scale from quant_config
629 + if quant_config is not None and (
630 + scale_name := quant_config.get_cache_scale(name)
631 + ):
632 + param = params_dict[scale_name]
633 + weight_loader = getattr(
634 + param, "weight_loader", default_weight_loader
635 + )
636 + if weight.numel() != 1:
637 + raise ValueError(
638 + f"KV cache scale '{scale_name}' is expected to be a "
639 + f"scalar, but got a tensor of shape {weight.shape}."
640 + )
641 + # Ensure weight is a scalar before passing to loader.
642 + weight_loader(param, weight.flatten()[0])
643 + loaded_params.add(scale_name)
644 + return True, loaded_params
645 +
646 + return False, loaded_params
647 +
648 + load_kv_cache_scale_completed, loaded_params = kv_cache_scale_loader(
649 + self.quant_config,
650 + name,
651 + params_dict,
652 + loaded_weight,
653 + default_weight_loader,
654 + loaded_params,
655 + )
656 + if load_kv_cache_scale_completed:
657 + continue
658 +
659 + if (
660 + all(key in name for key in ["input_scale", "mlp.experts"])
661 + and expert_id is not None
662 + ):
663 + assert loaded_weight.numel() == 1
664 + expert_data = params_dict[fused_name].data[expert_id]
665 + expert_data.copy_(loaded_weight)
666 + loaded_params.add(fused_name)
667 + continue
668 +
669 + # Unified handler for mxfp4 weights and scales
670 + elif moe_quant_method == "mxfp4" and any(
671 + name.endswith(suffix)
672 + for suffix in [
673 + ".w13_weight_scale",
674 + ".w2_weight_scale",
675 + ".w13_weight",
676 + ".w2_weight",
677 + ]
678 + ):
679 + is_w13 = ".w13_" in name
680 + is_scale = "_scale" in name
681 +
682 + # Reshape weight for mxfp4 if needed (not for scales)
683 + if not is_scale and expert_id is None:
684 + if is_w13:
685 + if loaded_weight.dim() < 3:
686 + raise ValueError(
687 + f"Expected w13_weight to have at least 3 "
688 + f"dimensions, got shape "
689 + f"{loaded_weight.shape}"
690 + )
691 + if loaded_weight.shape[0] != num_experts:
692 + raise ValueError(
693 + f"Expected w13_weight first dimension to be "
694 + f"{num_experts}, got "
695 + f"{loaded_weight.shape[0]}"
696 + )
697 + loaded_weight = loaded_weight.view(
698 + num_experts, 2 * intermediate_size, -1
699 + ).contiguous()
700 + else:
701 + if loaded_weight.dim() < 3:
702 + raise ValueError(
703 + f"Expected w2_weight to have at least 3 "
704 + f"dimensions, got shape "
705 + f"{loaded_weight.shape}"
706 + )
707 + if loaded_weight.shape[0] != num_experts:
708 + raise ValueError(
709 + f"Expected w2_weight first dimension to be "
710 + f"{num_experts}, got "
711 + f"{loaded_weight.shape[0]}"
712 + )
713 + loaded_weight = loaded_weight.view(
714 + num_experts, -1, intermediate_size // 2
715 + ).contiguous()
716 +
717 + if use_ep:
718 + sliced_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
719 + else:
720 + if is_w13:
721 + if expert_id is None:
722 + sliced_weight = loaded_weight[
723 + :, 2 * tp_rank_start : 2 * tp_rank_end, ...
724 + ]
725 + else:
726 + sliced_weight = loaded_weight[
727 + 2 * tp_rank_start : 2 * tp_rank_end, ...
728 + ]
729 + else:
730 + if is_scale:
731 + sliced_weight = loaded_weight[
732 + ...,
733 + tp_rank_start // OCP_MX_BLOCK_SIZE : tp_rank_end
734 + // OCP_MX_BLOCK_SIZE,
735 + ]
736 + else:
737 + sliced_weight = loaded_weight[
738 + ..., tp_rank_start // 2 : tp_rank_end // 2
739 + ]
740 +
741 + # NOTE(rob): because gpt-oss ckpt has "unique" structure with
742 + # fused gate_up_proj fused on disk, we cannot use the existing
743 + # weight loaders without added complexity, so just do the
744 + # direct load here.
745 + param = params_dict[fused_name]
746 + expert_data = param.data[expert_id]
747 + dim1 = sliced_weight.shape[0]
748 + dim2 = sliced_weight.shape[1]
749 + expert_data.data[:dim1, :dim2].copy_(sliced_weight)
750 + loaded_params.add(fused_name)
751 + continue
752 +
753 + elif name.endswith(".w13_weight") and moe_quant_method == "fp8":
754 + if use_ep:
755 + narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
756 + else:
757 + if expert_id is None:
758 + narrow_weight = loaded_weight[
759 + :, 2 * tp_rank_start : 2 * tp_rank_end, :
760 + ]
761 + else:
762 + narrow_weight = loaded_weight[
763 + 2 * tp_rank_start : 2 * tp_rank_end, :
764 + ]
765 +
766 + assert fused_name is not None
767 + param = params_dict[fused_name]
768 +
769 + if expert_id is None:
770 + param.data.copy_(narrow_weight)
771 + else:
772 + param.data[expert_id].copy_(narrow_weight)
773 +
774 + loaded_params.add(fused_name)
775 + continue
776 +
777 + elif name.endswith(".w13_weight_scale") and moe_quant_method == "fp8":
778 + assert fused_name is not None
779 + param = params_dict[fused_name]
780 +
781 + # Check if this is per-channel or per-tensor scale
782 + if loaded_weight.numel() > 1 and loaded_weight.dim() == 1:
783 + if use_ep:
784 + narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
785 + else:
786 + narrow_weight = loaded_weight[
787 + 2 * tp_rank_start : 2 * tp_rank_end
788 + ]
789 + else:
790 + narrow_weight = loaded_weight
791 +
792 + if expert_id is None:
793 + param.data.copy_(narrow_weight)
794 + else:
795 + param.data[expert_id].copy_(narrow_weight)
796 +
797 + loaded_params.add(fused_name)
798 + continue
799 +
800 + elif name.endswith(".w13_input_scale") and moe_quant_method == "fp8":
801 + assert fused_name is not None
802 + param = params_dict[fused_name]
803 +
804 + if expert_id is None:
805 + param.data.copy_(loaded_weight)
806 + else:
807 + param.data[expert_id].copy_(loaded_weight)
808 +
809 + loaded_params.add(fused_name)
810 + continue
811 +
812 + elif name.endswith(".w2_weight") and moe_quant_method == "fp8":
813 + if use_ep:
814 + narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
815 + else:
816 + if expert_id is None:
817 + narrow_weight = loaded_weight[..., tp_rank_start:tp_rank_end]
818 + else:
819 + narrow_weight = loaded_weight[..., tp_rank_start:tp_rank_end]
820 +
821 + assert fused_name is not None
822 + param = params_dict[fused_name]
823 +
824 + if expert_id is None:
825 + param.data.copy_(narrow_weight)
826 + else:
827 + param.data[expert_id].copy_(narrow_weight)
828 +
829 + loaded_params.add(fused_name)
830 + continue
831 +
832 + elif name.endswith(".w2_weight_scale") and moe_quant_method == "fp8":
833 + assert fused_name is not None
834 + param = params_dict[fused_name]
835 +
836 + if use_ep:
837 + narrow_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
838 + else:
839 + narrow_weight = loaded_weight
840 +
841 + if expert_id is None:
842 + param.data.copy_(narrow_weight)
843 + else:
844 + param.data[expert_id].copy_(narrow_weight)
845 +
846 + loaded_params.add(fused_name)
847 + continue
848 +
849 + # Unified handler for bias loading (w13_bias and w2_bias)
850 + elif name.endswith(".w13_bias") or name.endswith(".w2_bias"):
851 + is_w13_bias = name.endswith(".w13_bias")
852 +
853 + if use_ep:
854 + sliced_weight = loaded_weight[ep_rank_start:ep_rank_end, ...]
855 + else:
856 + if is_w13_bias:
857 + if expert_id is None:
858 + sliced_weight = loaded_weight[
859 + :, 2 * tp_rank_start : 2 * tp_rank_end
860 + ]
861 + else:
862 + sliced_weight = loaded_weight[
863 + 2 * tp_rank_start : 2 * tp_rank_end
864 + ]
865 + else:
866 + sliced_weight = loaded_weight
867 + if tp_rank != 0:
868 + sliced_weight = sliced_weight.zero_()
869 +
870 + # NOTE(rob): because gpt-oss ckpt has "unique" structure with
871 + # fused gate_up_proj fused on disk, we cannot use the existing
872 + # weight loaders without added complexity, so just do the
873 + # direct load here.
874 + assert fused_name is not None
875 + param = params_dict[fused_name]
876 + expert_data = param.data[expert_id]
877 + dim1 = sliced_weight.shape[0]
878 + expert_data.data[:dim1].copy_(sliced_weight)
879 + loaded_params.add(fused_name)
880 + continue
881 +
882 + elif "sinks" in name:
883 + # Handle attention sinks (distributed across ranks)
884 + param = params_dict[name]
885 + narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank)
886 + param.data.copy_(narrow_weight)
887 + loaded_params.add(name)
888 + continue
889 +
890 + for param_name, weight_name, shard_id in stacked_params_mapping:
891 + # Skip non-stacked layers and experts (experts handled below).
892 + if weight_name not in name:
893 + continue
894 + # We have mlp.experts[0].gate_proj in the checkpoint.
895 + # Since we handle the experts below in expert_params_mapping,
896 + # we need to skip here BEFORE we update the name, otherwise
897 + # name will be updated to mlp.experts[0].gate_up_proj, which
898 + # will then be updated below in expert_params_mapping
899 + # for mlp.experts[0].gate_gate_up_proj, which breaks load.
900 + if ("mlp.experts." in name) and name not in params_dict:
901 + continue
902 + name = name.replace(weight_name, param_name)
903 +
904 + if name.endswith("scale"):
905 + # Remapping the name of FP8 kv-scale.
906 + name = maybe_remap_kv_scale_name(name, params_dict)
907 + if name is None:
908 + continue
909 +
910 + param = params_dict[name]
911 + weight_loader = param.weight_loader
912 +
913 + weight_loader(param, loaded_weight, shard_id)
914 + loaded_params.add(name)
915 + break
916 + else:
917 + for mapping in expert_params_mapping:
918 + # Anyway, this is an expert weight and should not be
919 + # attempted to load as other weights later
920 + param_name, weight_name, mapping_expert_id, shard_id = mapping
921 + weight_name = (
922 + weight_name[:-1] if weight_name.endswith(".") else weight_name
923 + )
924 +
925 + if weight_name not in name:
926 + continue
927 +
928 + param = params_dict[fused_name]
929 + # We should ask the weight loader to return success or not
930 + # here since otherwise we may skip experts with other
931 + # available replicas.
932 + weight_loader = typing.cast(
933 + Callable[..., bool], param.weight_loader
934 + )
935 + # Use checkpoint's expert_id for quark format (when expert_id
936 + # is extracted from weight name), otherwise use mapping's expert_id
937 + actual_expert_id = (
938 + expert_id if expert_id is not None else mapping_expert_id
939 + )
940 + success = weight_loader(
941 + param,
942 + loaded_weight,
943 + fused_name,
944 + shard_id=shard_id,
945 + expert_id=actual_expert_id,
946 + return_success=True,
947 + )
948 + if success:
949 + name = fused_name
950 + loaded_params.add(name)
951 + break
952 + else:
953 + if name not in params_dict:
954 + continue
955 + param = params_dict[name]
956 + weight_loader = getattr(
957 + param, "weight_loader", default_weight_loader
958 + )
959 + weight_loader(param, loaded_weight)
960 +
961 + loaded_params.add(name)
962 + return loaded_params
963 +
964 def _load_weights_other(
965 self,
966 ep_rank_end: int,
@@ -635,6 +1101,7 @@ class GptOssModel(nn.Module):
1101 if hasattr(self.config, "quantization_config")
1102 else None
1103 )
1104 +
1105 if quant_method == "mxfp4":
1106 return self._load_weights_mxfp4(
1107 ep_rank_end,
@@ -644,6 +1111,15 @@ class GptOssModel(nn.Module):
1111 weights,
1112 stacked_params_mapping,
1113 )
1114 + elif quant_method == "quark":
1115 + return self._load_weights_quark(
1116 + ep_rank_end,
1117 + ep_rank_start,
1118 + heads_per_rank,
1119 + head_start,
1120 + weights,
1121 + stacked_params_mapping,
1122 + )
1123 else:
1124 return self._load_weights_other(
1125 ep_rank_end,
@@ -676,6 +1152,15 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
1152 # MoE Bias
1153 ".gate_up_proj_bias": ".w13_bias",
1154 ".down_proj_bias": ".w2_bias",
1155 + # For quark format
1156 + ".gate_up_proj.weight": ".w13_weight",
1157 + ".gate_up_proj.weight_scale": ".w13_weight_scale",
1158 + ".gate_up_proj.bias": ".w13_bias",
1159 + ".gate_up_proj.input_scale": ".w13_input_scale",
1160 + ".down_proj.weight": ".w2_weight",
1161 + ".down_proj.weight_scale": ".w2_weight_scale",
1162 + ".down_proj.bias": ".w2_bias",
1163 + ".down_proj.input_scale": ".w2_input_scale",
1164 },
1165 )
1166
@@ -725,18 +1210,6 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
1210 logits = self.logits_processor(self.lm_head, hidden_states)
1211 return logits
1212
728 - def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
729 - # Params for weights, weight scales, activation scales
730 - # (param_name, weight_name, expert_id, shard_id)
731 - return FusedMoE.make_expert_params_mapping(
732 - self,
733 - ckpt_gate_proj_name="gate_proj",
734 - ckpt_down_proj_name="down_proj",
735 - ckpt_up_proj_name="up_proj",
736 - num_experts=self.config.num_local_experts,
737 - num_redundant_experts=0,
738 - )
739 -
1213 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1214 loader = AutoWeightsLoader(
1215 self,