@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1976356ee6
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, 17:45 UTC

[MoE Refactor] MXFP4 Cutlass Experts to MK (#34542)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>

Yongye Zhu committed Feb 25, 2026 at 17:32 UTC 1976356ee69750630189eb127fc9eeaa6f8e0c9e
19 files changed +453 -168
.buildkite/test_areas/lm_eval.yaml
+26
@@ -73,3 +73,29 @@ steps:
73 num_devices: 2
74 commands:
75 - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
76 +
77 +- label: GPQA Eval (GPT-OSS) (H100)
78 + timeout_in_minutes: 120
79 + device: h100
80 + optional: true
81 + num_devices: 2
82 + source_file_dependencies:
83 + - csrc/
84 + - vllm/model_executor/layers/quantization
85 + - tests/evals/gpt_oss/
86 + commands:
87 + - uv pip install --system 'gpt-oss[eval]==0.0.5'
88 + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt
89 +
90 +- label: GPQA Eval (GPT-OSS) (B200)
91 + timeout_in_minutes: 120
92 + device: b200
93 + optional: true
94 + num_devices: 2
95 + source_file_dependencies:
96 + - csrc/
97 + - vllm/model_executor/layers/quantization
98 + - tests/evals/gpt_oss/
99 + commands:
100 + - uv pip install --system 'gpt-oss[eval]==0.0.5'
101 + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt
.buildkite/test_areas/misc.yaml
-27
@@ -153,33 +153,6 @@ steps:
153 - pytest -v -s transformers_utils
154 - pytest -v -s config
155
156 -- label: GPT-OSS Eval (H100)
157 - timeout_in_minutes: 60
158 - working_dir: "/vllm-workspace/"
159 - device: h100
160 - optional: true
161 - source_file_dependencies:
162 - - tests/evals/gpt_oss
163 - - vllm/model_executor/models/gpt_oss.py
164 - - vllm/model_executor/layers/quantization/mxfp4.py
165 - commands:
166 - - uv pip install --system 'gpt-oss[eval]==0.0.5'
167 - - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
168 -
169 -- label: GPT-OSS Eval (B200)
170 - timeout_in_minutes: 60
171 - working_dir: "/vllm-workspace/"
172 - device: b200
173 - optional: true
174 - source_file_dependencies:
175 - - tests/evals/gpt_oss
176 - - vllm/model_executor/models/gpt_oss.py
177 - - vllm/model_executor/layers/quantization/mxfp4.py
178 - - vllm/v1/attention/backends/flashinfer.py
179 - commands:
180 - - uv pip install --system 'gpt-oss[eval]==0.0.5'
181 - - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58
182 -
156 - label: Batch Invariance (H100)
157 timeout_in_minutes: 25
158 device: h100
tests/evals/gpt_oss/README.md new
+49
@@ -0,0 +1,49 @@
1 +# GPQA Evaluation using GPT-OSS
2 +
3 +This directory contains GPQA evaluation tests using the GPT-OSS evaluation package and vLLM server.
4 +
5 +## Usage
6 +
7 +### Run tests with pytest (like buildkite)
8 +
9 +```bash
10 +# H200
11 +pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
12 + --config-list-file=configs/models-h200.txt
13 +
14 +# B200
15 +pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
16 + --config-list-file=configs/models-b200.txt
17 +```
18 +
19 +## Configuration Format
20 +
21 +Model configs in `configs/` directory use this YAML format:
22 +
23 +```yaml
24 +model_name: "openai/gpt-oss-20b"
25 +metric_threshold: 0.568 # Minimum expected accuracy
26 +reasoning_effort: "low" # Reasoning effort level (default: "low")
27 +server_args: "--tensor-parallel-size 2" # Server arguments
28 +startup_max_wait_seconds: 1800 # Max wait for server startup (default: 1800)
29 +env: # Environment variables (optional)
30 + SOME_VAR: "value"
31 +```
32 +
33 +The `server_args` field accepts any arguments that can be passed to `vllm serve`.
34 +
35 +The `env` field accepts a dictionary of environment variables to set for the server process.
36 +
37 +## Adding New Models
38 +
39 +1. Create a new YAML config file in the `configs/` directory
40 +2. Add the filename to the appropriate `models-*.txt` file
41 +
42 +## Tiktoken Encoding Files
43 +
44 +The tiktoken encoding files required by the vLLM server are automatically downloaded from OpenAI's public blob storage on first run:
45 +
46 +- `cl100k_base.tiktoken`
47 +- `o200k_base.tiktoken`
48 +
49 +Files are cached in the `data/` directory. The `TIKTOKEN_ENCODINGS_BASE` environment variable is automatically set to point to this directory when running evaluations.
tests/evals/gpt_oss/configs/gpt-oss-20b-baseline.yaml new
+6
@@ -0,0 +1,6 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +model_name: "openai/gpt-oss-20b"
4 +metric_threshold: 0.568
5 +reasoning_effort: "low"
6 +server_args: "--tensor-parallel-size 2"
tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-bf16.yaml new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +model_name: "openai/gpt-oss-20b"
4 +metric_threshold: 0.568
5 +reasoning_effort: "low"
6 +server_args: "--tensor-parallel-size 2"
7 +env:
8 + VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: "1"
tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +model_name: "openai/gpt-oss-20b"
4 +metric_threshold: 0.568
5 +reasoning_effort: "low"
6 +server_args: "--tensor-parallel-size 2"
7 +env:
8 + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: "1"
tests/evals/gpt_oss/configs/gpt-oss-20b-marlin.yaml new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +model_name: "openai/gpt-oss-20b"
4 +metric_threshold: 0.568
5 +reasoning_effort: "low"
6 +server_args: "--tensor-parallel-size 2"
7 +env:
8 + VLLM_MXFP4_USE_MARLIN: "1"
tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml new
+8
@@ -0,0 +1,8 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +model_name: "openai/gpt-oss-20b"
4 +metric_threshold: 0.568
5 +reasoning_effort: "low"
6 +server_args: "--tensor-parallel-size 2"
7 +env:
8 + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: "1"
tests/evals/gpt_oss/configs/models-b200.txt new
+5
@@ -0,0 +1,5 @@
1 +# B200 model configurations for GPQA evaluation
2 +# Tests different environment variable combinations
3 +gpt-oss-20b-flashinfer-mxfp4-bf16.yaml
4 +gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml
5 +gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml
\ No newline at end of file
tests/evals/gpt_oss/configs/models-h100.txt new
+5
@@ -0,0 +1,5 @@
1 +# H100 model configurations for GPQA evaluation
2 +# Tests different environment variable combinations
3 +gpt-oss-20b-baseline.yaml
4 +gpt-oss-20b-flashinfer-mxfp4-bf16.yaml
5 +gpt-oss-20b-marlin.yaml
tests/evals/gpt_oss/conftest.py
+54 -6
@@ -4,13 +4,61 @@
4 Pytest configuration for GPT-OSS evaluation tests.
5 """
6
7 +from pathlib import Path
8 +
9
10 def pytest_addoption(parser):
9 - """Add command line options for pytest."""
10 - parser.addoption("--model", action="store", help="Model name to evaluate")
11 - parser.addoption(
12 - "--metric", action="store", type=float, help="Expected metric threshold"
13 - )
11 + """Add custom command line options."""
12 parser.addoption(
15 - "--server-args", action="store", default="", help="Additional server arguments"
13 + "--config-list-file",
14 + required=True,
15 + help="File containing list of config files to test",
16 )
17 +
18 +
19 +def pytest_generate_tests(metafunc):
20 + """Generate test parameters from config files."""
21 + if "config_filename" in metafunc.fixturenames:
22 + config_list_file = metafunc.config.getoption("--config-list-file")
23 +
24 + # Handle both relative and absolute paths
25 + config_list_path = Path(config_list_file)
26 + if not config_list_path.is_absolute():
27 + # If relative, try relative to test directory first
28 + test_dir_path = Path(__file__).parent / config_list_file
29 + if test_dir_path.exists():
30 + config_list_path = test_dir_path
31 + else:
32 + # Try relative to current working directory
33 + config_list_path = Path.cwd() / config_list_file
34 +
35 + print(f"Looking for config list at: {config_list_path}")
36 +
37 + config_files = []
38 + if config_list_path.exists():
39 + # Determine config directory (same directory as the list file)
40 + config_dir = config_list_path.parent
41 +
42 + with open(config_list_path) as f:
43 + for line in f:
44 + line = line.strip()
45 + if line and not line.startswith("#"):
46 + config_path = config_dir / line
47 + print(f"Checking config file: {config_path}")
48 + if config_path.exists():
49 + config_files.append(config_path)
50 + print(f" Found: {config_path}")
51 + else:
52 + print(f" Missing: {config_path}")
53 + else:
54 + print(f"Config list file not found: {config_list_path}")
55 +
56 + # Generate test parameters
57 + if config_files:
58 + metafunc.parametrize(
59 + "config_filename",
60 + config_files,
61 + ids=[config_file.stem for config_file in config_files],
62 + )
63 + else:
64 + print("No config files found, test will be skipped")
tests/evals/gpt_oss/test_gpqa_correctness.py
+82 -28
@@ -5,22 +5,48 @@ GPQA evaluation using vLLM server and GPT-OSS evaluation package.
5
6 Usage:
7 pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \
8 - --model openai/gpt-oss-20b \
9 - --metric 0.58 \
10 - --server-args "--tensor-parallel-size 2"
8 + --config-list-file=configs/models-h200.txt
9 """
10
11 +import os
12 +import shlex
13 import subprocess
14 import sys
15 +import urllib.request
16 +from pathlib import Path
17
18 import regex as re
19 +import yaml
20
21 from tests.utils import RemoteOpenAIServer
22
23 TOL = 0.05 # Absolute tolerance for accuracy comparison
24
25 +# Path to tiktoken encoding files
26 +TIKTOKEN_DATA_DIR = Path(__file__).parent / "data"
27
23 -def run_gpqa_eval(model_name: str, base_url: str) -> float:
28 +# Tiktoken encoding files to download
29 +TIKTOKEN_FILES = {
30 + "cl100k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken",
31 + "o200k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken",
32 +}
33 +
34 +
35 +def ensure_tiktoken_files():
36 + """Download tiktoken encoding files if they don't exist."""
37 + TIKTOKEN_DATA_DIR.mkdir(parents=True, exist_ok=True)
38 +
39 + for filename, url in TIKTOKEN_FILES.items():
40 + filepath = TIKTOKEN_DATA_DIR / filename
41 + if not filepath.exists():
42 + print(f"Downloading {filename} from {url}...")
43 + urllib.request.urlretrieve(url, filepath)
44 + print(f" Downloaded to {filepath}")
45 + else:
46 + print(f" {filename} already exists.")
47 +
48 +
49 +def run_gpqa_eval(model_name: str, base_url: str, reasoning_effort: str) -> float:
50 """Run GPQA evaluation using the gpt-oss evaluation package."""
51
52 # Build the command to run the evaluation
@@ -33,7 +59,7 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float:
59 "--model",
60 model_name,
61 "--reasoning-effort",
36 - "low",
62 + reasoning_effort,
63 "--base-url",
64 base_url,
65 "--n-threads",
@@ -41,16 +67,29 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float:
67 ]
68
69 try:
70 + # Set up environment for the evaluation subprocess
71 + # Inherit current environment and add required variables
72 + eval_env = os.environ.copy()
73 + eval_env["OPENAI_API_KEY"] = "dummy"
74 +
75 # Run the evaluation
76 result = subprocess.run(
77 cmd,
78 text=True,
79 capture_output=True,
80 timeout=1800, # 30 minute timeout
50 - env={"OPENAI_API_KEY": "dummy"},
81 + env=eval_env,
82 )
83
53 - print("Evaluation process output:\n", result.stdout)
84 + print("Evaluation process stdout:\n", result.stdout)
85 + print("Evaluation process stderr:\n", result.stderr)
86 + print(f"Evaluation process return code: {result.returncode}")
87 +
88 + if result.returncode != 0:
89 + raise RuntimeError(
90 + f"Evaluation failed with exit code {result.returncode}:\n"
91 + f"stdout: {result.stdout}\nstderr: {result.stderr}"
92 + )
93
94 # Parse the output to extract the score
95 match = re.search(r"'metric':\s*([\d.]+)", result.stdout)
@@ -64,47 +103,62 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float:
103
104 except subprocess.TimeoutExpired as e:
105 raise RuntimeError("Evaluation timed out") from e
67 - except subprocess.CalledProcessError as e:
68 - raise RuntimeError(
69 - f"Evaluation failed with exit code {e.returncode}:\n"
70 - f"stdout: {e.stdout}\nstderr: {e.stderr}"
71 - ) from e
106
107
74 -def test_gpqa_correctness(request):
75 - """Test GPQA correctness for GPT-OSS model."""
108 +def test_gpqa_correctness(config_filename):
109 + """Test GPQA correctness for a given model configuration."""
110 + # Ensure tiktoken files are downloaded
111 + ensure_tiktoken_files()
112 +
113 + # Verify tiktoken files exist
114 + for filename in TIKTOKEN_FILES:
115 + filepath = TIKTOKEN_DATA_DIR / filename
116 + assert filepath.exists(), f"Tiktoken file not found: {filepath}"
117
77 - # Get command line arguments
78 - model_name = request.config.getoption("--model")
79 - expected_metric = request.config.getoption("--metric")
80 - server_args_str = request.config.getoption("--server-args")
118 + eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8"))
119
82 - # Parse server arguments
83 - server_args = []
84 - if server_args_str:
85 - server_args = server_args_str.split()
120 + # Parse server arguments from config (use shlex to handle quoted strings)
121 + server_args_str = eval_config.get("server_args", "")
122 + server_args = shlex.split(server_args_str) if server_args_str else []
123
124 # Add standard server arguments
125 server_args.extend(
126 [
127 "--trust-remote-code",
128 + "--enforce-eager",
129 + "--disable-uvicorn-access-log",
130 ]
131 )
132
94 - print(f"Starting GPQA evaluation for model: {model_name}")
95 - print(f"Expected metric threshold: {expected_metric}")
133 + # Build server environment with tiktoken path and any config-specified vars
134 + server_env = {"TIKTOKEN_ENCODINGS_BASE": str(TIKTOKEN_DATA_DIR)}
135 + if eval_config.get("env"):
136 + server_env.update(eval_config["env"])
137 +
138 + reasoning_effort = eval_config.get("reasoning_effort", "low")
139 +
140 + print(f"Starting GPQA evaluation for model: {eval_config['model_name']}")
141 + print(f"Expected metric threshold: {eval_config['metric_threshold']}")
142 + print(f"Reasoning effort: {reasoning_effort}")
143 print(f"Server args: {' '.join(server_args)}")
144 + print(f"Server environment variables: {server_env}")
145
146 # Launch server and run evaluation
147 with RemoteOpenAIServer(
100 - model_name, server_args, max_wait_seconds=1800
148 + eval_config["model_name"],
149 + server_args,
150 + env_dict=server_env,
151 + max_wait_seconds=eval_config.get("startup_max_wait_seconds", 1800),
152 ) as remote_server:
153 base_url = remote_server.url_for("v1")
154 print(f"Server started at: {base_url}")
155
105 - measured_metric = run_gpqa_eval(model_name, base_url)
156 + measured_metric = run_gpqa_eval(
157 + eval_config["model_name"], base_url, reasoning_effort
158 + )
159 + expected_metric = eval_config["metric_threshold"]
160
107 - print(f"GPQA Results for {model_name}:")
161 + print(f"GPQA Results for {eval_config['model_name']}:")
162 print(f" Measured metric: {measured_metric:.4f}")
163 print(f" Expected metric: {expected_metric:.4f}")
164 print(f" Tolerance: {TOL:.4f}")
@@ -115,4 +169,4 @@ def test_gpqa_correctness(request):
169 f"{expected_metric:.4f} - {TOL:.4f} = {expected_metric - TOL:.4f}"
170 )
171
118 - print(f"✅ GPQA test passed for {model_name}")
172 + print(f"GPQA test passed for {eval_config['model_name']}")
vllm/model_executor/layers/fused_moe/config.py
+4
@@ -242,6 +242,10 @@ class FusedMoEQuantConfig:
242 def quant_dtype(self) -> torch.dtype | str | None:
243 return self._a1.dtype
244
245 + @property
246 + def weight_quant_dtype(self) -> torch.dtype | str | None:
247 + return self._w1.dtype
248 +
249 @property
250 def is_quantized(self) -> bool:
251 return self.quant_dtype is not None
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+110 -8
@@ -4,6 +4,7 @@
4 import torch
5
6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm.config import get_current_vllm_config
8 from vllm.logger import init_logger
9 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10 from vllm.model_executor.layers.fused_moe.config import (
@@ -18,6 +19,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
19 kFp8Dynamic128Sym,
20 kFp8Static128BlockSym,
21 kFp8StaticTensorSym,
22 + kMxfp4Static,
23 + kMxfp8Dynamic,
24 kNvfp4Dynamic,
25 kNvfp4Static,
26 )
@@ -64,10 +67,18 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
67 quant_config: FusedMoEQuantConfig,
68 ):
69 super().__init__(moe_config, quant_config)
67 - assert quant_config.quant_dtype in ("nvfp4", torch.float8_e4m3fn, None), (
68 - "Only nvfp4, fp8, bfloat16 and"
70 +
71 + assert quant_config.weight_quant_dtype in (
72 + "mxfp4",
73 + "nvfp4",
74 + torch.float8_e4m3fn,
75 + None,
76 + ), (
77 + "Only mxfp4, nvfp4, fp8, bfloat16 and"
78 " float16 quantization are currently supported."
79 )
80 + self.device = moe_config.device
81 + self.num_experts = moe_config.num_local_experts
82 self.ep_rank = moe_config.moe_parallel_config.ep_rank
83 self.ep_size = moe_config.moe_parallel_config.ep_size
84 self.tp_rank = moe_config.moe_parallel_config.tp_rank
@@ -78,6 +89,28 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
89 # - pass per-block weight scales to the kernel
90 # - skip input activation quantization (kernel applies scaling)
91 self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized
92 + self.max_capture_size = (
93 + get_current_vllm_config().compilation_config.max_cudagraph_capture_size
94 + )
95 +
96 + if quant_config.weight_quant_dtype == "mxfp4":
97 + # This value is used specifically for gpt-oss,
98 + # Need to revisit this for other models
99 + self.gemm1_alpha = torch.tensor(
100 + [1.702] * self.num_experts, dtype=torch.float32, device=self.device
101 + )
102 + self.gemm1_beta = torch.tensor(
103 + [1.0] * self.num_experts, dtype=torch.float32, device=self.device
104 + )
105 + self.gemm1_clamp_limit = torch.tensor(
106 + [7.0] * self.num_experts, dtype=torch.float32, device=self.device
107 + )
108 + if quant_config.quant_dtype == "mxfp8":
109 + self.fake_input_scale = torch.ones(
110 + self.num_experts,
111 + device=self.device,
112 + dtype=torch.float32,
113 + )
114
115 @property
116 def expects_unquantized_inputs(self) -> bool:
@@ -119,20 +152,33 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
152 ]
153 and p.has_device_capability(90)
154 )
122 - # fp8 block-scale on 9.0
155 + # fp8 block-scale, wmxfp4a16 on 9.0
156 or (
124 - scheme == (kFp8Static128BlockSym, kFp8Dynamic128Sym)
157 + scheme
158 + in [
159 + (kMxfp4Static, None),
160 + (kFp8Static128BlockSym, kFp8Dynamic128Sym),
161 + ]
162 and p.is_device_capability(90)
163 )
127 - # nvfp4 on 10.0+
164 + # nvfp4, wmxfp4amxfp8 on 10.0+
165 or (
129 - scheme == (kNvfp4Static, kNvfp4Dynamic) and p.has_device_capability(100)
166 + scheme
167 + in [
168 + (kMxfp4Static, kMxfp8Dynamic),
169 + (kNvfp4Static, kNvfp4Dynamic),
170 + ]
171 + and p.has_device_capability(100)
172 )
173 )
174
175 @staticmethod
176 def _supports_activation(activation: MoEActivation) -> bool:
135 - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
177 + return activation in [
178 + MoEActivation.SILU,
179 + MoEActivation.RELU2_NO_MUL,
180 + MoEActivation.SWIGLUOAI,
181 + ]
182
183 @staticmethod
184 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -216,12 +262,23 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
262
263 activation_str_to_value_map = {
264 MoEActivation.SILU: ActivationType.Swiglu, # This is the default
265 + MoEActivation.SWIGLUOAI: ActivationType.Swiglu, # gpt-oss alias
266 MoEActivation.RELU2_NO_MUL: ActivationType.Relu2,
267 }
268 assert activation in activation_str_to_value_map, (
269 f"{activation=} missing from {activation_str_to_value_map.keys()=}"
270 )
271
272 + quant_scales = None
273 + fc1_expert_weights = None
274 + fc2_expert_weights = None
275 + fc1_expert_biases = None
276 + fc2_expert_biases = None
277 + swiglu_alpha = None
278 + swiglu_beta = None
279 + swiglu_limit = None
280 + use_mxfp8_act_scaling = False
281 + use_w4_group_scaling = False
282 # Select quantization metadata based on FP8 format/path
283 if (
284 self.quant_dtype == torch.float8_e4m3fn
@@ -256,6 +313,43 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
313 # FlashInfer API requires weight to be long for nvfp4
314 fc1_expert_weights = w1.view(torch.long)
315 fc2_expert_weights = w2.view(torch.long)
316 + elif self.weight_quant_dtype == "mxfp4":
317 + assert self.w1_scale is not None and self.w2_scale is not None
318 + assert w1.is_contiguous() and w2.is_contiguous()
319 + assert self.gemm1_alpha is not None
320 + assert self.gemm1_beta is not None
321 + assert self.gemm1_clamp_limit is not None
322 + assert topk_ids.is_contiguous()
323 +
324 + fc1_expert_biases = self.w1_bias
325 + fc2_expert_biases = self.w2_bias
326 + swiglu_alpha = self.gemm1_alpha
327 + swiglu_beta = self.gemm1_beta
328 + swiglu_limit = self.gemm1_clamp_limit
329 +
330 + if self.quant_dtype == "mxfp8":
331 + assert self.fake_input_scale is not None
332 + fc1_expert_weights = w1.view(torch.long)
333 + fc2_expert_weights = w2.view(torch.long)
334 +
335 + quant_scales = [
336 + self.w1_scale.view(torch.int32),
337 + self.fake_input_scale,
338 + self.w2_scale.view(torch.int32),
339 + self.fake_input_scale,
340 + ]
341 + use_mxfp8_act_scaling = True
342 + else:
343 + assert hidden_states.dtype == torch.bfloat16
344 + fc1_expert_weights = w1
345 + fc2_expert_weights = w2
346 + quant_scales = [
347 + self.w1_scale,
348 + self.w2_scale,
349 + ]
350 + a1q_scale = None
351 + use_w4_group_scaling = True
352 +
353 elif self.use_deepseek_fp8_block_scale:
354 # FP8 block-scale path: provide block-scale weights, omit a1q_scale
355 quant_scales = [
@@ -277,6 +371,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
371 token_final_scales=topk_weights,
372 fc1_expert_weights=fc1_expert_weights,
373 fc2_expert_weights=fc2_expert_weights,
374 + fc1_expert_biases=fc1_expert_biases,
375 + fc2_expert_biases=fc2_expert_biases,
376 + swiglu_alpha=swiglu_alpha,
377 + swiglu_beta=swiglu_beta,
378 + swiglu_limit=swiglu_limit,
379 + output=output,
380 output_dtype=self.out_dtype,
381 quant_scales=quant_scales,
382 input_sf=a1q_scale,
@@ -284,10 +384,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
384 tp_rank=self.tp_rank,
385 ep_size=self.ep_size,
386 ep_rank=self.ep_rank,
287 - output=output,
387 activation_type=activation_str_to_value_map[activation],
388 # Informs FlashInfer to use the block-scale decoding path when True
389 use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale,
390 + use_mxfp8_act_scaling=use_mxfp8_act_scaling,
391 + use_w4_group_scaling=use_w4_group_scaling,
392 + tune_max_num_tokens=max(self.max_capture_size, 1),
393 )
394
395 def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
vllm/model_executor/layers/fused_moe/modular_kernel.py
+5 -1
@@ -564,9 +564,13 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
564 #
565
566 @property
567 - def quant_dtype(self) -> torch.dtype | None:
567 + def quant_dtype(self) -> torch.dtype | str | None:
568 return self.quant_config.quant_dtype
569
570 + @property
571 + def weight_quant_dtype(self) -> torch.dtype | str | None:
572 + return self.quant_config.weight_quant_dtype
573 +
574 @property
575 def block_shape(self) -> list[int] | None:
576 return self.quant_config.block_shape
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+11 -6
@@ -25,15 +25,20 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
25 self,
26 moe_config: FusedMoEConfig,
27 quant_config: FusedMoEQuantConfig,
28 - gemm1_alpha,
29 - gemm1_beta,
30 - gemm1_clamp_limit,
28 max_capture_size,
29 ):
30 super().__init__(moe_config, quant_config)
34 - self.gemm1_alpha = gemm1_alpha
35 - self.gemm1_beta = gemm1_beta
36 - self.gemm1_clamp_limit = gemm1_clamp_limit
31 + self.device = torch.cuda.current_device()
32 + self.num_experts = moe_config.num_local_experts
33 + self.gemm1_alpha = torch.tensor(
34 + [1.702] * self.num_experts, dtype=torch.float32, device=self.device
35 + )
36 + self.gemm1_beta = torch.tensor(
37 + [1.0] * self.num_experts, dtype=torch.float32, device=self.device
38 + )
39 + self.gemm1_clamp_limit = torch.tensor(
40 + [7.0] * self.num_experts, dtype=torch.float32, device=self.device
41 + )
42 self.max_capture_size = max_capture_size
43
44 @staticmethod
vllm/model_executor/layers/fused_moe/utils.py
+9 -2
@@ -195,11 +195,12 @@ def _mxfp8_e4m3_quantize(
195 A_scale: torch.Tensor | None,
196 per_act_token_quant: bool,
197 block_shape: list[int] | None = None,
198 + is_sf_swizzled_layout: bool = False,
199 ) -> tuple[torch.Tensor, torch.Tensor]:
200 assert A_scale is None
201 assert not per_act_token_quant
202 assert block_shape is None
202 - return mxfp8_e4m3_quantize(A)
203 + return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout)
204
205
206 def _mxfp6_e3m2_quantize(
@@ -275,7 +276,13 @@ def moe_kernel_quantize_input(
276 elif quant_dtype == "mxfp8":
277 # TODO: `quant_dtype == "mxfp8"` is ambiguous,
278 # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`.
278 - return _mxfp8_e4m3_quantize(A, A_scale, per_act_token_quant, block_shape)
279 + return _mxfp8_e4m3_quantize(
280 + A,
281 + A_scale,
282 + per_act_token_quant,
283 + block_shape,
284 + is_sf_swizzled_layout=is_fp4_scale_swizzled,
285 + )
286 elif quant_dtype == "mxfp6_e3m2":
287 return _mxfp6_e3m2_quantize(A, A_scale, per_act_token_quant, block_shape)
288 elif quant_dtype == "mxfp6_e2m3":
vllm/model_executor/layers/quantization/mxfp4.py
+42 -90
@@ -256,6 +256,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
256 "Please check your environment and try again."
257 )
258 self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {}
259 + # Initialized in process_weights_after_loading for CUTLASS/SM90 backends
260 self.moe_mk: mk.FusedMoEModularKernel | None = None
261
262 def create_weights(
@@ -648,19 +649,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
649 self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
650 or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
651 ):
651 - layer.gemm1_alpha = Parameter(
652 - torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(),
653 - requires_grad=False,
654 - )
655 - layer.gemm1_beta = Parameter(
656 - torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(),
657 - requires_grad=False,
658 - )
659 - layer.gemm1_clamp_limit = Parameter(
660 - torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(),
661 - requires_grad=False,
662 - )
663 -
652 sf_block_size = 32 # mxfp4 block size
653
654 # Common shape assertions
@@ -772,6 +760,30 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
760 layer.w2_weight_scale = torch.nn.Parameter(
761 w2_scales_interleaved, requires_grad=False
762 )
763 +
764 + # theses two kernels go through the `flashinfer_cutlass_fused_moe` path
765 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
766 + FlashInferExperts,
767 + )
768 +
769 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
770 + assert self.moe_quant_config is not None
771 + prepare_finalize = maybe_make_prepare_finalize(
772 + moe=self.moe,
773 + quant_config=self.moe_quant_config,
774 + routing_tables=layer._maybe_init_expert_routing_tables(),
775 + allow_new_interface=True,
776 + )
777 + assert prepare_finalize is not None
778 +
779 + self.moe_mk = mk.FusedMoEModularKernel(
780 + prepare_finalize,
781 + FlashInferExperts(
782 + moe_config=self.moe,
783 + quant_config=self.moe_quant_config,
784 + ),
785 + shared_experts=None,
786 + )
787 elif self.mxfp4_backend == Mxfp4Backend.TRITON:
788 from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig
789
@@ -847,7 +859,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
859 w1_scale=layer.w13_weight_scale,
860 w2_scale=layer.w2_weight_scale,
861 )
850 - elif self.mxfp4_backend in [Mxfp4Backend.SM100_FI_MXFP4_BF16]:
862 + elif self.mxfp4_backend in [
863 + Mxfp4Backend.SM100_FI_MXFP4_BF16,
864 + Mxfp4Backend.SM90_FI_MXFP4_BF16,
865 + ]:
866 return mxfp4_w4a16_moe_quant_config(
867 w1_bias=layer.w13_bias,
868 w2_bias=layer.w2_bias,
@@ -897,9 +912,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
912 ):
913 # B200 code-path
914 kwargs = {
900 - "gemm1_alpha": layer.gemm1_alpha,
901 - "gemm1_beta": layer.gemm1_beta,
902 - "gemm1_clamp_limit": layer.gemm1_clamp_limit,
915 # TODO(bnell): part of quant_config
916 "max_capture_size": self.max_capture_size,
917 }
@@ -935,20 +947,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
947 if layer.enable_eplb:
948 raise NotImplementedError("EPLB is not supported for mxfp4")
949
938 - if self.mxfp4_backend == Mxfp4Backend.MARLIN:
939 - assert self.moe_mk is not None
940 -
941 - return self.moe_mk(
942 - hidden_states=x,
943 - w1=layer.w13_weight,
944 - w2=layer.w2_weight,
945 - topk_weights=topk_weights,
946 - topk_ids=topk_ids,
947 - activation=layer.activation,
948 - global_num_experts=layer.global_num_experts,
949 - expert_map=layer.expert_map,
950 - apply_router_weight_on_input=layer.apply_router_weight_on_input,
951 - )
950 assert _can_support_mxfp4(
951 layer.use_grouped_topk,
952 layer.topk_group,
@@ -967,69 +965,23 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
965 assert (
966 self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS
967 or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
968 + or self.mxfp4_backend == Mxfp4Backend.MARLIN
969 )
971 - from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe
972 -
973 - # Backend-specific preparation
974 - if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS:
975 - from flashinfer import mxfp8_quantize
976 -
977 - x_quant, x_scale = mxfp8_quantize(x, True, 32)
978 -
979 - fake_input_scale = torch.ones(self.num_experts, device=x.device)
980 - quant_scales = [
981 - layer.w13_weight_scale.contiguous().view(torch.int32),
982 - fake_input_scale,
983 - layer.w2_weight_scale.contiguous().view(torch.int32),
984 - fake_input_scale,
985 - ]
986 -
987 - fi_input = x_quant
988 - extra_kwargs = dict(
989 - use_mxfp8_act_scaling=True,
990 - input_sf=x_scale,
991 - fc1_expert_weights=layer.w13_weight.contiguous().view(torch.long),
992 - fc2_expert_weights=layer.w2_weight.contiguous().view(torch.long),
993 - )
994 - elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16:
995 - assert x.dtype == torch.bfloat16
996 -
997 - quant_scales = [
998 - layer.w13_weight_scale,
999 - layer.w2_weight_scale,
1000 - ]
1001 -
1002 - fi_input = x
1003 - extra_kwargs = dict(
1004 - use_w4_group_scaling=True,
1005 - fc1_expert_weights=layer.w13_weight,
1006 - fc2_expert_weights=layer.w2_weight,
1007 - )
970
1009 - output = torch.empty_like(x, dtype=torch.bfloat16)
1010 -
1011 - flashinfer_cutlass_fused_moe(
1012 - input=fi_input,
1013 - token_selected_experts=topk_ids.to(torch.int).contiguous(),
1014 - token_final_scales=topk_weights,
1015 - output_dtype=torch.bfloat16,
1016 - output=output,
1017 - quant_scales=quant_scales,
1018 - fc1_expert_biases=layer.w13_bias,
1019 - fc2_expert_biases=layer.w2_bias,
1020 - swiglu_alpha=layer.gemm1_alpha,
1021 - swiglu_beta=layer.gemm1_beta,
1022 - swiglu_limit=layer.gemm1_clamp_limit,
1023 - tp_size=self.moe.tp_size,
1024 - tp_rank=self.moe.tp_rank,
1025 - ep_size=self.moe.ep_size,
1026 - ep_rank=self.moe.ep_rank,
1027 - tune_max_num_tokens=max(self.max_capture_size, 1),
1028 - **extra_kwargs,
971 + assert self.moe_mk is not None
972 + return self.moe_mk(
973 + hidden_states=x,
974 + w1=layer.w13_weight,
975 + w2=layer.w2_weight,
976 + topk_weights=topk_weights,
977 + topk_ids=topk_ids,
978 + activation=layer.activation,
979 + global_num_experts=layer.global_num_experts,
980 + apply_router_weight_on_input=layer.apply_router_weight_on_input,
981 + expert_map=layer.expert_map,
982 + shared_experts_input=shared_experts_input,
983 )
984
1031 - return output
1032 -
985 def apply_monolithic(
986 self,
987 layer: FusedMoE,
vllm/model_executor/layers/quantization/utils/quant_utils.py
+13
@@ -19,6 +19,7 @@ if TYPE_CHECKING:
19
20 FP8_DTYPE = current_platform.fp8_dtype()
21 FP4_DTYPE = torch.uint8
22 +MXFP_SCALE_DTYPE = torch.uint8
23
24
25 def get_fp8_min_max() -> tuple[float, float]:
@@ -151,6 +152,18 @@ kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True
152 kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64))
153 kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True)
154
155 +# TODO (zyongye): Convert all the torch.dtype to scale_dtype
156 +# Changing that requires changing torch compile fused AR+Quant Quant key
157 +# to avoid assertion error
158 +kMxfp4DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32))
159 +kMxfp4Dynamic = QuantKey(FP4_DTYPE, scale=kMxfp4DynamicGroupScale, symmetric=True)
160 +
161 +kMxfp8DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32))
162 +kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=True)
163 +
164 +kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32))
165 +kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True)
166 +
167
168 # Normalize the group_shape to the full extent for any dims that are -1
169 def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape):