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 10, 12:45 UTC
[MoE Refactor][15/N] Apply Refactor to Fp8 (#31415)
Robert Shaw committed
Jan 7, 2026 at 19:42 UTC
5dcd7ef1f219068e6b6be5b614bc43978f028651
38 files changed
+1446
-1535
.buildkite/test-pipeline.yaml
+7
@@ -1422,3 +1422,10 @@ steps:
1422
num_gpus: 2
1423
commands:
1424
- pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor/config-b200.txt
1425
+
1426
+- label: MoE Refactor Integration Test (B200 DP - TEMPORARY) # optional
1427
+ gpu: b200
1428
+ optional: true
1429
+ num_gpus: 2
1430
+ commands:
1431
+ - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
benchmarks/kernels/benchmark_cutlass_moe_fp8.py
+24
-92
@@ -6,13 +6,16 @@ kernel. Both kernels take in fp8 quantized weights and 16-bit activations,
6
but use different quantization strategies and backends.
7
"""
8
9
-import nvtx
9
import torch
10
11
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm import _custom_ops as ops
13
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
14
-from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
14
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
15
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
16
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
17
+ MoEPrepareAndFinalizeNoEP,
18
+)
19
from vllm.platforms import current_platform
20
from vllm.utils.argparse_utils import FlexibleArgumentParser
21
from vllm.v1.worker.workspace import init_workspace_manager
@@ -59,6 +62,7 @@ def bench_run(
62
per_out_ch: bool,
63
mkn: tuple[int, int, int],
64
):
65
+ init_workspace_manager(torch.cuda.current_device())
66
(m, k, n) = mkn
67
68
dtype = torch.half
@@ -121,85 +125,6 @@ def bench_run(
125
# Force per-tensor quantization for all cases
126
per_act_token = False
127
124
- # Create stride tensors for CUTLASS
125
- ab_strides1 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
126
- ab_strides2 = torch.full((num_experts,), n, dtype=torch.int64, device=device)
127
- c_strides1 = torch.full((num_experts,), 2 * n, dtype=torch.int64, device=device)
128
- c_strides2 = torch.full((num_experts,), k, dtype=torch.int64, device=device)
129
-
130
- def run_triton_moe(
131
- a: torch.Tensor,
132
- w1: torch.Tensor,
133
- w2: torch.Tensor,
134
- topk_weights: torch.Tensor,
135
- topk_ids: torch.Tensor,
136
- w1_scale: torch.Tensor,
137
- w2_scale: torch.Tensor,
138
- a1_scale: torch.Tensor,
139
- a2_scale: torch.Tensor,
140
- num_repeats: int,
141
- ):
142
- quant_config = fp8_w8a8_moe_quant_config(
143
- w1_scale=w1_scale,
144
- w2_scale=w2_scale,
145
- a1_scale=a1_scale,
146
- a2_scale=a2_scale,
147
- per_act_token_quant=per_act_token,
148
- per_out_ch_quant=per_out_ch,
149
- )
150
-
151
- for _ in range(num_repeats):
152
- fused_experts(
153
- a,
154
- w1,
155
- w2,
156
- topk_weights,
157
- topk_ids,
158
- quant_config=quant_config,
159
- )
160
-
161
- def run_cutlass_moe_fp8(
162
- a: torch.Tensor,
163
- w1: torch.Tensor,
164
- w2: torch.Tensor,
165
- topk_weights: torch.Tensor,
166
- topk_ids: torch.Tensor,
167
- ab_strides1: torch.Tensor,
168
- ab_strides2: torch.Tensor,
169
- c_strides1: torch.Tensor,
170
- c_strides2: torch.Tensor,
171
- w1_scale: torch.Tensor,
172
- w2_scale: torch.Tensor,
173
- a1_scale: torch.Tensor,
174
- a2_scale: torch.Tensor,
175
- num_repeats: int,
176
- ):
177
- quant_config = fp8_w8a8_moe_quant_config(
178
- w1_scale=w1_scale,
179
- w2_scale=w2_scale,
180
- a1_scale=a1_scale,
181
- a2_scale=a2_scale,
182
- per_act_token_quant=per_act_token,
183
- per_out_ch_quant=per_out_ch,
184
- )
185
-
186
- for _ in range(num_repeats):
187
- with nvtx.annotate("cutlass_moe_fp8", color="blue"):
188
- cutlass_moe_fp8(
189
- a=a,
190
- w1_q=w1,
191
- w2_q=w2,
192
- topk_weights=topk_weights,
193
- topk_ids=topk_ids,
194
- ab_strides1=ab_strides1,
195
- ab_strides2=ab_strides2,
196
- c_strides1=c_strides1,
197
- c_strides2=c_strides2,
198
- quant_config=quant_config,
199
- activation="silu",
200
- global_num_experts=num_experts,
201
- )
202
-
128
# Pre-create quantization config to avoid creating it inside CUDA graph
129
quant_config = fp8_w8a8_moe_quant_config(
130
w1_scale=w1_scale,
@@ -210,23 +135,30 @@ def bench_run(
135
per_out_ch_quant=per_out_ch,
136
)
137
138
+ fn = mk.FusedMoEModularKernel(
139
+ MoEPrepareAndFinalizeNoEP(),
140
+ CutlassExpertsFp8(
141
+ out_dtype=a.dtype,
142
+ e=num_experts,
143
+ n=n,
144
+ k=k,
145
+ quant_config=quant_config,
146
+ device=w1.device,
147
+ ),
148
+ )
149
+
150
# Create CUDA graphs for CUTLASS (match benchmark_moe.py pattern exactly)
151
cutlass_stream = torch.cuda.Stream()
152
cutlass_graph = torch.cuda.CUDAGraph()
153
with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):
154
# Capture 10 invocations like benchmark_moe.py
155
for _ in range(10):
219
- cutlass_moe_fp8(
220
- a=a,
221
- w1_q=w1_fp8q_cutlass,
222
- w2_q=w2_fp8q_cutlass,
223
- topk_weights=topk_weights,
224
- topk_ids=topk_ids,
225
- ab_strides1=ab_strides1,
226
- ab_strides2=ab_strides2,
227
- c_strides1=c_strides1,
228
- c_strides2=c_strides2,
229
- quant_config=quant_config,
156
+ fn(
157
+ a,
158
+ w1_fp8q_cutlass,
159
+ w2_fp8q_cutlass,
160
+ topk_weights,
161
+ topk_ids,
162
activation="silu",
163
global_num_experts=num_experts,
164
)
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+37
-53
@@ -5,14 +5,18 @@ import torch
5
import torch.utils.benchmark as benchmark
6
from benchmark_shapes import WEIGHT_SHAPES_MOE
7
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm import _custom_ops as ops
10
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
11
from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config
11
-from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp8
12
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
13
from vllm.model_executor.layers.fused_moe.fused_moe import (
14
fused_experts,
15
fused_topk,
16
)
17
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
18
+ MoEPrepareAndFinalizeNoEP,
19
+)
20
from vllm.utils.argparse_utils import FlexibleArgumentParser
21
from vllm.v1.worker.workspace import init_workspace_manager
22
@@ -45,6 +49,7 @@ def bench_run(
49
per_out_ch: bool,
50
mkn: tuple[int, int, int],
51
):
52
+ init_workspace_manager(torch.cuda.current_device())
53
label = "Quant Matmul"
54
55
sub_label = (
@@ -82,11 +87,6 @@ def bench_run(
87
a, score, topk, renormalize=False
88
)
89
85
- ab_strides1 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
86
- ab_strides2 = torch.full((num_experts,), n, device="cuda", dtype=torch.int64)
87
- c_strides1 = torch.full((num_experts,), 2 * n, device="cuda", dtype=torch.int64)
88
- c_strides2 = torch.full((num_experts,), k, device="cuda", dtype=torch.int64)
89
-
90
def run_triton_moe(
91
a: torch.Tensor,
92
w1: torch.Tensor,
@@ -120,10 +120,6 @@ def bench_run(
120
w2: torch.Tensor,
121
w1_scale: torch.Tensor,
122
w2_scale: torch.Tensor,
123
- ab_strides1: torch.Tensor,
124
- ab_strides2: torch.Tensor,
125
- c_strides1: torch.Tensor,
126
- c_strides2: torch.Tensor,
123
topk_weights: torch.Tensor,
124
topk_ids: torch.Tensor,
125
per_act_token: bool,
@@ -135,31 +131,29 @@ def bench_run(
131
per_act_token_quant=per_act_token,
132
)
133
138
- for _ in range(num_repeats):
139
- cutlass_moe_fp8(
140
- a,
141
- w1,
142
- w2,
143
- topk_weights,
144
- topk_ids,
145
- ab_strides1,
146
- ab_strides2,
147
- c_strides1,
148
- c_strides2,
134
+ fn = mk.FusedMoEModularKernel(
135
+ MoEPrepareAndFinalizeNoEP(),
136
+ CutlassExpertsFp8(
137
+ out_dtype=a.dtype,
138
+ # NOTE(rob): w2 is shaped as [E, hidden, intermediate]
139
+ e=w2.shape[0],
140
+ n=w2.shape[2],
141
+ k=w2.shape[1],
142
quant_config=quant_config,
150
- )
143
+ device=w1.device,
144
+ ),
145
+ )
146
+
147
+ for _ in range(num_repeats):
148
+ fn(a, w1, w2, topk_weights, topk_ids)
149
150
def run_cutlass_from_graph(
151
a: torch.Tensor,
152
a_scale: torch.Tensor,
155
- w1_q: torch.Tensor,
156
- w2_q: torch.Tensor,
153
+ w1: torch.Tensor,
154
+ w2: torch.Tensor,
155
w1_scale: torch.Tensor,
156
w2_scale: torch.Tensor,
159
- ab_strides1: torch.Tensor,
160
- ab_strides2: torch.Tensor,
161
- c_strides1: torch.Tensor,
162
- c_strides2: torch.Tensor,
157
topk_weights: torch.Tensor,
158
topk_ids: torch.Tensor,
159
):
@@ -169,21 +163,23 @@ def bench_run(
163
per_act_token_quant=per_act_token,
164
)
165
166
+ fn = mk.FusedMoEModularKernel(
167
+ MoEPrepareAndFinalizeNoEP(),
168
+ CutlassExpertsFp8(
169
+ out_dtype=a.dtype,
170
+ # NOTE(rob): w2 is shaped as [E, hidden, intermediate]
171
+ e=w2.shape[0],
172
+ n=w2.shape[2],
173
+ k=w2.shape[1],
174
+ quant_config=quant_config,
175
+ device=w1.device,
176
+ ),
177
+ )
178
+
179
with set_current_vllm_config(
180
VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
181
):
175
- return cutlass_moe_fp8(
176
- a,
177
- w1_q,
178
- w2_q,
179
- topk_weights,
180
- topk_ids,
181
- ab_strides1,
182
- ab_strides2,
183
- c_strides1,
184
- c_strides2,
185
- quant_config=quant_config,
186
- )
182
+ return fn(a, w1, w2, topk_weights, topk_ids)
183
184
def run_triton_from_graph(
185
a: torch.Tensor,
@@ -227,10 +223,6 @@ def bench_run(
223
w2_q,
224
w1_scale,
225
w2_scale,
230
- ab_strides1,
231
- ab_strides2,
232
- c_strides1,
233
- c_strides2,
226
topk_weights,
227
topk_ids,
228
)
@@ -268,10 +260,6 @@ def bench_run(
260
"w1_scale": w1_scale,
261
"w2_scale": w2_scale,
262
"per_act_token": per_act_token,
271
- "ab_strides1": ab_strides1,
272
- "ab_strides2": ab_strides2,
273
- "c_strides1": c_strides1,
274
- "c_strides2": c_strides2,
263
# cuda graph params
264
"cutlass_graph": cutlass_graph,
265
"triton_graph": triton_graph,
@@ -330,10 +318,6 @@ def bench_run(
318
w2_q,
319
w1_scale,
320
w2_scale,
333
- ab_strides1,
334
- ab_strides2,
335
- c_strides1,
336
- c_strides2,
321
topk_weights,
322
topk_ids,
323
per_act_token,
@@ -342,7 +326,7 @@ def bench_run(
326
327
results.append(
328
benchmark.Timer(
345
- stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, ab_strides1, ab_strides2, c_strides1, c_strides2, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
329
+ stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
330
globals=globals,
331
label=label,
332
sub_label=sub_label,
benchmarks/kernels/benchmark_moe.py
-2
@@ -48,8 +48,6 @@ def clear_triton_cache():
48
49
# Try to clear Triton's runtime cache
50
try:
51
- import triton
52
-
51
if (
52
hasattr(triton, "runtime")
53
and hasattr(triton.runtime, "cache")
docs/design/moe_kernel_features.md
+1
-1
@@ -87,7 +87,7 @@ To be used with a particular `FusedMoEPrepareAndFinalize` subclass, MoE kernels
87
| triton (batched) | batched | all<sup>1</sup> | G,A,T | silu, gelu | <sup>6</sup> | Y | [`BatchedTritonExperts`][vllm.model_executor.layers.fused_moe.fused_batched_moe.BatchedTritonExperts] |
88
| deep gemm | standard,</br>batched | fp8 | G(128),A,T | silu, gelu | <sup>6</sup> | Y | [`deep_gemm_moe_fp8`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.deep_gemm_moe_fp8],</br>[`DeepGemmExperts`][vllm.model_executor.layers.fused_moe.deep_gemm_moe.DeepGemmExperts],</br>[`BatchedDeepGemmExperts`][vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe.BatchedDeepGemmExperts] |
89
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp4],</br>[`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
90
-| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`cutlass_moe_fp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.cutlass_moe_fp8],</br>[`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
90
+| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
91
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`flashinfer_cutlass_moe_fp4`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.flashinfer_cutlass_moe_fp4],</br>[`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
92
| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
93
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Llama-4-Scout-Fp8-ModelOpt-triton.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8"
2
+accuracy_threshold: 0.92
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ht.yaml
new
+8
@@ -0,0 +1,8 @@
1
+model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_high_throughput"
6
+env:
7
+ VLLM_USE_DEEP_GEMM: "1"
8
+ VLLM_USE_DEEP_GEMM_MOE: "1"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml
new
+9
@@ -0,0 +1,9 @@
1
+model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --disable-uvicorn-access-log"
6
+env:
7
+ VLLM_USE_DEEP_GEMM: "1"
8
+ VLLM_USE_DEEP_GEMM_MOE: "1"
9
+ VLLM_USE_DEEP_GEMM_E8M0: "0"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm.yaml
new
+8
@@ -0,0 +1,8 @@
1
+model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_DEEP_GEMM: "1"
8
+ VLLM_USE_DEEP_GEMM_MOE: "1"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ht.yaml
new
+8
@@ -0,0 +1,8 @@
1
+model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block"
2
+accuracy_threshold: 0.85
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_high_throughput"
6
+env:
7
+ VLLM_USE_DEEP_GEMM: "1"
8
+ VLLM_USE_DEEP_GEMM_MOE: "1"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
new
+9
@@ -0,0 +1,9 @@
1
+model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block"
2
+accuracy_threshold: 0.85
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --disable-uvicorn-access-log"
6
+env:
7
+ VLLM_USE_DEEP_GEMM: "1"
8
+ VLLM_USE_DEEP_GEMM_MOE: "1"
9
+ VLLM_USE_DEEP_GEMM_E8M0: "0"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-Fp8-CT-Block-deepgemm.yaml
new
+8
@@ -0,0 +1,8 @@
1
+model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block"
2
+accuracy_threshold: 0.85
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_DEEP_GEMM: "1"
8
+ VLLM_USE_DEEP_GEMM_MOE: "1"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml
new
+8
@@ -0,0 +1,8 @@
1
+model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4"
2
+accuracy_threshold: 0.88
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel"
6
+env:
7
+ VLLM_USE_FLASHINFER_MOE_FP4: "1"
8
+ VLLM_FLASHINFER_MOE_BACKEND: "throughput"
tests/evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt
new
+8
@@ -0,0 +1,8 @@
1
+Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ht.yaml
2
+Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm-deepep-ll.yaml
3
+Qwen3-30B-A3B-Fp8-AutoFp8-deepgemm.yaml
4
+Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ht.yaml
5
+Qwen3-30B-A3B-Fp8-CT-Block-deepgemm-deepep-ll.yaml
6
+Qwen3-30B-A3B-Fp8-CT-Block-deepgemm.yaml
7
+Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml
8
+Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass-fi-a2av.yaml
tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-CT-vllm-cutlass.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "RedHatAI/Llama-4-Scout-17B-16E-Instruct-FP8-dynamic"
2
+accuracy_threshold: 0.92
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2"
tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml
renamed
tests/evals/gsm8k/configs/moe-refactor/config-b200.txt
+1
@@ -1,3 +1,4 @@
1
+Llama-4-Scout-Fp8-CT-vllm-cutlass.yaml
2
Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml
3
Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml
4
Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml
tests/evals/gsm8k/configs/moe-refactor/config-h100.txt
+1
-1
@@ -5,7 +5,7 @@ Qwen3-30B-A3B-Fp8-AutoFp8-marlin.yaml
5
Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml
6
Qwen3-30B-A3B-Fp8-CT-Block-deepgemm.yaml
7
Qwen3-30B-A3B-Fp8-CT-Block-marlin.yaml
8
-Qwen3-30B-A3B-Fp8-CT-Block-vllm-cutlass.yaml
8
+Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml
9
Qwen3-30B-A3B-Fp8-CT-Channel-marlin.yaml
10
Qwen3-30B-A3B-Fp8-CT-Channel-vllm-cutlass.yaml
11
Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml
tests/evals/gsm8k/test_gsm8k_correctness.py
+1
@@ -61,6 +61,7 @@ def test_gsm8k_correctness(config_filename):
61
server_args.extend(
62
[
63
"--trust-remote-code",
64
+ "--disable-uvicorn-access-log",
65
]
66
)
67
tests/kernels/moe/test_cutlass_moe.py
+46
-26
@@ -7,17 +7,22 @@ from math import prod
7
import pytest
8
import torch
9
10
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm import _custom_ops as ops
12
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
13
from vllm.model_executor.layers.fused_moe.config import (
14
FUSED_MOE_UNQUANTIZED_CONFIG,
15
+ FusedMoEQuantConfig,
16
fp8_w8a8_moe_quant_config,
17
)
18
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
17
- cutlass_moe_fp8,
19
+ CutlassExpertsFp8,
20
run_cutlass_moe_fp8,
21
)
22
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
23
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
24
+ MoEPrepareAndFinalizeNoEP,
25
+)
26
from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
27
from vllm.platforms import current_platform
28
from vllm.utils.torch_utils import set_random_seed
@@ -150,16 +155,15 @@ class MOETensors8Bit(MOETensors):
155
156
157
def run_with_expert_maps(
153
- num_experts: int, num_local_experts: int, **cutlass_moe_kwargs
158
+ num_experts: int,
159
+ num_local_experts: int,
160
+ quant_config: FusedMoEQuantConfig,
161
+ **cutlass_moe_kwargs,
162
):
163
def slice_experts():
164
slice_params = [
157
- "w1_q",
158
- "w2_q",
159
- "ab_strides1",
160
- "ab_strides2",
161
- "c_strides1",
162
- "c_strides2",
165
+ "w1",
166
+ "w2",
167
]
168
full_tensors = {
169
k: v
@@ -167,8 +171,6 @@ def run_with_expert_maps(
171
if k in slice_params and k in cutlass_moe_kwargs
172
}
173
170
- quant_config = cutlass_moe_kwargs["quant_config"]
171
-
174
for i in range(0, num_experts, num_local_experts):
175
s, e = i, i + num_local_experts
176
@@ -187,13 +189,23 @@ def run_with_expert_maps(
189
new_quant_config._w1.scale = quant_config.w1_scale[s:e]
190
new_quant_config._w2.scale = quant_config.w2_scale[s:e]
191
190
- cutlass_moe_kwargs["quant_config"] = new_quant_config
191
-
192
- yield cutlass_moe_kwargs
193
-
194
- out_tensor = torch.zeros_like(cutlass_moe_kwargs["a"])
195
- for kwargs in slice_experts():
196
- out_tensor = out_tensor + cutlass_moe_fp8(**kwargs)
192
+ yield cutlass_moe_kwargs, new_quant_config
193
+
194
+ out_tensor = torch.zeros_like(cutlass_moe_kwargs["hidden_states"])
195
+ for kwargs, new_quant_config in slice_experts():
196
+ kernel = mk.FusedMoEModularKernel(
197
+ MoEPrepareAndFinalizeNoEP(),
198
+ CutlassExpertsFp8(
199
+ out_dtype=kwargs["hidden_states"].dtype,
200
+ # NOTE(rob): w2 is shaped as [E, hidden, intermediate]
201
+ e=kwargs["w2"].shape[0], # type: ignore[union-attr]
202
+ n=kwargs["w2"].shape[2], # type: ignore[union-attr]
203
+ k=kwargs["w2"].shape[1], # type: ignore[union-attr]
204
+ quant_config=new_quant_config,
205
+ device="cuda",
206
+ ),
207
+ )
208
+ out_tensor = out_tensor + kernel(**kwargs)
209
210
return out_tensor
211
@@ -230,27 +242,35 @@ def run_8_bit(
242
)
243
244
kwargs = {
233
- "a": moe_tensors.a,
234
- "w1_q": moe_tensors.w1_q, # type: ignore[union-attr]
235
- "w2_q": moe_tensors.w2_q, # type: ignore[union-attr]
245
+ "hidden_states": moe_tensors.a,
246
+ "w1": moe_tensors.w1_q, # type: ignore[union-attr]
247
+ "w2": moe_tensors.w2_q, # type: ignore[union-attr]
248
"topk_weights": topk_weights,
249
"topk_ids": topk_ids,
238
- "ab_strides1": moe_tensors.ab_strides1,
239
- "ab_strides2": moe_tensors.ab_strides2,
240
- "c_strides1": moe_tensors.c_strides1,
241
- "c_strides2": moe_tensors.c_strides2,
242
- "quant_config": quant_config,
250
}
251
252
num_experts = moe_tensors.w1.size(0)
253
with_ep = num_local_experts is not None or num_local_experts == num_experts
254
if not with_ep:
248
- return cutlass_moe_fp8(**kwargs)
255
+ kernel = mk.FusedMoEModularKernel(
256
+ MoEPrepareAndFinalizeNoEP(),
257
+ CutlassExpertsFp8(
258
+ out_dtype=moe_tensors.a.dtype,
259
+ # NOTE(rob): w2 is shaped as [E, hidden, intermediate]
260
+ e=moe_tensors.w2_q.shape[0], # type: ignore[union-attr]
261
+ n=moe_tensors.w2_q.shape[2], # type: ignore[union-attr]
262
+ k=moe_tensors.w2_q.shape[1], # type: ignore[union-attr]
263
+ quant_config=quant_config,
264
+ device="cuda",
265
+ ),
266
+ )
267
+ return kernel(**kwargs)
268
269
assert num_local_experts is not None
270
return run_with_expert_maps(
271
num_experts,
272
num_local_experts, # type: ignore[arg-type]
273
+ quant_config,
274
**kwargs,
275
)
276
tests/kernels/moe/test_flashinfer.py
+37
-12
@@ -11,12 +11,17 @@ from vllm.model_executor.layers.fused_moe.config import (
11
FusedMoEQuantConfig,
12
fp8_w8a8_moe_quant_config,
13
)
14
+from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
15
+ FlashInferExperts,
16
+)
17
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
18
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
19
+ MoEPrepareAndFinalizeNoEP,
20
+)
21
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
16
- apply_flashinfer_per_tensor_scale_fp8,
17
- flashinfer_cutlass_moe_fp8,
22
+ apply_fi_trtllm_fp8_per_tensor_moe,
23
register_scales_for_trtllm_fp8_per_tensor_moe,
19
- rotate_flashinfer_fp8_moe_weights,
24
+ rotate_weights_for_fi_trtllm_fp8_per_tensor_moe,
25
swap_w13_to_w31,
26
)
27
from vllm.model_executor.layers.quantization.utils.fp8_utils import input_to_float8
@@ -103,6 +108,7 @@ class TestData:
108
w2_quantized, w2_weight_scale = quant_fp8_per_tensor_batches(w2)
109
110
layer = torch.nn.Module()
111
+ layer.orig_dtype = torch.bfloat16
112
layer.w13_weight = w13_quantized.clone()
113
layer.w2_weight = w2_quantized.clone()
114
layer.w13_input_scale = a1_scale
@@ -115,10 +121,10 @@ class TestData:
121
pcp_size=1,
122
dp_size=1,
123
ep_size=1,
118
- tp_rank=1,
119
- pcp_rank=1,
120
- dp_rank=1,
121
- ep_rank=1,
124
+ tp_rank=0,
125
+ pcp_rank=0,
126
+ dp_rank=0,
127
+ ep_rank=0,
128
use_ep=False,
129
all2all_backend="naive",
130
)
@@ -126,7 +132,9 @@ class TestData:
132
# flashinfer expects swapped rows for w13
133
layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)
134
if is_trtllm:
129
- rotate_flashinfer_fp8_moe_weights(layer.w13_weight, layer.w2_weight)
135
+ rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(
136
+ layer.w13_weight, layer.w2_weight
137
+ )
138
register_scales_for_trtllm_fp8_per_tensor_moe(
139
layer,
140
layer.w13_weight_scale,
@@ -199,7 +207,7 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
207
quant_config=quant_config,
208
)
209
202
- flashinfer_output = apply_flashinfer_per_tensor_scale_fp8(
210
+ flashinfer_output = apply_fi_trtllm_fp8_per_tensor_moe(
211
layer=td.layer,
212
hidden_states=td.hidden_states,
213
router_logits=score,
@@ -277,17 +285,34 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
285
td.layer.get_fused_moe_quant_config = get_fused_moe_quant_config
286
td.layer.quant_method = td.layer
287
280
- flashinfer_cutlass_output = flashinfer_cutlass_moe_fp8(
288
+ kernel = mk.FusedMoEModularKernel(
289
+ MoEPrepareAndFinalizeNoEP(
290
+ defer_input_quant=quant_config.is_block_quantized
291
+ ),
292
+ FlashInferExperts(
293
+ out_dtype=td.layer.orig_dtype,
294
+ quant_config=quant_config,
295
+ ep_rank=td.layer.moe_parallel_config.ep_rank,
296
+ ep_size=td.layer.moe_parallel_config.ep_size,
297
+ tp_rank=td.layer.moe_parallel_config.tp_rank,
298
+ tp_size=td.layer.moe_parallel_config.tp_size,
299
+ use_dp=False,
300
+ use_deepseek_fp8_block_scale=False,
301
+ ),
302
+ )
303
+
304
+ flashinfer_cutlass_output = kernel(
305
td.hidden_states,
282
- td.layer,
306
+ td.layer.w13_weight,
307
+ td.layer.w2_weight,
308
topk_weights,
309
topk_ids,
310
+ inplace=False,
311
activation=activation,
312
global_num_experts=e,
313
expert_map=None,
314
apply_router_weight_on_input=True,
315
)
290
-
316
torch.testing.assert_close(
317
output, flashinfer_cutlass_output, atol=5.5e-2, rtol=1e-2
318
)
tests/quantization/test_fp8.py
+12
-7
@@ -15,7 +15,6 @@ from vllm.model_executor.layers.quantization.fp8 import (
15
Fp8Config,
16
Fp8KVCacheMethod,
17
Fp8LinearMethod,
18
- Fp8MoeBackend,
18
Fp8MoEMethod,
19
)
20
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
@@ -278,8 +277,18 @@ def test_scaled_fp8_quant(dtype) -> None:
277
# this is the case for marlin as well as per-tensor Fp8MoEMethod
278
@pytest.mark.parametrize("use_marlin", [False]) # skip True
279
def test_fp8_reloading(
281
- method_cls, is_checkpoint_fp8_serialized, weight_block_size, use_marlin, dist_init
280
+ method_cls,
281
+ is_checkpoint_fp8_serialized,
282
+ weight_block_size,
283
+ use_marlin,
284
+ dist_init,
285
+ monkeypatch,
286
):
287
+ # NOTE(rob): this test fails when using DeepGEMM because the
288
+ # shapes are invalid. Previously the test was passing because
289
+ # we set fp8_backend to None, which sidestepped the issue.
290
+ monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "0")
291
+
292
if is_checkpoint_fp8_serialized is False:
293
pytest.skip("FP8 weight reloading does not support online quantization")
294
@@ -307,6 +316,7 @@ def test_fp8_reloading(
316
params_dtype=torch.bfloat16,
317
weight_loader=default_weight_loader,
318
)
319
+ method.use_marlin = use_marlin
320
321
else:
322
layer = FusedMoE(
@@ -325,11 +335,6 @@ def test_fp8_reloading(
335
weight_loader=default_weight_loader,
336
)
337
328
- # Fp8LinearMethod uses use_marlin
329
- # Fp8MoEMethod uses fp8_backend
330
- method.use_marlin = use_marlin
331
- method.fp8_backend = Fp8MoeBackend.MARLIN if use_marlin else None
332
-
338
# capture weights format during loading
339
original_metadata = [
340
(name, param.shape, getattr(param, "weight_loader", default_weight_loader))
vllm/model_executor/layers/fused_moe/__init__.py
-2
@@ -73,7 +73,6 @@ if HAS_TRITON:
73
CutlassExpertsFp8,
74
CutlassExpertsW4A8Fp8,
75
cutlass_moe_fp4,
76
- cutlass_moe_fp8,
76
cutlass_moe_w4a8_fp8,
77
)
78
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
@@ -96,7 +95,6 @@ if HAS_TRITON:
95
"fused_experts",
96
"get_config_file_name",
97
"GroupedTopk",
99
- "cutlass_moe_fp8",
98
"cutlass_moe_fp4",
99
"cutlass_moe_w4a8_fp8",
100
"CutlassExpertsFp8",
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+17
-145
@@ -249,20 +249,28 @@ def run_cutlass_moe_fp8(
249
class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):
250
def __init__(
251
self,
252
+ e: int,
253
+ n: int,
254
+ k: int,
255
out_dtype: torch.dtype | None,
253
- ab_strides1: torch.Tensor,
254
- ab_strides2: torch.Tensor,
255
- c_strides1: torch.Tensor,
256
- c_strides2: torch.Tensor,
256
quant_config: FusedMoEQuantConfig,
257
+ device: torch.dtype,
258
):
259
assert quant_config.use_fp8_w8a8
260
super().__init__(quant_config)
261
+
262
+ # E: num_experts
263
+ # N: intermediate size per partition
264
+ # K: hidden dim
265
+ ab_strides1_c_strides2 = torch.full((e,), k, device=device, dtype=torch.int64)
266
+ ab_strides2 = torch.full((e,), n, device=device, dtype=torch.int64)
267
+ c_strides1 = torch.full((e,), 2 * n, device=device, dtype=torch.int64)
268
+
269
self.out_dtype = out_dtype
262
- self.ab_strides1 = ab_strides1
270
+ self.ab_strides1 = ab_strides1_c_strides2
271
self.ab_strides2 = ab_strides2
272
self.c_strides1 = c_strides1
265
- self.c_strides2 = c_strides2
273
+ self.c_strides2 = ab_strides1_c_strides2
274
275
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
276
# Let PrepareAndFinalize::finalize() decide the impl.
@@ -329,24 +337,6 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):
337
338
339
class CutlassExpertsFp8(CutlassExpertsFp8Base):
332
- def __init__(
333
- self,
334
- out_dtype: torch.dtype | None,
335
- ab_strides1: torch.Tensor,
336
- ab_strides2: torch.Tensor,
337
- c_strides1: torch.Tensor,
338
- c_strides2: torch.Tensor,
339
- quant_config: FusedMoEQuantConfig,
340
- ):
341
- super().__init__(
342
- out_dtype,
343
- ab_strides1,
344
- ab_strides2,
345
- c_strides1,
346
- c_strides2,
347
- quant_config,
348
- )
349
-
340
@property
341
def activation_formats(
342
self,
@@ -390,21 +380,10 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
380
self,
381
max_experts_per_worker: int,
382
num_dispatchers: int,
393
- out_dtype: torch.dtype | None,
394
- ab_strides1: torch.Tensor,
395
- ab_strides2: torch.Tensor,
396
- c_strides1: torch.Tensor,
397
- c_strides2: torch.Tensor,
398
- quant_config: FusedMoEQuantConfig,
383
+ *args,
384
+ **kwargs,
385
):
400
- super().__init__(
401
- out_dtype,
402
- ab_strides1,
403
- ab_strides2,
404
- c_strides1,
405
- c_strides2,
406
- quant_config,
407
- )
386
+ super().__init__(*args, **kwargs)
387
assert max_experts_per_worker > 0
388
self.max_experts_per_worker = max_experts_per_worker
389
self.num_dispatchers = num_dispatchers
@@ -445,113 +424,6 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
424
return (workspace1, workspace2, output)
425
426
448
-def cutlass_moe_fp8(
449
- a: torch.Tensor,
450
- w1_q: torch.Tensor,
451
- w2_q: torch.Tensor,
452
- topk_weights: torch.Tensor,
453
- topk_ids: torch.Tensor,
454
- ab_strides1: torch.Tensor,
455
- ab_strides2: torch.Tensor,
456
- c_strides1: torch.Tensor,
457
- c_strides2: torch.Tensor,
458
- quant_config: FusedMoEQuantConfig,
459
- activation: str = "silu",
460
- expert_map: torch.Tensor | None = None,
461
- apply_router_weight_on_input: bool = False,
462
- global_num_experts: int = -1,
463
-) -> torch.Tensor:
464
- """
465
- This function computes a a8w8-quantized Mixture of Experts (MoE) layer
466
- using two sets of quantized weights, w1_q and w2_q, and top-k gating
467
- mechanism. The matrix multiplications are implemented with CUTLASS
468
- grouped gemm.
469
-
470
- Parameters:
471
- - a (torch.Tensor): The input tensor to the MoE layer.
472
- Shape: [M, K]
473
- - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
474
- Shape: [num_experts, K, 2N] (the weights are passed transposed)
475
- - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
476
- Shape: [num_experts, N, K] (the weights are passed transposed)
477
- - topk_weights (torch.Tensor): The weights of each token->expert mapping.
478
- - topk_ids (torch.Tensor): The token->expert mappings.
479
- - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
480
- Shape: [num_experts] or [num_experts, 2N]
481
- - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
482
- Shape: [num_experts] or [num_experts, K]
483
- - ab_strides1 (torch.Tensor): The input/weight strides for the first gemm.
484
- Shape: [num_experts]
485
- - ab_strides2 (torch.Tensor): The input/weight strides for the second gemm.
486
- Shape: [num_experts]
487
- - c_strides1 (torch.Tensor): The output strides for the first gemm.
488
- Shape: [num_experts]
489
- - c_strides2 (torch.Tensor): The output strides for the second gemm.
490
- Shape: [num_experts]
491
- - per_act_token (Optional[bool]): Whether the scale is per-token or
492
- per-tensor.
493
- - activation (str): The activation function to use.
494
- - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
495
- Shape: scalar or [M]
496
- - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
497
- quantize the intermediate result between the gemms.
498
- Shape: scalar or [M]
499
- - expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
500
- every Rank is responsible for a subset of experts. expert_map is a
501
- mapping from global expert-id to local expert-id. When expert_map[i]
502
- is -1, it means that this Rank is not responsible for global
503
- expert-id i.
504
- - apply_router_weight_on_input (bool): When true, the topk weights are
505
- applied directly on the inputs. This is only applicable when topk is 1.
506
- - global_num_experts (int): The total number of experts.
507
-
508
- Returns:
509
- - torch.Tensor: The fp16 output tensor after applying the MoE layer.
510
- """
511
- assert quant_config is not None
512
-
513
- if quant_config.a1_scale is not None:
514
- assert quant_config.per_act_token_quant == (quant_config.a1_scale.numel() != 1)
515
- if quant_config.a2_scale is not None:
516
- assert quant_config.per_act_token_quant == (quant_config.a2_scale.numel() != 1)
517
-
518
- if quant_config.w1_scale is not None:
519
- if quant_config.per_out_ch_quant:
520
- assert quant_config.w1_scale.dim() > 1 and quant_config.w1_scale.size(
521
- 1
522
- ) == w1_q.size(1)
523
- else:
524
- assert (
525
- quant_config.w1_scale.dim() == 1 or quant_config.w1_scale.size(1) == 1
526
- )
527
-
528
- num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(0)
529
-
530
- fn = mk.FusedMoEModularKernel(
531
- MoEPrepareAndFinalizeNoEP(),
532
- CutlassExpertsFp8(
533
- out_dtype=a.dtype,
534
- ab_strides1=ab_strides1,
535
- ab_strides2=ab_strides2,
536
- c_strides1=c_strides1,
537
- c_strides2=c_strides2,
538
- quant_config=quant_config,
539
- ),
540
- )
541
-
542
- return fn(
543
- a,
544
- w1_q,
545
- w2_q,
546
- topk_weights,
547
- topk_ids,
548
- activation=activation,
549
- global_num_experts=num_experts,
550
- expert_map=expert_map,
551
- apply_router_weight_on_input=apply_router_weight_on_input,
552
- )
553
-
554
-
427
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
428
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
429
vllm/model_executor/layers/fused_moe/fallback.py
new
+126
@@ -0,0 +1,126 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+
6
+import torch
7
+
8
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+
10
+
11
+class FallbackExperts(mk.FusedMoEPermuteExpertsUnpermute, ABC):
12
+ """Base class for runtime dispatching of expert implementations."""
13
+
14
+ def __init__(
15
+ self,
16
+ experts: mk.FusedMoEPermuteExpertsUnpermute,
17
+ fallback_experts: mk.FusedMoEPermuteExpertsUnpermute,
18
+ ):
19
+ super().__init__(experts.quant_config)
20
+ self.fallback_experts = fallback_experts
21
+ self.experts = experts
22
+
23
+ @property
24
+ def activation_formats(
25
+ self,
26
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
27
+ assert (
28
+ self.fallback_experts.activation_formats == self.experts.activation_formats
29
+ )
30
+ return self.fallback_experts.activation_formats
31
+
32
+ def supports_chunking(self) -> bool:
33
+ assert (
34
+ self.experts.supports_chunking()
35
+ == self.fallback_experts.supports_chunking()
36
+ )
37
+ return (
38
+ self.experts.supports_chunking()
39
+ and self.fallback_experts.supports_chunking()
40
+ )
41
+
42
+ def supports_expert_map(self) -> bool:
43
+ assert (
44
+ self.experts.supports_expert_map()
45
+ == self.fallback_experts.supports_expert_map()
46
+ )
47
+ return (
48
+ self.experts.supports_expert_map()
49
+ and self.fallback_experts.supports_expert_map()
50
+ )
51
+
52
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
53
+ e_war = self.experts.finalize_weight_and_reduce_impl()
54
+ fbe_war = self.fallback_experts.finalize_weight_and_reduce_impl()
55
+ is_dge_war = e_war is not None
56
+ is_fbe_war = fbe_war is not None
57
+
58
+ if is_dge_war and is_fbe_war:
59
+ assert e_war == fbe_war, (
60
+ "Both implementations should agree on WeightAndReduce impls. "
61
+ f"Got e_war: {e_war}, and fbe_war: {fbe_war}"
62
+ )
63
+
64
+ if e_war is not None:
65
+ return e_war
66
+ assert fbe_war is not None
67
+ return fbe_war
68
+
69
+ @abstractmethod
70
+ def workspace_shapes(
71
+ self,
72
+ M: int,
73
+ N: int,
74
+ K: int,
75
+ topk: int,
76
+ global_num_experts: int,
77
+ local_num_experts: int,
78
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
79
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
80
+ raise NotImplementedError
81
+
82
+ @abstractmethod
83
+ def _select_experts_impl(
84
+ self,
85
+ hidden_states: torch.Tensor,
86
+ w1: torch.Tensor,
87
+ w2: torch.Tensor,
88
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
89
+ raise NotImplementedError
90
+
91
+ def apply(
92
+ self,
93
+ output: torch.Tensor,
94
+ hidden_states: torch.Tensor,
95
+ w1: torch.Tensor,
96
+ w2: torch.Tensor,
97
+ topk_weights: torch.Tensor,
98
+ topk_ids: torch.Tensor,
99
+ activation: str,
100
+ global_num_experts: int,
101
+ expert_map: torch.Tensor | None,
102
+ a1q_scale: torch.Tensor | None,
103
+ a2_scale: torch.Tensor | None,
104
+ workspace13: torch.Tensor,
105
+ workspace2: torch.Tensor,
106
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
107
+ apply_router_weight_on_input: bool,
108
+ ):
109
+ experts = self._select_experts_impl(hidden_states, w1, w2)
110
+ experts.apply(
111
+ output,
112
+ hidden_states,
113
+ w1,
114
+ w2,
115
+ topk_weights,
116
+ topk_ids,
117
+ activation,
118
+ global_num_experts,
119
+ expert_map,
120
+ a1q_scale,
121
+ a2_scale,
122
+ workspace13,
123
+ workspace2,
124
+ expert_tokens_meta,
125
+ apply_router_weight_on_input,
126
+ )
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py
+5
-5
@@ -100,7 +100,7 @@ direct_register_custom_op(
100
)
101
102
103
-def flashinfer_fused_moe_per_tensor_scale_fp8(
103
+def fi_trtllm_fp8_per_tensor_moe(
104
routing_logits: torch.Tensor,
105
routing_bias: torch.Tensor | None,
106
hidden_states: torch.Tensor,
@@ -158,7 +158,7 @@ def flashinfer_fused_moe_per_tensor_scale_fp8(
158
)
159
160
161
-def flashinfer_fused_moe_per_tensor_scale_fp8_fake(
161
+def fi_trtllm_fp8_per_tensor_moe_fake(
162
routing_logits: torch.Tensor,
163
routing_bias: torch.Tensor | None,
164
hidden_states: torch.Tensor,
@@ -184,9 +184,9 @@ def flashinfer_fused_moe_per_tensor_scale_fp8_fake(
184
185
# TODO(bnell): Does this really need to be a torch.op?
186
direct_register_custom_op(
187
- op_name="flashinfer_fused_moe_per_tensor_scale_fp8",
188
- op_func=flashinfer_fused_moe_per_tensor_scale_fp8,
187
+ op_name="fi_trtllm_fp8_per_tensor_moe",
188
+ op_func=fi_trtllm_fp8_per_tensor_moe,
189
mutates_args=["hidden_states"],
190
- fake_impl=flashinfer_fused_moe_per_tensor_scale_fp8_fake,
190
+ fake_impl=fi_trtllm_fp8_per_tensor_moe_fake,
191
tags=(torch.Tag.needs_fixed_stride_order,),
192
)
vllm/model_executor/layers/fused_moe/oracle/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/layers/fused_moe/oracle/fp8.py
new
+358
@@ -0,0 +1,358 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from enum import Enum
4
+
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm import envs
9
+from vllm._aiter_ops import rocm_aiter_ops
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.fused_moe.config import (
12
+ FusedMoEConfig,
13
+ FusedMoEQuantConfig,
14
+ fp8_w8a8_moe_quant_config,
15
+ fp8_w8a16_moe_quant_config,
16
+)
17
+from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
18
+ FlashinferMoeBackend,
19
+ get_flashinfer_moe_backend,
20
+ make_fp8_moe_alpha_scales_for_fi,
21
+ prepare_fp8_moe_layer_for_fi,
22
+)
23
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
24
+ prepare_fp8_moe_layer_for_deepgemm,
25
+)
26
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
27
+ prepare_fp8_moe_layer_for_marlin,
28
+)
29
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
30
+ cutlass_group_gemm_supported,
31
+)
32
+from vllm.platforms import current_platform
33
+from vllm.utils.deep_gemm import is_deep_gemm_supported
34
+from vllm.utils.flashinfer import has_flashinfer_moe
35
+from vllm.utils.import_utils import has_deep_gemm
36
+
37
+logger = init_logger(__name__)
38
+
39
+
40
+class Fp8MoeBackend(Enum):
41
+ NONE = 0
42
+ FLASHINFER_TRTLLM = 1
43
+ FLASHINFER_CUTLASS = 2
44
+ DEEPGEMM = 3
45
+ MARLIN = 4
46
+ TRITON = 5
47
+ AITER = 6
48
+ VLLM_CUTLASS = 7
49
+
50
+
51
+def select_fp8_moe_backend(
52
+ block_quant: bool,
53
+ tp_size: int,
54
+ with_lora_support: bool,
55
+ is_act_and_mul: bool = True,
56
+ allow_vllm_cutlass: bool = False,
57
+) -> Fp8MoeBackend:
58
+ """
59
+ Select the primary FP8 MoE backend
60
+ Note: Shape-specific fallbacks may still occur at runtime.
61
+ """
62
+ # TODO(rob): in a future PR, we will query each mk for
63
+ # supported features and return the mk directly, just like
64
+ # we do for the Attention Backend.
65
+
66
+ if with_lora_support:
67
+ return Fp8MoeBackend.TRITON
68
+
69
+ def _make_log_backend(backend_name: str):
70
+ return f"Using {backend_name} backend for FP8 MoE"
71
+
72
+ # Prefer FlashInfer backends on supported GPUs; allow SM90 and SM100.
73
+ if (
74
+ current_platform.is_cuda()
75
+ and (
76
+ current_platform.is_device_capability_family(100)
77
+ or current_platform.is_device_capability(90)
78
+ )
79
+ and envs.VLLM_USE_FLASHINFER_MOE_FP8
80
+ and has_flashinfer_moe()
81
+ ):
82
+ backend = get_flashinfer_moe_backend()
83
+ if backend == FlashinferMoeBackend.TENSORRT_LLM:
84
+ logger.info_once(_make_log_backend("FlashInfer TRTLLM"))
85
+ if not is_act_and_mul:
86
+ raise ValueError(
87
+ "FlashInfer TRTLLM FP8 MoE backend only supports "
88
+ "act_and_mul gate_up_project fusion. Please set "
89
+ "VLLM_USE_FLASHINFER_MOE_FP8=throughput to use the "
90
+ "FlashInfer CUTLASS backend instead."
91
+ )
92
+ return Fp8MoeBackend.FLASHINFER_TRTLLM
93
+ else:
94
+ if block_quant and current_platform.is_device_capability_family(100):
95
+ raise ValueError(
96
+ "FlashInfer FP8 MoE throughput backend does not "
97
+ "support block quantization on SM100. Please use "
98
+ "VLLM_FLASHINFER_MOE_BACKEND=latency to use the "
99
+ "FlashInfer TRTLLM backend instead."
100
+ )
101
+ logger.info_once(_make_log_backend("FlashInfer CUTLASS"))
102
+ return Fp8MoeBackend.FLASHINFER_CUTLASS
103
+
104
+ # weight-only path for older GPUs without native FP8
105
+ if (
106
+ current_platform.is_cuda() and not current_platform.has_device_capability(89)
107
+ ) or envs.VLLM_TEST_FORCE_FP8_MARLIN:
108
+ logger.info_once(_make_log_backend("Marlin"), scope="local")
109
+ return Fp8MoeBackend.MARLIN
110
+
111
+ # Determine if we should use DeepGEMM with block-quantized weights:
112
+ # - If explicitly set by user, respect their choice
113
+ # - If not explicitly set (default), disable when TP size is >= 8
114
+ moe_use_deep_gemm = envs.VLLM_MOE_USE_DEEP_GEMM
115
+ if not envs.is_set("VLLM_MOE_USE_DEEP_GEMM") and tp_size >= 8:
116
+ moe_use_deep_gemm = False
117
+ logger.info_once(
118
+ "DeepGEMM MoE is disabled by default when TP size is >= 8. "
119
+ "Set VLLM_MOE_USE_DEEP_GEMM=1 to enable it.",
120
+ scope="local",
121
+ )
122
+
123
+ use_deep_gemm = envs.VLLM_USE_DEEP_GEMM
124
+ if not is_deep_gemm_supported():
125
+ use_deep_gemm = False
126
+ logger.info_once(
127
+ "DeepGEMM is disabled because the platform does not support it.",
128
+ scope="local",
129
+ )
130
+
131
+ if use_deep_gemm and moe_use_deep_gemm and block_quant:
132
+ if not has_deep_gemm():
133
+ logger.warning_once(
134
+ "DeepGEMM backend requested but not available.", scope="local"
135
+ )
136
+ elif is_deep_gemm_supported():
137
+ logger.info_once(_make_log_backend("DeepGEMM"), scope="local")
138
+ return Fp8MoeBackend.DEEPGEMM
139
+
140
+ if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MOE:
141
+ logger.info_once(_make_log_backend("ROCm AITER"), scope="local")
142
+ return Fp8MoeBackend.AITER
143
+
144
+ if allow_vllm_cutlass and not block_quant and cutlass_group_gemm_supported():
145
+ logger.info_once(_make_log_backend("vLLM CUTLASS"), scope="local")
146
+ return Fp8MoeBackend.VLLM_CUTLASS
147
+
148
+ # default to Triton
149
+ logger.info_once(_make_log_backend("Triton"), scope="local")
150
+ return Fp8MoeBackend.TRITON
151
+
152
+
153
+def convert_to_fp8_moe_kernel_format(
154
+ fp8_backend: Fp8MoeBackend,
155
+ layer: torch.nn.Module,
156
+ w13: torch.Tensor,
157
+ w2: torch.Tensor,
158
+ w13_scale: torch.Tensor,
159
+ w2_scale: torch.Tensor,
160
+ w13_input_scale: torch.Tensor | None,
161
+ w2_input_scale: torch.Tensor | None,
162
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
163
+ block_quant = hasattr(layer, "weight_block_size")
164
+ if fp8_backend == Fp8MoeBackend.DEEPGEMM:
165
+ assert block_quant
166
+ w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_deepgemm(
167
+ w13,
168
+ w2,
169
+ w13_scale,
170
+ w2_scale,
171
+ tuple(layer.weight_block_size),
172
+ )
173
+ elif fp8_backend == Fp8MoeBackend.AITER:
174
+ w13, w2 = rocm_aiter_ops.shuffle_weights(w13, w2)
175
+ elif fp8_backend == Fp8MoeBackend.MARLIN:
176
+ w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_marlin(
177
+ layer,
178
+ w13,
179
+ w2,
180
+ w13_scale,
181
+ w2_scale,
182
+ )
183
+ elif fp8_backend in [
184
+ Fp8MoeBackend.FLASHINFER_CUTLASS,
185
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
186
+ ]:
187
+ w13, w2, w13_scale = prepare_fp8_moe_layer_for_fi(
188
+ layer=layer,
189
+ w13=w13,
190
+ w2=w2,
191
+ w13_scale=w13_scale,
192
+ w13_input_scale=w13_input_scale,
193
+ w2_scale=w2_scale,
194
+ w2_input_scale=w2_input_scale,
195
+ is_trtllm=(fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM),
196
+ )
197
+
198
+ return w13, w2, w13_scale, w2_scale
199
+
200
+
201
+def make_fp8_moe_quant_config(
202
+ fp8_backend: Fp8MoeBackend,
203
+ w1_scale: torch.Tensor,
204
+ w2_scale: torch.Tensor,
205
+ a1_scale: torch.Tensor | None,
206
+ a2_scale: torch.Tensor | None,
207
+ block_shape: list[int] | None = None,
208
+) -> FusedMoEQuantConfig | None:
209
+ """
210
+ Create FusedMoEQuantConfig for the specifed FP8 Backend.
211
+ The FusedMoEQuantConfig holds the scales that are used
212
+ at runtime by the Modular Kernel abstraction.
213
+
214
+ Note that certain kernels (e.g. Flashinfer CUTLASS) need
215
+ special Quant configs to handle non-standard inputs to
216
+ their kernel interfaces.
217
+
218
+ In a future PR, we will have this function should be
219
+ a method of the modular kernel itself.
220
+ """
221
+ # TRTLLM does not use Modular Kernel abstraction yet.
222
+ if fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
223
+ return None
224
+
225
+ # MARLIN is mixed precision W8A16 config.
226
+ if fp8_backend == Fp8MoeBackend.MARLIN:
227
+ return fp8_w8a16_moe_quant_config(
228
+ w1_scale=w1_scale,
229
+ w2_scale=w2_scale,
230
+ block_shape=block_shape,
231
+ )
232
+
233
+ # Flashinfer CUTLASS per-tensor uses single dq scale
234
+ # (alpha = w_scale * a_scale) and inverse a2 scale.
235
+ if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS and block_shape is None:
236
+ assert a1_scale is not None and a2_scale is not None
237
+ g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(
238
+ w1_scale,
239
+ a1_scale,
240
+ w2_scale,
241
+ a2_scale,
242
+ )
243
+ return fp8_w8a8_moe_quant_config(
244
+ w1_scale=w1_scale,
245
+ w2_scale=w2_scale,
246
+ a1_scale=a1_scale,
247
+ a2_scale=a2_scale,
248
+ a1_gscale=(1.0 / a1_scale),
249
+ a2_gscale=(1.0 / a2_scale),
250
+ g1_alphas=g1_alphas,
251
+ g2_alphas=g2_alphas,
252
+ )
253
+ # All other backends use normal config.
254
+ return fp8_w8a8_moe_quant_config(
255
+ w1_scale=w1_scale,
256
+ w2_scale=w2_scale,
257
+ a1_scale=a1_scale,
258
+ a2_scale=a2_scale,
259
+ block_shape=block_shape,
260
+ )
261
+
262
+
263
+def make_fp8_moe_kernel(
264
+ layer: torch.nn.Module,
265
+ moe_quant_config: FusedMoEQuantConfig,
266
+ moe_config: FusedMoEConfig,
267
+ fp8_backend: Fp8MoeBackend,
268
+) -> tuple[mk.FusedMoEModularKernel, bool]:
269
+ # Delayed import is required since the oracle is imported
270
+ # by CPU backends which cannot import all of these experts.
271
+ # TODO: update the experts to make this not happen.
272
+ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
273
+ MoEPrepareAndFinalizeNoEP,
274
+ )
275
+
276
+ # NOTE(rob): this is a WIP refactor. We are first migrating
277
+ # all of the kernels in the TP case to use mk. Once this is
278
+ # done, then we will initialzie the TP case and DP/EP case
279
+ # via the same code path (i.e. via maybe_init_modular_kernel).
280
+ # NOTE(rob): in progress migrating all into this format.
281
+ use_inplace = True
282
+ if fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
283
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
284
+ FlashInferExperts,
285
+ )
286
+
287
+ kernel = mk.FusedMoEModularKernel(
288
+ MoEPrepareAndFinalizeNoEP(
289
+ defer_input_quant=moe_quant_config.is_block_quantized
290
+ ),
291
+ FlashInferExperts(
292
+ out_dtype=layer.orig_dtype,
293
+ quant_config=moe_quant_config,
294
+ ep_rank=moe_config.ep_rank,
295
+ ep_size=moe_config.ep_size,
296
+ tp_rank=moe_config.tp_rank,
297
+ tp_size=moe_config.tp_size,
298
+ use_dp=(moe_config.dp_size > 1),
299
+ use_deepseek_fp8_block_scale=moe_quant_config.is_block_quantized,
300
+ ),
301
+ )
302
+ use_inplace = False
303
+
304
+ elif fp8_backend == Fp8MoeBackend.AITER:
305
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
306
+ AiterExperts,
307
+ )
308
+
309
+ kernel = mk.FusedMoEModularKernel(
310
+ # TODO: make defer_input_quant an attr of the AiterExperts
311
+ MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
312
+ AiterExperts(quant_config=moe_quant_config),
313
+ )
314
+ elif fp8_backend == Fp8MoeBackend.MARLIN:
315
+ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
316
+ MarlinExperts,
317
+ )
318
+
319
+ kernel = mk.FusedMoEModularKernel(
320
+ MoEPrepareAndFinalizeNoEP(),
321
+ MarlinExperts(quant_config=moe_quant_config),
322
+ )
323
+ elif fp8_backend == Fp8MoeBackend.VLLM_CUTLASS:
324
+ from vllm.model_executor.layers.fused_moe.triton_cutlass_moe import (
325
+ TritonOrCutlassExperts,
326
+ )
327
+
328
+ kernel = mk.FusedMoEModularKernel(
329
+ MoEPrepareAndFinalizeNoEP(),
330
+ TritonOrCutlassExperts(
331
+ out_dtype=moe_config.in_dtype,
332
+ e=layer.local_num_experts,
333
+ n=layer.intermediate_size_per_partition,
334
+ k=layer.hidden_size,
335
+ device=layer.w13_weight.device,
336
+ quant_config=moe_quant_config,
337
+ ),
338
+ )
339
+ elif fp8_backend == Fp8MoeBackend.DEEPGEMM:
340
+ from vllm.model_executor.layers.fused_moe import (
341
+ TritonOrDeepGemmExperts,
342
+ )
343
+
344
+ kernel = mk.FusedMoEModularKernel(
345
+ MoEPrepareAndFinalizeNoEP(),
346
+ TritonOrDeepGemmExperts(quant_config=moe_quant_config),
347
+ )
348
+ else:
349
+ from vllm.model_executor.layers.fused_moe.fused_moe import (
350
+ TritonExperts,
351
+ )
352
+
353
+ assert fp8_backend == Fp8MoeBackend.TRITON
354
+ kernel = mk.FusedMoEModularKernel(
355
+ MoEPrepareAndFinalizeNoEP(),
356
+ TritonExperts(quant_config=moe_quant_config),
357
+ )
358
+ return kernel, use_inplace
vllm/model_executor/layers/fused_moe/triton_cutlass_moe.py
new
+75
@@ -0,0 +1,75 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9
+from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
10
+from vllm.model_executor.layers.fused_moe.fallback import FallbackExperts
11
+from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
12
+from vllm.platforms import current_platform
13
+
14
+
15
+class TritonOrCutlassExperts(FallbackExperts):
16
+ """Cutlass with fallback to Triton for low latency shapes on SM100."""
17
+
18
+ def __init__(
19
+ self,
20
+ e: int,
21
+ n: int,
22
+ k: int,
23
+ out_dtype: torch.dtype | None,
24
+ quant_config: FusedMoEQuantConfig,
25
+ device: torch.dtype,
26
+ ):
27
+ self.is_sm100 = current_platform.has_device_capability(100)
28
+ super().__init__(
29
+ experts=CutlassExpertsFp8(e, n, k, out_dtype, quant_config, device),
30
+ fallback_experts=TritonExperts(quant_config),
31
+ )
32
+
33
+ def workspace_shapes(
34
+ self,
35
+ M: int,
36
+ N: int,
37
+ K: int,
38
+ topk: int,
39
+ global_num_experts: int,
40
+ local_num_experts: int,
41
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
42
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
43
+ # Small batch fallback for sm100.
44
+ if self.is_sm100 and M <= 8:
45
+ return self.fallback_experts.workspace_shapes(
46
+ M,
47
+ N,
48
+ K,
49
+ topk,
50
+ global_num_experts,
51
+ local_num_experts,
52
+ expert_tokens_meta,
53
+ )
54
+ else:
55
+ return self.experts.workspace_shapes(
56
+ M,
57
+ N,
58
+ K,
59
+ topk,
60
+ global_num_experts,
61
+ local_num_experts,
62
+ expert_tokens_meta,
63
+ )
64
+
65
+ def _select_experts_impl(
66
+ self,
67
+ hidden_states: torch.Tensor,
68
+ w1: torch.Tensor,
69
+ w2: torch.Tensor,
70
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
71
+ # Small batch fallback for sm100.
72
+ if self.is_sm100 and hidden_states.shape[0] <= 8:
73
+ return self.fallback_experts
74
+ else:
75
+ return self.experts
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+16
-107
@@ -10,78 +10,22 @@ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
10
_valid_deep_gemm,
11
_valid_deep_gemm_shape,
12
)
13
+from vllm.model_executor.layers.fused_moe.fallback import FallbackExperts
14
from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
15
from vllm.utils.deep_gemm import (
15
- get_mk_alignment_for_contiguous_layout,
16
is_deep_gemm_e8m0_used,
17
)
18
19
20
-class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
21
- def __init__(
22
- self,
23
- quant_config: FusedMoEQuantConfig,
24
- allow_deep_gemm: bool = False,
25
- ):
26
- super().__init__(quant_config)
27
-
28
- self.triton_expert = TritonExperts(quant_config)
29
-
30
- self.allow_deep_gemm = (
31
- allow_deep_gemm
32
- and self.quant_config.use_fp8_w8a8
33
- and self.block_shape == get_mk_alignment_for_contiguous_layout()
34
- )
35
-
36
- self.deep_gemm_expert = (
37
- DeepGemmExperts(self.quant_config) if self.allow_deep_gemm else None
38
- )
20
+class TritonOrDeepGemmExperts(FallbackExperts):
21
+ """DeepGemm with fallback to Triton for low latency shapes."""
22
40
- @property
41
- def activation_formats(
42
- self,
43
- ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
44
- assert (
45
- self.deep_gemm_expert is None
46
- or self.triton_expert.activation_formats
47
- == self.deep_gemm_expert.activation_formats
48
- )
49
- return self.triton_expert.activation_formats
50
-
51
- def supports_chunking(self) -> bool:
52
- dge = self.deep_gemm_expert
53
- te = self.triton_expert
54
- return (dge is None or dge.supports_chunking()) and (
55
- te is None or te.supports_chunking()
23
+ def __init__(self, quant_config: FusedMoEQuantConfig):
24
+ super().__init__(
25
+ experts=DeepGemmExperts(quant_config),
26
+ fallback_experts=TritonExperts(quant_config),
27
)
28
58
- def supports_expert_map(self) -> bool:
59
- dge = self.deep_gemm_expert
60
- te = self.triton_expert
61
- return (dge is None or dge.supports_expert_map()) and (
62
- te is None or te.supports_expert_map()
63
- )
64
-
65
- def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
66
- dge = self.deep_gemm_expert
67
- te = self.triton_expert
68
- dge_war = dge.finalize_weight_and_reduce_impl() if dge else None
69
- te_war = te.finalize_weight_and_reduce_impl() if te else None
70
- is_dge_war = dge_war is not None
71
- is_te_war = te_war is not None
72
-
73
- if is_dge_war and is_te_war:
74
- assert dge_war == te_war, (
75
- "Both implementations should agree on WeightAndReduce impls. "
76
- f"Got dge_war: {dge_war}, and te_war: {te_war}"
77
- )
78
-
79
- if dge_war is not None:
80
- return dge_war
81
-
82
- assert te_war is not None
83
- return te_war
84
-
29
def workspace_shapes(
30
self,
31
M: int,
@@ -95,11 +39,8 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
39
# Note: the deep gemm workspaces are strictly larger than the triton
40
# workspaces so we can be pessimistic here and allocate for DeepGemm
41
# even if we fall back to triton later, e.g. if expert maps are set.
98
- if self.allow_deep_gemm and (
99
- is_deep_gemm_e8m0_used() or _valid_deep_gemm_shape(M, N, K)
100
- ):
101
- assert self.deep_gemm_expert is not None
102
- return self.deep_gemm_expert.workspace_shapes(
42
+ if is_deep_gemm_e8m0_used() or _valid_deep_gemm_shape(M, N, K):
43
+ return self.experts.workspace_shapes(
44
M,
45
N,
46
K,
@@ -109,7 +50,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
50
expert_tokens_meta,
51
)
52
else:
112
- return self.triton_expert.workspace_shapes(
53
+ return self.fallback_experts.workspace_shapes(
54
M,
55
N,
56
K,
@@ -119,45 +60,13 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
60
expert_tokens_meta,
61
)
62
122
- def apply(
63
+ def _select_experts_impl(
64
self,
124
- output: torch.Tensor,
65
hidden_states: torch.Tensor,
66
w1: torch.Tensor,
67
w2: torch.Tensor,
128
- topk_weights: torch.Tensor,
129
- topk_ids: torch.Tensor,
130
- activation: str,
131
- global_num_experts: int,
132
- expert_map: torch.Tensor | None,
133
- a1q_scale: torch.Tensor | None,
134
- a2_scale: torch.Tensor | None,
135
- workspace13: torch.Tensor,
136
- workspace2: torch.Tensor,
137
- expert_tokens_meta: mk.ExpertTokensMetadata | None,
138
- apply_router_weight_on_input: bool,
139
- ):
140
- use_deep_gemm = self.allow_deep_gemm and (
141
- is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2)
142
- )
143
-
144
- experts = self.deep_gemm_expert if use_deep_gemm else self.triton_expert
145
- assert experts is not None
146
-
147
- experts.apply(
148
- output,
149
- hidden_states,
150
- w1,
151
- w2,
152
- topk_weights,
153
- topk_ids,
154
- activation,
155
- global_num_experts,
156
- expert_map,
157
- a1q_scale,
158
- a2_scale,
159
- workspace13,
160
- workspace2,
161
- expert_tokens_meta,
162
- apply_router_weight_on_input,
163
- )
68
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
69
+ if is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2):
70
+ return self.experts
71
+ else:
72
+ return self.fallback_experts
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+135
-355
@@ -13,10 +13,8 @@ from compressed_tensors.quantization import (
13
)
14
from torch.nn.parameter import Parameter
15
16
-import vllm.envs as envs
16
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
17
from vllm import _custom_ops as ops
19
-from vllm._aiter_ops import rocm_aiter_ops
18
from vllm.distributed import get_tensor_model_parallel_world_size
19
from vllm.logger import init_logger
20
from vllm.model_executor.layers.fused_moe import (
@@ -31,6 +29,7 @@ from vllm.model_executor.layers.fused_moe import (
29
from vllm.model_executor.layers.fused_moe.config import (
30
FusedMoEQuantConfig,
31
fp8_w8a8_moe_quant_config,
32
+ fp8_w8a16_moe_quant_config,
33
int4_w4a16_moe_quant_config,
34
int4_w4afp8_moe_quant_config,
35
int8_w8a8_moe_quant_config,
@@ -46,11 +45,16 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
45
MarlinExperts,
46
fused_marlin_moe,
47
)
48
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
49
+ Fp8MoeBackend,
50
+ convert_to_fp8_moe_kernel_format,
51
+ make_fp8_moe_kernel,
52
+ select_fp8_moe_backend,
53
+)
54
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
55
WNA16_SUPPORTED_BITS,
56
WNA16_SUPPORTED_TYPES_MAP,
57
)
53
-from vllm.model_executor.layers.quantization.utils import replace_parameter
58
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
59
build_flashinfer_fp4_cutlass_moe_prepare_finalize,
60
flashinfer_trtllm_fp4_moe,
@@ -63,8 +67,8 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
67
get_flashinfer_moe_backend,
68
)
69
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
66
- expert_weight_is_col_major,
67
- requant_weight_ue8m0_inplace,
70
+ process_fp8_input_tensor_strategy_moe,
71
+ process_fp8_weight_tensor_strategy_moe,
72
)
73
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
74
check_moe_marlin_supports_layer,
@@ -76,29 +80,17 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
80
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
81
prepare_moe_fp4_layer_for_marlin,
82
)
79
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
80
- prepare_moe_fp8_layer_for_marlin,
81
-)
83
from vllm.model_executor.layers.quantization.utils.quant_utils import (
84
convert_bf16_scales_to_fp8,
85
convert_packed_uint4b8_to_signed_int4_inplace,
86
swizzle_blockscale,
87
)
88
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
88
- all_close_1d,
89
normalize_e4m3fn_to_e4m3fnuz,
90
- per_tensor_dequantize,
90
)
92
-from vllm.model_executor.utils import set_weight_attrs
91
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
92
from vllm.platforms import CpuArchEnum, current_platform
93
from vllm.scalar_type import scalar_types
95
-from vllm.utils.deep_gemm import (
96
- get_col_major_tma_aligned_tensor,
97
- get_mk_alignment_for_contiguous_layout,
98
- is_deep_gemm_e8m0_used,
99
- is_deep_gemm_supported,
100
-)
101
-from vllm.utils.import_utils import has_deep_gemm
94
95
logger = init_logger(__name__)
96
@@ -657,10 +649,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
649
moe: FusedMoEConfig,
650
layer_name: str | None = None,
651
):
660
- from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
661
- CompressedTensorsConfig,
662
- )
663
-
652
super().__init__(moe)
653
self.weight_quant = weight_quant
654
self.input_quant = input_quant
@@ -687,42 +675,31 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
675
"For FP8 Fused MoE layer, we require either per tensor or "
676
"channelwise, dynamic per token quantization."
677
)
690
-
691
- # For GPUs that lack FP8 hardware support, we can leverage the Marlin
692
- # kernel for fast weight-only FP8 quantization
693
- self.use_marlin = (
694
- not current_platform.has_device_capability(89)
695
- or envs.VLLM_TEST_FORCE_FP8_MARLIN
696
- and not self.block_quant
697
- )
698
- # Disable marlin for rocm
699
- if current_platform.is_rocm():
700
- self.use_marlin = False
701
-
702
- self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
703
-
704
- # cutlass path
705
- self.is_fp8_w8a8_sm100 = CompressedTensorsConfig._is_fp8_w8a8_sm100(
706
- self.weight_quant, self.input_quant
678
+ self.fp8_backend = select_fp8_moe_backend(
679
+ block_quant=self.block_quant,
680
+ tp_size=moe.tp_size,
681
+ with_lora_support=moe.is_lora_enabled,
682
+ # TODO(rob): enable selecting this externally.
683
+ allow_vllm_cutlass=True,
684
)
708
- self.use_cutlass = not self.block_quant and (
709
- CompressedTensorsConfig._is_fp8_w8a8_sm90(
710
- self.weight_quant, self.input_quant
685
+ if self.fp8_backend != Fp8MoeBackend.MARLIN:
686
+ per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
687
+ per_channel_quant = (
688
+ self.weight_quant.strategy == QuantizationStrategy.CHANNEL
689
+ )
690
+ if per_act_token != per_channel_quant:
691
+ raise NotImplementedError(
692
+ "For FP8 Fused MoE layers, per-token and per-channel must be "
693
+ "used together."
694
+ )
695
+ # TODO(rob): hook this up in a follow up PR.
696
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
697
+ raise NotImplementedError(
698
+ "FlashInfer TRTLLM backend not supported for compressed-tensors yet."
699
)
712
- or self.is_fp8_w8a8_sm100
713
- )
700
self.disable_expert_map = False
715
- self.layer_name = layer_name
716
- self.marlin_input_dtype = (
717
- get_marlin_input_dtype(layer_name) if self.use_marlin else None
718
- )
701
720
- self.allow_deep_gemm = (
721
- self.block_quant
722
- and envs.VLLM_MOE_USE_DEEP_GEMM
723
- and is_deep_gemm_supported()
724
- and list(self.weight_block_size) == get_mk_alignment_for_contiguous_layout()
725
- )
702
+ self.kernel: mk.FusedMoEModularKernel | None = None
703
704
def create_weights(
705
self,
@@ -880,163 +857,75 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
857
layer.w2_input_scale = None
858
859
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
883
- # Fp8 moe kernels require a single activation scale.
884
- # We take the max of all the scales in case they differ.
885
- if self.static_input_scales:
886
- assert self.input_quant.strategy == QuantizationStrategy.TENSOR
887
- if layer.w13_input_scale is None or layer.w2_input_scale is None:
888
- raise ValueError(
889
- "QuantConfig has static quantization, but found "
890
- "activation scales are None."
891
- )
892
- if not all_close_1d(layer.w13_input_scale) or not all_close_1d(
893
- layer.w2_input_scale
894
- ):
895
- logger.warning_once(
896
- "Found input_scales that are not equal for "
897
- "fp8 MoE layer. Using the maximum across experts "
898
- "for each layer."
899
- )
900
- layer.w13_input_scale = torch.nn.Parameter(
901
- layer.w13_input_scale.max(), requires_grad=False
902
- )
903
- layer.w2_input_scale = torch.nn.Parameter(
904
- layer.w2_input_scale.max(), requires_grad=False
905
- )
906
-
860
+ # Allow for accessing weights and scales in standard way.
861
+ w13 = layer.w13_weight
862
+ w2 = layer.w2_weight
863
+ w13_scale = layer.w13_weight_scale
864
+ w2_scale = layer.w2_weight_scale
865
+ w13_input_scale = layer.w13_input_scale
866
+ w2_input_scale = layer.w2_input_scale
867
+
868
+ # MI300x and MI325x use FNUZ format for FP8. Convert if needed.
869
if current_platform.is_fp8_fnuz():
908
- # Normalize the weights and scales
909
- w13_weight, w13_weight_scale, w13_input_scale = (
910
- normalize_e4m3fn_to_e4m3fnuz(
911
- layer.w13_weight, layer.w13_weight_scale, layer.w13_input_scale
912
- )
870
+ w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
871
+ w13, w13_scale, w13_input_scale
872
)
914
- w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
915
- layer.w2_weight, layer.w2_weight_scale, layer.w2_input_scale
916
- )
917
- # Reset the parameter
918
- layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
919
- layer.w13_weight_scale = torch.nn.Parameter(
920
- w13_weight_scale, requires_grad=False
873
+ w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
874
+ w2, w2_scale, w2_input_scale
875
)
922
- if w13_input_scale is not None:
923
- layer.w13_input_scale = torch.nn.Parameter(
924
- w13_input_scale, requires_grad=False
925
- )
926
- layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
927
- layer.w2_weight_scale = torch.nn.Parameter(
928
- w2_weight_scale, requires_grad=False
929
- )
930
- if w2_input_scale is not None:
931
- layer.w2_input_scale = torch.nn.Parameter(
932
- w2_input_scale, requires_grad=False
933
- )
876
935
- # For Per-TENSOR case, Fp8 moe kernel needs single weight scale
936
- # for w13 per expert. Use max then dequant and requant each expert.
937
- if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
938
- assert layer.w13_weight_scale is not None
939
- shard_size = layer.intermediate_size_per_partition
940
- max_w13_scales = layer.w13_weight_scale.max(dim=1).values
941
- for expert_id in range(layer.local_num_experts):
942
- start = 0
943
- for shard_id in range(2):
944
- dq_weight = per_tensor_dequantize(
945
- layer.w13_weight[expert_id][start : start + shard_size, :],
946
- layer.w13_weight_scale[expert_id][shard_id],
947
- )
948
- layer.w13_weight[expert_id][start : start + shard_size, :], _ = (
949
- ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
950
- )
951
- start += shard_size
952
- layer.w13_weight_scale = torch.nn.Parameter(
953
- max_w13_scales, requires_grad=False
954
- )
955
-
956
- # Property to determine if AITER is used
957
- if self.rocm_aiter_moe_enabled:
958
- # reshaping weights is required for aiter moe kernel.
959
- shuffled_w13, shuffled_w2 = rocm_aiter_ops.shuffle_weights(
960
- layer.w13_weight.data, layer.w2_weight.data
877
+ # Per tensor kernels require single activation scale. Use the max.
878
+ if self.static_input_scales:
879
+ assert self.input_quant.strategy == QuantizationStrategy.TENSOR
880
+ assert w13_input_scale is not None and w2_input_scale is not None
881
+ w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
882
+ w13_input_scale, w2_input_scale
883
)
884
+ replace_parameter(layer, "w13_input_scale", w13_input_scale)
885
+ replace_parameter(layer, "w2_input_scale", w2_input_scale)
886
963
- layer.w13_weight = torch.nn.Parameter(shuffled_w13, requires_grad=False)
964
- layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
965
-
966
- elif self.use_marlin:
967
- (
968
- workspace,
969
- w13_weight,
970
- w2_weight,
971
- w13_weight_scale,
972
- w2_weight_scale,
973
- ) = prepare_moe_fp8_layer_for_marlin(
974
- layer,
975
- layer.w13_weight,
976
- layer.w2_weight,
977
- layer.w13_weight_scale,
978
- layer.w2_weight_scale,
979
- input_dtype=self.marlin_input_dtype,
887
+ # Per-tensor kernels use a single scale, for W13, but on disk there
888
+ # is a separate scale for W1 and W3. Requantize with the max scale.
889
+ if self.weight_quant.strategy == QuantizationStrategy.TENSOR:
890
+ process_fp8_weight_tensor_strategy_moe(
891
+ w13,
892
+ w13_scale,
893
+ shard_size=layer.intermediate_size_per_partition,
894
+ num_experts=layer.num_local_experts,
895
+ )
896
+
897
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
898
+ fp8_backend=self.fp8_backend,
899
+ layer=layer,
900
+ w13=w13,
901
+ w2=w2,
902
+ w13_scale=w13_scale,
903
+ w2_scale=w2_scale,
904
+ w13_input_scale=w13_input_scale,
905
+ w2_input_scale=w2_input_scale,
906
+ )
907
+
908
+ # Replace parameters with updated versions. Note that this helper
909
+ # function ensures the replacement is compatible with RL weight reloads.
910
+ replace_parameter(layer, "w13_weight", w13)
911
+ replace_parameter(layer, "w2_weight", w2)
912
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
913
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
914
+
915
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
916
+ if self.moe_quant_config:
917
+ self.kernel, self.use_inplace = make_fp8_moe_kernel(
918
+ layer=layer,
919
+ moe_quant_config=self.moe_quant_config,
920
+ moe_config=self.moe,
921
+ fp8_backend=self.fp8_backend,
922
)
981
- layer.workspace = workspace
982
- replace_parameter(layer, "w13_weight", w13_weight)
983
- replace_parameter(layer, "w2_weight", w2_weight)
984
- replace_parameter(layer, "w13_weight_scale", w13_weight_scale)
985
- replace_parameter(layer, "w2_weight_scale", w2_weight_scale)
986
-
987
- if self.use_cutlass:
988
- assert self.weight_quant.strategy != QuantizationStrategy.BLOCK
989
- device = layer.w13_weight.device
990
- # ab_strides1 and c_strides2 are the same
991
- self.ab_strides1_c_strides2 = torch.full(
992
- (layer.local_num_experts,),
993
- layer.hidden_size,
994
- device=device,
995
- dtype=torch.int64,
996
- )
997
- self.ab_strides2 = torch.full(
998
- (layer.local_num_experts,),
999
- layer.intermediate_size_per_partition,
1000
- device=device,
1001
- dtype=torch.int64,
1002
- )
1003
- self.c_strides1 = torch.full(
1004
- (layer.local_num_experts,),
1005
- 2 * layer.intermediate_size_per_partition,
1006
- device=device,
1007
- dtype=torch.int64,
1008
- )
1009
-
1010
- if is_deep_gemm_e8m0_used() and self.block_quant:
1011
- assert layer.weight_block_size is not None
1012
- # Re-quantise the expert weights so their scales are UE8M0.
1013
- block_sz = tuple(layer.weight_block_size)
1014
- requant_weight_ue8m0_inplace(
1015
- layer.w13_weight.data,
1016
- layer.w13_weight_scale.data,
1017
- block_sz,
1018
- )
1019
- requant_weight_ue8m0_inplace(
1020
- layer.w2_weight.data,
1021
- layer.w2_weight_scale.data,
1022
- block_sz,
1023
- )
1024
-
1025
- # Ensure column-major TMA alignment expected by DeepGEMM.
1026
- if expert_weight_is_col_major(layer.w13_weight_scale):
1027
- layer.w13_weight_scale = get_col_major_tma_aligned_tensor(
1028
- layer.w13_weight_scale
1029
- )
1030
- if expert_weight_is_col_major(layer.w2_weight_scale):
1031
- layer.w2_weight_scale = get_col_major_tma_aligned_tensor(
1032
- layer.w2_weight_scale
1033
- )
923
924
def maybe_make_prepare_finalize(
925
self,
926
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
927
) -> mk.FusedMoEPrepareAndFinalize | None:
1039
- if self.use_marlin or self.rocm_aiter_moe_enabled:
928
+ if self.fp8_backend in [Fp8MoeBackend.MARLIN, Fp8MoeBackend.AITER]:
929
return None
930
else:
931
return super().maybe_make_prepare_finalize(routing_tables)
@@ -1048,7 +937,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
937
) -> FusedMoEPermuteExpertsUnpermute:
938
# cutlass path
939
assert self.moe_quant_config is not None
1051
- if self.use_cutlass:
940
+ if self.fp8_backend == Fp8MoeBackend.VLLM_CUTLASS:
941
from vllm.model_executor.layers.fused_moe import (
942
CutlassBatchedExpertsFp8,
943
CutlassExpertsFp8,
@@ -1064,26 +953,27 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
953
):
954
logger.debug("CutlassBatchedExpertsFp8(%s)", self.__class__.__name__)
955
experts = CutlassBatchedExpertsFp8(
1067
- self.moe.num_local_experts,
1068
- num_dispatchers,
1069
- self.moe.in_dtype,
1070
- ab_strides1=self.ab_strides1_c_strides2,
1071
- ab_strides2=self.ab_strides2,
1072
- c_strides1=self.c_strides1,
1073
- c_strides2=self.ab_strides1_c_strides2,
956
+ max_experts_per_worker=self.moe.num_local_experts,
957
+ num_dispatchers=num_dispatchers,
958
+ out_dtype=self.moe.in_dtype,
959
+ e=layer.local_num_experts,
960
+ n=layer.intermediate_size_per_partition,
961
+ k=layer.hidden_size,
962
+ device=layer.w13_weight.device,
963
quant_config=self.moe_quant_config,
964
)
965
else:
966
logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
967
experts = CutlassExpertsFp8(
1079
- self.moe.in_dtype,
1080
- ab_strides1=self.ab_strides1_c_strides2,
1081
- ab_strides2=self.ab_strides2,
1082
- c_strides1=self.c_strides1,
1083
- c_strides2=self.ab_strides1_c_strides2,
968
+ out_dtype=self.moe.in_dtype,
969
+ e=layer.local_num_experts,
970
+ n=layer.intermediate_size_per_partition,
971
+ k=layer.hidden_size,
972
+ device=layer.w13_weight.device,
973
quant_config=self.moe_quant_config,
974
)
975
976
+ # TODO(rob): investigate disable_expert_map
977
self.disable_expert_map = (
978
num_dispatchers > 1 or not experts.supports_expert_map()
979
)
@@ -1096,13 +986,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
986
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
987
BatchedTritonExperts,
988
)
989
+ from vllm.model_executor.layers.fused_moe.fused_moe import (
990
+ TritonExperts,
991
+ )
992
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
993
TritonOrDeepGemmExperts,
994
)
995
1103
- assert not self.rocm_aiter_moe_enabled and not self.use_marlin
1104
-
1105
- use_deep_gemm = envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
996
+ assert self.fp8_backend not in [Fp8MoeBackend.AITER, Fp8MoeBackend.MARLIN]
997
998
if (
999
prepare_finalize.activation_format
@@ -1111,28 +1002,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1002
max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
1003
assert max_num_tokens_per_rank is not None
1004
1114
- if use_deep_gemm and not has_deep_gemm():
1115
- raise RuntimeError(
1116
- "DeepGEMM requested for MoE layer but not installed."
1117
- )
1118
-
1119
- compatible_with_deep_gemm = (
1120
- self.moe_quant_config.use_fp8_w8a8
1121
- and self.moe_quant_config.block_shape
1122
- == get_mk_alignment_for_contiguous_layout()
1123
- )
1124
-
1125
- # If this MoE layer is compatible with DeepGEMM, the proper env
1126
- # vars are set and DeepGEMM is not installed, throw an error.
1127
- if use_deep_gemm and compatible_with_deep_gemm and not has_deep_gemm():
1128
- raise RuntimeError(
1129
- f"MoE layer incompatible with DeepGEMM, expected "
1130
- f"fp8==True, got {self.moe_quant_config.use_fp8_w8a8}"
1131
- f"or block_shape {self.moe_quant_config.block_shape}"
1132
- f"=={get_mk_alignment_for_contiguous_layout()}."
1133
- )
1134
-
1135
- if use_deep_gemm and compatible_with_deep_gemm and has_deep_gemm():
1005
+ if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:
1006
logger.debug("BatchedDeepGemmExperts(%s)", self.__class__.__name__)
1007
return BatchedDeepGemmExperts(
1008
max_num_tokens=max_num_tokens_per_rank,
@@ -1148,17 +1018,22 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1018
)
1019
1020
else:
1151
- logger.debug("TritonOrDeepGemmExperts(%s)", self.__class__.__name__)
1152
- return TritonOrDeepGemmExperts(
1153
- self.moe_quant_config,
1154
- allow_deep_gemm=use_deep_gemm,
1155
- )
1021
+ if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:
1022
+ logger.debug("TritonOrDeepGemmExperts(%s)", self.__class__.__name__)
1023
+ return TritonOrDeepGemmExperts(self.moe_quant_config)
1024
+ else:
1025
+ logger.debug("TritonExperts(%s)", self.__class__.__name__)
1026
+ return TritonExperts(self.moe_quant_config)
1027
1028
def get_fused_moe_quant_config(
1029
self, layer: torch.nn.Module
1030
) -> FusedMoEQuantConfig | None:
1160
- if self.use_marlin:
1161
- return None
1031
+ if self.fp8_backend == Fp8MoeBackend.MARLIN:
1032
+ return fp8_w8a16_moe_quant_config(
1033
+ w1_scale=layer.w13_weight_scale,
1034
+ w2_scale=layer.w2_weight_scale,
1035
+ block_shape=self.weight_block_size,
1036
+ )
1037
1038
per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
1039
per_channel_quant = self.weight_quant.strategy == QuantizationStrategy.CHANNEL
@@ -1184,118 +1059,23 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1059
router_logits=router_logits,
1060
)
1061
1187
- per_act_token = self.input_quant.strategy == QuantizationStrategy.TOKEN
1188
- per_channel_quant = self.weight_quant.strategy == QuantizationStrategy.CHANNEL
1189
-
1190
- if self.use_marlin:
1191
- assert layer.activation == "silu", (
1192
- f"{layer.activation} not supported for Marlin MoE."
1193
- )
1194
- return fused_marlin_moe(
1195
- x,
1196
- layer.w13_weight,
1197
- layer.w2_weight,
1198
- None,
1199
- None,
1200
- layer.w13_weight_scale,
1201
- layer.w2_weight_scale,
1202
- router_logits,
1203
- topk_weights,
1204
- topk_ids,
1205
- quant_type_id=scalar_types.float8_e4m3fn.id,
1206
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1207
- global_num_experts=layer.global_num_experts,
1208
- expert_map=layer.expert_map,
1209
- input_dtype=self.marlin_input_dtype,
1210
- workspace=layer.workspace,
1211
- )
1212
-
1213
- elif self.rocm_aiter_moe_enabled:
1214
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501
1215
- rocm_aiter_fused_experts,
1216
- )
1217
-
1218
- assert per_act_token == per_channel_quant
1219
- assert self.moe_quant_config is not None
1220
- return rocm_aiter_fused_experts(
1221
- hidden_states=x,
1222
- w1=layer.w13_weight,
1223
- w2=layer.w2_weight,
1224
- topk_weights=topk_weights,
1225
- topk_ids=topk_ids,
1226
- activation=layer.activation,
1227
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1228
- expert_map=layer.expert_map,
1229
- quant_config=self.moe_quant_config,
1230
- )
1231
-
1232
- # cutlass path
1233
- elif self.use_cutlass:
1234
- assert self.moe_quant_config is not None
1235
-
1236
- # small-batch fallback on SM100
1237
- if self.is_fp8_w8a8_sm100 and topk_ids.shape[0] <= 8:
1238
- from vllm.model_executor.layers.fused_moe import fused_experts
1239
-
1240
- assert per_act_token == per_channel_quant
1241
- return fused_experts(
1242
- hidden_states=x,
1243
- w1=layer.w13_weight,
1244
- w2=layer.w2_weight,
1245
- topk_weights=topk_weights,
1246
- topk_ids=topk_ids,
1247
- inplace=True,
1248
- activation=layer.activation,
1249
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1250
- global_num_experts=layer.global_num_experts,
1251
- expert_map=None
1252
- if self.disable_expert_map
1253
- else layer.expert_map, # ???
1254
- quant_config=self.moe_quant_config,
1255
- allow_deep_gemm=self.allow_deep_gemm,
1256
- )
1257
- else:
1258
- from vllm.model_executor.layers.fused_moe.cutlass_moe import (
1259
- cutlass_moe_fp8,
1260
- )
1261
-
1262
- assert per_act_token == per_channel_quant
1263
- assert self.moe_quant_config is not None
1264
- return cutlass_moe_fp8(
1265
- x,
1266
- layer.w13_weight,
1267
- layer.w2_weight,
1268
- topk_weights,
1269
- topk_ids,
1270
- quant_config=self.moe_quant_config,
1271
- activation=layer.activation,
1272
- global_num_experts=layer.global_num_experts,
1273
- expert_map=None if self.disable_expert_map else layer.expert_map,
1274
- ab_strides1=self.ab_strides1_c_strides2,
1275
- ab_strides2=self.ab_strides2,
1276
- c_strides1=self.c_strides1,
1277
- c_strides2=self.ab_strides1_c_strides2,
1278
- )
1279
-
1280
- else:
1281
- from vllm.model_executor.layers.fused_moe import fused_experts
1062
+ assert self.kernel is not None
1063
+ result = self.kernel(
1064
+ x,
1065
+ layer.w13_weight,
1066
+ layer.w2_weight,
1067
+ topk_weights,
1068
+ topk_ids,
1069
+ inplace=self.use_inplace,
1070
+ activation=layer.activation,
1071
+ global_num_experts=layer.global_num_experts,
1072
+ # TODO(rob): investigate the disable_expert_map introduced by:
1073
+ # https://github.com/vllm-project/vllm/commit/84166fee9770e6fba71a96978b3e7d149392fb28 # noqa: E501
1074
+ expert_map=None if self.disable_expert_map else layer.expert_map,
1075
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
1076
+ )
1077
1283
- assert per_act_token == per_channel_quant
1284
- assert self.moe_quant_config is not None
1285
- return fused_experts(
1286
- hidden_states=x,
1287
- w1=layer.w13_weight,
1288
- w2=layer.w2_weight,
1289
- topk_weights=topk_weights,
1290
- topk_ids=topk_ids,
1291
- inplace=True,
1292
- activation=layer.activation,
1293
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1294
- global_num_experts=layer.global_num_experts,
1295
- expert_map=layer.expert_map,
1296
- quant_config=self.moe_quant_config,
1297
- allow_deep_gemm=self.allow_deep_gemm,
1298
- )
1078
+ return result
1079
1080
@property
1081
def supports_eplb(self) -> bool:
vllm/model_executor/layers/quantization/fp8.py
+112
-384
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from enum import Enum
4
from typing import TYPE_CHECKING, Any, Optional
5
6
import torch
@@ -27,13 +26,17 @@ from vllm.model_executor.layers.fused_moe import (
26
FusedMoeWeightScaleSupported,
27
)
28
from vllm.model_executor.layers.fused_moe.config import (
30
- FusedMoEParallelConfig,
29
FusedMoEQuantConfig,
30
RoutingMethodType,
33
- fp8_w8a8_moe_quant_config,
34
- fp8_w8a16_moe_quant_config,
31
)
32
from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod
33
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
34
+ Fp8MoeBackend,
35
+ convert_to_fp8_moe_kernel_format,
36
+ make_fp8_moe_kernel,
37
+ make_fp8_moe_quant_config,
38
+ select_fp8_moe_backend,
39
+)
40
from vllm.model_executor.layers.linear import (
41
LinearBase,
42
LinearMethodBase,
@@ -46,25 +49,20 @@ from vllm.model_executor.layers.quantization.base_config import (
49
)
50
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
51
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
49
- FlashinferMoeBackend,
50
- apply_flashinfer_per_tensor_scale_fp8,
52
+ apply_fi_trtllm_fp8_per_tensor_moe,
53
build_flashinfer_fp8_cutlass_moe_prepare_finalize,
52
- get_flashinfer_moe_backend,
53
- make_fp8_moe_alpha_scales_for_fi,
54
- register_scales_for_trtllm_fp8_per_tensor_moe,
55
- rotate_flashinfer_fp8_moe_weights,
54
select_cutlass_fp8_gemm_impl,
57
- swap_w13_to_w31,
55
)
56
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
57
W8A8BlockFp8LinearOp,
58
create_fp8_input_scale,
59
create_fp8_scale_parameter,
60
create_fp8_weight_parameter,
64
- deepgemm_post_process_fp8_weight_block,
61
maybe_post_process_fp8_weight_block,
62
+ process_fp8_input_tensor_strategy_moe,
63
process_fp8_weight_block_strategy,
64
process_fp8_weight_tensor_strategy,
65
+ process_fp8_weight_tensor_strategy_moe,
66
validate_fp8_block_shape,
67
)
68
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
@@ -73,7 +71,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
71
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
72
apply_fp8_marlin_linear,
73
prepare_fp8_layer_for_marlin,
76
- prepare_moe_fp8_layer_for_marlin,
74
)
75
from vllm.model_executor.layers.quantization.utils.quant_utils import (
76
GroupShape,
@@ -81,12 +78,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
78
)
79
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
80
Fp8LinearOp,
84
- all_close_1d,
81
cutlass_block_fp8_supported,
82
cutlass_fp8_supported,
83
maybe_create_device_identity,
84
normalize_e4m3fn_to_e4m3fnuz,
89
- per_tensor_dequantize,
85
)
86
from vllm.model_executor.parameter import (
87
BlockQuantScaleParameter,
@@ -96,11 +91,8 @@ from vllm.model_executor.parameter import (
91
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
92
from vllm.platforms import current_platform
93
from vllm.utils.deep_gemm import (
99
- is_deep_gemm_e8m0_used,
94
is_deep_gemm_supported,
95
)
102
-from vllm.utils.flashinfer import has_flashinfer_moe
103
-from vllm.utils.import_utils import has_deep_gemm
96
97
if TYPE_CHECKING:
98
from vllm.model_executor.models.utils import WeightsMapper
@@ -110,107 +102,6 @@ ACTIVATION_SCHEMES = ["static", "dynamic"]
102
logger = init_logger(__name__)
103
104
113
-class Fp8MoeBackend(Enum):
114
- NONE = 0
115
- FLASHINFER_TRTLLM = 1
116
- FLASHINFER_CUTLASS = 2
117
- DEEPGEMM = 3
118
- MARLIN = 4
119
- TRITON = 5
120
- AITER = 6
121
-
122
-
123
-def get_fp8_moe_backend(
124
- block_quant: bool,
125
- moe_parallel_config: FusedMoEParallelConfig,
126
- with_lora_support: bool,
127
-) -> Fp8MoeBackend | None:
128
- """
129
- Select the primary FP8 MoE backend
130
- Note: Shape-specific fallbacks may still occur at runtime.
131
- """
132
- if current_platform.is_xpu():
133
- return None
134
- if with_lora_support:
135
- return Fp8MoeBackend.TRITON
136
- # Prefer FlashInfer backends on supported GPUs; allow SM90 and SM100.
137
- if (
138
- current_platform.is_cuda()
139
- and (
140
- current_platform.is_device_capability_family(100)
141
- or current_platform.is_device_capability(90)
142
- )
143
- and envs.VLLM_USE_FLASHINFER_MOE_FP8
144
- and has_flashinfer_moe()
145
- ):
146
- backend = get_flashinfer_moe_backend()
147
- if backend == FlashinferMoeBackend.TENSORRT_LLM:
148
- logger.info_once("Using FlashInfer FP8 MoE TRTLLM backend for SM100")
149
- return Fp8MoeBackend.FLASHINFER_TRTLLM
150
- else:
151
- if block_quant and current_platform.is_device_capability_family(100):
152
- raise ValueError(
153
- "FlashInfer FP8 MoE throughput backend does not "
154
- "support block quantization on SM100. Please use "
155
- "VLLM_FLASHINFER_MOE_BACKEND=latency "
156
- "instead."
157
- )
158
- logger.info_once("Using FlashInfer FP8 MoE CUTLASS backend for SM90/SM100")
159
- return Fp8MoeBackend.FLASHINFER_CUTLASS
160
-
161
- # weight-only path for older GPUs without native FP8
162
- use_marlin = (
163
- not current_platform.has_device_capability(89)
164
- or envs.VLLM_TEST_FORCE_FP8_MARLIN
165
- )
166
- if current_platform.is_rocm():
167
- use_marlin = False
168
- if use_marlin:
169
- logger.info_once("Using Marlin backend for FP8 MoE")
170
- return Fp8MoeBackend.MARLIN
171
-
172
- # Determine if we should use DeepGEMM with block-quantized weights:
173
- # - If explicitly set by user, respect their choice
174
- # - If not explicitly set (default), disable when TP size is >= 8
175
- moe_use_deep_gemm = envs.VLLM_MOE_USE_DEEP_GEMM
176
- if not envs.is_set("VLLM_MOE_USE_DEEP_GEMM") and moe_parallel_config.tp_size >= 8:
177
- moe_use_deep_gemm = False
178
- logger.info_once(
179
- "DeepGEMM MoE is disabled by default when TP size is >= 8. "
180
- "Set VLLM_MOE_USE_DEEP_GEMM=1 to enable it.",
181
- scope="local",
182
- )
183
-
184
- # Determine if we should use DeepGEMM (top-level enable switch)
185
- # - If explicitly set by user, respect their choice
186
- # - If not platform supports DeepGEMM, disable it
187
- # This helps avoid warning messages on unsupported platforms.
188
- use_deep_gemm = envs.VLLM_USE_DEEP_GEMM
189
- if not is_deep_gemm_supported():
190
- use_deep_gemm = False
191
- logger.info_once(
192
- "DeepGEMM is disabled because the platform does not support it.",
193
- scope="local",
194
- )
195
-
196
- if use_deep_gemm and moe_use_deep_gemm and block_quant:
197
- if not has_deep_gemm():
198
- logger.warning_once(
199
- "DeepGEMM backend requested but not available.", scope="local"
200
- )
201
- elif is_deep_gemm_supported():
202
- logger.info_once("Using DeepGEMM backend for FP8 MoE", scope="local")
203
- return Fp8MoeBackend.DEEPGEMM
204
-
205
- if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MOE:
206
- logger.info_once("Using ROCm AITER backend for FP8 MoE", scope="local")
207
- return Fp8MoeBackend.AITER
208
-
209
- # default to Triton
210
- logger.info_once("Using Triton backend for FP8 MoE")
211
- return Fp8MoeBackend.TRITON
212
-
213
-
105
class Fp8Config(QuantizationConfig):
106
"""Config class for FP8."""
107
@@ -348,7 +239,6 @@ class Fp8Config(QuantizationConfig):
239
moe_quant_method = Fp8MoEMethod(self, layer)
240
else:
241
moe_quant_method = Fp8OnlineMoEMethod(self, layer)
351
- moe_quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
242
return moe_quant_method
243
elif isinstance(layer, Attention):
244
return Fp8KVCacheMethod(self)
@@ -736,40 +626,24 @@ class Fp8MoEMethod(FusedMoEMethodBase):
626
627
def __init__(self, quant_config: Fp8Config, layer: torch.nn.Module):
628
super().__init__(layer.moe_config)
739
- self.layer = layer
629
self.quant_config = quant_config
630
self.weight_block_size = self.quant_config.weight_block_size
631
self.block_quant: bool = self.weight_block_size is not None
632
self.weight_scale_name = (
633
"weight_scale_inv" if self.block_quant else "weight_scale"
634
)
746
- self.fp8_backend = get_fp8_moe_backend(
747
- self.block_quant, layer.moe_parallel_config, self.moe.is_lora_enabled
635
+ self.fp8_backend = select_fp8_moe_backend(
636
+ block_quant=self.block_quant,
637
+ tp_size=layer.moe_parallel_config.tp_size,
638
+ with_lora_support=self.moe.is_lora_enabled,
639
)
640
750
- self.marlin_input_dtype = None
751
- self.flashinfer_moe_backend: FlashinferMoeBackend | None = None
752
- if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
753
- self.flashinfer_moe_backend = FlashinferMoeBackend.TENSORRT_LLM
754
- elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
755
- self.flashinfer_moe_backend = FlashinferMoeBackend.CUTLASS
641
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
642
if self.block_quant and self.weight_block_size != [128, 128]:
643
raise NotImplementedError(
644
"FlashInfer CUTLASS FP8 MoE backend only supports block "
645
"size [128, 128]."
646
)
761
- if not self.block_quant:
762
- if layer.renormalize or layer.custom_routing_function is not None:
763
- raise NotImplementedError(
764
- "FlashInfer CUTLASS FP8 MoE backend does custom routing "
765
- f"function or renormalization, but got {layer.renormalize} and "
766
- f"{layer.custom_routing_function}."
767
- )
768
- if layer.scoring_func != "sigmoid":
769
- raise NotImplementedError(
770
- "FlashInfer CUTLASS FP8 MoE backend only supports "
771
- f"'sigmoid' scoring function, but got {layer.scoring_func}."
772
- )
647
if layer.activation != "silu":
648
raise NotImplementedError(
649
"FlashInfer CUTLASS FP8 MoE backend only supports SiLU "
@@ -778,12 +652,17 @@ class Fp8MoEMethod(FusedMoEMethodBase):
652
dynamic_per_token = (
653
not self.block_quant and self.quant_config.activation_scheme != "static"
654
)
781
- if self.flashinfer_moe_backend is not None and dynamic_per_token:
655
+ if dynamic_per_token and self.fp8_backend in [
656
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
657
+ Fp8MoeBackend.FLASHINFER_CUTLASS,
658
+ ]:
659
raise NotImplementedError(
660
"FlashInfer FP8 MoE backend does not support dynamic per token "
661
"activation quantization."
662
)
663
664
+ self.kernel: mk.FusedMoEModularKernel | None = None
665
+
666
def create_weights(
667
self,
668
layer: Module,
@@ -907,148 +786,43 @@ class Fp8MoEMethod(FusedMoEMethodBase):
786
layer.w13_input_scale = None
787
layer.w2_input_scale = None
788
910
- def _convert_weights_to_kernel_format(
789
+ def _setup_kernel(
790
self,
791
layer: Module,
913
- w13_weight: torch.Tensor,
914
- w2_weight: torch.Tensor,
915
- w13_weight_scale: torch.Tensor,
916
- w2_weight_scale: torch.Tensor,
792
+ w13: torch.Tensor,
793
+ w2: torch.Tensor,
794
+ w13_scale: torch.Tensor,
795
+ w2_scale: torch.Tensor,
796
w13_input_scale: torch.Tensor | None,
797
w2_input_scale: torch.Tensor | None,
798
) -> None:
920
- if self.fp8_backend == Fp8MoeBackend.DEEPGEMM:
921
- assert self.block_quant
922
- w13_weight, w13_weight_scale = deepgemm_post_process_fp8_weight_block(
923
- wq=w13_weight,
924
- ws=w13_weight_scale,
925
- quant_block_shape=tuple(layer.weight_block_size),
926
- use_e8m0=is_deep_gemm_e8m0_used(),
927
- )
928
- w2_weight, w2_weight_scale = deepgemm_post_process_fp8_weight_block(
929
- wq=w2_weight,
930
- ws=w2_weight_scale,
931
- quant_block_shape=tuple(layer.weight_block_size),
932
- use_e8m0=is_deep_gemm_e8m0_used(),
933
- )
934
- elif self.fp8_backend == Fp8MoeBackend.AITER:
935
- w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(
936
- w13_weight, w2_weight
937
- )
938
- elif self.fp8_backend == Fp8MoeBackend.MARLIN:
939
- (
940
- workspace,
941
- w13_weight,
942
- w2_weight,
943
- w13_weight_scale,
944
- w2_weight_scale,
945
- ) = prepare_moe_fp8_layer_for_marlin(
946
- layer,
947
- w13_weight,
948
- w2_weight,
949
- w13_weight_scale,
950
- w2_weight_scale,
951
- input_dtype=self.marlin_input_dtype,
952
- )
953
- layer.workspace = workspace
954
-
955
- elif self.fp8_backend in [
956
- Fp8MoeBackend.FLASHINFER_CUTLASS,
957
- Fp8MoeBackend.FLASHINFER_TRTLLM,
958
- ]:
959
- w13_weight = swap_w13_to_w31(w13_weight)
960
- if self.block_quant:
961
- w13_weight_scale = swap_w13_to_w31(w13_weight_scale)
962
- else:
963
- if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
964
- rotate_flashinfer_fp8_moe_weights(w13_weight, w2_weight)
965
- register_scales_for_trtllm_fp8_per_tensor_moe(
966
- layer=layer,
967
- w13_weight_scale=w13_weight,
968
- w13_input_scale=w13_input_scale,
969
- w2_weight_scale=w2_weight,
970
- w2_input_scale=w2_input_scale,
971
- )
972
-
973
- elif self.fp8_backend == Fp8MoeBackend.AITER:
974
- w13_weight, w2_weight = rocm_aiter_ops.shuffle_weights(
975
- w13_weight, w2_weight
976
- )
799
+ # Shuffle weights to runtime format.
800
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
801
+ fp8_backend=self.fp8_backend,
802
+ layer=layer,
803
+ w13=w13,
804
+ w2=w2,
805
+ w13_scale=w13_scale,
806
+ w2_scale=w2_scale,
807
+ w13_input_scale=w13_input_scale,
808
+ w2_input_scale=w2_input_scale,
809
+ )
810
811
# Replace parameters with updated versions. Note that this helper
812
# function ensures the replacement is compatible with RL weight reloads.
980
- replace_parameter(layer, "w13_weight", w13_weight)
981
- replace_parameter(layer, "w2_weight", w2_weight)
982
- replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_weight_scale)
983
- replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_weight_scale)
984
-
985
- def _setup_kernel(self, layer: Module) -> None:
986
- """Setup Modular Kernel for TP Case"""
987
- # NOTE(rob): this is a WIP refactor. We are first migrating
988
- # all of the kernels in the TP case to use mk. Once this is
989
- # done, then we will initialzie the TP case and DP/EP case
990
- # via the same code path (i.e. via maybe_init_modular_kernel).
991
- # NOTE(rob): in progress migrating all into this format.
992
-
993
- from vllm.model_executor.layers.fused_moe import (
994
- TritonOrDeepGemmExperts,
995
- )
996
- from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
997
- FlashInferExperts,
998
- )
999
- from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
1000
- MarlinExperts,
1001
- )
1002
- from vllm.model_executor.layers.fused_moe.prepare_finalize import (
1003
- MoEPrepareAndFinalizeNoEP,
1004
- )
1005
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
1006
- AiterExperts,
1007
- )
1008
-
1009
- # Flashinfer TRTLLM does not use the modular kernel abstraction.
1010
- if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
1011
- return
813
+ replace_parameter(layer, "w13_weight", w13)
814
+ replace_parameter(layer, "w2_weight", w2)
815
+ replace_parameter(layer, f"w13_{self.weight_scale_name}", w13_scale)
816
+ replace_parameter(layer, f"w2_{self.weight_scale_name}", w2_scale)
817
818
+ # Setup modular kernel for TP case.
819
self.moe_quant_config = self.get_fused_moe_quant_config(layer)
1014
- assert self.moe_quant_config is not None
1015
- self.use_inplace = True
1016
-
1017
- if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
1018
- self.kernel = mk.FusedMoEModularKernel(
1019
- # TODO: make defer_input_quant an attr of the FlashInferExperts
1020
- MoEPrepareAndFinalizeNoEP(defer_input_quant=self.block_quant),
1021
- FlashInferExperts(
1022
- out_dtype=layer.orig_dtype,
1023
- quant_config=self.moe_quant_config,
1024
- ep_rank=self.moe.ep_rank,
1025
- ep_size=self.moe.ep_size,
1026
- tp_rank=self.moe.tp_rank,
1027
- tp_size=self.moe.tp_size,
1028
- use_dp=(self.moe.dp_size > 1),
1029
- use_deepseek_fp8_block_scale=self.block_quant,
1030
- ),
1031
- )
1032
- self.use_inplace = False
1033
-
1034
- elif self.fp8_backend == Fp8MoeBackend.AITER:
1035
- self.kernel = mk.FusedMoEModularKernel(
1036
- # TODO: make defer_input_quant an attr of the AiterExperts
1037
- MoEPrepareAndFinalizeNoEP(defer_input_quant=True),
1038
- AiterExperts(quant_config=self.moe_quant_config),
1039
- )
1040
- elif self.fp8_backend == Fp8MoeBackend.MARLIN:
1041
- self.kernel = mk.FusedMoEModularKernel(
1042
- MoEPrepareAndFinalizeNoEP(),
1043
- MarlinExperts(quant_config=self.moe_quant_config),
1044
- )
1045
- else:
1046
- self.kernel = mk.FusedMoEModularKernel(
1047
- MoEPrepareAndFinalizeNoEP(),
1048
- TritonOrDeepGemmExperts(
1049
- quant_config=self.moe_quant_config,
1050
- allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),
1051
- ),
820
+ if self.moe_quant_config:
821
+ self.kernel, self.use_inplace = make_fp8_moe_kernel(
822
+ layer=layer,
823
+ moe_quant_config=self.moe_quant_config,
824
+ moe_config=self.moe,
825
+ fp8_backend=self.fp8_backend,
826
)
827
828
def process_weights_after_loading(self, layer: Module) -> None:
@@ -1056,78 +830,58 @@ class Fp8MoEMethod(FusedMoEMethodBase):
830
return
831
832
# Allow for accessing weights and scales in standard way.
1059
- w13_weight = layer.w13_weight
1060
- w2_weight = layer.w2_weight
1061
- w13_weight_scale = getattr(layer, f"w13_{self.weight_scale_name}")
1062
- w2_weight_scale = getattr(layer, f"w2_{self.weight_scale_name}")
833
+ w13 = layer.w13_weight
834
+ w2 = layer.w2_weight
835
+ w13_scale = getattr(layer, f"w13_{self.weight_scale_name}")
836
+ w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
837
w13_input_scale = layer.w13_input_scale
838
w2_input_scale = layer.w2_input_scale
839
840
# MI300x and MI325x use FNUZ format for FP8. Convert if needed.
841
if current_platform.is_fp8_fnuz():
1068
- w13_weight, w13_weight_scale, w13_input_scale = (
1069
- normalize_e4m3fn_to_e4m3fnuz(
1070
- w13_weight, w13_weight_scale, w13_input_scale
1071
- )
842
+ w13, w13_scale, w13_input_scale = normalize_e4m3fn_to_e4m3fnuz(
843
+ w13,
844
+ w13_scale,
845
+ w13_input_scale,
846
)
1073
- w2_weight, w2_weight_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
1074
- w2_weight, w2_weight_scale, w2_input_scale
847
+ w2, w2_scale, w2_input_scale = normalize_e4m3fn_to_e4m3fnuz(
848
+ w2,
849
+ w2_scale,
850
+ w2_input_scale,
851
)
852
853
# Per tensor kernels require single activation scale. Use the max.
854
if self.quant_config.activation_scheme == "static":
855
assert not self.block_quant
856
assert w13_input_scale is not None and w2_input_scale is not None
1081
- if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
1082
- logger.warning_once(
1083
- "Found input_scales that are not equal for "
1084
- "fp8 MoE layer. Using the maximum across experts "
1085
- "for each layer."
1086
- )
1087
- replace_parameter(layer, "w13_input_scale", w13_input_scale.max())
1088
- replace_parameter(layer, "w2_input_scale", w2_input_scale.max())
857
+ w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
858
+ w13_input_scale, w2_input_scale
859
+ )
860
+ replace_parameter(layer, "w13_input_scale", w13_input_scale)
861
+ replace_parameter(layer, "w2_input_scale", w2_input_scale)
862
863
# Per tensor kernels require single weight scale for w13 per expert, but
864
# on disk there is a scale for w1 and w3. Use the max to requantize.
865
if not self.block_quant:
866
shard_size = layer.intermediate_size_per_partition
1094
- max_w13_scales = w13_weight_scale.max(dim=1).values
1095
- for expert_id in range(layer.local_num_experts):
1096
- start = 0
1097
- for shard_id in range(2):
1098
- dq_weight = per_tensor_dequantize(
1099
- w13_weight[expert_id][start : start + shard_size, :],
1100
- w13_weight_scale[expert_id][shard_id],
1101
- )
1102
- w13_weight[expert_id][start : start + shard_size, :], _ = (
1103
- ops.scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
1104
- )
1105
- start += shard_size
1106
- w13_weight_scale = max_w13_scales
867
+ w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
868
+ w13, w13_scale, shard_size, layer.local_num_experts
869
+ )
870
1108
- # Shuffle weights into the runtime format.
1109
- self._convert_weights_to_kernel_format(
1110
- layer=layer,
1111
- w13_weight=w13_weight,
1112
- w2_weight=w2_weight,
1113
- w13_weight_scale=w13_weight_scale,
1114
- w2_weight_scale=w2_weight_scale,
1115
- w13_input_scale=w13_input_scale,
1116
- w2_input_scale=w2_input_scale,
871
+ # Shuffle weights to runtime format and setup kernel.
872
+ self._setup_kernel(
873
+ layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
874
)
875
1119
- # Setup modular kernel for TP case.
1120
- self._setup_kernel(layer)
1121
-
876
def maybe_make_prepare_finalize(
877
self,
878
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
879
) -> mk.FusedMoEPrepareAndFinalize | None:
1126
- if (
1127
- self.fp8_backend == Fp8MoeBackend.AITER
1128
- or self.fp8_backend == Fp8MoeBackend.MARLIN
1129
- or self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
1130
- ):
880
+ if self.fp8_backend in [
881
+ Fp8MoeBackend.AITER,
882
+ Fp8MoeBackend.MARLIN,
883
+ Fp8MoeBackend.FLASHINFER_TRTLLM,
884
+ ]:
885
return None
886
elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
887
prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
@@ -1184,7 +938,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
938
)
939
elif self.moe.is_lora_enabled:
940
return TritonExperts(quant_config=self.moe_quant_config)
1187
- elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
941
+ elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
942
# Select GEMM experts with block-scale when weights are block-quantized
943
experts = select_cutlass_fp8_gemm_impl(
944
self.moe,
@@ -1193,17 +947,23 @@ class Fp8MoEMethod(FusedMoEMethodBase):
947
)
948
logger.debug_once("Using %s", experts.__class__.__name__)
949
return experts
1196
- else:
950
+ elif self.fp8_backend == Fp8MoeBackend.DEEPGEMM:
951
logger.debug(
952
"TritonOrDeepGemmExperts(%s): block_size=%s, per_act_token=%s",
953
self.__class__.__name__,
954
self.weight_block_size,
955
False,
956
)
1203
- return TritonOrDeepGemmExperts(
1204
- quant_config=self.moe_quant_config,
1205
- allow_deep_gemm=(self.fp8_backend == Fp8MoeBackend.DEEPGEMM),
957
+ return TritonOrDeepGemmExperts(self.moe_quant_config)
958
+ else:
959
+ assert self.fp8_backend == Fp8MoeBackend.TRITON
960
+ logger.debug(
961
+ "TritonExperts(%s): block_size=%s, per_act_token=%s",
962
+ self.__class__.__name__,
963
+ self.weight_block_size,
964
+ False,
965
)
966
+ return TritonExperts(self.moe_quant_config)
967
968
def get_fused_moe_quant_config(
969
self, layer: torch.nn.Module
@@ -1212,42 +972,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
972
if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
973
return None
974
1215
- # MARLIN uses mixed precision W8A16 config.
1216
- if self.fp8_backend == Fp8MoeBackend.MARLIN:
1217
- return fp8_w8a16_moe_quant_config(
1218
- w1_scale=getattr(layer, f"w13_{self.weight_scale_name}"),
1219
- w2_scale=getattr(layer, f"w2_{self.weight_scale_name}"),
1220
- block_shape=self.weight_block_size,
1221
- )
1222
-
975
w1_scale = getattr(layer, f"w13_{self.weight_scale_name}")
976
w2_scale = getattr(layer, f"w2_{self.weight_scale_name}")
977
a1_scale = layer.w13_input_scale
978
a2_scale = layer.w2_input_scale
979
1228
- # Flashinfer CUTLASS per-tensor uses single dq scale
1229
- # (alpha = w_scale * a_scale) and inverse a2 scale.
1230
- if (
1231
- self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS
1232
- and not self.block_quant
1233
- ):
1234
- g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(
1235
- w1_scale,
1236
- a1_scale,
1237
- w2_scale,
1238
- a2_scale,
1239
- )
1240
- return fp8_w8a8_moe_quant_config(
1241
- w1_scale=w1_scale,
1242
- w2_scale=w2_scale,
1243
- a1_scale=a1_scale,
1244
- a2_scale=(1.0 / a2_scale),
1245
- g1_alphas=g1_alphas,
1246
- g2_alphas=g2_alphas,
1247
- )
1248
-
1249
- # All other backends use normal config.
1250
- return fp8_w8a8_moe_quant_config(
980
+ return make_fp8_moe_quant_config(
981
+ fp8_backend=self.fp8_backend,
982
w1_scale=w1_scale,
983
w2_scale=w2_scale,
984
a1_scale=a1_scale,
@@ -1269,7 +1000,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1000
x: torch.Tensor,
1001
router_logits: torch.Tensor,
1002
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1272
- if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1003
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
1004
# TODO(rob): convert this to MK.
1005
if layer.enable_eplb:
1006
raise NotImplementedError("EPLB not supported for `Fp8MoEMethod` yet.")
@@ -1308,10 +1039,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1039
routed_scaling=layer.routed_scaling_factor,
1040
)
1041
else:
1311
- assert (
1312
- not layer.renormalize and layer.custom_routing_function is not None
1313
- )
1314
- result = apply_flashinfer_per_tensor_scale_fp8(
1042
+ result = apply_fi_trtllm_fp8_per_tensor_moe(
1043
layer=layer,
1044
hidden_states=x,
1045
router_logits=router_logits,
@@ -1327,6 +1055,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1055
hidden_states=x,
1056
router_logits=router_logits,
1057
)
1058
+
1059
+ assert self.kernel is not None
1060
result = self.kernel(
1061
x,
1062
layer.w13_weight,
@@ -1358,7 +1088,6 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1088
assert not quant_config.is_checkpoint_fp8_serialized
1089
assert quant_config.activation_scheme == "dynamic"
1090
assert quant_config.weight_block_size is None
1361
- assert self.flashinfer_moe_backend is None
1091
1092
def create_weights(
1093
self,
@@ -1447,6 +1176,8 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1176
)
1177
layer.register_parameter("w13_weight_scale", w13_weight_scale)
1178
layer.register_parameter("w2_weight_scale", w2_weight_scale)
1179
+ set_weight_attrs(w13_weight_scale, extra_weight_attrs)
1180
+ set_weight_attrs(w2_weight_scale, extra_weight_attrs)
1181
1182
layer.w13_input_scale = None
1183
layer.w2_input_scale = None
@@ -1457,33 +1188,30 @@ class Fp8OnlineMoEMethod(Fp8MoEMethod):
1188
1189
# If checkpoint is fp16, quantize in place.
1190
fp8_dtype = current_platform.fp8_dtype()
1460
- w13_weight = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
1461
- w2_weight = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
1191
+ w13 = torch.empty_like(layer.w13_weight, dtype=fp8_dtype)
1192
+ w2 = torch.empty_like(layer.w2_weight, dtype=fp8_dtype)
1193
+ w13_scale = layer.w13_weight_scale
1194
+ w2_scale = layer.w2_weight_scale
1195
1196
for expert in range(layer.local_num_experts):
1464
- w13_weight[expert, :, :], layer.w13_weight_scale[expert] = (
1465
- ops.scaled_fp8_quant(layer.w13_weight[expert, :, :])
1197
+ w13[expert, :, :], w13_scale[expert] = ops.scaled_fp8_quant(
1198
+ layer.w13_weight[expert, :, :]
1199
)
1467
- w2_weight[expert, :, :], layer.w2_weight_scale[expert] = (
1468
- ops.scaled_fp8_quant(layer.w2_weight[expert, :, :])
1200
+ w2[expert, :, :], w2_scale[expert] = ops.scaled_fp8_quant(
1201
+ layer.w2_weight[expert, :, :]
1202
)
1470
- replace_parameter(layer, "w13_weight", w13_weight)
1471
- replace_parameter(layer, "w2_weight", w2_weight)
1203
1473
- # Shuffle weights into the runtime format.
1474
- self._convert_weights_to_kernel_format(
1475
- layer=layer,
1476
- w13_weight=w13_weight,
1477
- w2_weight=w2_weight,
1478
- w13_weight_scale=layer.w13_weight_scale,
1479
- w2_weight_scale=layer.w2_weight_scale,
1480
- w13_input_scale=None,
1481
- w2_input_scale=None,
1204
+ # Shuffle weights to runtime format and setup kernel.
1205
+ self._setup_kernel(
1206
+ layer,
1207
+ w13,
1208
+ w2,
1209
+ w13_scale,
1210
+ w2_scale,
1211
+ layer.w13_input_scale,
1212
+ layer.w2_input_scale,
1213
)
1214
1484
- # Setup modular kernel for TP case.
1485
- self._setup_kernel(layer)
1486
-
1215
1216
class Fp8KVCacheMethod(BaseKVCacheMethod):
1217
"""
vllm/model_executor/layers/quantization/modelopt.py
+147
-264
@@ -15,7 +15,6 @@ from vllm.attention.layer import Attention
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEQuantConfig,
18
- fp8_w8a8_moe_quant_config,
18
nvfp4_moe_quant_config,
19
)
20
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
@@ -24,6 +23,13 @@ from vllm.model_executor.layers.fused_moe.layer import (
23
FusedMoEMethodBase,
24
FusedMoeWeightScaleSupported,
25
)
26
+from vllm.model_executor.layers.fused_moe.oracle.fp8 import (
27
+ Fp8MoeBackend,
28
+ convert_to_fp8_moe_kernel_format,
29
+ make_fp8_moe_kernel,
30
+ make_fp8_moe_quant_config,
31
+ select_fp8_moe_backend,
32
+)
33
from vllm.model_executor.layers.linear import (
34
LinearBase,
35
LinearMethodBase,
@@ -45,19 +51,16 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
51
)
52
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
53
FlashinferMoeBackend,
48
- apply_flashinfer_per_tensor_scale_fp8,
54
+ apply_fi_trtllm_fp8_per_tensor_moe,
55
build_flashinfer_fp8_cutlass_moe_prepare_finalize,
50
- flashinfer_cutlass_moe_fp8,
56
get_flashinfer_moe_backend,
57
is_flashinfer_supporting_global_sf,
53
- make_fp8_moe_alpha_scales_for_fi,
54
- register_scales_for_trtllm_fp8_per_tensor_moe,
55
- rotate_flashinfer_fp8_moe_weights,
58
select_cutlass_fp8_gemm_impl,
57
- swap_w13_to_w31,
59
)
60
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
61
W8A8BlockFp8LinearOp,
62
+ process_fp8_input_tensor_strategy_moe,
63
+ process_fp8_weight_tensor_strategy_moe,
64
)
65
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
66
get_marlin_input_dtype,
@@ -85,13 +88,12 @@ from vllm.model_executor.parameter import (
88
ModelWeightParameter,
89
PerTensorScaleParameter,
90
)
91
+from vllm.model_executor.utils import replace_parameter
92
from vllm.scalar_type import scalar_types
93
from vllm.utils.flashinfer import (
94
flashinfer_scaled_fp4_mm,
95
has_flashinfer,
92
- has_flashinfer_moe,
96
)
94
-from vllm.utils.math_utils import round_up
97
98
if TYPE_CHECKING:
99
from vllm.model_executor.models.utils import WeightsMapper
@@ -721,38 +723,23 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
723
layer: FusedMoE,
724
) -> None:
725
super().__init__(layer.moe_config)
724
- self.layer = layer
726
self.quant_config = quant_config
726
- from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
727
- cutlass_fp8_supported,
727
+ assert self.quant_config.is_checkpoint_fp8_serialized
728
+ self.fp8_backend = select_fp8_moe_backend(
729
+ block_quant=False,
730
+ tp_size=layer.moe_parallel_config.tp_size,
731
+ with_lora_support=self.moe.is_lora_enabled,
732
)
729
-
730
- self.cutlass_fp8_supported = cutlass_fp8_supported()
731
- self.flashinfer_moe_backend: FlashinferMoeBackend | None = None
732
- if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
733
- self.flashinfer_moe_backend = get_flashinfer_moe_backend()
734
- if (
735
- self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
736
- and not self.moe.is_act_and_mul
737
- ):
738
- logger.info_once(
739
- "Non-gated MoE is not supported for min-latency mode,"
740
- "falling back to high-throughput mode"
741
- )
742
- self.flashinfer_moe_backend = FlashinferMoeBackend.CUTLASS
743
-
744
- logger.info_once(
745
- f"Using FlashInfer {self.flashinfer_moe_backend.value} kernels"
746
- )
733
+ self.kernel: mk.FusedMoEModularKernel | None = None
734
735
def maybe_make_prepare_finalize(
736
self,
737
routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
738
) -> mk.FusedMoEPrepareAndFinalize | None:
739
# TRT LLM not supported with all2all yet.
753
- if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
740
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
741
return None
755
- elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
742
+ elif self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
743
# TP case: avoid convert to ModularKernelMethod - to be refactored.
744
if self.moe.dp_size == 1:
745
return None
@@ -787,6 +774,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
774
params_dtype: torch.dtype,
775
**extra_weight_attrs,
776
):
777
+ layer.orig_dtype = params_dtype
778
+ layer.num_experts = num_experts
779
+
780
# Use FP8 dtype if checkpoint is serialized
781
weight_dtype = (
782
torch.float8_e4m3fn
@@ -826,217 +816,121 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
816
)
817
layer.register_parameter("w2_weight", w2_weight)
818
829
- if self.quant_config.is_checkpoint_fp8_serialized:
830
- # WEIGHT SCALES - Per-tensor scaling for ModelOpts
831
- # For gated MoE, allocate 2 scales for w1 and w3 respectively.
832
- # They will be combined to a single scale after weight loading.
833
- # For non-gated MoE, allocate 1 scale for w13.
834
- if self.moe.is_act_and_mul:
835
- w13_weight_scale_shape = (num_experts, 2)
836
- else:
837
- w13_weight_scale_shape = (num_experts, 1)
838
- w13_weight_scale = PerTensorScaleParameter(
839
- data=torch.full(
840
- w13_weight_scale_shape,
841
- 1.0,
842
- dtype=torch.float32,
843
- ),
844
- weight_loader=weight_loader,
845
- )
846
- w2_weight_scale = PerTensorScaleParameter(
847
- data=torch.full((num_experts,), 1.0, dtype=torch.float32),
848
- weight_loader=weight_loader,
849
- )
850
- layer.register_parameter("w13_weight_scale", w13_weight_scale)
851
- layer.register_parameter("w2_weight_scale", w2_weight_scale)
852
-
853
- # Set weight loader attributes for scales
854
- extra_weight_attrs.update(
855
- {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value}
856
- )
857
-
858
- # INPUT SCALES - Per-tensor scaling for ModelOpt
859
- w13_input_scale = PerTensorScaleParameter(
860
- data=torch.full((num_experts,), 1.0, dtype=torch.float32),
861
- weight_loader=weight_loader,
862
- )
863
- w2_input_scale = PerTensorScaleParameter(
864
- data=torch.full((num_experts,), 1.0, dtype=torch.float32),
865
- weight_loader=weight_loader,
866
- )
867
- layer.register_parameter("w13_input_scale", w13_input_scale)
868
- layer.register_parameter("w2_input_scale", w2_input_scale)
869
-
870
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
871
- """Process FP8 MoE weights after loading from serialized checkpoint.
872
- Only supports pre-quantized checkpoints with FP8 weights and scales.
873
- """
874
-
875
- if self.flashinfer_moe_backend is not None:
876
- self._maybe_pad_intermediate_for_flashinfer(layer)
877
-
878
- layer.w13_weight = Parameter(layer.w13_weight.data, requires_grad=False)
879
- layer.w2_weight = Parameter(layer.w2_weight.data, requires_grad=False)
819
+ # WEIGHT SCALES - Per-tensor scaling for ModelOpts
820
+ # For gated MoE, allocate 2 scales for w1 and w3 respectively.
821
+ # They will be combined to a single scale after weight loading.
822
+ # For non-gated MoE, allocate 1 scale for w13.
823
+ w13_weight_scale = PerTensorScaleParameter(
824
+ data=torch.full(
825
+ (num_experts, 2 if self.moe.is_act_and_mul else 1),
826
+ 1.0,
827
+ dtype=torch.float32,
828
+ ),
829
+ weight_loader=weight_loader,
830
+ )
831
+ w2_weight_scale = PerTensorScaleParameter(
832
+ data=torch.full((num_experts,), 1.0, dtype=torch.float32),
833
+ weight_loader=weight_loader,
834
+ )
835
+ layer.register_parameter("w13_weight_scale", w13_weight_scale)
836
+ layer.register_parameter("w2_weight_scale", w2_weight_scale)
837
881
- from vllm._custom_ops import scaled_fp8_quant
882
- from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
883
- per_tensor_dequantize,
838
+ # INPUT SCALES - Per-tensor scaling for ModelOpt
839
+ w13_input_scale = PerTensorScaleParameter(
840
+ data=torch.full((num_experts,), 1.0, dtype=torch.float32),
841
+ weight_loader=weight_loader,
842
+ )
843
+ w2_input_scale = PerTensorScaleParameter(
844
+ data=torch.full((num_experts,), 1.0, dtype=torch.float32),
845
+ weight_loader=weight_loader,
846
)
847
+ layer.register_parameter("w13_input_scale", w13_input_scale)
848
+ layer.register_parameter("w2_input_scale", w2_input_scale)
849
886
- # Handle scale parameters
887
- if hasattr(layer, "w13_weight_scale") and layer.w13_weight_scale is not None:
888
- # Fp8 moe kernel needs single weight scale for w13 per expert.
889
- # We take the max of the w1 and w3 scales
890
- # then dequant and requant each expert.
891
- if (
892
- layer.w13_weight_scale.dim() == 2
893
- and layer.w13_weight_scale.shape[1] == 2
894
- ):
895
- assert self.moe.is_act_and_mul, (
896
- "w13_weight_scale should have 2 elements per expert "
897
- "only for gated MoE"
898
- )
899
- # Get the maximum scale across w1 and w3 for each expert
900
- max_w13_scales = layer.w13_weight_scale.max(dim=1).values
901
-
902
- # Requantize each expert's weights using the combined scale
903
- # w13_weight (num_experts, 2 * intermediate_size, hidden_size)
904
- # where the first intermediate_size rows are w1, the next are w3
905
- intermediate_size = layer.w13_weight.shape[1] // 2
906
- for expert_id in range(layer.w13_weight.shape[0]):
907
- start = 0
908
- for shard_id in range(2): # w1 and w3
909
- # Dequantize using the original scale for this shard
910
- dq_weight = per_tensor_dequantize(
911
- layer.w13_weight[expert_id][
912
- start : start + intermediate_size, :
913
- ],
914
- layer.w13_weight_scale[expert_id][shard_id],
915
- )
916
- # Requantize using the combined max scale
917
-
918
- (
919
- layer.w13_weight[expert_id][
920
- start : start + intermediate_size, :
921
- ],
922
- _,
923
- ) = scaled_fp8_quant(dq_weight, max_w13_scales[expert_id])
924
-
925
- start += intermediate_size
926
-
927
- # Update the scale parameter to be per-expert
928
- layer.w13_weight_scale = Parameter(max_w13_scales, requires_grad=False)
929
- else:
930
- layer.w13_weight_scale = Parameter(
931
- layer.w13_weight_scale.data, requires_grad=False
932
- )
850
+ def _setup_kernel(
851
+ self,
852
+ layer: torch.nn.Module,
853
+ w13: torch.Tensor,
854
+ w2: torch.Tensor,
855
+ w13_scale: torch.Tensor,
856
+ w2_scale: torch.Tensor,
857
+ w13_input_scale: torch.Tensor,
858
+ w2_input_scale: torch.Tensor,
859
+ ):
860
+ w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(
861
+ fp8_backend=self.fp8_backend,
862
+ layer=layer,
863
+ w13=w13,
864
+ w2=w2,
865
+ w13_scale=w13_scale,
866
+ w2_scale=w2_scale,
867
+ w13_input_scale=w13_input_scale,
868
+ w2_input_scale=w2_input_scale,
869
+ )
870
934
- if hasattr(layer, "w2_weight_scale") and layer.w2_weight_scale is not None:
935
- layer.w2_weight_scale = Parameter(
936
- layer.w2_weight_scale.data, requires_grad=False
937
- )
938
- # Input scales must be equal for each expert in fp8 MoE layers.
939
- if hasattr(layer, "w13_input_scale") and layer.w13_input_scale is not None:
940
- layer.w13_input_scale = Parameter(
941
- layer.w13_input_scale.max(), requires_grad=False
942
- )
943
- if hasattr(layer, "w2_input_scale") and layer.w2_input_scale is not None:
944
- layer.w2_input_scale = Parameter(
945
- layer.w2_input_scale.max(), requires_grad=False
871
+ # Replace parameters with updated versions. Note that this helper
872
+ # function ensures the replacement is compatible with RL weight reloads.
873
+ replace_parameter(layer, "w13_weight", w13)
874
+ replace_parameter(layer, "w2_weight", w2)
875
+ replace_parameter(layer, "w13_weight_scale", w13_scale)
876
+ replace_parameter(layer, "w2_weight_scale", w2_scale)
877
+
878
+ # Setup modular kernel for TP case.
879
+ self.moe_quant_config = self.get_fused_moe_quant_config(layer)
880
+ if self.moe_quant_config:
881
+ self.kernel, self.use_inplace = make_fp8_moe_kernel(
882
+ layer=layer,
883
+ moe_quant_config=self.moe_quant_config,
884
+ moe_config=self.moe,
885
+ fp8_backend=self.fp8_backend,
886
)
887
948
- if self.flashinfer_moe_backend is not None:
949
- if self.moe.is_act_and_mul:
950
- layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)
951
-
952
- # NOTE: this adds some attributes used by the trtllm kernel,
953
- # which does not conform to the modular kernels abstraction (yet).
954
- if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
955
- rotate_flashinfer_fp8_moe_weights(layer.w13_weight, layer.w2_weight)
956
- register_scales_for_trtllm_fp8_per_tensor_moe(
957
- layer=layer,
958
- w13_weight_scale=layer.w13_weight_scale,
959
- w13_input_scale=layer.w13_input_scale,
960
- w2_weight_scale=layer.w2_weight_scale,
961
- w2_input_scale=layer.w2_input_scale,
962
- )
963
-
964
- def _maybe_pad_intermediate_for_flashinfer(self, layer: torch.nn.Module) -> None:
965
- """Pad intermediate size so FlashInfer kernels' alignment constraints hold.
966
-
967
- Some FlashInfer FP8 MoE kernels require the (gated) intermediate size
968
- used for GEMM to be divisible by a small alignment value. When this is
969
- not satisfied (e.g. with certain tensor-parallel sizes), we pad the
970
- gate/up and down projection weights along the intermediate dim.
971
- """
972
- if not hasattr(layer, "w13_weight") or not hasattr(layer, "w2_weight"):
973
- return
974
-
975
- # Current local intermediate size (per partition) is the K dimension of
976
- # the down projection.
977
- num_experts, hidden_size, intermediate = layer.w2_weight.shape
978
-
979
- min_alignment = 16
980
- padded_intermediate = round_up(intermediate, min_alignment)
981
-
982
- if padded_intermediate == intermediate:
983
- return
984
-
985
- logger.info(
986
- "Padding intermediate size from %d to %d for up/down projection weights.",
987
- intermediate,
988
- padded_intermediate,
888
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
889
+ w13 = layer.w13_weight
890
+ w2 = layer.w2_weight
891
+ w13_scale = layer.w13_weight_scale
892
+ w2_scale = layer.w2_weight_scale
893
+ w13_input_scale = layer.w13_input_scale
894
+ w2_input_scale = layer.w2_input_scale
895
+
896
+ # Per tensor kernels require single activation scale. Use the max.
897
+ w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe(
898
+ w13_input_scale, w2_input_scale
899
+ )
900
+ replace_parameter(layer, "w13_input_scale", w13_input_scale)
901
+ replace_parameter(layer, "w2_input_scale", w2_input_scale)
902
+
903
+ # Per tensor kernels require single weight scale for w13 per expert, but
904
+ # on disk there is a scale for w1 and w3. Use the max to requantize.
905
+ shard_size = layer.intermediate_size_per_partition
906
+ w13, w13_scale = process_fp8_weight_tensor_strategy_moe(
907
+ w13,
908
+ w13_scale,
909
+ shard_size,
910
+ num_experts=layer.w13_weight.shape[0],
911
+ is_act_and_mul=self.moe.is_act_and_mul,
912
)
913
991
- up_mult = 2 if self.moe.is_act_and_mul else 1
992
- padded_gate_up_dim = up_mult * padded_intermediate
993
-
994
- # Pad w13 and w12 along its intermediate dimension.
995
- w13 = layer.w13_weight.data
996
- padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size))
997
- padded_w13[:, : w13.shape[1], :] = w13
998
- layer.w13_weight.data = padded_w13
999
-
1000
- w2 = layer.w2_weight.data
1001
- padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate))
1002
- padded_w2[:, :, :intermediate] = w2
1003
- layer.w2_weight.data = padded_w2
1004
-
1005
- if hasattr(layer, "intermediate_size_per_partition"):
1006
- layer.intermediate_size_per_partition = padded_intermediate
914
+ # Shuffle weights to runtime format and setup kernel.
915
+ self._setup_kernel(
916
+ layer, w13, w2, w13_scale, w2_scale, w13_input_scale, w2_input_scale
917
+ )
918
919
def get_fused_moe_quant_config(
920
self, layer: torch.nn.Module
921
) -> FusedMoEQuantConfig | None:
1011
- if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1012
- # TRTLLM does not use modular kernels
1013
- return None
1014
-
1015
- elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1016
- g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(
1017
- layer.w13_weight_scale,
1018
- layer.w13_input_scale,
1019
- layer.w2_weight_scale,
1020
- layer.w2_input_scale,
1021
- )
1022
- return fp8_w8a8_moe_quant_config(
1023
- w1_scale=layer.w13_weight_scale,
1024
- w2_scale=layer.w2_weight_scale,
1025
- a1_scale=layer.w13_input_scale,
1026
- a2_scale=layer.w2_input_scale,
1027
- a1_gscale=(1.0 / layer.w13_input_scale),
1028
- a2_gscale=(1.0 / layer.w2_input_scale),
1029
- g1_alphas=g1_alphas,
1030
- g2_alphas=g2_alphas,
1031
- )
1032
- else:
1033
- assert self.flashinfer_moe_backend is None
1034
- return fp8_w8a8_moe_quant_config(
1035
- w1_scale=layer.w13_weight_scale,
1036
- w2_scale=layer.w2_weight_scale,
1037
- a1_scale=layer.w13_input_scale,
1038
- a2_scale=layer.w2_input_scale,
1039
- )
922
+ w1_scale = layer.w13_weight_scale
923
+ w2_scale = layer.w2_weight_scale
924
+ a1_scale = layer.w13_input_scale
925
+ a2_scale = layer.w2_input_scale
926
+
927
+ return make_fp8_moe_quant_config(
928
+ fp8_backend=self.fp8_backend,
929
+ w1_scale=w1_scale,
930
+ w2_scale=w2_scale,
931
+ a1_scale=a1_scale,
932
+ a2_scale=a2_scale,
933
+ )
934
935
def apply(
936
self,
@@ -1044,17 +938,18 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
938
x: torch.Tensor,
939
router_logits: torch.Tensor,
940
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
1047
- if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
941
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM:
942
if layer.enable_eplb:
943
raise NotImplementedError(
1050
- "EPLB not supported for `ModelOptFp8MoEMethod` yet."
944
+ "EPLB not supported for FlashInfer TRTLLM FP8 MoE Backend."
945
)
946
+ # TODO(rob): this validation should happen at kernel selection
947
+ # time in the oracle rather than here.
948
assert layer.activation == "silu", (
949
f"Expected 'silu' activation but got {layer.activation}"
950
)
1055
-
951
assert not layer.renormalize
1057
- return apply_flashinfer_per_tensor_scale_fp8(
952
+ return apply_fi_trtllm_fp8_per_tensor_moe(
953
layer=layer,
954
hidden_states=x,
955
router_logits=router_logits,
@@ -1066,46 +961,34 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
961
apply_router_weight_on_input=layer.apply_router_weight_on_input,
962
)
963
1069
- # Expert selection
964
topk_weights, topk_ids = layer.select_experts(
965
hidden_states=x,
966
router_logits=router_logits,
967
)
968
1075
- if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
969
+ # TODO(rob): this validation should happen at kernel selection
970
+ # time in the oracle rather than here.
971
+ if self.fp8_backend == Fp8MoeBackend.FLASHINFER_CUTLASS:
972
assert layer.activation in ("silu", "relu2_no_mul"), (
973
"Expected activation to be in ('silu', 'relu2_no_mul'),"
974
f"but got {layer.activation}"
975
)
1080
- return flashinfer_cutlass_moe_fp8(
1081
- x,
1082
- layer,
1083
- topk_weights,
1084
- topk_ids,
1085
- inplace=False,
1086
- activation=layer.activation,
1087
- global_num_experts=layer.global_num_experts,
1088
- expert_map=layer.expert_map,
1089
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1090
- )
1091
- else:
1092
- from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
976
1094
- assert self.moe_quant_config is not None
977
+ assert self.kernel is not None
978
+ result = self.kernel(
979
+ x,
980
+ layer.w13_weight,
981
+ layer.w2_weight,
982
+ topk_weights,
983
+ topk_ids,
984
+ inplace=self.use_inplace,
985
+ activation=layer.activation,
986
+ global_num_experts=layer.global_num_experts,
987
+ expert_map=layer.expert_map,
988
+ apply_router_weight_on_input=layer.apply_router_weight_on_input,
989
+ )
990
1096
- return fused_experts(
1097
- x,
1098
- layer.w13_weight,
1099
- layer.w2_weight,
1100
- topk_weights=topk_weights,
1101
- topk_ids=topk_ids,
1102
- inplace=True,
1103
- activation=layer.activation,
1104
- quant_config=self.moe_quant_config,
1105
- global_num_experts=layer.global_num_experts,
1106
- expert_map=layer.expert_map,
1107
- apply_router_weight_on_input=layer.apply_router_weight_on_input,
1108
- )
991
+ return result
992
993
994
ModelOptFp8Config.LinearMethodCls = ModelOptFp8LinearMethod
vllm/model_executor/layers/quantization/quark/quark_moe.py
+3
-4
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
22
)
23
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe
24
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
25
- prepare_moe_fp8_layer_for_marlin,
25
+ prepare_fp8_moe_layer_for_marlin,
26
)
27
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
28
OCP_MX_BLOCK_SIZE,
@@ -315,8 +315,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
315
layer.w2_weight = torch.nn.Parameter(shuffled_w2, requires_grad=False)
316
317
elif self.use_marlin:
318
- (workspace, w13_weight, w2_weight, w13_weight_scale, w2_weight_scale) = (
319
- prepare_moe_fp8_layer_for_marlin(
318
+ w13_weight, w2_weight, w13_weight_scale, w2_weight_scale = (
319
+ prepare_fp8_moe_layer_for_marlin(
320
layer,
321
layer.w13_weight,
322
layer.w2_weight,
@@ -324,7 +324,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
324
layer.w2_weight_scale,
325
)
326
)
327
- layer.workspace = workspace
327
# TODO(rob): once we apply refactor to Quark, switch to using
328
# replace_parameter for compatibility with reloading in RL.
329
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+119
-53
@@ -18,6 +18,7 @@ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize im
18
create_flashinfer_prepare_finalize,
19
)
20
from vllm.platforms import current_platform
21
+from vllm.utils.math_utils import round_up
22
23
logger = init_logger(__name__)
24
@@ -58,9 +59,10 @@ def swap_w13_to_w31(x: torch.Tensor) -> torch.Tensor:
59
)
60
61
61
-def rotate_flashinfer_fp8_moe_weights(
62
+def rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(
63
gemm1_weights: torch.Tensor, gemm2_weights: torch.Tensor
64
):
65
+ """Shuffle weights for for FI TRT-LLM Format"""
66
from flashinfer import reorder_rows_for_gated_act_gemm, shuffle_matrix_a
67
68
epilogue_tile_m = 128
@@ -105,16 +107,16 @@ def rotate_flashinfer_fp8_moe_weights(
107
108
def register_scales_for_trtllm_fp8_per_tensor_moe(
109
layer: torch.nn.Module,
108
- w13_weight_scale: torch.Tensor,
110
+ w13_scale: torch.Tensor,
111
w13_input_scale: torch.Tensor,
110
- w2_weight_scale: torch.Tensor,
112
+ w2_scale: torch.Tensor,
113
w2_input_scale: torch.Tensor,
114
) -> None:
115
"""Register necessary scales for FlashInfer TRTLLM FP8 MoE kernel"""
116
g1_alphas, g2_alphas = make_fp8_moe_alpha_scales_for_fi(
115
- w13_scale=w13_weight_scale,
117
+ w13_scale=w13_scale,
118
w13_input_scale=w13_input_scale,
117
- w2_scale=w2_weight_scale,
119
+ w2_scale=w2_scale,
120
w2_input_scale=w2_input_scale,
121
)
122
layer.w2_input_scale_inv = 1.0 / w2_input_scale
@@ -123,7 +125,7 @@ def register_scales_for_trtllm_fp8_per_tensor_moe(
125
layer.output2_scales_scalar = g2_alphas
126
127
126
-def apply_flashinfer_per_tensor_scale_fp8(
128
+def apply_fi_trtllm_fp8_per_tensor_moe(
129
layer: torch.nn.Module,
130
hidden_states: torch.Tensor,
131
router_logits: torch.Tensor,
@@ -139,16 +141,23 @@ def apply_flashinfer_per_tensor_scale_fp8(
141
import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
142
from vllm.model_executor.models.llama4 import Llama4MoE
143
144
+ # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe
145
assert (
146
hasattr(layer, "output1_scales_scalar")
147
and hasattr(layer, "output1_scales_gate_scalar")
148
and hasattr(layer, "output2_scales_scalar")
149
)
150
148
- assert layer.custom_routing_function == Llama4MoE.custom_routing_function, (
149
- "FusedMoE flashinfer kernels are only supported for Llama4"
151
+ # Added to the layer by: register_scales_for_trtllm_fp8_per_tensor_moe
152
+ assert (
153
+ hasattr(layer, "output1_scales_scalar")
154
+ and hasattr(layer, "output1_scales_gate_scalar")
155
+ and hasattr(layer, "output2_scales_scalar")
156
)
151
- return torch.ops.vllm.flashinfer_fused_moe_per_tensor_scale_fp8(
157
+
158
+ is_llama4 = layer.custom_routing_function == Llama4MoE.custom_routing_function
159
+ assert is_llama4, "FusedMoE flashinfer kernels are only supported for Llama4"
160
+ return torch.ops.vllm.fi_trtllm_fp8_per_tensor_moe(
161
routing_logits=router_logits,
162
routing_bias=routing_bias,
163
hidden_states=hidden_states,
@@ -221,50 +230,6 @@ def select_cutlass_fp8_gemm_impl(
230
)
231
232
224
-def flashinfer_cutlass_moe_fp8(
225
- hidden_states: torch.Tensor,
226
- layer: torch.nn.Module,
227
- topk_weights: torch.Tensor,
228
- topk_ids: torch.Tensor,
229
- inplace: bool = False,
230
- activation: str = "silu",
231
- global_num_experts: int = -1,
232
- expert_map: torch.Tensor | None = None,
233
- apply_router_weight_on_input: bool = False,
234
- use_deepseek_fp8_block_scale: bool = False,
235
- moe: FusedMoEConfig | None = None,
236
-) -> torch.Tensor:
237
- quant_config = layer.quant_method.get_fused_moe_quant_config(layer)
238
- assert quant_config is not None
239
-
240
- # Construct modular kernel with block-scale support when requested.
241
- fused_experts = mk.FusedMoEModularKernel(
242
- build_flashinfer_fp8_cutlass_moe_prepare_finalize(
243
- moe=moe, use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale
244
- ),
245
- select_cutlass_fp8_gemm_impl(
246
- moe=moe,
247
- quant_config=quant_config,
248
- out_dtype=hidden_states.dtype,
249
- use_deepseek_fp8_block_scale=use_deepseek_fp8_block_scale,
250
- ),
251
- moe_parallel_config=layer.moe_parallel_config,
252
- )
253
-
254
- return fused_experts(
255
- hidden_states,
256
- layer.w13_weight,
257
- layer.w2_weight,
258
- topk_weights,
259
- topk_ids,
260
- inplace=inplace,
261
- activation=activation,
262
- global_num_experts=global_num_experts,
263
- expert_map=expert_map,
264
- apply_router_weight_on_input=apply_router_weight_on_input,
265
- )
266
-
267
-
233
def get_flashinfer_moe_backend() -> FlashinferMoeBackend:
234
backend_map = {
235
"throughput": FlashinferMoeBackend.CUTLASS,
@@ -301,3 +266,104 @@ def is_flashinfer_supporting_global_sf(backend: FlashinferMoeBackend | None) ->
266
FlashinferMoeBackend.TENSORRT_LLM,
267
)
268
return backend in backends_supporting_global_sf
269
+
270
+
271
+def align_fp8_moe_weights_for_fi(
272
+ w13: torch.Tensor, w2: torch.Tensor, is_act_and_mul: bool
273
+) -> tuple[torch.Tensor, torch.Tensor, int]:
274
+ """Pad intermediate size so FlashInfer kernels' alignment constraints hold.
275
+
276
+ Some FlashInfer FP8 MoE kernels require the (gated) intermediate size
277
+ used for GEMM to be divisible by a small alignment value. When this is
278
+ not satisfied (e.g. with certain tensor-parallel sizes), we pad the
279
+ gate/up and down projection weights along the intermediate dim.
280
+ """
281
+
282
+ # Current local intermediate size (per partition) is the K dimension of
283
+ # the down projection.
284
+ num_experts, hidden_size, intermediate = w2.shape
285
+
286
+ min_alignment = 16
287
+ padded_intermediate = round_up(intermediate, min_alignment)
288
+
289
+ if padded_intermediate == intermediate:
290
+ return w13, w2, intermediate
291
+
292
+ logger.info_once(
293
+ "Padding intermediate size from %d to %d for up/down projection weights.",
294
+ intermediate,
295
+ padded_intermediate,
296
+ scope="local",
297
+ )
298
+
299
+ up_mult = 2 if is_act_and_mul else 1
300
+ padded_gate_up_dim = up_mult * padded_intermediate
301
+
302
+ # Pad w13 and w2 along its intermediate dimension.
303
+ padded_w13 = w13.new_zeros((num_experts, padded_gate_up_dim, hidden_size))
304
+ padded_w13[:, : w13.shape[1], :] = w13
305
+
306
+ padded_w2 = w2.new_zeros((num_experts, hidden_size, padded_intermediate))
307
+ padded_w2[:, :, :intermediate] = w2
308
+
309
+ return padded_w13, padded_w2, padded_intermediate
310
+
311
+
312
+def prepare_fp8_moe_layer_for_fi(
313
+ layer: torch.nn.Module,
314
+ w13: torch.Tensor,
315
+ w2: torch.Tensor,
316
+ w13_scale: torch.Tensor,
317
+ w13_input_scale: torch.Tensor | None,
318
+ w2_scale: torch.Tensor,
319
+ w2_input_scale: torch.Tensor | None,
320
+ is_trtllm: bool = False,
321
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
322
+ """
323
+ Convert Fp8 MoE weights to flashinfer kernel format
324
+
325
+ Note that for trtllm we update the model state dict
326
+ with the scale format needed for these kernels.
327
+
328
+ Note that for per-tensor, we update the layer's
329
+ intermediate size if the weights needed padding.
330
+ """
331
+
332
+ assert hasattr(layer.moe_config, "is_act_and_mul")
333
+ block_quant = (
334
+ hasattr(layer, "weight_block_size") and layer.weight_block_size is not None
335
+ )
336
+
337
+ # Some FI MoE kernels require internal alignment of 16
338
+ # for the gate-up proj. Pad the weights to respect this.
339
+ if not block_quant:
340
+ w13, w2, new_intermediate = align_fp8_moe_weights_for_fi(
341
+ w13,
342
+ w2,
343
+ layer.moe_config.is_act_and_mul,
344
+ )
345
+ layer.intermediate_size_per_partition = new_intermediate
346
+
347
+ # FI kernels require W31 layout rather than W13.
348
+ if layer.moe_config.is_act_and_mul:
349
+ w13 = swap_w13_to_w31(w13)
350
+ if block_quant:
351
+ w13_scale = swap_w13_to_w31(w13_scale)
352
+
353
+ # FI TRT-LLM FP8 per-tensor MoE kernel requires weight shuffle
354
+ # and registration of alpha scales. Note that we do not register
355
+ # as nn.Parameters since they are not needed for weight-reloading.
356
+ if is_trtllm and not block_quant:
357
+ assert w13_input_scale is not None
358
+ assert w2_input_scale is not None
359
+
360
+ rotate_weights_for_fi_trtllm_fp8_per_tensor_moe(w13, w2)
361
+ register_scales_for_trtllm_fp8_per_tensor_moe(
362
+ layer,
363
+ w13_scale=w13_scale,
364
+ w13_input_scale=w13_input_scale,
365
+ w2_scale=w2_scale,
366
+ w2_input_scale=w2_input_scale,
367
+ )
368
+
369
+ return w13, w2, w13_scale
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+71
-4
@@ -21,6 +21,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
21
)
22
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
23
CUTLASS_BLOCK_FP8_SUPPORTED,
24
+ all_close_1d,
25
+ per_tensor_dequantize,
26
)
27
from vllm.model_executor.parameter import (
28
BlockQuantScaleParameter,
@@ -1350,6 +1352,29 @@ def deepgemm_post_process_fp8_weight_block(
1352
return wq, dg_ws
1353
1354
1355
+def prepare_fp8_moe_layer_for_deepgemm(
1356
+ w13: torch.Tensor,
1357
+ w2: torch.Tensor,
1358
+ w13_scale: torch.Tensor,
1359
+ w2_scale: torch.Tensor,
1360
+ block_shape: tuple[int],
1361
+):
1362
+ w13, w13_scale = deepgemm_post_process_fp8_weight_block(
1363
+ wq=w13,
1364
+ ws=w13_scale,
1365
+ quant_block_shape=block_shape,
1366
+ use_e8m0=is_deep_gemm_e8m0_used(),
1367
+ )
1368
+ w2, w2_scale = deepgemm_post_process_fp8_weight_block(
1369
+ wq=w2,
1370
+ ws=w2_scale,
1371
+ quant_block_shape=block_shape,
1372
+ use_e8m0=is_deep_gemm_e8m0_used(),
1373
+ )
1374
+
1375
+ return w13, w2, w13_scale, w2_scale
1376
+
1377
+
1378
def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
1379
"""Pad the weight tensor. This is an optimization on ROCm platform, which
1380
can benefit from tensors located far enough from one another in memory"""
@@ -1584,7 +1609,49 @@ def maybe_post_process_fp8_weight_block(layer: torch.nn.Module):
1609
replace_parameter(layer, scale_attr, dg_weight_scale)
1610
1611
1587
-def expert_weight_is_col_major(x: torch.Tensor) -> bool:
1588
- assert x.dim() == 3
1589
- b, m, n = x.shape
1590
- return x.stride(0) == m * n and x.stride(1) == 1 and x.stride(2) == m
1612
+def process_fp8_weight_tensor_strategy_moe(
1613
+ weight: torch.Tensor,
1614
+ weight_scales: torch.Tensor,
1615
+ shard_size: int,
1616
+ num_experts: int,
1617
+ is_act_and_mul: bool = True,
1618
+) -> tuple[torch.Tensor, torch.Tensor]:
1619
+ """Process moe weights for tensor-wise quantization strategy."""
1620
+ max_scales = weight_scales.max(dim=1).values
1621
+
1622
+ # For w1 case (i.e. not w13): just collapse the last dim since
1623
+ # there is already just one scale per expert in this case.
1624
+ if not is_act_and_mul:
1625
+ assert weight_scales.shape[1] == 1
1626
+ return weight, weight_scales.max()
1627
+
1628
+ # For w13 case (common): require single scale for w13 per expert, but
1629
+ # on disk there is a scale for w1 and w3. Use the max to requantize.
1630
+ for expert_id in range(num_experts):
1631
+ start = 0
1632
+ for shard_id in range(2):
1633
+ dq_weight = per_tensor_dequantize(
1634
+ weight[expert_id][start : start + shard_size, :],
1635
+ weight_scales[expert_id][shard_id],
1636
+ )
1637
+ weight[expert_id][start : start + shard_size, :], _ = ops.scaled_fp8_quant(
1638
+ dq_weight, max_scales[expert_id]
1639
+ )
1640
+ start += shard_size
1641
+ return weight, max_scales
1642
+
1643
+
1644
+def process_fp8_input_tensor_strategy_moe(
1645
+ w13_input_scale: torch.Tensor,
1646
+ w2_input_scale: torch.Tensor,
1647
+) -> tuple[torch.Tensor, torch.Tensor]:
1648
+ """Process moe input scales for tensor-wise quantization strategy."""
1649
+
1650
+ if not all_close_1d(w13_input_scale) or not all_close_1d(w2_input_scale):
1651
+ logger.info_once(
1652
+ "Found input_scales that are not equal for "
1653
+ "fp8 MoE layer. Using the maximum across experts "
1654
+ "for each layer."
1655
+ )
1656
+
1657
+ return w13_input_scale.max(), w2_input_scale.max()
vllm/model_executor/layers/quantization/utils/marlin_utils.py
+1
-1
@@ -496,7 +496,7 @@ def get__quant_fp8_method() -> QuantFP8:
496
return _quant_fp8_method
497
498
499
-def get_marlin_input_dtype(prefix):
499
+def get_marlin_input_dtype(prefix: str | None = None):
500
if envs.VLLM_MARLIN_INPUT_DTYPE is None:
501
return
502
elif envs.VLLM_MARLIN_INPUT_DTYPE.lower() == "int8":
vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py
+16
-17
@@ -8,6 +8,7 @@ import vllm._custom_ops as ops
8
from vllm.logger import init_logger
9
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
10
USE_FP32_REDUCE_DEFAULT,
11
+ get_marlin_input_dtype,
12
marlin_make_workspace_new,
13
marlin_permute_bias,
14
marlin_permute_scales,
@@ -197,26 +198,28 @@ def prepare_fp8_layer_for_marlin(
198
replace_parameter(layer, "bias", bias)
199
200
200
-def prepare_moe_fp8_layer_for_marlin(
201
+def prepare_fp8_moe_layer_for_marlin(
202
layer: torch.nn.Module,
203
w13_weight: torch.Tensor,
204
w2_weight: torch.Tensor,
205
w13_weight_scale: torch.Tensor,
206
w2_weight_scale: torch.Tensor,
206
- input_dtype: torch.dtype | None = None,
207
-) -> tuple[
208
- torch.Tensor, # workspace
209
- torch.Tensor, # w13_weight
210
- torch.Tensor, # w2_weight
211
- torch.Tensor, # w13_weight_scale
212
- torch.Tensor, # w2_weight_scale
213
-]:
207
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
208
+ """
209
+ Shuffle weights and scales into marlin format.
210
+
211
+ Note that this function has the side effect of adding a `workspace`
212
+ attribute to the layer. This `workspace` does not need to be
213
+ registered as a Parameter as it is not used during weight reloading.
214
+ """
215
+
216
logger.warning_once(
217
"Your GPU does not have native support for FP8 computation but "
218
"FP8 quantization is being used. Weight-only FP8 compression will "
219
"be used leveraging the Marlin kernel. This may degrade "
220
"performance for compute-heavy workloads."
221
)
222
+ input_dtype = get_marlin_input_dtype()
223
if input_dtype is not None and input_dtype.itemsize == 1:
224
raise NotImplementedError("Marlin W8A8 is not supported.")
225
@@ -227,7 +230,9 @@ def prepare_moe_fp8_layer_for_marlin(
230
231
# WORKSPACE
232
device = layer.w13_weight.device
230
- workspace = marlin_make_workspace_new(device, 4)
233
+ # NOTE(rob): we do not need to register the workspace as a param
234
+ # because it is not used as part of the weight reloading process.
235
+ layer.workspace = marlin_make_workspace_new(device, 4)
236
perm = torch.empty(0, dtype=torch.int, device=device)
237
238
# WEIGHT
@@ -310,13 +315,7 @@ def prepare_moe_fp8_layer_for_marlin(
315
w13_weight_scale = permute_scales(w13_weight_scale, "w13")
316
w2_weight_scale = permute_scales(w2_weight_scale, "w2")
317
313
- return (
314
- workspace,
315
- w13_weight,
316
- w2_weight,
317
- w13_weight_scale,
318
- w2_weight_scale,
319
- )
318
+ return w13_weight, w2_weight, w13_weight_scale, w2_weight_scale
319
320
321
def pack_fp8_to_int32(