@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5963b98b46
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Kernel] Delegate construction of FusedMoEQuantConfig to FusedMoEMethodBase subclasses (#22537)

Signed-off-by: Bill Nell <bnell@redhat.com>

bnellnm committed Sep 17, 2025 at 19:43 UTC 5963b98b465007e3cfb0d39447e4459a8afa96dc
68 files changed +2675 -2503
benchmarks/kernels/benchmark_cutlass_fp4_moe.py
+36 -22
@@ -13,6 +13,10 @@ import torch.utils.benchmark as benchmark
13
14 from vllm import _custom_ops as ops
15 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
16 +from vllm.model_executor.layers.fused_moe.config import (
17 + fp8_w8a8_moe_quant_config,
18 + nvfp4_moe_quant_config,
19 +)
20 from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
21 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts, fused_topk
22 from vllm.scalar_type import scalar_types
@@ -140,6 +144,12 @@ def bench_run(
144 a_fp8_scale: torch.Tensor,
145 num_repeats: int,
146 ):
147 + quant_config = fp8_w8a8_moe_quant_config(
148 + w1_scale=w1_scale,
149 + w2_scale=w2_scale,
150 + a1_scale=a_fp8_scale,
151 + )
152 +
153 for _ in range(num_repeats):
154 fused_experts(
155 a,
@@ -147,10 +157,7 @@ def bench_run(
157 w2,
158 topk_weights,
159 topk_ids,
150 - use_fp8_w8a8=True,
151 - w1_scale=w1_scale,
152 - w2_scale=w2_scale,
153 - a1_scale=a_fp8_scale,
160 + quant_config=quant_config,
161 )
162
163 def run_cutlass_moe_fp4(
@@ -172,25 +179,27 @@ def bench_run(
179 device: torch.device,
180 num_repeats: int,
181 ):
182 + quant_config = nvfp4_moe_quant_config(
183 + a1_gscale=a1_gs,
184 + a2_gscale=a2_gs,
185 + w1_scale=w1_blockscale,
186 + w2_scale=w2_blockscale,
187 + g1_alphas=w1_gs,
188 + g2_alphas=w2_gs,
189 + )
190 for _ in range(num_repeats):
191 with nvtx.annotate("cutlass_moe_fp4", color="green"):
192 cutlass_moe_fp4(
193 a=a,
179 - a1_gscale=a1_gs,
180 - a2_gscale=a2_gs,
194 w1_fp4=w1_fp4,
182 - w1_blockscale=w1_blockscale,
183 - w1_alphas=w1_gs,
195 w2_fp4=w2_fp4,
185 - w2_blockscale=w2_blockscale,
186 - w2_alphas=w2_gs,
196 topk_weights=topk_weights,
197 topk_ids=topk_ids,
198 m=m,
199 n=n,
200 k=k,
201 e=num_experts,
193 - device=device,
202 + quant_config=quant_config,
203 )
204
205 def run_cutlass_from_graph(
@@ -211,26 +220,29 @@ def bench_run(
220 e: int,
221 device: torch.device,
222 ):
223 + quant_config = nvfp4_moe_quant_config(
224 + a1_gscale=a1_gs,
225 + a2_gscale=a2_gs,
226 + w1_scale=w1_blockscale,
227 + w2_scale=w2_blockscale,
228 + g1_alphas=w1_gs,
229 + g2_alphas=w2_gs,
230 + )
231 +
232 with set_current_vllm_config(
233 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
234 ):
235 return cutlass_moe_fp4(
236 a=a,
219 - a1_gscale=a1_gs,
237 w1_fp4=w1_fp4,
221 - w1_blockscale=w1_blockscale,
222 - w1_alphas=w1_alphas,
223 - a2_gscale=a2_gs,
238 w2_fp4=w2_fp4,
225 - w2_blockscale=w2_blockscale,
226 - w2_alphas=w2_alphas,
239 topk_weights=topk_weights,
240 topk_ids=topk_ids,
241 m=m,
242 n=n,
243 k=k,
244 e=num_experts,
233 - device=device,
245 + quant_config=quant_config,
246 )
247
248 def run_triton_from_graph(
@@ -246,16 +258,18 @@ def bench_run(
258 with set_current_vllm_config(
259 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
260 ):
261 + quant_config = fp8_w8a8_moe_quant_config(
262 + w1_scale=w1_scale,
263 + w2_scale=w2_scale,
264 + a1_scale=a_fp8_scale,
265 + )
266 return fused_experts(
267 a,
268 w1,
269 w2,
270 topk_weights,
271 topk_ids,
255 - use_fp8_w8a8=True,
256 - w1_scale=w1_scale,
257 - w2_scale=w2_scale,
258 - a1_scale=a_fp8_scale,
272 + quant_config=quant_config,
273 )
274
275 def replay_graph(graph, num_repeats):
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+27 -16
@@ -7,6 +7,7 @@ from benchmark_shapes import WEIGHT_SHAPES_MOE
7
8 from vllm import _custom_ops as ops
9 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
10 +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.fused_moe import (
13 fused_experts,
@@ -96,6 +97,11 @@ def bench_run(
97 a_scale: torch.Tensor,
98 num_repeats: int,
99 ):
100 + quant_config = fp8_w8a8_moe_quant_config(
101 + w1_scale=w1_scale,
102 + w2_scale=w2_scale,
103 + a1_scale=a_scale,
104 + )
105 for _ in range(num_repeats):
106 fused_experts(
107 a,
@@ -103,10 +109,7 @@ def bench_run(
109 w2,
110 topk_weights,
111 topk_ids,
106 - use_fp8_w8a8=True,
107 - w1_scale=w1_scale,
108 - w2_scale=w2_scale,
109 - a1_scale=a_scale,
112 + quant_config=quant_config,
113 )
114
115 def run_cutlass_moe(
@@ -125,6 +128,12 @@ def bench_run(
128 per_act_token: bool,
129 num_repeats: int,
130 ):
131 + quant_config = fp8_w8a8_moe_quant_config(
132 + w1_scale=w1_scale,
133 + w2_scale=w2_scale,
134 + per_act_token_quant=per_act_token,
135 + )
136 +
137 for _ in range(num_repeats):
138 cutlass_moe_fp8(
139 a,
@@ -132,14 +141,11 @@ def bench_run(
141 w2,
142 topk_weights,
143 topk_ids,
135 - w1_scale,
136 - w2_scale,
144 ab_strides1,
145 ab_strides2,
146 c_strides1,
147 c_strides2,
141 - per_act_token,
142 - a1_scale=None,
148 + quant_config=quant_config,
149 )
150
151 def run_cutlass_from_graph(
@@ -156,6 +162,12 @@ def bench_run(
162 topk_weights: torch.Tensor,
163 topk_ids: torch.Tensor,
164 ):
165 + quant_config = fp8_w8a8_moe_quant_config(
166 + w1_scale=w1_scale,
167 + w2_scale=w2_scale,
168 + per_act_token_quant=per_act_token,
169 + )
170 +
171 with set_current_vllm_config(
172 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
173 ):
@@ -165,14 +177,11 @@ def bench_run(
177 w2_q,
178 topk_weights,
179 topk_ids,
168 - w1_scale,
169 - w2_scale,
180 ab_strides1,
181 ab_strides2,
182 c_strides1,
183 c_strides2,
174 - per_act_token,
175 - a1_scale=None,
184 + quant_config=quant_config,
185 )
186
187 def run_triton_from_graph(
@@ -185,6 +194,11 @@ def bench_run(
194 w2_scale: torch.Tensor,
195 a_scale: torch.Tensor,
196 ):
197 + quant_config = fp8_w8a8_moe_quant_config(
198 + w1_scale=w1_scale,
199 + w2_scale=w2_scale,
200 + a1_scale=a_scale,
201 + )
202 with set_current_vllm_config(
203 VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
204 ):
@@ -194,10 +208,7 @@ def bench_run(
208 w2,
209 topk_weights,
210 topk_ids,
197 - use_fp8_w8a8=True,
198 - w1_scale=w1_scale,
199 - w2_scale=w2_scale,
200 - a1_scale=a_scale,
211 + quant_config=quant_config,
212 )
213
214 def replay_graph(graph, num_repeats):
benchmarks/kernels/benchmark_moe.py
+35 -38
@@ -14,6 +14,10 @@ import ray
14 import torch
15 from ray.experimental.tqdm_ray import tqdm
16
17 +from vllm.model_executor.layers.fused_moe.config import (
18 + FusedMoEQuantConfig,
19 + _get_config_dtype_str,
20 +)
21 from vllm.model_executor.layers.fused_moe.fused_moe import *
22 from vllm.platforms import current_platform
23 from vllm.transformers_utils.config import get_config
@@ -134,43 +138,36 @@ def benchmark_config(
138 def run():
139 from vllm.model_executor.layers.fused_moe import override_config
140
141 + if use_fp8_w8a8:
142 + quant_dtype = torch.float8_e4m3fn
143 + elif use_int8_w8a16:
144 + quant_dtype = torch.int8
145 + else:
146 + quant_dtype = None
147 +
148 + quant_config = FusedMoEQuantConfig.make(
149 + quant_dtype=quant_dtype,
150 + w1_scale=w1_scale,
151 + w2_scale=w2_scale,
152 + a1_scale=a1_scale,
153 + a2_scale=a2_scale,
154 + block_shape=block_quant_shape,
155 + )
156 +
157 with override_config(config):
138 - if use_deep_gemm:
139 - topk_weights, topk_ids, token_expert_indices = fused_topk(
140 - x, input_gating, topk, False
141 - )
142 - return fused_experts(
143 - x,
144 - w1,
145 - w2,
146 - topk_weights,
147 - topk_ids,
148 - inplace=True,
149 - use_fp8_w8a8=use_fp8_w8a8,
150 - w1_scale=w1_scale,
151 - w2_scale=w2_scale,
152 - a1_scale=a1_scale,
153 - a2_scale=a2_scale,
154 - block_shape=block_quant_shape,
155 - allow_deep_gemm=True,
156 - )
157 - else:
158 - fused_moe(
159 - x,
160 - w1,
161 - w2,
162 - input_gating,
163 - topk,
164 - renormalize=True,
165 - inplace=True,
166 - use_fp8_w8a8=use_fp8_w8a8,
167 - use_int8_w8a16=use_int8_w8a16,
168 - w1_scale=w1_scale,
169 - w2_scale=w2_scale,
170 - a1_scale=a1_scale,
171 - a2_scale=a2_scale,
172 - block_shape=block_quant_shape,
173 - )
158 + topk_weights, topk_ids, token_expert_indices = fused_topk(
159 + x, input_gating, topk, renormalize=not use_deep_gemm
160 + )
161 + return fused_experts(
162 + x,
163 + w1,
164 + w2,
165 + topk_weights,
166 + topk_ids,
167 + inplace=True,
168 + quant_config=quant_config,
169 + allow_deep_gemm=use_deep_gemm,
170 + )
171
172 # JIT compilation & warmup
173 run()
@@ -414,7 +411,7 @@ class BenchmarkWorker:
411 use_deep_gemm: bool = False,
412 ) -> tuple[dict[str, int], float]:
413 current_platform.seed_everything(self.seed)
417 - dtype_str = get_config_dtype_str(
414 + dtype_str = _get_config_dtype_str(
415 dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
416 )
417 # NOTE(woosuk): The current naming convention uses w2.shape[2], which
@@ -547,7 +544,7 @@ def save_configs(
544 block_quant_shape: list[int],
545 save_dir: str,
546 ) -> None:
550 - dtype_str = get_config_dtype_str(
547 + dtype_str = _get_config_dtype_str(
548 dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8
549 )
550
tests/kernels/moe/modular_kernel_tools/common.py
+62 -48
@@ -20,7 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (
20 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
21 from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
22
23 -from .mk_objects import (expert_info, make_fused_experts,
23 +from .mk_objects import (TestMoEQuantConfig, expert_info, make_fused_experts,
24 make_prepare_finalize, prepare_finalize_info)
25 from .parallel_utils import ProcessGroupInfo
26
@@ -40,7 +40,7 @@ class Config:
40 E: int
41 topks: Union[list[int], int]
42 dtype: torch.dtype
43 - quant_config: Optional[FusedMoEQuantConfig]
43 + quant_config: Optional[TestMoEQuantConfig]
44
45 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
46 fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute
@@ -52,7 +52,7 @@ class Config:
52
53 def __post_init__(self):
54 if self.quant_config is None:
55 - self.quant_config = FusedMoEQuantConfig()
55 + self.quant_config = TestMoEQuantConfig(None, False, False, None)
56
57 def describe(self) -> str:
58 s = ""
@@ -275,21 +275,19 @@ class WeightTensors:
275 or self.w1.dtype == torch.uint8 or self.w1.dtype == torch.int8)
276
277 def to_current_device(self):
278 - self.w1 = self.w1.to(device=torch.cuda.current_device())
279 - self.w2 = self.w2.to(device=torch.cuda.current_device())
278 + device = torch.cuda.current_device()
279 + self.w1 = self.w1.to(device=device)
280 + self.w2 = self.w2.to(device=device)
281
281 - if self.is_quantized():
282 - assert self.w1_scale is not None
283 - assert self.w2_scale is not None
284 - self.w1_scale = self.w1_scale.to(
285 - device=torch.cuda.current_device())
286 - self.w2_scale = self.w2_scale.to(
287 - device=torch.cuda.current_device())
282 + if self.w1_scale is not None:
283 + self.w1_scale = self.w1_scale.to(device=device)
284 + if self.w2_scale is not None:
285 + self.w2_scale = self.w2_scale.to(device=device)
286
287 if self.w1_gs is not None:
290 - assert self.w2_gs is not None
291 - self.w1_gs = self.w1_gs.to(device=torch.cuda.current_device())
292 - self.w2_gs = self.w2_gs.to(device=torch.cuda.current_device())
288 + self.w1_gs = self.w1_gs.to(device=device)
289 + if self.w2_gs is not None:
290 + self.w2_gs = self.w2_gs.to(device=device)
291
292 def slice_weights(self, rank: int,
293 num_local_experts: int) -> "WeightTensors":
@@ -297,20 +295,12 @@ class WeightTensors:
295 e = s + num_local_experts
296 w1 = self.w1[s:e, :, :]
297 w2 = self.w2[s:e, :, :]
300 -
301 - w1_scale, w2_scale = (None, None)
302 - if self.is_quantized():
303 - assert self.w1_scale is not None
304 - assert self.w2_scale is not None
305 - w1_scale = self.w1_scale[s:e, :, :]
306 - w2_scale = self.w2_scale[s:e, :, :]
307 -
308 - w1_gs = self.w1_gs
309 - w2_gs = self.w2_gs
310 - if w1_gs is not None:
311 - assert w2_gs is not None
312 - w1_gs = w1_gs[s:e]
313 - w2_gs = w2_gs[s:e]
298 + w1_scale = self.w1_scale[
299 + s:e, :, :] if self.w1_scale is not None else None
300 + w2_scale = self.w2_scale[
301 + s:e, :, :] if self.w2_scale is not None else None
302 + w1_gs = self.w1_gs[s:e] if self.w1_gs is not None else None
303 + w2_gs = self.w2_gs[s:e] if self.w2_gs is not None else None
304
305 return WeightTensors(w1, w2, w1_scale, w2_scale, w1_gs, w2_gs)
306
@@ -323,7 +313,8 @@ class WeightTensors:
313 in_dtype=config.dtype,
314 quant_dtype=config.quant_dtype,
315 block_shape=config.quant_block_shape,
326 - per_act_token_quant=config.is_per_out_ch_quant,
316 + per_out_ch_quant=config.
317 + is_per_act_token_quant, # or config.is_per_out_ch_quant
318 )
319 return WeightTensors(w1=w1,
320 w2=w2,
@@ -342,8 +333,6 @@ class RankTensors:
333 topk_ids: torch.Tensor
334 expert_map: Optional[torch.Tensor]
335
345 - quant_config: Optional[FusedMoEQuantConfig]
346 -
336 def describe(self):
337 s = ""
338 s += "== Rank Tensors: \n"
@@ -426,7 +415,6 @@ class RankTensors:
415 topk_weights=topk_weights,
416 topk_ids=topk_ids,
417 expert_map=expert_map,
429 - quant_config=config.quant_config,
418 )
419
420
@@ -522,10 +510,16 @@ def reference_moe_impl(config: Config, weights: WeightTensors,
510 and config.supports_apply_weight_on_input())
511
512
513 +def _make_gscale(num_experts: int) -> torch.Tensor:
514 + return torch.ones((num_experts, ),
515 + device=torch.cuda.current_device(),
516 + dtype=torch.float32)
517 +
518 +
519 def make_modular_kernel(
520 config: Config,
521 vllm_config: VllmConfig,
528 - weights: WeightTensors,
522 + quant_config: FusedMoEQuantConfig,
523 ) -> mk.FusedMoEModularKernel:
524
525 def next_power_of_2(x):
@@ -548,20 +542,20 @@ def make_modular_kernel(
542 num_local_experts=config.num_local_experts,
543 moe_parallel_config=moe_parallel_config,
544 in_dtype=config.dtype,
551 - quant_config=config.quant_config,
545 max_num_tokens=next_power_of_2(config.M),
546 )
547
548 # make modular kernel
549 prepare_finalize = make_prepare_finalize(config.prepare_finalize_type,
557 - config.all2all_backend(), moe)
550 + config.all2all_backend(), moe,
551 + quant_config)
552
553 fused_experts = make_fused_experts(
554 config.fused_experts_type,
555 moe,
556 + quant_config,
557 prepare_finalize.num_dispatchers(),
563 - weights.w1_gs,
564 - weights.w2_gs,
558 + config.N,
559 )
560
561 modular_kernel = mk.FusedMoEModularKernel(
@@ -583,12 +577,38 @@ def run_modular_kernel(
577 # weights for rank
578 rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
579
586 - mk = make_modular_kernel(config, vllm_config, weights)
580 + if config.quant_dtype == "nvfp4":
581 + gscale = _make_gscale(config.num_local_experts)
582 + else:
583 + gscale = None
584 +
585 + quant_config = FusedMoEQuantConfig.make(
586 + config.quant_dtype,
587 + w1_scale=rank_weights.w1_scale,
588 + w2_scale=rank_weights.w2_scale,
589 + a1_scale=rank_tensors.hidden_states_scale,
590 + g1_alphas=(1 / rank_weights.w1_gs)
591 + if rank_weights.w1_gs is not None else None,
592 + g2_alphas=(1 / rank_weights.w2_gs)
593 + if rank_weights.w2_gs is not None else None,
594 + a1_gscale=gscale,
595 + a2_gscale=gscale,
596 + block_shape=config.quant_block_shape,
597 + per_act_token_quant=config.is_per_act_token_quant,
598 + per_out_ch_quant=config.is_per_out_ch_quant,
599 + )
600 +
601 + mk = make_modular_kernel(config, vllm_config, quant_config)
602 +
603 + # impls might update the tensor in place
604 + hidden_states = rank_tensors.hidden_states.clone()
605 +
606 + topk_ids = rank_tensors.topk_ids.to(
607 + mk.prepare_finalize.topk_indices_dtype())
608
609 mk_kwargs = {
610 "hidden_states":
590 - rank_tensors.hidden_states.clone(
591 - ), # impls might update the tensor in place
611 + hidden_states,
612 "w1":
613 rank_weights.w1,
614 "w2":
@@ -596,15 +616,9 @@ def run_modular_kernel(
616 "topk_weights":
617 rank_tensors.topk_weights,
618 "topk_ids":
599 - rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype()),
619 + topk_ids,
620 "expert_map":
621 rank_tensors.expert_map,
602 - "w1_scale":
603 - rank_weights.w1_scale,
604 - "w2_scale":
605 - rank_weights.w2_scale,
606 - "a1_scale":
607 - rank_tensors.hidden_states_scale,
622 "global_num_experts":
623 config.E,
624 "apply_router_weight_on_input":
tests/kernels/moe/modular_kernel_tools/make_feature_matrix.py
+3 -2
@@ -10,7 +10,8 @@ import torch
10 from tqdm import tqdm
11
12 from vllm.config import VllmConfig, set_current_vllm_config
13 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
13 +from vllm.model_executor.layers.fused_moe.config import (
14 + FUSED_MOE_UNQUANTIZED_CONFIG)
15 from vllm.platforms import current_platform
16
17 from .common import (Config, RankTensors, WeightTensors, reference_moe_impl,
@@ -86,7 +87,7 @@ def make_feature_matrix(csv_file_path: str):
87 quant_config_dict = config_dict['quant_config']
88 del config_dict['quant_config']
89 if quant_config_dict is None:
89 - quant_config = FusedMoEQuantConfig(None)
90 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
91 quant_config_dict = asdict(quant_config)
92
93 config_dict |= quant_config_dict
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+79 -84
@@ -32,6 +32,14 @@ from vllm.utils.deep_gemm import is_deep_gemm_supported
32 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
33
34
35 +@dataclass
36 +class TestMoEQuantConfig:
37 + quant_dtype: Union[torch.dtype, str, None]
38 + per_out_ch_quant: bool
39 + per_act_token_quant: bool
40 + block_shape: Optional[list[int]]
41 +
42 +
43 @dataclass
44 class PrepareFinalizeInfo:
45 activation_format: mk.FusedMoEActivationFormat
@@ -66,7 +74,7 @@ common_float_types: list[Union[torch.dtype, str]] = [
74 torch.float8_e4m3fn, torch.bfloat16, torch.float16, torch.float32
75 ]
76 common_float_and_int_types = common_float_types + [torch.int8]
69 -nv_fp4_types = ["nvfp4"]
77 +nvfp4_types = ["nvfp4"]
78 fp8_types = [torch.float8_e4m3fn]
79
80
@@ -219,7 +227,7 @@ if (has_flashinfer_cutlass_fused_moe()
227 register_prepare_and_finalize(
228 FlashInferCutlassMoEPrepareAndFinalize,
229 standard_format,
222 - nv_fp4_types,
230 + nvfp4_types,
231 blocked_quantization_support=True,
232 backend=None,
233 force_multigpu=True,
@@ -229,7 +237,7 @@ if (has_flashinfer_cutlass_fused_moe()
237 register_experts(
238 FlashInferExperts,
239 standard_format,
232 - nv_fp4_types,
240 + nvfp4_types,
241 blocked_quantization_support=True,
242 supports_chunking=True,
243 # Note: this is a hack to get it to run for now
@@ -306,39 +314,39 @@ if cutlass_fp4_supported():
314 register_experts(
315 CutlassExpertsFp4,
316 standard_format,
309 - nv_fp4_types,
317 + nvfp4_types,
318 blocked_quantization_support=True,
319 supports_chunking=True,
320 supports_expert_map=False,
321 )
322
315 -MK_QUANT_CONFIGS = [
323 +MK_QUANT_CONFIGS: list[Optional[TestMoEQuantConfig]] = [
324 None,
325 # per-channel / per-column weights and per-tensor activations
318 - FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
319 - per_out_ch_quant=True,
320 - per_act_token_quant=False,
321 - block_shape=None),
326 + TestMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
327 + per_out_ch_quant=True,
328 + per_act_token_quant=False,
329 + block_shape=None),
330 # per-channel / per-column weights and per-token activations
323 - FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
324 - per_out_ch_quant=True,
325 - per_act_token_quant=True,
326 - block_shape=None),
331 + TestMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
332 + per_out_ch_quant=True,
333 + per_act_token_quant=True,
334 + block_shape=None),
335 # per-tensor weights and per-tensor activations
328 - FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
329 - per_out_ch_quant=False,
330 - per_act_token_quant=False,
331 - block_shape=None),
336 + TestMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
337 + per_out_ch_quant=False,
338 + per_act_token_quant=False,
339 + block_shape=None),
340 # per-tensor weights and per-token activations
333 - FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
334 - per_out_ch_quant=False,
335 - per_act_token_quant=True,
336 - block_shape=None),
341 + TestMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
342 + per_out_ch_quant=False,
343 + per_act_token_quant=True,
344 + block_shape=None),
345 # block-quantized weights and 128 block per-token activations
338 - FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
339 - per_out_ch_quant=False,
340 - per_act_token_quant=False,
341 - block_shape=[128, 128]),
346 + TestMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
347 + per_out_ch_quant=False,
348 + per_act_token_quant=False,
349 + block_shape=[128, 128]),
350 # TODO (varun) : Should we test the following combinations ?
351 # block-quantized weights and per-token activations
352 # block-quantized weights and per-tensor activations
@@ -346,33 +354,27 @@ MK_QUANT_CONFIGS = [
354
355 if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():
356 MK_QUANT_CONFIGS += [
349 - FusedMoEQuantConfig(quant_dtype="nvfp4",
350 - per_out_ch_quant=False,
351 - per_act_token_quant=False,
352 - block_shape=None),
357 + TestMoEQuantConfig(quant_dtype="nvfp4",
358 + per_out_ch_quant=False,
359 + per_act_token_quant=False,
360 + block_shape=None),
361 ]
362
363
356 -def _make_gscale(num_experts: int) -> torch.Tensor:
357 - return torch.ones((num_experts, ),
358 - device=torch.cuda.current_device(),
359 - dtype=torch.float32)
360 -
361 -
364 def make_prepare_finalize(
365 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
366 backend: Optional[str],
367 moe: FusedMoEConfig,
368 + quant_config: FusedMoEQuantConfig,
369 ) -> mk.FusedMoEPrepareAndFinalize:
370 if backend != "naive" and backend is not None:
368 - prepare_finalize = FusedMoEMethodBase._maybe_make_prepare_finalize(moe)
371 + prepare_finalize = FusedMoEMethodBase._maybe_make_prepare_finalize(
372 + moe, quant_config)
373 assert prepare_finalize is not None
374 return prepare_finalize
375 elif prepare_finalize_type == FlashInferCutlassMoEPrepareAndFinalize:
376 return FlashInferCutlassMoEPrepareAndFinalize(
373 - use_dp=moe.moe_parallel_config.dp_size > 1,
374 - a1_gscale=_make_gscale(moe.num_local_experts),
375 - )
377 + use_dp=moe.moe_parallel_config.dp_size > 1)
378 else:
379 return MoEPrepareAndFinalizeNoEP()
380
@@ -383,34 +385,39 @@ def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:
385 return t[s:e]
386
387
388 +def make_cutlass_strides(
389 + e: int,
390 + n: int,
391 + k: int,
392 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
393 + ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
394 + ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
395 + c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
396 + c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
397 + return ab_strides1, ab_strides2, c_strides1, c_strides2
398 +
399 +
400 def make_fused_experts(
401 fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
402 moe: FusedMoEConfig,
403 + quant_config: FusedMoEQuantConfig,
404 num_dispatchers: int,
390 - w1_gs: Optional[torch.Tensor],
391 - w2_gs: Optional[torch.Tensor],
405 + N: int,
406 ) -> mk.FusedMoEPermuteExpertsUnpermute:
407
394 - use_fp8 = moe.quant_dtype == torch.float8_e4m3fn
408 batch_kwargs = {
409 "max_num_tokens": moe.max_num_tokens,
410 "num_dispatchers": num_dispatchers,
411 }
412 quant_kwargs = {
400 - "use_fp8_w8a8": use_fp8,
401 - "use_int8_w8a8": False,
402 - "use_int8_w8a16": False,
403 - "use_int4_w4a16": False,
404 - "block_shape": moe.block_shape,
405 - "per_act_token_quant": moe.per_act_token_quant,
413 + "quant_config": quant_config,
414 }
415 deepgemm_kwargs = {"allow_deep_gemm": has_deep_gemm()}
416
417 + torch.set_printoptions(threshold=0, edgeitems=0, linewidth=10000)
418 +
419 if fused_experts_type == BatchedDeepGemmExperts:
410 - kwargs = batch_kwargs | {
411 - "block_shape": moe.block_shape,
412 - "per_act_token_quant": moe.per_act_token_quant,
413 - }
420 + kwargs = batch_kwargs | quant_kwargs
421 print(f"Making BatchedDeepGemmExperts {kwargs} ...")
422 experts = BatchedDeepGemmExperts(**kwargs)
423 elif fused_experts_type == BatchedTritonExperts:
@@ -422,8 +429,8 @@ def make_fused_experts(
429 print(f"Making BatchedTritonOrDeepGemmExperts {kwargs} ...")
430 experts = BatchedTritonOrDeepGemmExperts(**kwargs)
431 elif fused_experts_type == DeepGemmExperts:
425 - print("Making DeepGemmExperts () ...")
426 - experts = DeepGemmExperts()
432 + print("Making DeepGemmExperts {quant_config} ...")
433 + experts = DeepGemmExperts(quant_config)
434 elif fused_experts_type == TritonExperts:
435 kwargs = quant_kwargs
436 print(f"Making TritonExperts {kwargs} ...")
@@ -437,62 +444,50 @@ def make_fused_experts(
444 print(f"Making NaiveBatchedExperts {kwargs} ...")
445 experts = NaiveBatchedExperts(**kwargs)
446 elif fused_experts_type == CutlassExpertsFp8:
447 + strides = make_cutlass_strides(moe.num_experts, N, moe.hidden_dim)
448 kwargs = {
449 "out_dtype": moe.in_dtype,
442 - "per_act_token_quant": moe.per_act_token_quant,
443 - "per_out_ch_quant": moe.per_out_ch_quant,
444 - "block_shape": moe.block_shape,
445 - }
450 + "ab_strides1": strides[0],
451 + "ab_strides2": strides[1],
452 + "c_strides1": strides[2],
453 + "c_strides2": strides[3],
454 + } | quant_kwargs
455 print(f"Making CutlassExpertsFp8 {kwargs} ...")
456 experts = CutlassExpertsFp8(**kwargs)
457 elif fused_experts_type == CutlassBatchedExpertsFp8:
458 + strides = make_cutlass_strides(moe.num_experts, N, moe.hidden_dim)
459 kwargs = {
460 "max_experts_per_worker": moe.num_local_experts,
461 "num_dispatchers": num_dispatchers,
462 "out_dtype": moe.in_dtype,
453 - "per_act_token_quant": moe.per_act_token_quant,
454 - "per_out_ch_quant": moe.per_out_ch_quant,
455 - "block_shape": moe.block_shape,
456 - }
463 + "ab_strides1": strides[0],
464 + "ab_strides2": strides[1],
465 + "c_strides1": strides[2],
466 + "c_strides2": strides[3],
467 + } | quant_kwargs
468 print(f"Making CutlassBatchedExpertsFp8 {kwargs} ...")
469 experts = CutlassBatchedExpertsFp8(**kwargs)
470 elif fused_experts_type == CutlassExpertsFp4:
460 - assert w1_gs is not None and w2_gs is not None
461 - num_experts = moe.num_local_experts
462 - rank = moe.moe_parallel_config.dp_rank
471 kwargs = {
464 - "g1_alphas": _slice(rank, num_experts, (1 / w1_gs)),
465 - "g2_alphas": _slice(rank, num_experts, (1 / w2_gs)),
466 - "a1_gscale": _make_gscale(num_experts),
467 - "a2_gscale": _make_gscale(num_experts),
468 - "max_experts_per_worker": num_experts,
469 - "out_dtype": moe.in_dtype,
470 - "per_act_token_quant": moe.per_act_token_quant,
471 - "per_out_ch_quant": moe.per_out_ch_quant,
472 - "block_shape": moe.block_shape,
472 + "max_experts_per_worker": moe.num_local_experts,
473 "num_dispatchers": num_dispatchers,
474 - }
474 + "out_dtype": moe.in_dtype,
475 + } | quant_kwargs
476 print(f"Making CutlassExpertsFp4 {kwargs} ...")
477 experts = CutlassExpertsFp4(**kwargs)
478 elif fused_experts_type == FlashInferExperts:
478 - assert w1_gs is not None and w2_gs is not None
479 - num_experts = moe.num_local_experts
480 - rank = moe.moe_parallel_config.dp_rank
479 kwargs = {
482 - "g1_alphas": _slice(rank, num_experts, (1 / w1_gs)),
483 - "g2_alphas": _slice(rank, num_experts, (1 / w2_gs)),
484 - "a1_gscale": _make_gscale(num_experts),
485 - "a2_gscale": _make_gscale(num_experts),
480 "out_dtype": moe.in_dtype,
487 - "quant_dtype": "nvfp4",
481 "ep_rank": moe.ep_rank,
482 "ep_size": moe.ep_size,
483 "tp_rank": moe.tp_rank,
484 "tp_size": moe.tp_size,
492 - }
485 + } | quant_kwargs
486 print(f"Making FlashInferExperts {kwargs} ...")
487 experts = FlashInferExperts(**kwargs)
488 else:
489 raise RuntimeError(f"Unknown fused experts type: {fused_experts_type}")
490
491 + torch.set_printoptions(threshold=1000, edgeitems=5, linewidth=80)
492 +
493 return experts
tests/kernels/moe/test_batched_deepgemm.py
+11 -9
@@ -6,6 +6,8 @@ import torch
6
7 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8 BatchedDeepGemmExperts)
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + fp8_w8a8_moe_quant_config)
11 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
12 BatchedPrepareAndFinalize, BatchedTritonExperts)
13 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -56,13 +58,18 @@ def test_batched_deepgemm_vs_triton(E: int, T: int, K: int, N: int, topk: int,
58 rank=0,
59 )
60
61 + quant_config = fp8_w8a8_moe_quant_config(
62 + w1_scale=w1_s,
63 + w2_scale=w2_s,
64 + per_act_token_quant=False,
65 + block_shape=BLOCK_SIZE,
66 + )
67 +
68 # triton (reference)
69 triton_experts = BatchedTritonExperts(
70 max_num_tokens=max_num_tokens,
71 num_dispatchers=1,
63 - use_fp8_w8a8=True,
64 - per_act_token_quant=False,
65 - block_shape=BLOCK_SIZE,
72 + quant_config=quant_config,
73 )
74 mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts)
75
@@ -73,8 +80,6 @@ def test_batched_deepgemm_vs_triton(E: int, T: int, K: int, N: int, topk: int,
80 topk_weights=topk_weights,
81 topk_ids=topk_ids,
82 inplace=False,
76 - w1_scale=w1_s,
77 - w2_scale=w2_s,
83 global_num_experts=E,
84 )
85
@@ -82,8 +87,7 @@ def test_batched_deepgemm_vs_triton(E: int, T: int, K: int, N: int, topk: int,
87 deepgemm_experts = BatchedDeepGemmExperts(
88 max_num_tokens=max_num_tokens,
89 num_dispatchers=1,
85 - block_shape=BLOCK_SIZE,
86 - per_act_token_quant=False,
90 + quant_config=quant_config,
91 )
92 mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts)
93
@@ -94,8 +98,6 @@ def test_batched_deepgemm_vs_triton(E: int, T: int, K: int, N: int, topk: int,
98 topk_weights=topk_weights,
99 topk_ids=topk_ids,
100 inplace=False,
97 - w1_scale=w1_s,
98 - w2_scale=w2_s,
101 global_num_experts=E,
102 )
103
tests/kernels/moe/test_batched_moe.py
+2 -2
@@ -140,7 +140,7 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
140 in_dtype=act_dtype,
141 quant_dtype=quant_dtype,
142 block_shape=block_shape,
143 - per_act_token_quant=per_act_token_quant,
143 + per_out_ch_quant=per_act_token_quant,
144 )
145
146 out_shape = (num_experts, max_tokens_per_expert, N)
@@ -250,7 +250,7 @@ def test_fused_moe_batched_experts(
250 block_shape=block_shape,
251 in_dtype=act_dtype,
252 quant_dtype=quant_dtype,
253 - per_act_token_quant=per_act_token_quant,
253 + per_out_ch_quant=per_act_token_quant,
254 )
255
256 if input_scales and quant_dtype is not None:
tests/kernels/moe/test_block_fp8.py
+32 -49
@@ -4,7 +4,7 @@
4 import pytest
5 import torch
6
7 -from tests.kernels.moe.utils import make_test_weights
7 +from tests.kernels.moe.utils import make_test_quant_config, make_test_weights
8 from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
9 native_w8a8_block_matmul)
10 from vllm.config import VllmConfig, set_current_vllm_config
@@ -161,22 +161,17 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed,
161 a = torch.randn((M, K), dtype=dtype) / 10
162 score = torch.randn((M, E), dtype=dtype)
163
164 - (_, w1, w1_s, _), (_, w2, w2_s,
165 - _) = make_test_weights(E,
166 - N,
167 - K,
168 - dtype,
169 - torch.float8_e4m3fn,
170 - per_act_token_quant=False,
171 - block_shape=block_size)
172 -
173 - m_fused_moe = modular_triton_fused_moe(use_fp8_w8a8=True,
174 - use_int8_w8a8=False,
175 - use_int8_w8a16=False,
176 - use_int4_w4a16=False,
177 - use_mxfp4_w4a4=False,
178 - per_act_token_quant=False,
179 - block_shape=block_size)
164 + w1, w2, quant_config = make_test_quant_config(
165 + E,
166 + N,
167 + K,
168 + dtype,
169 + quant_dtype=torch.float8_e4m3fn,
170 + per_act_token_quant=False,
171 + block_shape=block_size,
172 + )
173 +
174 + m_fused_moe = modular_triton_fused_moe(quant_config)
175
176 topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
177
@@ -186,37 +181,24 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed,
181 a,
182 w1,
183 w2,
189 - w1_s,
190 - w2_s,
184 + quant_config.w1_scale,
185 + quant_config.w2_scale,
186 topk_weights,
187 topk_ids,
188 block_size,
189 )
190
196 - out = fused_experts(
197 - a,
198 - w1,
199 - w2,
200 - topk_weights,
201 - topk_ids,
202 - use_fp8_w8a8=True,
203 - w1_scale=w1_s,
204 - w2_scale=w2_s,
205 - block_shape=block_size,
206 - )
191 + out = fused_experts(a,
192 + w1,
193 + w2,
194 + topk_weights,
195 + topk_ids,
196 + quant_config=quant_config)
197
208 - m_out = m_fused_moe(
209 - a,
210 - w1,
211 - w2,
212 - topk_weights,
213 - topk_ids,
214 - w1_scale=w1_s,
215 - w2_scale=w2_s,
216 - )
198 + m_out = m_fused_moe(a, w1, w2, topk_weights, topk_ids)
199
218 - # 0.039 only needed for [40000-4608-7168-2-1-block_size852-dtype852-0]
219 - tol = 0.035 if M < 40000 else 0.039
200 + # 0.039 only needed for M >= 8192
201 + tol = 0.035 if M < 8192 else 0.039
202 torch.testing.assert_close(out, ref_out, atol=tol, rtol=tol)
203 torch.testing.assert_close(m_out, ref_out, atol=tol, rtol=tol)
204
@@ -248,14 +230,15 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
230 a = torch.randn((M, K), dtype=dtype) / 10
231 score = torch.randn((M, E), dtype=dtype)
232
251 - (_, w1, w1_s, _), (_, w2, w2_s,
252 - _) = make_test_weights(E,
253 - N,
254 - K,
255 - dtype,
256 - torch.float8_e4m3fn,
257 - per_act_token_quant=False,
258 - block_shape=block_size)
233 + (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
234 + E,
235 + N,
236 + K,
237 + dtype,
238 + torch.float8_e4m3fn,
239 + per_out_ch_quant=False,
240 + block_shape=block_size,
241 + )
242
243 # Note: for now use_compile will error out if the problem size is
244 # large enough to trigger chunking. I'm leaving the flag and
tests/kernels/moe/test_block_int8.py
+22 -25
@@ -4,12 +4,12 @@
4 import pytest
5 import torch
6
7 -from tests.kernels.moe.utils import make_test_weights
7 +from tests.kernels.moe.utils import make_test_quant_config
8 from tests.kernels.quant_utils import (native_per_token_group_quant_int8,
9 native_w8a8_block_matmul)
10 from vllm.config import VllmConfig, set_current_vllm_config
11 from vllm.model_executor.layers.activation import SiluAndMul
12 -from vllm.model_executor.layers.fused_moe import fused_moe
12 +from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
13 from vllm.platforms import current_platform
14
15 if current_platform.get_device_capability() < (7, 0):
@@ -50,7 +50,7 @@ MNK_FACTORS = [
50 (2048, 128, 128),
51 (2048, 1024, 7168),
52 (2048, 4096, 512),
53 - (2048, 4096, 7168),
53 + (2048, 4096, 4096),
54 ]
55
56 E = [8, 24]
@@ -117,31 +117,28 @@ def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
117
118 a = torch.randn((M, K), dtype=dtype) / 10
119 score = torch.randn((M, E), dtype=dtype)
120 -
121 - (_, w1, w1_s, _), (_, w2, w2_s,
122 - _) = make_test_weights(E,
123 - N,
124 - K,
125 - dtype,
126 - torch.int8,
127 - per_act_token_quant=False,
128 - block_shape=block_size)
120 + topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
121 +
122 + w1, w2, quant_config = make_test_quant_config(
123 + E,
124 + N,
125 + K,
126 + dtype,
127 + quant_dtype=torch.int8,
128 + per_act_token_quant=False,
129 + block_shape=block_size,
130 + )
131
132 # Set the context to avoid lots of warning spam.
133 with set_current_vllm_config(vllm_config):
132 - out = fused_moe(
133 - a,
134 - w1,
135 - w2,
136 - score,
137 - topk,
138 - renormalize=False,
139 - use_int8_w8a8=True,
140 - w1_scale=w1_s,
141 - w2_scale=w2_s,
142 - block_shape=block_size,
143 - )
144 - ref_out = torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk,
134 + out = fused_experts(a,
135 + w1,
136 + w2,
137 + topk_weights,
138 + topk_ids,
139 + quant_config=quant_config)
140 + ref_out = torch_w8a8_block_int8_moe(a, w1, w2, quant_config.w1_scale,
141 + quant_config.w2_scale, score, topk,
142 block_size)
143
144 # Check results
tests/kernels/moe/test_cutlass_moe.py
+42 -11
@@ -1,5 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import copy
4 import dataclasses
5 from math import prod
6 from typing import Optional
@@ -9,6 +10,8 @@ import torch
10
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, fp8_w8a8_moe_quant_config)
15 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
16 cutlass_moe_fp8, run_cutlass_moe_fp8)
17 from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts,
@@ -154,7 +157,7 @@ def run_with_expert_maps(num_experts: int, num_local_experts: int,
157 def slice_experts():
158 slice_params = [
159 "w1_q", "w2_q", "ab_strides1", "ab_strides2", "c_strides1",
157 - "c_strides2", "w1_scale", "w2_scale"
160 + "c_strides2"
161 ]
162 full_tensors = {
163 k: v
@@ -162,6 +165,8 @@ def run_with_expert_maps(num_experts: int, num_local_experts: int,
165 if k in slice_params and k in cutlass_moe_kwargs
166 }
167
168 + quant_config = cutlass_moe_kwargs["quant_config"]
169 +
170 for i in range(0, num_experts, num_local_experts):
171 s, e = i, i + num_local_experts
172
@@ -178,6 +183,12 @@ def run_with_expert_maps(num_experts: int, num_local_experts: int,
183 for k, t in full_tensors.items():
184 cutlass_moe_kwargs[k] = t[s:e]
185
186 + new_quant_config = copy.deepcopy(quant_config)
187 + new_quant_config._w1.scale = quant_config.w1_scale[s:e]
188 + new_quant_config._w2.scale = quant_config.w2_scale[s:e]
189 +
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"])
@@ -191,6 +202,7 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
202 topk_weights: torch.Tensor,
203 topk_ids: torch.Tensor,
204 per_act_token: bool,
205 + per_out_ch: bool,
206 num_local_experts: Optional[int] = None) -> torch.Tensor:
207 assert not any([
208 t is None for t in [
@@ -199,20 +211,27 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
211 ]
212 ])
213
214 + quant_config = fp8_w8a8_moe_quant_config(
215 + w1_scale=moe_tensors.w1_scale,
216 + w2_scale=moe_tensors.w2_scale,
217 + per_act_token_quant=per_act_token,
218 + per_out_ch_quant=per_out_ch,
219 + # Set to moe_tensors.a_scale iff static scales + per tensor.
220 + # This is not currently being tested.
221 + a1_scale=None,
222 + )
223 +
224 kwargs = {
225 'a': moe_tensors.a,
226 'w1_q': moe_tensors.w1_q, # type: ignore[union-attr]
227 'w2_q': moe_tensors.w2_q, # type: ignore[union-attr]
228 'topk_weights': topk_weights,
229 'topk_ids': topk_ids,
208 - 'w1_scale': moe_tensors.w1_scale,
209 - 'w2_scale': moe_tensors.w2_scale,
230 'ab_strides1': moe_tensors.ab_strides1,
231 'ab_strides2': moe_tensors.ab_strides2,
232 'c_strides1': moe_tensors.c_strides1,
233 'c_strides2': moe_tensors.c_strides2,
214 - 'per_act_token': per_act_token,
215 - 'a1_scale': None #moe_tensors.a_scale
234 + 'quant_config': quant_config,
235 }
236
237 num_experts = moe_tensors.w1.size(0)
@@ -261,16 +280,23 @@ def test_cutlass_moe_8_bit_no_graph(
280
281 # Note that we are using the dequantized versions of the tensors.
282 # Using a, w1 and w2 directly results in minor output differences.
264 - triton_output = fused_experts(mt.a_d, mt.w1_d, mt.w2_d, topk_weights,
265 - topk_ids)
283 +
284 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
285 + triton_output = fused_experts(mt.a_d,
286 + mt.w1_d,
287 + mt.w2_d,
288 + topk_weights,
289 + topk_ids,
290 + quant_config=quant_config)
291
292 if ep_size is not None:
293 assert e % ep_size == 0, "Cannot distribute experts evenly"
294 number_local_experts = e // ep_size
295 else:
296 number_local_experts = None
297 +
298 cutlass_output = run_8_bit(mt, topk_weights, topk_ids, per_act_token,
273 - number_local_experts)
299 + per_out_ch, number_local_experts)
300
301 # Note 5.5 only needed for larger problem sizes, 5 works ok for
302 # the rest.
@@ -315,14 +341,19 @@ def test_cutlass_moe_8_bit_cuda_graph(
341
342 # Note that we are using the dequantized versions of the tensors.
343 # Using a, w1 and w2 directly results in minor output differences.
318 - triton_output = fused_experts(mt.a_d, mt.w1_d, mt.w2_d, topk_weights,
319 - topk_ids)
344 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
345 + triton_output = fused_experts(mt.a_d,
346 + mt.w1_d,
347 + mt.w2_d,
348 + topk_weights,
349 + topk_ids,
350 + quant_config=quant_config)
351
352 stream = torch.cuda.Stream()
353 graph = torch.cuda.CUDAGraph()
354 with torch.cuda.graph(graph, stream=stream):
355 cutlass_output = run_8_bit(mt, topk_weights, topk_ids,
325 - per_act_token)
356 + per_act_token, per_out_ch)
357
358 torch.cuda.synchronize()
359 graph.replay()
tests/kernels/moe/test_deepep_deepgemm_moe.py
+52 -35
@@ -15,6 +15,8 @@ from torch.distributed import ProcessGroup
15 from typing_extensions import ParamSpec
16
17 from vllm.config import VllmConfig, set_current_vllm_config
18 +from vllm.model_executor.layers.fused_moe.config import (
19 + FusedMoEQuantConfig, fp8_w8a8_moe_quant_config)
20 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
21 from vllm.model_executor.layers.fused_moe.modular_kernel import (
22 FusedMoEModularKernel)
@@ -71,9 +73,12 @@ def make_block_quant_fp8_weights(
73 Return weights w1q, w2q, w1_scale, w2_scale
74 """
75 (_, w1q, w1_scale, _), (_, w2q, w2_scale,
74 - _) = make_test_weights(e, n, k, torch.bfloat16,
76 + _) = make_test_weights(e,
77 + n,
78 + k,
79 + torch.bfloat16,
80 torch.float8_e4m3fn,
76 - block_size)
81 + block_shape=block_size)
82 return w1q, w2q, w1_scale, w2_scale
83
84
@@ -130,10 +135,11 @@ class TestTensors:
135 config=config)
136
137
133 -def make_ll_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
134 - max_tokens_per_rank: int, dp_size: int,
135 - hidden_size: int, q_dtype: Optional[torch.dtype],
136 - test_config: TestConfig) -> FusedMoEModularKernel:
138 +def make_ll_modular_kernel(
139 + pg: ProcessGroup, pgi: ProcessGroupInfo, max_tokens_per_rank: int,
140 + dp_size: int, hidden_size: int, q_dtype: Optional[torch.dtype],
141 + test_config: TestConfig,
142 + quant_config: FusedMoEQuantConfig) -> FusedMoEModularKernel:
143
144 assert test_config.low_latency
145 assert test_config.use_fp8_dispatch is not None
@@ -154,17 +160,18 @@ def make_ll_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
160 fused_experts = BatchedDeepGemmExperts(
161 max_num_tokens=max_tokens_per_rank,
162 num_dispatchers=pgi.world_size // dp_size,
157 - block_shape=test_config.block_size,
158 - per_act_token_quant=test_config.per_act_token_quant)
163 + quant_config=quant_config,
164 + )
165 mk = FusedMoEModularKernel(prepare_finalize=a2a,
166 fused_experts=fused_experts)
167 return mk
168
169
164 -def make_ht_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
165 - dp_size: int, num_local_experts: int,
166 - q_dtype: Optional[torch.dtype],
167 - test_config: TestConfig) -> FusedMoEModularKernel:
170 +def make_ht_modular_kernel(
171 + pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
172 + num_local_experts: int, q_dtype: Optional[torch.dtype],
173 + test_config: TestConfig,
174 + quant_config: FusedMoEQuantConfig) -> FusedMoEModularKernel:
175
176 assert not test_config.low_latency
177 assert test_config.use_fp8_dispatch is None
@@ -178,15 +185,16 @@ def make_ht_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
185 q_dtype=q_dtype,
186 block_shape=test_config.block_size)
187
181 - fused_experts = DeepGemmExperts()
188 + fused_experts = DeepGemmExperts(quant_config)
189 mk = FusedMoEModularKernel(prepare_finalize=a2a,
190 fused_experts=fused_experts)
191 return mk
192
193
187 -def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
188 - num_local_experts: int,
189 - test_tensors: TestTensors) -> FusedMoEModularKernel:
194 +def make_modular_kernel(
195 + pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
196 + num_local_experts: int, test_tensors: TestTensors,
197 + quant_config: FusedMoEQuantConfig) -> FusedMoEModularKernel:
198
199 q_dtype = torch.float8_e4m3fn
200 test_config = test_tensors.config
@@ -204,10 +212,16 @@ def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
212 dp_size=dp_size,
213 hidden_size=hidden_size,
214 q_dtype=q_dtype,
207 - test_config=test_config)
215 + test_config=test_config,
216 + quant_config=quant_config)
217 else:
209 - mk = make_ht_modular_kernel(pg, pgi, dp_size, num_local_experts,
210 - q_dtype, test_config)
218 + mk = make_ht_modular_kernel(pg,
219 + pgi,
220 + dp_size,
221 + num_local_experts,
222 + q_dtype,
223 + test_config,
224 + quant_config=quant_config)
225
226 return mk
227
@@ -233,17 +247,23 @@ def deepep_deepgemm_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
247 return expert_map.to(device=torch.cuda.current_device(),
248 dtype=torch.int32)
249
250 + quant_config = fp8_w8a8_moe_quant_config(
251 + w1_scale=w1_scale,
252 + w2_scale=w2_scale,
253 + # Low-Latency kernels can't dispatch scales.
254 + a1_scale=(None if test_config.low_latency else
255 + test_tensors.rank_token_scales),
256 + block_shape=test_config.block_size,
257 + )
258 +
259 # Make modular kernel
260 mk: FusedMoEModularKernel = make_modular_kernel(
261 pg=pg,
262 pgi=pgi,
263 dp_size=dp_size,
264 num_local_experts=num_local_experts,
242 - test_tensors=test_tensors)
243 -
244 - # Low-Latency kernels can't dispatch scales.
245 - a1_scale = (None
246 - if test_config.low_latency else test_tensors.rank_token_scales)
265 + test_tensors=test_tensors,
266 + quant_config=quant_config)
267
268 out = mk.forward(hidden_states=test_tensors.rank_tokens,
269 w1=w1,
@@ -254,12 +274,6 @@ def deepep_deepgemm_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
274 activation="silu",
275 global_num_experts=num_experts,
276 expert_map=build_expert_map(),
257 - w1_scale=w1_scale,
258 - w2_scale=w2_scale,
259 - w1_zp=None,
260 - w2_zp=None,
261 - a1_scale=a1_scale,
262 - a2_scale=None,
277 apply_router_weight_on_input=False)
278 return out
279
@@ -269,6 +283,13 @@ def triton_impl(a: torch.Tensor, topk_ids: torch.Tensor,
283 w1_scale: torch.Tensor, w2_scale: torch.Tensor,
284 a1_scale: torch.Tensor, block_shape: list[int]):
285
286 + quant_config = fp8_w8a8_moe_quant_config(
287 + w1_scale=w1_scale,
288 + w2_scale=w2_scale,
289 + a1_scale=a1_scale,
290 + block_shape=block_shape,
291 + )
292 +
293 return fused_experts(
294 hidden_states=a,
295 w1=w1,
@@ -276,11 +297,7 @@ def triton_impl(a: torch.Tensor, topk_ids: torch.Tensor,
297 topk_weights=topk_weights,
298 topk_ids=topk_ids,
299 inplace=False,
279 - use_fp8_w8a8=True,
280 - w1_scale=w1_scale,
281 - w2_scale=w2_scale,
282 - a1_scale=a1_scale,
283 - block_shape=block_shape,
300 + quant_config=quant_config,
301 # Make sure this is set to False so we
302 # don't end up comparing the same implementation.
303 allow_deep_gemm=False)
tests/kernels/moe/test_deepep_moe.py
+33 -37
@@ -15,6 +15,7 @@ from vllm import _custom_ops as ops
15 from vllm.config import VllmConfig, set_current_vllm_config
16 from vllm.model_executor.layers.activation import SiluAndMul
17 from vllm.model_executor.layers.fused_moe import TritonExperts
18 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
19 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
20 BatchedTritonExperts)
21 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -129,11 +130,9 @@ def make_modular_kernel(
130 num_local_experts: int,
131 q_dtype: Optional[torch.dtype],
132 use_fp8_dispatch: bool,
132 - per_act_token_quant: bool,
133 + quant_config: FusedMoEQuantConfig,
134 ) -> FusedMoEModularKernel:
135
135 - is_quantized = q_dtype is not None
136 -
136 ht_args: Optional[DeepEPHTArgs] = None
137 ll_args: Optional[DeepEPLLArgs] = None
138
@@ -159,24 +158,14 @@ def make_modular_kernel(
158 num_dispatchers = pgi.world_size // dp_size
159
160 if low_latency_mode:
162 - assert not per_act_token_quant, "not supported in ll mode"
161 + assert not quant_config.per_act_token_quant, "not supported in ll mode"
162 fused_experts = BatchedTritonExperts(
163 max_num_tokens=MAX_TOKENS_PER_RANK,
164 num_dispatchers=num_dispatchers,
166 - use_fp8_w8a8=is_quantized,
167 - use_int8_w8a8=False,
168 - use_int8_w8a16=False,
169 - use_int4_w4a16=False,
170 - per_act_token_quant=False,
165 + quant_config=quant_config,
166 )
167 else:
173 - fused_experts = TritonExperts(
174 - use_fp8_w8a8=is_quantized,
175 - use_int8_w8a8=False,
176 - use_int8_w8a16=False,
177 - use_int4_w4a16=False,
178 - per_act_token_quant=per_act_token_quant,
179 - )
168 + fused_experts = TritonExperts(quant_config=quant_config)
169
170 mk = FusedMoEModularKernel(prepare_finalize=a2a,
171 fused_experts=fused_experts)
@@ -217,11 +206,6 @@ def deep_ep_moe_impl(
206 if is_quantized:
207 q_dtype = torch.float8_e4m3fn
208
220 - # Make modular kernel
221 - mk: FusedMoEModularKernel = make_modular_kernel(
222 - pg, pgi, low_latency_mode, hidden_size, dp_size, num_experts,
223 - num_local_experts, q_dtype, use_fp8_dispatch, per_act_token_quant)
224 -
209 out_hidden_states = torch.empty_like(test_tensors.rank_tokens)
210 total_num_tokens = test_tensors.rank_tokens.size(0)
211
@@ -236,6 +220,19 @@ def deep_ep_moe_impl(
220 rank_token_scales_chunk = rank_token_scales_chunk[
221 chunk_start:chunk_end]
222
223 + quant_config = FusedMoEQuantConfig.make(
224 + q_dtype,
225 + w1_scale=w1_scale,
226 + w2_scale=w2_scale,
227 + per_act_token_quant=per_act_token_quant,
228 + a1_scale=rank_token_scales_chunk,
229 + )
230 +
231 + # Make modular kernel
232 + mk: FusedMoEModularKernel = make_modular_kernel(
233 + pg, pgi, low_latency_mode, hidden_size, dp_size, num_experts,
234 + num_local_experts, q_dtype, use_fp8_dispatch, quant_config)
235 +
236 out = mk.forward(hidden_states=rank_tokens_chunk,
237 w1=w1,
238 w2=w2,
@@ -245,12 +242,6 @@ def deep_ep_moe_impl(
242 activation="silu",
243 global_num_experts=num_experts,
244 expert_map=build_expert_map(),
248 - w1_scale=w1_scale,
249 - w2_scale=w2_scale,
250 - w1_zp=None,
251 - w2_zp=None,
252 - a1_scale=rank_token_scales_chunk,
253 - a2_scale=None,
245 apply_router_weight_on_input=False)
246
247 if not skip_result_store:
@@ -407,7 +398,7 @@ DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
398
399
400 @pytest.mark.parametrize("dtype", DTYPES)
410 -@pytest.mark.parametrize("mnk", MNKs)
401 +@pytest.mark.parametrize("m,n,k", MNKs)
402 @pytest.mark.parametrize("num_experts", [32])
403 @pytest.mark.parametrize("topk", [6])
404 @pytest.mark.parametrize("world_dp_size", [(2, 1)])
@@ -416,7 +407,9 @@ DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
407 @requires_deep_ep
408 def test_deep_ep_moe(
409 dtype: torch.dtype,
419 - mnk: tuple[int, int, int],
410 + m: int,
411 + n: int,
412 + k: int,
413 num_experts: int,
414 topk: int,
415 world_dp_size: tuple[int, int],
@@ -424,7 +417,6 @@ def test_deep_ep_moe(
417 ):
418 low_latency_mode = False
419 use_fp8_dispatch = False
427 - m, n, k = mnk
420
421 current_platform.seed_everything(7)
422 world_size, dp_size = world_dp_size
@@ -456,20 +448,24 @@ USE_FP8_DISPATCH = [True, False]
448
449
450 @pytest.mark.parametrize("dtype", DTYPES)
459 -@pytest.mark.parametrize("mnk", MNKs)
451 +@pytest.mark.parametrize("m,n,k", MNKs)
452 @pytest.mark.parametrize("num_experts", [32])
453 @pytest.mark.parametrize("topk", [6])
454 @pytest.mark.parametrize("world_dp_size", [(2, 1)])
455 @pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH)
456 @multi_gpu_test(num_gpus=2)
457 @requires_deep_ep
466 -def test_low_latency_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],
467 - num_experts: int, topk: int,
468 - world_dp_size: tuple[int, int],
469 - use_fp8_dispatch: bool):
470 -
458 +def test_low_latency_deep_ep_moe(
459 + dtype: torch.dtype,
460 + m: int,
461 + n: int,
462 + k: int,
463 + num_experts: int,
464 + topk: int,
465 + world_dp_size: tuple[int, int],
466 + use_fp8_dispatch: bool,
467 +):
468 low_latency_mode = True
472 - m, n, k = mnk
469
470 if (low_latency_mode
471 and k not in DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES):
tests/kernels/moe/test_deepgemm.py
+16 -18
@@ -11,6 +11,8 @@ import math
11 import pytest
12 import torch
13
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + fp8_w8a8_moe_quant_config)
16 # vLLM fused-expert reference (Triton fallback + DeepGEMM option)
17 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
18 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
@@ -94,6 +96,13 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
96 topk_weights, topk_ids = torch.topk(router_logits, k=topk, dim=-1)
97 topk_weights = torch.nn.functional.softmax(topk_weights, dim=-1)
98
99 + quant_config = fp8_w8a8_moe_quant_config(
100 + w1_scale=w1_s,
101 + w2_scale=w2_s,
102 + a1_scale=a1_scale,
103 + block_shape=block_size,
104 + )
105 +
106 # triton reference
107 out_triton = fused_experts(
108 hidden_states=tokens_bf16,
@@ -102,11 +111,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
111 topk_weights=topk_weights,
112 topk_ids=topk_ids,
113 inplace=False,
105 - use_fp8_w8a8=True,
106 - w1_scale=w1_s,
107 - w2_scale=w2_s,
108 - a1_scale=a1_scale,
109 - block_shape=block_size,
114 + quant_config=quant_config,
115 allow_deep_gemm=False,
116 )
117
@@ -118,19 +123,14 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
123 topk_weights=topk_weights,
124 topk_ids=topk_ids,
125 inplace=False,
121 - use_fp8_w8a8=True,
122 - w1_scale=w1_s,
123 - w2_scale=w2_s,
124 - a1_scale=a1_scale,
125 - block_shape=block_size,
126 + quant_config=quant_config,
127 allow_deep_gemm=True,
128 )
129 diff = calc_diff(out_deepgemm, out_triton)
130 assert diff < 0.001, f"Diff exceeded 1%: {diff}"
131
132
132 -# Note: W1 has shape (E, 2N, K), so N = 512
133 -# can trigger the deepgemm path.
133 +# Note: N <= 512 will disable the deepgemm path due to performance issues.
134 MNKs = [
135 (1024, 768, 128),
136 (1024, 768, 512),
@@ -144,15 +144,15 @@ TOPKS = [2, 6]
144 NUM_EXPERTS = [32]
145
146
147 -@pytest.mark.parametrize("mnk", MNKs)
147 +@pytest.mark.parametrize(("m", "n", "k"), MNKs)
148 @pytest.mark.parametrize("topk", TOPKS)
149 @pytest.mark.parametrize("num_experts", NUM_EXPERTS)
150 @pytest.mark.skipif(not is_deep_gemm_supported(),
151 reason="Requires deep_gemm kernels")
152 -def test_deepgemm_vs_triton(mnk, topk, num_experts, monkeypatch):
152 +def test_deepgemm_vs_triton(m, n, k, topk, num_experts, monkeypatch):
153
154 - with monkeypatch.context() as m:
155 - m.setenv("VLLM_USE_DEEP_GEMM", "1")
154 + with monkeypatch.context() as mp:
155 + mp.setenv("VLLM_USE_DEEP_GEMM", "1")
156
157 _fused_moe_mod = importlib.import_module(
158 "vllm.model_executor.layers.fused_moe.fused_moe")
@@ -168,8 +168,6 @@ def test_deepgemm_vs_triton(mnk, topk, num_experts, monkeypatch):
168 monkeypatch.setattr(_fused_moe_mod, "deep_gemm_moe_fp8",
169 _spy_deep_gemm_moe_fp8)
170
171 - m, n, k = mnk
172 -
171 if topk > num_experts:
172 pytest.skip(f"topk={topk} > num_experts={num_experts}")
173
tests/kernels/moe/test_flashinfer.py
+20 -12
@@ -6,6 +6,8 @@ import pytest
6 import torch
7
8 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + fp8_w8a8_moe_quant_config)
11 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
12 from vllm.model_executor.layers.fused_moe.layer import FusedMoE
13 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
@@ -145,6 +147,14 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
147 custom_routing_function=Llama4MoE.custom_routing_function,
148 scoring_func="softmax")
149
150 + quant_config = fp8_w8a8_moe_quant_config(
151 + w1_scale=td.w13_weight_scale,
152 + w2_scale=td.w2_weight_scale,
153 + a1_scale=td.a1_scale,
154 + a2_scale=td.a2_scale,
155 + per_act_token_quant=False,
156 + )
157 +
158 output = fused_experts(
159 td.hidden_states,
160 td.w13_quantized,
@@ -153,15 +163,10 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
163 topk_ids=topk_ids,
164 inplace=False,
165 activation="silu",
156 - use_fp8_w8a8=True,
157 - per_channel_quant=False,
166 global_num_experts=e,
167 expert_map=None,
160 - w1_scale=td.w13_weight_scale,
161 - w2_scale=td.w2_weight_scale,
162 - a1_scale=td.a1_scale,
163 - a2_scale=td.a2_scale,
168 apply_router_weight_on_input=True,
169 + quant_config=quant_config,
170 )
171
172 flashinfer_output = apply_flashinfer_per_tensor_scale_fp8(
@@ -210,6 +215,14 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
215 custom_routing_function=Llama4MoE.custom_routing_function,
216 scoring_func="softmax")
217
218 + quant_config = fp8_w8a8_moe_quant_config(
219 + w1_scale=td.w13_weight_scale,
220 + w2_scale=td.w2_weight_scale,
221 + a1_scale=td.a1_scale,
222 + a2_scale=td.a2_scale,
223 + per_act_token_quant=False,
224 + )
225 +
226 output = fused_experts(
227 td.hidden_states,
228 td.w13_quantized,
@@ -218,15 +231,10 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
231 topk_ids=topk_ids,
232 inplace=False,
233 activation="silu",
221 - use_fp8_w8a8=True,
222 - per_channel_quant=False,
234 global_num_experts=e,
235 expert_map=None,
225 - w1_scale=td.w13_weight_scale,
226 - w2_scale=td.w2_weight_scale,
227 - a1_scale=td.a1_scale,
228 - a2_scale=td.a2_scale,
236 apply_router_weight_on_input=True,
237 + quant_config=quant_config,
238 )
239
240 td.layer.dp_size = 1
tests/kernels/moe/test_flashinfer_moe.py
+24 -44
@@ -3,7 +3,7 @@
3 import pytest
4 import torch
5
6 -from tests.kernels.moe.utils import make_test_weights
6 +from tests.kernels.moe.utils import make_test_quant_config
7 from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
8 FLOAT8_E4M3_MAX,
9 dequantize_nvfp4_to_dtype)
@@ -41,7 +41,6 @@ MNK_FACTORS = [
41
42 @pytest.mark.parametrize("m,n,k", MNK_FACTORS)
43 @pytest.mark.parametrize("e", [40, 64, 256])
44 -#@pytest.mark.parametrize("e", [128, 256])
44 @pytest.mark.parametrize("topk", [1, 6, 8])
45 @pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
46 @torch.inference_mode()
@@ -56,16 +55,15 @@ def test_flashinfer_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
55
56 quant_blocksize = 16
57
59 - (_, w1_q, w1_blockscale,
60 - w1_gs), (_, w2_q, w2_blockscale, w2_gs) = make_test_weights(
61 - e,
62 - n,
63 - k,
64 - in_dtype=dtype,
65 - quant_dtype="nvfp4",
66 - block_shape=None, # use quant_blocksize?
67 - per_act_token_quant=False,
68 - )
58 + w1_q, w2_q, quant_config = make_test_quant_config(
59 + e,
60 + n,
61 + k,
62 + in_dtype=dtype,
63 + quant_dtype="nvfp4",
64 + block_shape=None,
65 + per_act_token_quant=False,
66 + )
67
68 score = torch.randn((m, e), device="cuda", dtype=dtype)
69 topk_weights, topk_ids, _ = fused_topk(a,
@@ -73,35 +71,17 @@ def test_flashinfer_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
71 topk,
72 renormalize=False)
73
76 - a1_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
77 - a2_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
78 -
74 assert is_valid_flashinfer_cutlass_fused_moe(a, w1_q, w2_q)
75
81 - assert w1_gs is not None
82 - assert w2_gs is not None
83 - assert w1_blockscale is not None
84 - assert w2_blockscale is not None
85 -
76 flashinfer_experts = FusedMoEModularKernel(
77 MoEPrepareAndFinalizeNoEP(),
88 - FlashInferExperts(
89 - a1_gscale=a1_gs,
90 - g1_alphas=(1 / w1_gs),
91 - a2_gscale=a2_gs,
92 - g2_alphas=(1 / w2_gs),
93 - out_dtype=dtype,
94 - quant_dtype="nvfp4",
95 - ))
78 + FlashInferExperts(out_dtype=dtype, quant_config=quant_config),
79 + )
80
81 flashinfer_output = flashinfer_experts(
82 hidden_states=a,
83 w1=w1_q,
100 - w1_scale=w1_blockscale,
84 w2=w2_q,
102 - w2_scale=w2_blockscale,
103 - a1_scale=a1_gs,
104 - a2_scale=a2_gs,
85 topk_weights=topk_weights,
86 topk_ids=topk_ids,
87 )
@@ -122,18 +102,18 @@ def test_flashinfer_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
102 w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
103
104 for idx in range(0, e):
125 - w1_d[idx] = dequantize_nvfp4_to_dtype(w1_q[idx],
126 - w1_blockscale[idx],
127 - w1_gs[idx],
128 - dtype=dtype,
129 - device=w1_q.device,
130 - block_size=quant_blocksize)
131 - w2_d[idx] = dequantize_nvfp4_to_dtype(w2_q[idx],
132 - w2_blockscale[idx],
133 - w2_gs[idx],
134 - dtype=dtype,
135 - device=w2_q.device,
136 - block_size=quant_blocksize)
105 + w1_d[idx] = dequantize_nvfp4_to_dtype(
106 + w1_q[idx],
107 + quant_config.w1_scale[idx], (1 / quant_config.g1_alphas[idx]),
108 + dtype=dtype,
109 + device=w1_q.device,
110 + block_size=quant_blocksize)
111 + w2_d[idx] = dequantize_nvfp4_to_dtype(
112 + w2_q[idx],
113 + quant_config.w2_scale[idx], (1 / quant_config.g2_alphas[idx]),
114 + dtype=dtype,
115 + device=w2_q.device,
116 + block_size=quant_blocksize)
117
118 torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk)
119
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+17 -13
@@ -23,6 +23,7 @@ from triton_kernels.tensor import FP4, convert_layout, wrap_torch_tensor
23 from triton_kernels.tensor_details import layout
24 from triton_kernels.testing import assert_close
25
26 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
27 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
28 BatchedPrepareAndFinalize)
29 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
@@ -293,6 +294,13 @@ def test_equiv(num_token, a_dtype, w_dtype, tp):
294 pc2,
295 ) = init_compute_data(M, K, N, E, a_dtype, w_dtype, num_warps=8)
296
297 + quant_config = FusedMoEQuantConfig.make(
298 + w1_bias=w1_bias_tri,
299 + w2_bias=w2_bias_tri,
300 + w1_precision=pc1,
301 + w2_precision=pc2,
302 + )
303 +
304 out_triton_monolithic = triton_kernel_moe_forward(
305 hidden_states=x_tri,
306 w1=w1_tri,
@@ -300,10 +308,7 @@ def test_equiv(num_token, a_dtype, w_dtype, tp):
308 gating_output=exp_data_tri,
309 topk=topk,
310 renormalize=True,
303 - w1_bias=w1_bias_tri,
304 - w2_bias=w2_bias_tri,
305 - w1_precision=pc1,
306 - w2_precision=pc2,
311 + quant_config=quant_config,
312 )
313 out_triton_monolithic = out_triton_monolithic[..., :K]
314
@@ -336,6 +341,13 @@ def batched_moe(
341 ) -> torch.Tensor:
342 max_num_tokens = round_up(a.shape[0], 64)
343
344 + quant_config = FusedMoEQuantConfig.make(
345 + w1_precision=w1_precision,
346 + w2_precision=w2_precision,
347 + w1_bias=w1_bias,
348 + w2_bias=w2_bias,
349 + )
350 +
351 fused_experts = FusedMoEModularKernel(
352 BatchedPrepareAndFinalize(
353 max_num_tokens,
@@ -344,19 +356,12 @@ def batched_moe(
356 rank=0,
357 ),
358 BatchedOAITritonExperts(
347 - None,
359 max_num_tokens=max_num_tokens,
360 num_dispatchers=1,
350 - w1_precision=w1_precision,
351 - w2_precision=w2_precision,
361 + quant_config=quant_config,
362 ),
363 )
364
355 - extra_expert_args = {
356 - "w1_bias": w1_bias,
357 - "w2_bias": w2_bias,
358 - }
359 -
365 topk_weight, topk_ids, _ = fused_topk(a, gating_output, topk, renormalize)
366
367 return fused_experts(
@@ -365,7 +370,6 @@ def batched_moe(
370 w2,
371 topk_weight,
372 topk_ids,
368 - extra_expert_args=extra_expert_args,
373 )
374
375
tests/kernels/moe/test_modular_kernel_combinations.py
+21 -19
@@ -12,7 +12,6 @@ import torch
12
13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 from vllm.config import VllmConfig, current_platform, set_current_vllm_config
15 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
15 from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
16 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
17
@@ -22,7 +21,8 @@ from .modular_kernel_tools.common import (Config, RankTensors, WeightTensors,
21 run_modular_kernel)
22 from .modular_kernel_tools.mk_objects import (
23 MK_FUSED_EXPERT_TYPES, MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
25 - MK_QUANT_CONFIGS, MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, expert_info)
24 + MK_QUANT_CONFIGS, MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, TestMoEQuantConfig,
25 + expert_info)
26 from .modular_kernel_tools.parallel_utils import (ProcessGroupInfo,
27 parallel_launch_with_config)
28
@@ -55,7 +55,7 @@ def rank_worker(
55 pgi: ProcessGroupInfo,
56 vllm_config: VllmConfig,
57 cpu_group,
58 - config: Config,
58 + base_config: Config,
59 weights: WeightTensors,
60 verbose: bool,
61 ):
@@ -63,42 +63,44 @@ def rank_worker(
63
64 # sanity check
65 from vllm import envs
66 - if config.fused_moe_chunk_size is not None:
67 - assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
66 + if base_config.fused_moe_chunk_size is not None:
67 + assert (
68 + base_config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
69
70 # get weights to this device
71 weights.to_current_device()
72
72 - Ms = config.Ms
73 + Ms = base_config.Ms
74 assert isinstance(Ms, list)
74 - TOPKs = config.topks
75 + TOPKs = base_config.topks
76 assert isinstance(TOPKs, list)
77
78 exceptions = []
79 count = 0
80
81 for m, topk in product(Ms, TOPKs):
82 + # override m and topk
83 + config = copy.deepcopy(base_config)
84 + config.Ms = m
85 + config.topks = topk
86 +
87 try:
88 print(f"Running[{pgi.rank}]: m={m}, topk={topk} ...")
89 count = count + 1
84 - # override m and topk
85 - cfgx = copy.deepcopy(config)
86 - cfgx.Ms = m
87 - cfgx.topks = topk
90
91 # inputs for rank
90 - rank_tensors = RankTensors.make(cfgx, pgi)
92 + rank_tensors = RankTensors.make(config, pgi)
93
94 # modular kernel out
93 - mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights,
95 + mk_out = run_modular_kernel(pgi, vllm_config, config, weights,
96 rank_tensors)
97
98 with set_current_vllm_config(vllm_config):
97 - ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
99 + ref_out = reference_moe_impl(config, weights, rank_tensors)
100
101 if config.quant_dtype == "nvfp4":
100 - atol = 1e-1
101 - rtol = 1e-1
102 + atol = 1e-1 if config.K < 4096 else 2e-1
103 + rtol = 1e-1 if config.K < 4096 else 2e-1
104 else:
105 atol = 3e-2
106 rtol = 3e-2
@@ -132,7 +134,7 @@ Ms = [32, 64]
134 # hidden sizes, making this too large will cause fp4 tests to fail.
135 # Also needs to be a multiple of 1024 for deep_gemm.
136 Ks = [2048]
135 -Ns = [2048]
137 +Ns = [1024]
138 TOPKs = [4, 1]
139 Es = [32]
140 DTYPEs = [torch.bfloat16]
@@ -167,7 +169,7 @@ def is_nyi_config(config: Config) -> bool:
169 @meets_multi_gpu_requirements
170 def test_modular_kernel_combinations_multigpu(
171 k: int, n: int, e: int, dtype: torch.dtype,
170 - quant_config: Optional[FusedMoEQuantConfig],
172 + quant_config: Optional[TestMoEQuantConfig],
173 combination: tuple[mk.FusedMoEPrepareAndFinalize,
174 mk.FusedMoEPermuteExpertsUnpermute],
175 fused_moe_chunk_size: Optional[int], world_size: int, pytestconfig):
@@ -208,7 +210,7 @@ def test_modular_kernel_combinations_multigpu(
210 @pytest.mark.parametrize("world_size", [1])
211 def test_modular_kernel_combinations_singlegpu(
212 k: int, n: int, e: int, dtype: torch.dtype,
211 - quant_config: Optional[FusedMoEQuantConfig],
213 + quant_config: Optional[TestMoEQuantConfig],
214 combination: tuple[mk.FusedMoEPrepareAndFinalize,
215 mk.FusedMoEPermuteExpertsUnpermute],
216 fused_moe_chunk_size: Optional[int], world_size: int, pytestconfig):
tests/kernels/moe/test_moe.py
+19 -15
@@ -15,11 +15,14 @@ from transformers import MixtralConfig
15 from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
16
17 import vllm.model_executor.layers.fused_moe # noqa
18 +from tests.kernels.moe.utils import fused_moe
19 from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
20 from vllm.config import VllmConfig, set_current_vllm_config
21 from vllm.distributed.parallel_state import init_distributed_environment
22 from vllm.forward_context import set_forward_context
22 -from vllm.model_executor.layers.fused_moe import fused_moe
23 +from vllm.model_executor.layers.fused_moe.config import (
24 + FUSED_MOE_UNQUANTIZED_CONFIG, int4_w4a16_moe_quant_config,
25 + int8_w8a16_moe_quant_config)
26 from vllm.model_executor.layers.fused_moe.fused_moe import (
27 fused_topk, modular_triton_fused_moe)
28 from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
@@ -187,14 +190,9 @@ def test_fused_moe(
190 #
191 # Setup test functions
192 #
193 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
194
191 - m_fused_moe_fn = modular_triton_fused_moe(use_fp8_w8a8=False,
192 - use_int8_w8a8=False,
193 - use_int8_w8a16=False,
194 - use_int4_w4a16=False,
195 - use_mxfp4_w4a4=False,
196 - per_act_token_quant=False,
197 - block_shape=None)
195 + m_fused_moe_fn = modular_triton_fused_moe(quant_config)
196
197 def m_fused_moe(
198 a: torch.Tensor,
@@ -340,6 +338,18 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
338 else:
339 e_map = None
340
341 + if weight_bits == 4:
342 + quant_config_builder = int4_w4a16_moe_quant_config
343 + else:
344 + assert weight_bits == 8
345 + quant_config_builder = int8_w8a16_moe_quant_config
346 +
347 + quant_config = quant_config_builder(w1_scale=w1_scales,
348 + w2_scale=w2_scales,
349 + w1_zp=w1_qzeros if has_zp else None,
350 + w2_zp=w2_qzeros if has_zp else None,
351 + block_shape=[0, group_size])
352 +
353 with set_current_vllm_config(vllm_config):
354 triton_output = fused_moe(a,
355 w1_qweight,
@@ -347,15 +357,9 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
357 score,
358 topk,
359 renormalize=False,
350 - use_int4_w4a16=weight_bits == 4,
351 - use_int8_w8a16=weight_bits == 8,
360 global_num_experts=e,
361 expert_map=e_map,
354 - w1_scale=w1_scales,
355 - w2_scale=w2_scales,
356 - w1_zp=w1_qzeros if has_zp else None,
357 - w2_zp=w2_qzeros if has_zp else None,
358 - block_shape=[0, group_size])
362 + quant_config=quant_config)
363 torch_output = torch_moe(a,
364 w1_ref,
365 w2_ref,
tests/kernels/moe/test_nvfp4_moe.py
+12 -7
@@ -10,6 +10,7 @@ from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
10 from tests.kernels.utils import torch_moe
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 nvfp4_moe_quant_config
14 from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
15 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
16 from vllm.platforms import current_platform
@@ -56,7 +57,7 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
57 in_dtype=dtype,
58 quant_dtype="nvfp4",
59 block_shape=None, # use quant_blocksize?
59 - per_act_token_quant=False,
60 + per_out_ch_quant=False,
61 )
62
63 score = torch.randn((m, e), device="cuda", dtype=dtype)
@@ -73,18 +74,22 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
74 assert w1_blockscale is not None
75 assert w2_blockscale is not None
76
77 + quant_config = nvfp4_moe_quant_config(
78 + g1_alphas=(1 / w1_gs),
79 + g2_alphas=(1 / w2_gs),
80 + a1_gscale=a1_gs,
81 + a2_gscale=a2_gs,
82 + w1_scale=w1_blockscale,
83 + w2_scale=w2_blockscale,
84 + )
85 +
86 cutlass_output = cutlass_moe_fp4(
87 a=a,
78 - a1_gscale=a1_gs,
88 w1_fp4=w1_q,
80 - w1_blockscale=w1_blockscale,
81 - g1_alphas=(1 / w1_gs),
82 - a2_gscale=a2_gs,
89 w2_fp4=w2_q,
84 - w2_blockscale=w2_blockscale,
85 - g2_alphas=(1 / w2_gs),
90 topk_weights=topk_weights,
91 topk_ids=topk_ids,
92 + quant_config=quant_config,
93 m=m,
94 n=n,
95 k=k,
tests/kernels/moe/test_pplx_cutlass_moe.py
+13 -8
@@ -9,6 +9,8 @@ import torch
9 from tests.kernels.utils import torch_experts
10 from vllm import _custom_ops as ops
11 from vllm.config import VllmConfig, set_current_vllm_config
12 +from vllm.model_executor.layers.fused_moe.config import (
13 + fp8_w8a8_moe_quant_config)
14 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
15 CutlassBatchedExpertsFp8)
16 from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
@@ -143,10 +145,16 @@ def pplx_cutlass_moe(
145 device="cuda",
146 dtype=torch.int64)
147
146 - experts = CutlassBatchedExpertsFp8(num_local_experts, num_dispatchers,
147 - out_dtype, per_act_token, per_out_ch,
148 - ab_strides1, ab_strides2, c_strides1,
149 - c_strides2)
148 + experts = CutlassBatchedExpertsFp8(
149 + num_local_experts, num_dispatchers, out_dtype, ab_strides1,
150 + ab_strides2, c_strides1, c_strides2,
151 + fp8_w8a8_moe_quant_config(
152 + per_act_token_quant=per_act_token,
153 + per_out_ch_quant=per_out_ch,
154 + w1_scale=chunk_by_rank(w1_scale, rank, world_size),
155 + w2_scale=chunk_by_rank(w2_scale, rank, world_size),
156 + a1_scale=chunk_by_rank(a1_scale, rank, world_size)
157 + if per_act_token else a1_scale[rank]))
158
159 fused_cutlass_experts = FusedMoEModularKernel(
160 prepare_finalize,
@@ -167,10 +175,7 @@ def pplx_cutlass_moe(
175 chunk_topk_ids,
176 global_num_experts=num_experts,
177 expert_map=None, #TODO
170 - w1_scale=chunk_by_rank(w1_scale, rank, world_size),
171 - w2_scale=chunk_by_rank(w2_scale, rank, world_size),
172 - a1_scale=chunk_by_rank(a1_scale, rank, world_size)
173 - if per_act_token else a1_scale[rank])
178 + )
179
180 torch.cuda.synchronize()
181
tests/kernels/moe/test_pplx_moe.py
+31 -31
@@ -58,7 +58,7 @@ BATCHED_MOE_MNK_FACTORS = [
58 ]
59
60 PPLX_COMBOS = [
61 - # TODO: figure out why this fails, seems to be test problem
61 + # TODO(bnell): figure out why this fails, seems to be test problem
62 #(1, 128, 128),
63 (2, 128, 512),
64 (3, 1024, 2048),
@@ -360,18 +360,18 @@ def pplx_prepare_finalize(
360
361 b_a, b_a_scale, expert_num_tokens, _, _ = prepare_finalize.prepare(
362 a_chunk,
363 - a1_scale,
364 - a2_scale,
363 chunk_topk_weight,
364 chunk_topk_ids,
365 num_experts,
366 None,
367 False,
370 - FusedMoEQuantConfig(
368 + FusedMoEQuantConfig.make(
369 quant_dtype,
372 - per_act_token_quant,
373 - False,
374 - block_shape,
370 + per_act_token_quant=per_act_token_quant,
371 + per_out_ch_quant=False,
372 + block_shape=block_shape,
373 + a1_scale=a1_scale,
374 + a2_scale=a2_scale,
375 ),
376 )
377
@@ -540,20 +540,6 @@ def pplx_moe(
540
541 topk_ids = topk_ids.to(dtype=torch.uint32)
542
543 - experts = BatchedTritonExperts(
544 - max_num_tokens=max_num_tokens,
545 - num_dispatchers=prepare_finalize.num_dispatchers(),
546 - use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
547 - block_shape=block_shape,
548 - per_act_token_quant=per_act_token_quant,
549 - )
550 -
551 - fused_experts = FusedMoEModularKernel(
552 - prepare_finalize,
553 - experts,
554 - shared_experts,
555 - )
556 -
543 # Note: workers with the same dp_rank must use the exact same inputs.
544 a_chunk = chunk_by_rank(a, rank, world_size)
545 chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size)
@@ -567,6 +553,28 @@ def pplx_moe(
553 a1_scale_chunk = chunk_scales_by_rank(a1_scale, rank, world_size)
554 a2_scale_chunk = chunk_scales_by_rank(a2_scale, rank, world_size)
555
556 + quant_config = FusedMoEQuantConfig.make(
557 + quant_dtype,
558 + block_shape=block_shape,
559 + per_act_token_quant=per_act_token_quant,
560 + w1_scale=w1_scale_chunk,
561 + w2_scale=w2_scale_chunk,
562 + a1_scale=a1_scale_chunk,
563 + a2_scale=a2_scale_chunk,
564 + )
565 +
566 + experts = BatchedTritonExperts(
567 + max_num_tokens=max_num_tokens,
568 + num_dispatchers=prepare_finalize.num_dispatchers(),
569 + quant_config=quant_config,
570 + )
571 +
572 + fused_experts = FusedMoEModularKernel(
573 + prepare_finalize,
574 + experts,
575 + shared_experts,
576 + )
577 +
578 # Note: for now use_compile will error out if the problem size is
579 # large enough to trigger chunking. I'm leaving the flag and
580 # setup code in case we are able to revisit this later.
@@ -585,10 +593,6 @@ def pplx_moe(
593 w2_chunk,
594 chunk_topk_weight,
595 chunk_topk_ids,
588 - w1_scale=w1_scale_chunk,
589 - w2_scale=w2_scale_chunk,
590 - a1_scale=a1_scale_chunk,
591 - a2_scale=a2_scale_chunk,
596 global_num_experts=num_experts)
597
598 if use_cudagraphs:
@@ -605,10 +609,6 @@ def pplx_moe(
609 w2_chunk,
610 chunk_topk_weight,
611 chunk_topk_ids,
608 - w1_scale=w1_scale_chunk,
609 - w2_scale=w2_scale_chunk,
610 - a1_scale=a1_scale_chunk,
611 - a2_scale=a2_scale_chunk,
612 global_num_experts=num_experts)
613
614 torch.cuda.synchronize()
@@ -820,7 +820,7 @@ def test_pplx_moe_slow(
820 k,
821 quant_dtype=quant_dtype,
822 block_shape=block_shape,
823 - per_act_token_quant=per_act_token_quant,
823 + per_out_ch_quant=per_act_token_quant,
824 )
825
826 parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk, e,
@@ -897,7 +897,7 @@ def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
897 k,
898 quant_dtype=quant_dtype,
899 block_shape=block_shape,
900 - per_act_token_quant=per_act_token_quant,
900 + per_out_ch_quant=per_act_token_quant,
901 )
902 args["w1"] = w1
903 args["w2"] = w2
tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+9 -6
@@ -7,10 +7,12 @@ import itertools
7 import pytest
8 import torch
9
10 +from tests.kernels.moe.utils import fused_moe
11 from vllm import _custom_ops as ops
12 from vllm.config import VllmConfig, set_current_vllm_config
13 from vllm.model_executor.layers.activation import SiluAndMul
13 -from vllm.model_executor.layers.fused_moe import fused_moe
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + fp8_w8a8_moe_quant_config)
16 from vllm.platforms import current_platform
17
18 if current_platform.get_device_capability() < (9, 0):
@@ -152,11 +154,12 @@ def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):
154 score,
155 topk,
156 renormalize=False,
155 - use_fp8_w8a8=True, # using fp8
156 - per_channel_quant=True,
157 - w1_scale=w1_s,
158 - w2_scale=w2_s,
159 - block_shape=None, # Not using block quantization
157 + quant_config=fp8_w8a8_moe_quant_config(
158 + per_act_token_quant=True,
159 + w1_scale=w1_s,
160 + w2_scale=w2_s,
161 + block_shape=None, # Not using block quantization
162 + ),
163 )
164
165 # Check results
tests/kernels/moe/utils.py
+112 -37
@@ -9,7 +9,8 @@ from tests.kernels.quant_utils import per_block_cast_to_int8
9 from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
10 FLOAT8_E4M3_MAX)
11 from vllm.model_executor.layers.activation import SiluAndMul
12 -from vllm.model_executor.layers.fused_moe import fused_experts
12 +from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
13 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
14 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
15 BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
16 from vllm.model_executor.layers.fused_moe.modular_kernel import (
@@ -34,18 +35,22 @@ def triton_moe(
35 per_act_token_quant=False,
36 block_shape: Optional[list[int]] = None,
37 ) -> torch.Tensor:
38 + quant_config = FusedMoEQuantConfig.make(
39 + quant_dtype,
40 + per_act_token_quant=per_act_token_quant,
41 + block_shape=block_shape,
42 + w1_scale=w1_scale,
43 + w2_scale=w2_scale,
44 + a1_scale=a1_scale,
45 + a2_scale=a2_scale,
46 + )
47 +
48 return fused_experts(a,
49 w1,
50 w2,
51 topk_weight,
52 topk_ids,
42 - w1_scale=w1_scale,
43 - w2_scale=w2_scale,
44 - a1_scale=a1_scale,
45 - a2_scale=a2_scale,
46 - per_channel_quant=per_act_token_quant,
47 - use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
48 - block_shape=block_shape)
53 + quant_config=quant_config)
54
55
56 def batched_moe(
@@ -64,6 +69,16 @@ def batched_moe(
69 ) -> torch.Tensor:
70 max_num_tokens = round_up(a.shape[0], 64)
71
72 + quant_config = FusedMoEQuantConfig.make(
73 + quant_dtype,
74 + per_act_token_quant=per_act_token_quant,
75 + block_shape=block_shape,
76 + w1_scale=w1_scale,
77 + w2_scale=w2_scale,
78 + a1_scale=a1_scale,
79 + a2_scale=a2_scale,
80 + )
81 +
82 fused_experts = FusedMoEModularKernel(
83 BatchedPrepareAndFinalize(max_num_tokens,
84 num_dispatchers=1,
@@ -72,21 +87,11 @@ def batched_moe(
87 BatchedTritonExperts(
88 max_num_tokens=max_num_tokens,
89 num_dispatchers=1,
75 - use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
76 - per_act_token_quant=per_act_token_quant,
77 - block_shape=block_shape,
90 + quant_config=quant_config,
91 ),
92 )
93
81 - return fused_experts(a,
82 - w1,
83 - w2,
84 - topk_weight,
85 - topk_ids,
86 - w1_scale=w1_scale,
87 - w2_scale=w2_scale,
88 - a1_scale=a1_scale,
89 - a2_scale=a2_scale)
94 + return fused_experts(a, w1, w2, topk_weight, topk_ids)
95
96
97 def naive_batched_moe(
@@ -105,6 +110,16 @@ def naive_batched_moe(
110 ) -> torch.Tensor:
111 max_num_tokens = round_up(a.shape[0], 64)
112
113 + quant_config = FusedMoEQuantConfig.make(
114 + quant_dtype,
115 + per_act_token_quant=per_act_token_quant,
116 + block_shape=block_shape,
117 + w1_scale=w1_scale,
118 + w2_scale=w2_scale,
119 + a1_scale=a1_scale,
120 + a2_scale=a2_scale,
121 + )
122 +
123 fused_experts = FusedMoEModularKernel(
124 BatchedPrepareAndFinalize(max_num_tokens,
125 num_dispatchers=1,
@@ -113,21 +128,11 @@ def naive_batched_moe(
128 NaiveBatchedExperts(
129 max_num_tokens=max_num_tokens,
130 num_dispatchers=1,
116 - use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
117 - per_act_token_quant=per_act_token_quant,
118 - block_shape=block_shape,
131 + quant_config=quant_config,
132 ),
133 )
134
122 - return fused_experts(a,
123 - w1,
124 - w2,
125 - topk_weight,
126 - topk_ids,
127 - w1_scale=w1_scale,
128 - w2_scale=w2_scale,
129 - a1_scale=a1_scale,
130 - a2_scale=a2_scale)
135 + return fused_experts(a, w1, w2, topk_weight, topk_ids)
136
137
138 def chunk_scales(scales: Optional[torch.Tensor], start: int,
@@ -216,7 +221,7 @@ def make_test_weight(
221 in_dtype: torch.dtype = torch.bfloat16,
222 quant_dtype: Union[torch.dtype, str, None] = None,
223 block_shape: Optional[list[int]] = None,
219 - per_act_token_quant: bool = False,
224 + per_out_ch_quant: bool = False,
225 ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
226 Optional[torch.Tensor]]:
227 w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15
@@ -228,7 +233,7 @@ def make_test_weight(
233 w_gs_l = [None] * e
234 for idx in range(e):
235 w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights(
231 - w_16[idx], None, quant_dtype, per_act_token_quant, block_shape)
236 + w_16[idx], None, quant_dtype, per_out_ch_quant, block_shape)
237
238 w = torch.stack(w_l)
239 w_s = torch.stack(w_s_l)
@@ -258,16 +263,16 @@ def make_test_weights(
263 in_dtype: torch.dtype = torch.bfloat16,
264 quant_dtype: Union[torch.dtype, str, None] = None,
265 block_shape: Optional[list[int]] = None,
261 - per_act_token_quant: bool = False,
266 + per_out_ch_quant: bool = False,
267 ) -> tuple[tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
268 Optional[torch.Tensor]],
269 tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
270 Optional[torch.Tensor]]]:
271 return (
272 make_test_weight(e, 2 * n, k, in_dtype, quant_dtype, block_shape,
268 - per_act_token_quant),
273 + per_out_ch_quant),
274 make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape,
270 - per_act_token_quant),
275 + per_out_ch_quant),
276 )
277
278
@@ -285,6 +290,76 @@ def per_token_cast_to_fp8(
290 return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
291
292
293 +def make_test_quant_config(
294 + e: int,
295 + n: int,
296 + k: int,
297 + in_dtype: torch.dtype,
298 + quant_dtype: Union[torch.dtype, str, None] = None,
299 + per_act_token_quant: bool = False,
300 + block_shape: Optional[list[int]] = None,
301 +) -> tuple[torch.Tensor, torch.Tensor, FusedMoEQuantConfig]:
302 + (_, w1, w1_s, w1_gs), (_, w2, w2_s, w2_gs) = make_test_weights(
303 + e,
304 + n,
305 + k,
306 + in_dtype,
307 + quant_dtype,
308 + per_out_ch_quant=per_act_token_quant,
309 + block_shape=block_shape,
310 + )
311 +
312 + # Hacky/trivial scales for nvfp4.
313 + a1_gscale: Optional[torch.Tensor] = None
314 + a2_gscale: Optional[torch.Tensor] = None
315 + if quant_dtype == "nvfp4":
316 + a1_gscale = torch.ones((e, ), device="cuda", dtype=torch.float32)
317 + a2_gscale = torch.ones((e, ), device="cuda", dtype=torch.float32)
318 + a1_scale = a1_gscale
319 + a2_scale = a2_gscale
320 + else:
321 + a1_scale = None
322 + a2_scale = None
323 +
324 + return w1, w2, FusedMoEQuantConfig.make(
325 + quant_dtype,
326 + per_act_token_quant=per_act_token_quant,
327 + block_shape=block_shape,
328 + w1_scale=w1_s,
329 + w2_scale=w2_s,
330 + a1_gscale=a1_gscale,
331 + a2_gscale=a2_gscale,
332 + a1_scale=a1_scale,
333 + a2_scale=a2_scale,
334 + # TODO: make sure this is handled properly
335 + g1_alphas=(1 / w1_gs) if w1_gs is not None else None,
336 + g2_alphas=(1 / w2_gs) if w2_gs is not None else None,
337 + )
338 +
339 +
340 +def fused_moe(
341 + hidden_states: torch.Tensor,
342 + w1: torch.Tensor,
343 + w2: torch.Tensor,
344 + score: torch.Tensor,
345 + topk: int,
346 + renormalize: bool = False,
347 + quant_config: Optional[FusedMoEQuantConfig] = None,
348 + global_num_experts: int = -1,
349 + expert_map: Optional[torch.Tensor] = None,
350 +) -> torch.Tensor:
351 + topk_weights, topk_ids, _ = fused_topk(hidden_states, score.float(), topk,
352 + renormalize)
353 + return fused_experts(hidden_states,
354 + w1,
355 + w2,
356 + topk_weights,
357 + topk_ids,
358 + global_num_experts=global_num_experts,
359 + expert_map=expert_map,
360 + quant_config=quant_config)
361 +
362 +
363 # CustomOp?
364 class BaselineMM(torch.nn.Module):
365
tests/kernels/quantization/test_int8_kernel.py
+21 -14
@@ -8,7 +8,8 @@ import pytest
8 import torch
9
10 from vllm.model_executor.layers.activation import SiluAndMul
11 -from vllm.model_executor.layers.fused_moe import fused_moe
11 +from vllm.model_executor.layers.fused_moe import fused_experts
12 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
13 from vllm.model_executor.layers.quantization.utils.int8_utils import (
14 per_token_quant_int8)
15 from vllm.platforms import current_platform
@@ -42,7 +43,8 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
43 return C.reshape(origin_C_shape).to(output_dtype)
44
45
45 -def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
46 +def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight,
47 + topk_ids):
48 """This function performs fused moe with per-column int8 quantization
49 using native torch."""
50
@@ -57,8 +59,6 @@ def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk):
59 out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
60
61 # Calculate routing
60 - score = torch.softmax(score, dim=-1, dtype=torch.float32)
61 - topk_weight, topk_ids = torch.topk(score, topk)
62 topk_weight = topk_weight.view(-1)
63 topk_ids = topk_ids.view(-1)
64 # Process each expert
@@ -127,20 +127,27 @@ def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):
127 w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
128 w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
129 score = torch.randn((M, E), dtype=dtype)
130 + score = torch.softmax(score, dim=-1, dtype=torch.float32)
131 + topk_weights, topk_ids = torch.topk(score, topk)
132 +
133 + ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk,
134 + topk_weights, topk_ids)
135
131 - ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk)
132 - out = fused_moe(
136 + quant_config = FusedMoEQuantConfig.make(
137 + torch.int8,
138 + per_act_token_quant=True,
139 + block_shape=None,
140 + w1_scale=w1_s,
141 + w2_scale=w2_s,
142 + )
143 +
144 + out = fused_experts(
145 a,
146 w1,
147 w2,
136 - score,
137 - topk,
138 - renormalize=False,
139 - use_int8_w8a8=True, # Using int8-w8a8
140 - per_channel_quant=True,
141 - w1_scale=w1_s,
142 - w2_scale=w2_s,
143 - block_shape=None, # Not using block quantization
148 + topk_weights,
149 + topk_ids,
150 + quant_config=quant_config,
151 )
152
153 # Check results
vllm/model_executor/layers/fused_moe/__init__.py
+4 -4
@@ -10,6 +10,7 @@ from vllm.model_executor.layers.fused_moe.layer import (
10 from vllm.model_executor.layers.fused_moe.modular_kernel import (
11 FusedMoEActivationFormat, FusedMoEPermuteExpertsUnpermute,
12 FusedMoEPrepareAndFinalize)
13 +from vllm.model_executor.layers.fused_moe.utils import activation_without_mul
14 from vllm.triton_utils import HAS_TRITON
15
16 _config: Optional[dict[str, Any]] = None
@@ -36,6 +37,7 @@ __all__ = [
37 "FusedMoEPermuteExpertsUnpermute",
38 "FusedMoEActivationFormat",
39 "FusedMoEPrepareAndFinalize",
40 + "activation_without_mul",
41 "override_config",
42 "get_config",
43 ]
@@ -43,7 +45,6 @@ __all__ = [
45 if HAS_TRITON:
46 # import to register the custom ops
47 import vllm.model_executor.layers.fused_moe.fused_marlin_moe # noqa
46 - import vllm.model_executor.layers.fused_moe.fused_moe # noqa
48 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
49 BatchedDeepGemmExperts)
50 from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
@@ -56,13 +57,12 @@ if HAS_TRITON:
57 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
58 BatchedTritonExperts)
59 from vllm.model_executor.layers.fused_moe.fused_moe import (
59 - TritonExperts, fused_experts, fused_moe, fused_topk,
60 - get_config_file_name, grouped_topk)
60 + TritonExperts, fused_experts, fused_topk, get_config_file_name,
61 + grouped_topk)
62 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
63 TritonOrDeepGemmExperts)
64
65 __all__ += [
65 - "fused_moe",
66 "fused_topk",
67 "fused_experts",
68 "get_config_file_name",
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+15 -25
@@ -8,6 +8,8 @@ import torch
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 from vllm.logger import init_logger
10 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11 +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
12 + deep_gemm_block_shape)
13 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
14 TopKWeightAndReduceDelegate)
15 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
@@ -212,27 +214,20 @@ def silu_mul_fp8_quant_deep_gemm_cuda(
214
215
216 class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
215 - # The Deep Gemm kernels only support block size of 128
216 - DEEPGEMM_BLOCK_SHAPE: list[int] = [128, 128]
217 -
218 - def __init__(self,
219 - max_num_tokens: int,
220 - num_dispatchers: int,
221 - block_shape: list[int],
222 - per_act_token_quant=False):
217 +
218 + def __init__(
219 + self,
220 + max_num_tokens: int,
221 + num_dispatchers: int,
222 + quant_config: FusedMoEQuantConfig,
223 + ):
224 """
225 max_num_tokens: Maximum number of tokens from a DP Rank
226 num_dispatchers: The number of DP dispatchers.
226 - block_shape: Block quantization block shape.
227 - per_act_token_quant: Per activation token quantization flag.
227 + quant_config: Quantization configuration
228 """
229 - super().__init__(
230 - FusedMoEQuantConfig(
231 - quant_dtype=torch.float8_e4m3fn,
232 - per_act_token_quant=per_act_token_quant,
233 - block_shape=block_shape,
234 - ))
235 - assert self.block_shape == self.DEEPGEMM_BLOCK_SHAPE
229 + super().__init__(quant_config)
230 + assert self.block_shape == deep_gemm_block_shape()
231 self.max_num_tokens = max_num_tokens
232 self.num_dispatchers = num_dispatchers
233
@@ -290,12 +285,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
285 activation: str,
286 global_num_experts: int,
287 expert_map: Optional[torch.Tensor],
293 - w1_scale: Optional[torch.Tensor],
294 - w2_scale: Optional[torch.Tensor],
295 - w1_zp: Optional[torch.Tensor],
296 - w2_zp: Optional[torch.Tensor],
288 a1q_scale: Optional[torch.Tensor],
298 - a2_scale: Optional[torch.Tensor],
289 workspace13: torch.Tensor,
290 workspace2: torch.Tensor,
291 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -321,11 +311,11 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
311 # for the M expectation of each batch, correctly setting this value
312 # may lead to better performance.
313 expected_m = max_num_tokens
324 - fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, w1_scale),
314 + fp8_m_grouped_gemm_nt_masked((a1q, a1q_scale), (w1, self.w1_scale),
315 workspace1, expert_num_tokens, expected_m)
316
317 a2q, a2q_scale = silu_mul_fp8_quant_deep_gemm_cuda(
318 workspace1, expert_num_tokens)
319
330 - fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, w2_scale), output,
331 - expert_num_tokens, expected_m)
320 + fp8_m_grouped_gemm_nt_masked((a2q, a2q_scale), (w2, self.w2_scale),
321 + output, expert_num_tokens, expected_m)
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+17 -41
@@ -8,55 +8,37 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
9 BatchedDeepGemmExperts)
10 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11 +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
12 + deep_gemm_block_shape)
13 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
14 BatchedTritonExperts)
15
16
17 class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
18
17 - def __init__(self,
18 - max_num_tokens: int,
19 - num_dispatchers: int,
20 - use_fp8_w8a8: bool = False,
21 - use_int8_w8a8: bool = False,
22 - use_int8_w8a16: bool = False,
23 - use_int4_w4a16: bool = False,
24 - block_shape: Optional[list[int]] = None,
25 - per_act_token_quant: bool = False,
26 - allow_deep_gemm: bool = False):
27 - assert not use_int8_w8a8, "NYI"
28 - assert not use_int8_w8a16, "NYI"
29 - assert not use_int4_w4a16, "NYI"
30 -
31 - super().__init__(
32 - FusedMoEQuantConfig.make(
33 - use_fp8_w8a8=use_fp8_w8a8,
34 - use_int8_w8a8=use_int8_w8a8,
35 - use_int8_w8a16=use_int8_w8a16,
36 - use_int4_w4a16=use_int4_w4a16,
37 - block_shape=block_shape,
38 - per_act_token_quant=per_act_token_quant,
39 - ))
19 + def __init__(
20 + self,
21 + max_num_tokens: int,
22 + num_dispatchers: int,
23 + quant_config: FusedMoEQuantConfig,
24 + allow_deep_gemm: bool = False,
25 + ):
26 + super().__init__(quant_config)
27
28 self.batched_triton_experts = BatchedTritonExperts(
29 max_num_tokens=max_num_tokens,
30 num_dispatchers=num_dispatchers,
44 - use_fp8_w8a8=use_fp8_w8a8,
45 - use_int8_w8a8=use_int8_w8a8,
46 - use_int8_w8a16=use_int8_w8a16,
47 - use_int4_w4a16=use_int4_w4a16,
48 - per_act_token_quant=self.per_act_token_quant,
49 - block_shape=self.block_shape,
31 + quant_config=self.quant_config,
32 )
33
52 - self.allow_deep_gemm = (allow_deep_gemm and use_fp8_w8a8
53 - and self.block_shape
54 - == BatchedDeepGemmExperts.DEEPGEMM_BLOCK_SHAPE)
34 + self.allow_deep_gemm = (allow_deep_gemm
35 + and self.quant_config.use_fp8_w8a8 and
36 + self.block_shape == deep_gemm_block_shape())
37
38 self.batched_deep_gemm_experts = BatchedDeepGemmExperts(
39 max_num_tokens=max_num_tokens,
40 num_dispatchers=num_dispatchers,
59 - block_shape=self.block_shape, # type: ignore[arg-type]
41 + quant_config=self.quant_config,
42 ) if self.allow_deep_gemm else None
43
44 assert (self.batched_deep_gemm_experts is not None
@@ -143,12 +125,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
125 activation: str,
126 global_num_experts: int,
127 expert_map: Optional[torch.Tensor],
146 - w1_scale: Optional[torch.Tensor],
147 - w2_scale: Optional[torch.Tensor],
148 - w1_zp: Optional[torch.Tensor],
149 - w2_zp: Optional[torch.Tensor],
128 a1q_scale: Optional[torch.Tensor],
151 - a2_scale: Optional[torch.Tensor],
129 workspace13: torch.Tensor,
130 workspace2: torch.Tensor,
131 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -158,7 +135,6 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
135 if self.allow_deep_gemm else self.batched_triton_experts)
136 assert experts is not None
137 experts.apply(output, hidden_states, w1, w2, topk_weights, topk_ids,
161 - activation, global_num_experts, expert_map, w1_scale,
162 - w2_scale, w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,
163 - workspace2, expert_tokens_meta,
138 + activation, global_num_experts, expert_map, a1q_scale,
139 + workspace13, workspace2, expert_tokens_meta,
140 apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/config.py
+482 -200
@@ -1,103 +1,322 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from dataclasses import dataclass
4 -from typing import Optional, Union
4 +from typing import TYPE_CHECKING, Optional, Union
5
6 import torch
7 -from compressed_tensors.quantization import (QuantizationArgs,
8 - QuantizationStrategy,
9 - QuantizationType)
7
8 import vllm.envs as envs
9 from vllm.config import ParallelConfig
10 from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
11 from vllm.logger import init_logger
15 -from vllm.model_executor.layers.quantization.base_config import (
16 - QuantizationConfig)
17 -from vllm.utils import cdiv
12 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
13 + GroupShape)
14 +from vllm.utils import cdiv, has_triton_kernels
15 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
16
17 +if TYPE_CHECKING and has_triton_kernels:
18 + from triton_kernels.matmul_ogs import PrecisionConfig
19 +
20 logger = init_logger(__name__)
21
22
23 -def _get_quant_config_quantization_args(
24 - quant_config: Optional[QuantizationConfig],
25 - prop_name: str,
26 -) -> Optional[QuantizationArgs]:
27 - if (quant_config is not None and hasattr(quant_config, 'target_scheme_map')
28 - and "Linear" in quant_config.target_scheme_map and
29 - "input_activations" in quant_config.target_scheme_map["Linear"]):
30 - return quant_config.target_scheme_map["Linear"].get(prop_name)
23 +def _get_config_dtype_str(
24 + dtype: torch.dtype,
25 + use_fp8_w8a8: bool = False,
26 + use_int8_w8a16: bool = False,
27 + use_int4_w4a16: bool = False,
28 + use_mxfp4_w4a4: bool = False,
29 +) -> Optional[str]:
30 + """
31 + Return a string used to construct the filename that contains the
32 + tuning info for a particular quantization scheme. See
33 + try_get_optimal_moe_config in fused_moe.py.
34 + """
35 + if use_fp8_w8a8:
36 + return "fp8_w8a8"
37 + elif use_int8_w8a16:
38 + return "int8_w8a16"
39 + elif use_int4_w4a16:
40 + return "int4_w4a16"
41 + elif use_mxfp4_w4a4:
42 + return "mxfp4_w4a4"
43 + elif dtype == torch.float:
44 + # avoiding cases where kernel fails when float32 MoE
45 + # use fp16/bfloat16 configs
46 + return "float32"
47 + return None
48 +
49 +
50 +def _quant_flags_to_group_shape(
51 + quant_dtype: Union[torch.dtype, str, None],
52 + per_act_token_quant: bool,
53 + per_out_ch_quant: bool,
54 + block_shape: Optional[list[int]],
55 +) -> tuple[Optional[GroupShape], Optional[GroupShape]]:
56 + """
57 + Convert MoE quantization flags into more generic GroupShapes.
58 + """
59 + a_shape: Optional[GroupShape]
60 + w_shape: Optional[GroupShape]
61 + if block_shape is not None:
62 + assert not per_act_token_quant
63 + assert not per_out_ch_quant
64 + # TODO(bnell): this is not quite right for activations since first
65 + # dim should be 1.
66 + a_shape = GroupShape(row=block_shape[0], col=block_shape[1])
67 + w_shape = GroupShape(row=block_shape[0], col=block_shape[1])
68 else:
32 - return None
69 + w_shape = None
70 + a_shape = None if quant_dtype is None else GroupShape.PER_TENSOR
71
72 + if per_act_token_quant:
73 + a_shape = GroupShape.PER_TOKEN
74
35 -def get_quant_config_input_quant(
36 - quant_config: Optional[QuantizationConfig]
37 -) -> Optional[QuantizationArgs]:
38 - return _get_quant_config_quantization_args(quant_config,
39 - "input_activations")
75 + if per_out_ch_quant:
76 + w_shape = GroupShape.PER_TOKEN
77
78 + return a_shape, w_shape
79
42 -def get_quant_config_weight_quant(
43 - quant_config: Optional[QuantizationConfig]
44 -) -> Optional[QuantizationArgs]:
45 - return _get_quant_config_quantization_args(quant_config, "weights")
80
81 +@dataclass
82 +class FusedMoEQuantDesc:
83 + """
84 + A quantization descriptor for fused MoE ops. This class can describe
85 + either activations or weights.
86 + """
87 +
88 + # The quantized type of this parameters. None means unquantized or
89 + # already quantized.
90 + # TODO (bnell): use scalar_type instead of Union.
91 + dtype: Union[torch.dtype, str, None] = None
92
48 -def get_config_quant_dtype(
49 - use_fp8_w8a8: bool,
50 - use_int8_w8a8: bool,
51 - use_int8_w8a16: bool,
52 - use_int4_w4a16: bool,
53 - use_mxfp4_w4a4: bool,
54 -) -> Union[None, torch.dtype, str]:
55 - if use_fp8_w8a8:
56 - return torch.float8_e4m3fn
57 - elif use_int8_w8a8:
58 - return torch.int8
59 - elif use_mxfp4_w4a4:
60 - return "mxfp4"
61 - return None
93 + # A field that describes the quantization group shape, from quant_utils.py.
94 + # * (-1, -1) for per-tensor quantization
95 + # * (1, -1) for per-row quantization
96 + # * (-1, 1) for per-column quantization
97 + # * (128, 128) for 128x128 deepseek style block quantization
98 + # * (1, 128) for deepseek style activation quantization
99 + # (i.e. per-token-per-group)
100 + shape: Optional[GroupShape] = None
101 +
102 + # Quantization scales.
103 + # TODO(bnell): maybe put PrecisionConfigs in subclass of QuantDesc?
104 + scale: Union[torch.Tensor, "PrecisionConfig", None] = None
105 +
106 + # Quantization alphas or gscales, used for nvfp4 types.
107 + # TODO(bnell): put some of these in subclasses
108 + alpha_or_gscale: Optional[torch.Tensor] = None
109 +
110 + # Zero points for int4/int8 types
111 + zp: Optional[torch.Tensor] = None
112 +
113 + # Biases for GPT triton MoE
114 + bias: Optional[torch.Tensor] = None
115
116
117 +# TODO(bnell): have subclasses for specific moe methods?
118 +# e.g. for specific arguments bias, precision, etc.
119 @dataclass
120 class FusedMoEQuantConfig:
66 - # The post quantization activation type.
67 - # TODO (bnell): use scalar_type instead of Union.
68 - quant_dtype: Union[torch.dtype, str, None] = None
69 - per_act_token_quant: bool = False
70 - per_out_ch_quant: bool = False
71 - block_shape: Optional[list[int]] = None
72 -
73 - # TODO: add col major flag?
74 - # add detailed quant info for input, intermediates, weights, etc?
121 + """
122 + The FusedMoEQuantConfig contains all the quantization parameters for
123 + a single FusedMoEMethodBase operation. It consists of four
124 + FusedMoEQuantDescs, one for each activation and set of weights.
125 +
126 + Each FusedMoEMethodBase must implement a get_fused_moe_quant_config
127 + method to construct a FusedMoEQuantConfig for use with that class.
128 +
129 + FusedMoEQuant configs are only used for modular kernels, fused_experts
130 + (from fused_moe.py), cutlass_moe_fp[48], rocm_aiter_fused_experts and
131 + triton_kernel_moe_forward. Other MoE methods can ignore the
132 + FusedMoEQuantConfig (for now) and hardcode it to None.
133 +
134 + There are currently some restrictions on what can be expressed:
135 + - Most MoE ops only support similar quantization strategies for
136 + each parameter, e.g. both weights must have the same GroupShape
137 + and both activations must share the same GroupShape. One exception to
138 + this is the cutlass moe which allows per channel quantization on the
139 + outputs. Note: this restrictions are not always rigorously checked.
140 + - Not all fused MoE functions support all the parameters, e.g. zero points,
141 + global scales, alphas and biases are not universally supported.
142 + - Fully general GroupShapes are not allowed. Activations only support
143 + per token, per tensor or K-blocked.
144 + - Weights are not required to have a GroupShape since they have already
145 + been quantized.
146 +
147 + Other notes:
148 + - PrecisionConfigs are specific to GPT OSS Triton.
149 + - As a follow up it would probably make sense to subclass FusedMoEQuantDesc
150 + or FusedMoEQuantConfig for particular FusedMoEMethodBase subclasses
151 + so that only the required quantization parameters are used/stored.
152 + """
153 +
154 + # TODO(bnell) make sure a1_scales/a2_scales don't interfere with chunking
155 + _a1: FusedMoEQuantDesc
156 + _a2: FusedMoEQuantDesc
157 + _w1: FusedMoEQuantDesc
158 + _w2: FusedMoEQuantDesc
159
160 def __post_init__(self):
161 assert (not self.per_act_token_quant
162 or self.block_shape is None), "illegal quantization"
163
164 + #
165 + # Convenience accessors for various properties.
166 + #
167 +
168 + @property
169 + def quant_dtype(self) -> Union[torch.dtype, str, None]:
170 + return self._a1.dtype
171 +
172 @property
173 def is_quantized(self) -> bool:
174 return self.quant_dtype is not None
175
176 @property
177 def is_per_act_token(self) -> bool:
86 - return self.per_act_token_quant
178 + return self._a1.shape == GroupShape.PER_TOKEN
179 +
180 + @property
181 + def per_act_token_quant(self) -> bool:
182 + return self._a1.shape == GroupShape.PER_TOKEN
183 +
184 + @property
185 + def per_out_ch_quant(self) -> bool:
186 + return self._w1.shape == GroupShape.PER_TOKEN
187 +
188 + @property
189 + def is_per_tensor(self) -> bool:
190 + return self._a1.shape == GroupShape.PER_TENSOR
191 +
192 + @property
193 + def block_shape(self) -> Optional[list[int]]:
194 + if (self._a1.shape is not None
195 + and self._a1.shape != GroupShape.PER_TENSOR
196 + and self._a1.shape != GroupShape.PER_TOKEN):
197 + return [self._a1.shape.row, self._a1.shape.col]
198 + else:
199 + return None
200
201 @property
202 def is_block_quantized(self) -> bool:
203 return self.block_shape is not None
204
205 @property
93 - def is_per_tensor(self) -> bool:
94 - return not self.per_act_token_quant and self.block_shape is None
206 + def a1_scale(self) -> Optional[torch.Tensor]:
207 + assert self._a1.scale is None or isinstance(self._a1.scale,
208 + torch.Tensor)
209 + return self._a1.scale
210 +
211 + @property
212 + def a1_gscale(self) -> Optional[torch.Tensor]:
213 + return self._a1.alpha_or_gscale
214 +
215 + @property
216 + def a2_scale(self) -> Optional[torch.Tensor]:
217 + assert self._a2.scale is None or isinstance(self._a2.scale,
218 + torch.Tensor)
219 + return self._a2.scale
220 +
221 + @property
222 + def a2_gscale(self) -> Optional[torch.Tensor]:
223 + return self._a2.alpha_or_gscale
224 +
225 + @property
226 + def w1_scale(self) -> Optional[torch.Tensor]:
227 + assert self._w1.scale is None or isinstance(self._w1.scale,
228 + torch.Tensor)
229 + return self._w1.scale
230 +
231 + @property
232 + def w1_zp(self) -> Optional[torch.Tensor]:
233 + return self._w1.zp
234 +
235 + @property
236 + def w1_bias(self) -> Optional[torch.Tensor]:
237 + return self._w1.bias
238 +
239 + @property
240 + def w1_precision(self) -> Optional["PrecisionConfig"]:
241 + assert self._w1.scale is None or isinstance(self._w1.scale,
242 + PrecisionConfig)
243 + return self._w1.scale
244 +
245 + @property
246 + def g1_alphas(self) -> Optional[torch.Tensor]:
247 + return self._w1.alpha_or_gscale
248 +
249 + @property
250 + def w2_scale(self) -> Optional[torch.Tensor]:
251 + assert self._w2.scale is None or isinstance(self._w2.scale,
252 + torch.Tensor)
253 + return self._w2.scale
254 +
255 + @property
256 + def w2_zp(self) -> Optional[torch.Tensor]:
257 + return self._w2.zp
258 +
259 + @property
260 + def w2_bias(self) -> Optional[torch.Tensor]:
261 + return self._w2.bias
262 +
263 + @property
264 + def w2_precision(self) -> Optional["PrecisionConfig"]:
265 + assert self._w2.scale is None or isinstance(self._w2.scale,
266 + PrecisionConfig)
267 + return self._w2.scale
268 +
269 + @property
270 + def g2_alphas(self) -> Optional[torch.Tensor]:
271 + return self._w2.alpha_or_gscale
272 +
273 + @property
274 + def use_fp8_w8a8(self) -> bool:
275 + return self.quant_dtype == torch.float8_e4m3fn
276 +
277 + @property
278 + def use_int8_w8a8(self) -> bool:
279 + return self.quant_dtype == torch.int8
280 +
281 + @property
282 + def use_int8_w8a16(self) -> bool:
283 + return (self._a1.dtype is None and self._w1.dtype == torch.int8)
284 +
285 + @property
286 + def use_int4_w4a16(self) -> bool:
287 + return (self._a1.dtype is None and self._w1.dtype == "int4")
288 +
289 + @property
290 + def use_mxfp4_w4a4(self) -> bool:
291 + return self.quant_dtype == "mxfp4"
292 +
293 + @property
294 + def use_nvfp4_w4a4(self) -> bool:
295 + return self.quant_dtype == "nvfp4"
296 +
297 + def config_name(self, dtype: torch.dtype) -> Optional[str]:
298 + """
299 + Return a string used to construct the filename that contains the
300 + tuning info for a particular quantization scheme. See
301 + try_get_optimal_moe_config in fused_moe.py.
302 + """
303 + return _get_config_dtype_str(
304 + use_fp8_w8a8=self.use_fp8_w8a8,
305 + use_int8_w8a16=self.use_int8_w8a16,
306 + use_int4_w4a16=self.use_int4_w4a16,
307 + use_mxfp4_w4a4=self.use_mxfp4_w4a4,
308 + dtype=dtype,
309 + )
310
311 def scale_shape(
312 self,
313 max_tokens: int,
314 hidden_dim: int,
315 ) -> Optional[tuple[int, int]]:
316 + """
317 + Construct the proper activation scale shape for this
318 + config.
319 + """
320 if self.is_quantized:
321 if self.is_block_quantized:
322 assert self.block_shape is not None
@@ -117,6 +336,10 @@ class FusedMoEQuantConfig:
336 max_tokens: int,
337 hidden_dim: int,
338 ) -> Optional[tuple[int, int, int]]:
339 + """
340 + Construct the proper activation batched scale shape for this
341 + config, e.g. (num experts, *scale_shape).
342 + """
343 if self.is_quantized:
344 scale_shape = self.scale_shape(max_tokens, hidden_dim)
345 assert scale_shape is not None
@@ -126,38 +349,218 @@ class FusedMoEQuantConfig:
349
350 @staticmethod
351 def make(
129 - use_fp8_w8a8: bool = False,
130 - use_int8_w8a8: bool = False,
131 - use_int8_w8a16: bool = False,
132 - use_int4_w4a16: bool = False,
133 - use_mxfp4_w4a4: bool = False,
352 + quant_dtype: Union[torch.dtype, str, None] = None,
353 per_act_token_quant: bool = False,
354 per_out_ch_quant: bool = False,
355 block_shape: Optional[list[int]] = None,
356 + w1_scale: Union[torch.Tensor, "PrecisionConfig", None] = None,
357 + w2_scale: Union[torch.Tensor, "PrecisionConfig", None] = None,
358 + a1_scale: Optional[torch.Tensor] = None,
359 + a2_scale: Optional[torch.Tensor] = None,
360 + g1_alphas: Optional[torch.Tensor] = None,
361 + g2_alphas: Optional[torch.Tensor] = None,
362 + a1_gscale: Optional[torch.Tensor] = None,
363 + a2_gscale: Optional[torch.Tensor] = None,
364 + w1_bias: Optional[torch.Tensor] = None,
365 + w2_bias: Optional[torch.Tensor] = None,
366 + w1_zp: Optional[torch.Tensor] = None,
367 + w2_zp: Optional[torch.Tensor] = None,
368 ) -> "FusedMoEQuantConfig":
138 - assert sum([
139 - int(flag) for flag in [
140 - use_fp8_w8a8,
141 - use_int8_w8a8,
142 - use_int8_w8a16,
143 - use_int4_w4a16,
144 - use_mxfp4_w4a4,
145 - ]
146 - ]) <= 1, "Quantization flags are mutually exclusive."
147 -
148 - quant_dtype = get_config_quant_dtype(
149 - use_fp8_w8a8=use_fp8_w8a8,
150 - use_int8_w8a8=use_int8_w8a8,
151 - use_int8_w8a16=use_int8_w8a16,
152 - use_int4_w4a16=use_int4_w4a16,
153 - use_mxfp4_w4a4=use_mxfp4_w4a4,
154 - )
155 - return FusedMoEQuantConfig(
156 - quant_dtype,
157 - per_act_token_quant,
158 - per_out_ch_quant,
159 - block_shape,
369 + """
370 + General builder function for a FusedMoEQuantConfig.
371 + - quant_dtype: Optional quantization type. None if activations are
372 + unquantized or quantized prior to calling. Note: "nvfp4" and
373 + "mxfp4" are the only valid string values for quant_dtype.
374 + - per_act_token_quant: Activations have per token quantization.
375 + - per_out_ch_quant: Outputs have per channel quantization. (only
376 + for cutlass).
377 + - block_shape: Optional block size for block-wise quantization.
378 + Incompatible with per_act_token and per_out_ch quant.
379 + - w1_scale: Optional scale to be used for w1.
380 + - w2_scale: Optional scale to be used for w2.
381 + - a1_scale: Optional scale to be used for a1.
382 + - a2_scale: Optional scale to be used for a2.
383 + - g1_alphas: Optional global quantization scales for w1 (for nvfp4).
384 + - g2_alphas: Optional global quantization scales for w2 (for nvfp4).
385 + - a1_gscale: Optional global quantization scales for a1 (for nvfp4).
386 + - a2_gscale: Optional global quantization scales for a2 (for nvfp4).
387 + - w1_bias: Optional biases for w1 (GPT OSS Triton).
388 + - w2_bias: Optional biases for w1 (GPT OSS Triton).
389 + - w1_zp: Optional w1 zero points for int4/int8 quantization.
390 + - w2_zp: Optional w2 zero points for int4/int8 quantization.
391 + """
392 + assert (not isinstance(quant_dtype, str) or quant_dtype == "nvfp4"
393 + or quant_dtype == "mxfp4")
394 + a_shape, w_shape = _quant_flags_to_group_shape(quant_dtype,
395 + per_act_token_quant,
396 + per_out_ch_quant,
397 + block_shape)
398 + quant_config = FusedMoEQuantConfig(
399 + _a1=FusedMoEQuantDesc(quant_dtype, a_shape, a1_scale, a1_gscale),
400 + _a2=FusedMoEQuantDesc(quant_dtype, a_shape, a2_scale, a2_gscale),
401 + _w1=FusedMoEQuantDesc(quant_dtype, w_shape, w1_scale, g1_alphas,
402 + w1_zp, w1_bias),
403 + _w2=FusedMoEQuantDesc(quant_dtype, w_shape, w2_scale, g2_alphas,
404 + w2_zp, w2_bias),
405 )
406 + assert quant_config.per_act_token_quant == per_act_token_quant
407 + assert quant_config.per_out_ch_quant == per_out_ch_quant
408 + assert quant_config.block_shape == block_shape
409 + return quant_config
410 +
411 +
412 +def fp8_w8a8_moe_quant_config(
413 + w1_scale: torch.Tensor,
414 + w2_scale: torch.Tensor,
415 + a1_scale: Optional[torch.Tensor] = None,
416 + a2_scale: Optional[torch.Tensor] = None,
417 + per_act_token_quant: bool = False,
418 + per_out_ch_quant: bool = False,
419 + block_shape: Optional[list[int]] = None,
420 +) -> FusedMoEQuantConfig:
421 + """
422 + Construct a quant config for fp8 activations and fp8 weights.
423 + """
424 + return FusedMoEQuantConfig.make(torch.float8_e4m3fn,
425 + w1_scale=w1_scale,
426 + w2_scale=w2_scale,
427 + a1_scale=a1_scale,
428 + a2_scale=a2_scale,
429 + per_act_token_quant=per_act_token_quant,
430 + per_out_ch_quant=per_out_ch_quant,
431 + block_shape=block_shape)
432 +
433 +
434 +def int8_w8a8_moe_quant_config(
435 + w1_scale: torch.Tensor,
436 + w2_scale: torch.Tensor,
437 + a1_scale: Optional[torch.Tensor],
438 + a2_scale: Optional[torch.Tensor],
439 + per_act_token_quant: bool = False,
440 +) -> FusedMoEQuantConfig:
441 + """
442 + Construct a quant config for int8 activations and int8 weights.
443 + """
444 + return FusedMoEQuantConfig.make(
445 + torch.int8,
446 + w1_scale=w1_scale,
447 + w2_scale=w2_scale,
448 + a1_scale=a1_scale,
449 + a2_scale=a2_scale,
450 + per_act_token_quant=per_act_token_quant,
451 + per_out_ch_quant=False,
452 + block_shape=None,
453 + )
454 +
455 +
456 +def mxfp4_w4a4_moe_quant_config(
457 + w1_scale: Union[torch.Tensor, "PrecisionConfig"],
458 + w2_scale: Union[torch.Tensor, "PrecisionConfig"],
459 + a1_scale: Optional[torch.Tensor] = None,
460 + a2_scale: Optional[torch.Tensor] = None,
461 + w1_bias: Optional[torch.Tensor] = None,
462 + w2_bias: Optional[torch.Tensor] = None,
463 + block_shape: Optional[list[int]] = None,
464 +) -> FusedMoEQuantConfig:
465 + """
466 + Construct a quant config for mxfp4 activations and mxfp4 weights.
467 + """
468 + return FusedMoEQuantConfig.make(
469 + "mxfp4",
470 + w1_scale=w1_scale,
471 + w2_scale=w2_scale,
472 + a1_scale=a1_scale,
473 + a2_scale=a2_scale,
474 + w1_bias=w1_bias,
475 + w2_bias=w2_bias,
476 + per_act_token_quant=False,
477 + per_out_ch_quant=False,
478 + block_shape=block_shape,
479 + )
480 +
481 +
482 +def nvfp4_moe_quant_config(
483 + g1_alphas: torch.Tensor,
484 + g2_alphas: torch.Tensor,
485 + a1_gscale: torch.Tensor,
486 + a2_gscale: torch.Tensor,
487 + w1_scale: torch.Tensor,
488 + w2_scale: torch.Tensor,
489 +) -> FusedMoEQuantConfig:
490 + """
491 + Construct a quant config for mxfp4 activations and nvp4 weights.
492 + """
493 + return FusedMoEQuantConfig.make(
494 + "nvfp4",
495 + w1_scale=w1_scale,
496 + w2_scale=w2_scale,
497 + a1_gscale=a1_gscale,
498 + a2_gscale=a2_gscale,
499 + g1_alphas=g1_alphas,
500 + g2_alphas=g2_alphas,
501 + per_act_token_quant=False,
502 + per_out_ch_quant=False,
503 + block_shape=None,
504 + )
505 +
506 +
507 +def int4_w4a16_moe_quant_config(
508 + w1_scale: torch.Tensor,
509 + w2_scale: torch.Tensor,
510 + w1_zp: Optional[torch.Tensor],
511 + w2_zp: Optional[torch.Tensor],
512 + block_shape: Optional[list[int]] = None,
513 +) -> FusedMoEQuantConfig:
514 + """
515 + Construct a quant config for 16-bit float activations and int4 weights.
516 + Note: Activations are pre-quantized.
517 + """
518 + group_shape = GroupShape(*block_shape) if block_shape is not None else None
519 + return FusedMoEQuantConfig(
520 + _a1=FusedMoEQuantDesc(shape=group_shape),
521 + _a2=FusedMoEQuantDesc(shape=group_shape),
522 + _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp),
523 + _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp),
524 + )
525 +
526 +
527 +def int8_w8a16_moe_quant_config(
528 + w1_scale: torch.Tensor,
529 + w2_scale: torch.Tensor,
530 + w1_zp: Optional[torch.Tensor],
531 + w2_zp: Optional[torch.Tensor],
532 + block_shape: Optional[list[int]] = None,
533 +) -> FusedMoEQuantConfig:
534 + """
535 + Construct a quant config for 16-bit float activations and int8 weights.
536 + Note: Activations are pre-quantized.
537 + """
538 + group_shape = GroupShape(*block_shape) if block_shape is not None else None
539 + return FusedMoEQuantConfig(
540 + _a1=FusedMoEQuantDesc(shape=group_shape),
541 + _a2=FusedMoEQuantDesc(shape=group_shape),
542 + _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp),
543 + _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp),
544 + )
545 +
546 +
547 +def biased_moe_quant_config(
548 + w1_bias: Optional[torch.Tensor],
549 + w2_bias: Optional[torch.Tensor],
550 +) -> FusedMoEQuantConfig:
551 + """
552 + Construct a quant config for unquantized activations with biases.
553 + """
554 + return FusedMoEQuantConfig(
555 + _a1=FusedMoEQuantDesc(),
556 + _a2=FusedMoEQuantDesc(),
557 + _w1=FusedMoEQuantDesc(bias=w1_bias),
558 + _w2=FusedMoEQuantDesc(bias=w2_bias),
559 + )
560 +
561 +
562 +# A FusedMoEQuantConfig constant for an unquantized MoE op.
563 +FUSED_MOE_UNQUANTIZED_CONFIG: FusedMoEQuantConfig = FusedMoEQuantConfig.make()
564
565
566 @dataclass
@@ -315,8 +718,6 @@ class FusedMoEConfig:
718 # The activation type.
719 in_dtype: torch.dtype
720
318 - quant_config: Optional[FusedMoEQuantConfig] = None
319 -
721 max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
722
723 has_bias: bool = False
@@ -328,34 +729,6 @@ class FusedMoEConfig:
729
730 assert self.max_num_tokens > 0
731
331 - @property
332 - def quant_dtype(self) -> Union[torch.dtype, str, None]:
333 - if self.quant_config is not None:
334 - return self.quant_config.quant_dtype
335 - else:
336 - return None
337 -
338 - @property
339 - def block_shape(self) -> Optional[list[int]]:
340 - if self.quant_config is not None:
341 - return self.quant_config.block_shape
342 - else:
343 - return None
344 -
345 - @property
346 - def per_act_token_quant(self) -> bool:
347 - if self.quant_config is not None:
348 - return self.quant_config.per_act_token_quant
349 - else:
350 - return False
351 -
352 - @property
353 - def per_out_ch_quant(self) -> bool:
354 - if self.quant_config is not None:
355 - return self.quant_config.per_out_ch_quant
356 - else:
357 - return False
358 -
732 @property
733 def tp_size(self):
734 return self.moe_parallel_config.tp_size
@@ -401,97 +774,6 @@ class FusedMoEConfig:
774 """
775 Whether to use FlashInfer cutlass kernels for NVFP4 MoE.
776 """
404 - return (self.quant_config is not None
405 - and self.quant_config.quant_dtype == "nvfp4"
406 - and envs.VLLM_USE_FLASHINFER_MOE_FP4
777 + return (envs.VLLM_USE_FLASHINFER_MOE_FP4
778 and has_flashinfer_cutlass_fused_moe()
779 and envs.VLLM_FLASHINFER_MOE_BACKEND == "throughput")
409 -
410 - @staticmethod
411 - def make(
412 - num_experts: int,
413 - experts_per_token: int,
414 - hidden_dim: int,
415 - num_local_experts: int,
416 - moe_parallel_config: FusedMoEParallelConfig,
417 - in_dtype: torch.dtype,
418 - max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE,
419 - quant_config: Optional[Union[FusedMoEQuantConfig,
420 - QuantizationConfig]] = None,
421 - has_bias: bool = False,
422 - ) -> "FusedMoEConfig":
423 -
424 - _quant_config: Optional[FusedMoEQuantConfig] = None
425 -
426 - if quant_config is not None and isinstance(quant_config,
427 - QuantizationConfig):
428 - if hasattr(quant_config, 'weight_block_size'):
429 - block_shape = quant_config.weight_block_size
430 - else:
431 - block_shape = None
432 - per_act_token_quant = False
433 - per_out_ch_quant = False
434 - quant_dtype: Union[torch.dtype, str, None] = None
435 -
436 - input_quant = get_quant_config_input_quant(quant_config)
437 - weight_quant = get_quant_config_weight_quant(quant_config)
438 -
439 - if input_quant is not None:
440 - per_act_token_quant = (input_quant.strategy
441 - == QuantizationStrategy.TOKEN
442 - if input_quant is not None else False)
443 -
444 - if input_quant.num_bits == 8:
445 - if input_quant.type == QuantizationType.FLOAT:
446 - quant_dtype = torch.float8_e4m3fn
447 - elif input_quant.type == QuantizationType.INT:
448 - quant_dtype = torch.int8
449 -
450 - from vllm.model_executor.layers.quantization.fp8 import Fp8Config
451 - if quant_dtype is None and isinstance(quant_config, Fp8Config):
452 - quant_dtype = torch.float8_e4m3fn
453 -
454 - from vllm.model_executor.layers.quantization.mxfp4 import (
455 - Mxfp4Config)
456 - if (quant_dtype is None and isinstance(quant_config, Mxfp4Config)
457 - and envs.VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8):
458 - quant_dtype = "mxfp8"
459 -
460 - from vllm.model_executor.layers.quantization.modelopt import (
461 - ModelOptNvFp4Config)
462 - if quant_dtype is None and isinstance(quant_config,
463 - ModelOptNvFp4Config):
464 - quant_dtype = "nvfp4"
465 -
466 - if weight_quant is not None:
467 - per_out_ch_quant = (
468 - weight_quant.strategy == QuantizationStrategy.CHANNEL)
469 -
470 - if quant_dtype is not None:
471 - _quant_config = FusedMoEQuantConfig(
472 - quant_dtype=quant_dtype,
473 - per_act_token_quant=per_act_token_quant,
474 - per_out_ch_quant=per_out_ch_quant,
475 - block_shape=block_shape,
476 - )
477 - else:
478 - _quant_config = FusedMoEQuantConfig()
479 - if moe_parallel_config.dp_size > 1:
480 - logger.warning_once("MoE DP setup unable to determine "
481 - "quantization scheme or unsupported "
482 - "quantization type. This model will "
483 - "not run with DP enabled.")
484 - else:
485 - _quant_config = quant_config
486 -
487 - return FusedMoEConfig(
488 - num_experts=num_experts,
489 - experts_per_token=experts_per_token,
490 - hidden_dim=hidden_dim,
491 - num_local_experts=num_local_experts,
492 - moe_parallel_config=moe_parallel_config,
493 - in_dtype=in_dtype,
494 - quant_config=_quant_config,
495 - max_num_tokens=max_num_tokens,
496 - has_bias=has_bias,
497 - )
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+62 -103
@@ -211,21 +211,14 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):
211 def __init__(
212 self,
213 out_dtype: Optional[torch.dtype],
214 - per_act_token_quant: bool,
215 - per_out_ch_quant: bool,
214 ab_strides1: torch.Tensor,
215 ab_strides2: torch.Tensor,
216 c_strides1: torch.Tensor,
217 c_strides2: torch.Tensor,
220 - block_shape: Optional[list[int]] = None,
218 + quant_config: FusedMoEQuantConfig,
219 ):
222 - super().__init__(
223 - FusedMoEQuantConfig(
224 - quant_dtype=torch.float8_e4m3fn,
225 - per_act_token_quant=per_act_token_quant,
226 - per_out_ch_quant=per_out_ch_quant,
227 - block_shape=block_shape,
228 - ))
220 + assert quant_config.use_fp8_w8a8
221 + super().__init__(quant_config)
222 self.out_dtype = out_dtype
223 self.ab_strides1 = ab_strides1
224 self.ab_strides2 = ab_strides2
@@ -247,19 +240,14 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):
240 activation: str,
241 global_num_experts: int,
242 expert_map: Optional[torch.Tensor],
250 - w1_scale: Optional[torch.Tensor],
251 - w2_scale: Optional[torch.Tensor],
252 - w1_zp: Optional[torch.Tensor],
253 - w2_zp: Optional[torch.Tensor],
243 a1q_scale: Optional[torch.Tensor],
255 - a2_scale: Optional[torch.Tensor],
244 workspace13: torch.Tensor,
245 workspace2: torch.Tensor,
246 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
247 apply_router_weight_on_input: bool,
248 ):
261 - assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
262 - assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
249 + assert self.w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
250 + assert self.w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
251
252 expert_num_tokens = None
253 if expert_tokens_meta is not None:
@@ -273,9 +261,10 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):
261 in_dtype = hidden_states.dtype
262 run_cutlass_moe_fp8(
263 output, hidden_states, w1, w2, topk_ids, activation_callable,
276 - global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,
277 - a2_scale, self.ab_strides1, self.ab_strides2, self.c_strides1,
278 - self.c_strides2, workspace13, workspace2, expert_num_tokens,
264 + global_num_experts, expert_map, self.w1_scale, self.w2_scale,
265 + a1q_scale, self.a2_scale, self.ab_strides1, self.ab_strides2,
266 + self.c_strides1, self.c_strides2, workspace13, workspace2,
267 + expert_num_tokens,
268 self.out_dtype if self.out_dtype is not None else in_dtype,
269 self.per_act_token_quant, self.per_out_ch_quant,
270 use_batched_format, topk_weights)
@@ -286,23 +275,19 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
275 def __init__(
276 self,
277 out_dtype: Optional[torch.dtype],
289 - per_act_token_quant: bool,
290 - per_out_ch_quant: bool,
278 ab_strides1: torch.Tensor,
279 ab_strides2: torch.Tensor,
280 c_strides1: torch.Tensor,
281 c_strides2: torch.Tensor,
295 - block_shape: Optional[list[int]] = None,
282 + quant_config: FusedMoEQuantConfig,
283 ):
284 super().__init__(
285 out_dtype,
299 - per_act_token_quant,
300 - per_out_ch_quant,
286 ab_strides1,
287 ab_strides2,
288 c_strides1,
289 c_strides2,
305 - block_shape,
290 + quant_config,
291 )
292
293 @property
@@ -348,23 +333,19 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
333 max_experts_per_worker: int,
334 num_dispatchers: int,
335 out_dtype: Optional[torch.dtype],
351 - per_act_token_quant: bool,
352 - per_out_ch_quant: bool,
336 ab_strides1: torch.Tensor,
337 ab_strides2: torch.Tensor,
338 c_strides1: torch.Tensor,
339 c_strides2: torch.Tensor,
357 - block_shape: Optional[list[int]] = None,
340 + quant_config: FusedMoEQuantConfig,
341 ):
342 super().__init__(
343 out_dtype,
361 - per_act_token_quant,
362 - per_out_ch_quant,
344 ab_strides1,
345 ab_strides2,
346 c_strides1,
347 c_strides2,
367 - block_shape,
348 + quant_config,
349 )
350 assert max_experts_per_worker > 0
351 self.max_experts_per_worker = max_experts_per_worker
@@ -414,16 +395,12 @@ def cutlass_moe_fp8(
395 w2_q: torch.Tensor,
396 topk_weights: torch.Tensor,
397 topk_ids: torch.Tensor,
417 - w1_scale: torch.Tensor,
418 - w2_scale: torch.Tensor,
398 ab_strides1: torch.Tensor,
399 ab_strides2: torch.Tensor,
400 c_strides1: torch.Tensor,
401 c_strides2: torch.Tensor,
423 - per_act_token: Optional[bool] = None,
402 + quant_config: FusedMoEQuantConfig,
403 activation: str = "silu",
425 - a1_scale: Optional[torch.Tensor] = None,
426 - a2_scale: Optional[torch.Tensor] = None,
404 expert_map: Optional[torch.Tensor] = None,
405 apply_router_weight_on_input: bool = False,
406 global_num_experts: int = -1,
@@ -475,10 +452,18 @@ def cutlass_moe_fp8(
452 Returns:
453 - torch.Tensor: The fp16 output tensor after applying the MoE layer.
454 """
478 - if per_act_token is None:
479 - per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
480 - a2_scale.numel() != 1 if a2_scale is not None else False)
481 - per_out_ch = w1_scale.numel() != w1_q.size(0)
455 + assert quant_config is not None
456 +
457 + if quant_config.a1_scale is not None:
458 + assert (quant_config.per_act_token_quant ==
459 + quant_config.a1_scale.numel() != 1)
460 + if quant_config.a2_scale is not None:
461 + assert (quant_config.per_act_token_quant ==
462 + quant_config.a2_scale.numel() != 1)
463 +
464 + assert (quant_config.w1_scale is None
465 + or (quant_config.per_out_ch_quant == (quant_config.w1_scale.size(1)
466 + == w1_q.size(1))))
467
468 num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(
469 0)
@@ -487,12 +472,11 @@ def cutlass_moe_fp8(
472 MoEPrepareAndFinalizeNoEP(),
473 CutlassExpertsFp8(
474 out_dtype=a.dtype,
490 - per_act_token_quant=per_act_token,
491 - per_out_ch_quant=per_out_ch,
475 ab_strides1=ab_strides1,
476 ab_strides2=ab_strides2,
477 c_strides1=c_strides1,
478 c_strides2=c_strides2,
479 + quant_config=quant_config,
480 ),
481 )
482
@@ -502,14 +486,9 @@ def cutlass_moe_fp8(
486 w2_q,
487 topk_weights,
488 topk_ids,
505 - False,
506 - activation,
507 - num_experts,
508 - expert_map,
509 - w1_scale,
510 - w2_scale,
511 - a1_scale=a1_scale,
512 - a2_scale=a2_scale,
489 + activation=activation,
490 + global_num_experts=num_experts,
491 + expert_map=expert_map,
492 apply_router_weight_on_input=apply_router_weight_on_input,
493 )
494
@@ -542,7 +521,7 @@ def run_cutlass_moe_fp4(
521 ) -> None:
522 """
523 MoE implementation for FP4 Inputs
545 -
524 +
525 # Gemm 1
526 a: Input tensor: [m, k] (half/bfloat16)
527 a1_gscale: Activation scale per expert: [e] (float32)
@@ -552,16 +531,16 @@ def run_cutlass_moe_fp4(
531 full precision)
532 w1_blockscale: [e, 2 * n, k // block_size] (float8_e4m3)
533 (Block size = 16 for NVFP4)
555 -
534 +
535 # Gemm 2
536 a2_gscale: Activation scale per expert: [e]
537 w2(down projection) (not an argument to cutlass_moe_fp4): [e, k, n]
538 w2_fp4: [e, k, n // 2], dtype: torch.uint8 (stacked E2M1)
539 w2_blockscale: [e, k, n // block_size], dtype: float8_e4m3
561 -
540 +
541 topk_weights: [m, topk] dtype: float8
542 topk_ids: [m, topk] dtype: float8
564 -
543 +
544 m, n, k: Unquantized weight shapes, dtype: int
545 e: number of experts, dtype: int
546
@@ -652,42 +631,21 @@ def run_cutlass_moe_fp4(
631 return
632
633
634 +# Split into batched and non-batched
635 class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
636
637 def __init__(
638 self,
659 - g1_alphas: torch.Tensor,
660 - g2_alphas: torch.Tensor,
661 - a1_gscale: torch.Tensor,
662 - a2_gscale: torch.Tensor,
639 max_experts_per_worker: int,
640 out_dtype: torch.dtype,
665 - per_act_token_quant: bool,
666 - per_out_ch_quant: bool,
667 - block_shape: Optional[list[int]] = None,
641 + quant_config: FusedMoEQuantConfig,
642 use_batched_format: bool = False,
643 ):
670 - super().__init__(
671 - # NVFP4 requires two levels of quantization, which involves
672 - # computing some scaling factors dynamically. This makes it
673 - # incompatible with the typical prepare -> MoE -> finalize
674 - # pipeline. Move the quantization logic into the MoE body.
675 - FusedMoEQuantConfig(
676 - quant_dtype=None, # skip quantization in prepare/finalize
677 - per_act_token_quant=per_act_token_quant,
678 - per_out_ch_quant=per_out_ch_quant,
679 - block_shape=block_shape,
680 - ))
644 + super().__init__(quant_config)
645 self.max_experts_per_worker = max_experts_per_worker
646 self.out_dtype = out_dtype
647 self.use_batched_format = use_batched_format
648
685 - # TODO(bnell): put this stuff into quant config?
686 - self.g1_alphas = g1_alphas
687 - self.g2_alphas = g2_alphas
688 - self.a1_gscale = a1_gscale
689 - self.a2_gscale = a2_gscale
690 -
649 @property
650 def activation_formats(
651 self
@@ -746,12 +704,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
704 activation: str,
705 global_num_experts: int,
706 expert_map: Optional[torch.Tensor],
749 - w1_scale: torch.Tensor,
750 - w2_scale: torch.Tensor,
751 - w1_zp: Optional[torch.Tensor],
752 - w2_zp: Optional[torch.Tensor],
753 - a1q_scale: Optional[torch.Tensor],
754 - a2_scale: torch.Tensor,
707 + a1q_scale: Optional[torch.Tensor], # unused
708 workspace13: Optional[torch.Tensor],
709 workspace2: Optional[torch.Tensor],
710 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -765,11 +718,11 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
718 a=hidden_states,
719 a1_gscale=self.a1_gscale,
720 w1_fp4=w1,
768 - w1_blockscale=w1_scale,
721 + w1_blockscale=self.w1_scale,
722 w1_alphas=self.g1_alphas,
723 a2_gscale=self.a2_gscale,
724 w2_fp4=w2,
772 - w2_blockscale=w2_scale,
725 + w2_blockscale=self.w2_scale,
726 w2_alphas=self.g2_alphas,
727 topk_weights=topk_weights,
728 topk_ids=topk_ids,
@@ -788,14 +741,9 @@ def cutlass_moe_fp4(
741 a: torch.Tensor,
742 w1_fp4: torch.Tensor,
743 w2_fp4: torch.Tensor,
791 - w1_blockscale: torch.Tensor,
792 - w2_blockscale: torch.Tensor,
793 - g1_alphas: torch.Tensor,
794 - g2_alphas: torch.Tensor,
795 - a1_gscale: torch.Tensor,
796 - a2_gscale: torch.Tensor,
744 topk_weights: torch.Tensor,
745 topk_ids: torch.Tensor,
746 + quant_config: FusedMoEQuantConfig,
747 m: int,
748 n: int,
749 k: int,
@@ -805,17 +753,31 @@ def cutlass_moe_fp4(
753 assert expert_map is None, ("Expert Parallelism / expert_map "
754 "is currently not supported for "
755 "ModelOptNvFp4FusedMoE's cutlass_moe_fp4.")
756 +
757 + # TODO(bnell): this feels a bit hacky
758 + # NVFP4 requires two levels of quantization, which involves
759 + # computing some scaling factors dynamically. This makes it
760 + # incompatible with the typical prepare -> MoE -> finalize
761 + # pipeline. Move the quantization logic into the MoE body.
762 + quant_config = FusedMoEQuantConfig.make(
763 + quant_dtype=None, # skip quantization in prepare/finalize
764 + per_act_token_quant=quant_config.per_act_token_quant,
765 + per_out_ch_quant=quant_config.per_out_ch_quant,
766 + block_shape=quant_config.block_shape,
767 + g1_alphas=quant_config.g1_alphas,
768 + g2_alphas=quant_config.g2_alphas,
769 + a1_gscale=quant_config.a1_gscale,
770 + a2_gscale=quant_config.a2_gscale,
771 + w1_scale=quant_config.w1_scale,
772 + w2_scale=quant_config.w2_scale,
773 + )
774 +
775 fn = mk.FusedMoEModularKernel(
776 MoEPrepareAndFinalizeNoEP(),
777 CutlassExpertsFp4(
811 - g1_alphas,
812 - g2_alphas,
813 - a1_gscale,
814 - a2_gscale,
778 max_experts_per_worker=e,
779 out_dtype=a.dtype,
817 - per_act_token_quant=False,
818 - per_out_ch_quant=False,
780 + quant_config=quant_config,
781 use_batched_format=False,
782 ),
783 )
@@ -830,10 +792,6 @@ def cutlass_moe_fp4(
792 activation="silu",
793 global_num_experts=e,
794 expert_map=None,
833 - w1_scale=w1_blockscale,
834 - w2_scale=w2_blockscale,
835 - a1_scale=None,
836 - a2_scale=None,
795 apply_router_weight_on_input=apply_router_weight_on_input,
796 )
797
@@ -891,6 +849,7 @@ def _valid_cutlass_block_scaled_grouped_gemm(
849 return True
850
851
852 +# TODO(bnell): would be nice combine/integrate with regular cutlass_fp8.
853 def run_cutlass_block_scaled_fused_experts(
854 a: torch.Tensor,
855 w1: torch.Tensor,
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+28 -37
@@ -1,6 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import functools
3 from typing import Optional
4
5 import torch
@@ -9,9 +8,11 @@ from tqdm import tqdm
8 import vllm.envs as env
9 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10 from vllm.logger import init_logger
12 -from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11 +from vllm.model_executor.layers.fused_moe.config import (
12 + FusedMoEQuantConfig, fp8_w8a8_moe_quant_config)
13 from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
14 - compute_aligned_M, deepgemm_moe_permute, deepgemm_unpermute_and_reduce)
14 + compute_aligned_M, deep_gemm_block_shape, deepgemm_moe_permute,
15 + deepgemm_unpermute_and_reduce)
16 from vllm.model_executor.layers.fused_moe.prepare_finalize import (
17 MoEPrepareAndFinalizeNoEP)
18 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
@@ -25,14 +26,6 @@ from vllm.utils.deep_gemm import m_grouped_fp8_gemm_nt_contiguous
26 logger = init_logger(__name__)
27
28
28 -@functools.cache
29 -def deep_gemm_block_shape() -> list[int]:
30 - # Lazy import to avoid CUDA initialization problems.
31 - import deep_gemm as dg
32 - block = dg.get_m_alignment_for_contiguous_layout()
33 - return [block, block]
34 -
35 -
29 def _valid_deep_gemm_shape(M: int, N: int, K: int) -> bool:
30 align = deep_gemm_block_shape()[0]
31 return align <= M and N % align == 0 and K % align == 0
@@ -163,13 +156,12 @@ def warmup_deepgemm_gg_contiguous_kernels(w1: torch.Tensor, w2: torch.Tensor,
156
157 class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
158
166 - def __init__(self):
167 - super().__init__(
168 - FusedMoEQuantConfig(
169 - quant_dtype=torch.float8_e4m3fn,
170 - per_act_token_quant=False,
171 - block_shape=deep_gemm_block_shape(),
172 - ))
159 + def __init__(self, quant_config: FusedMoEQuantConfig):
160 + super().__init__(quant_config)
161 + assert quant_config.block_shape == deep_gemm_block_shape()
162 + assert quant_config.quant_dtype == torch.float8_e4m3fn
163 + assert not quant_config.per_act_token_quant
164 + assert not quant_config.per_out_ch_quant
165
166 @property
167 def activation_formats(
@@ -221,21 +213,17 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
213 activation: str,
214 global_num_experts: int,
215 expert_map: Optional[torch.Tensor],
224 - w1_scale: Optional[torch.Tensor],
225 - w2_scale: Optional[torch.Tensor],
226 - w1_zp: Optional[torch.Tensor],
227 - w2_zp: Optional[torch.Tensor],
216 a1q_scale: Optional[torch.Tensor],
229 - a2_scale: Optional[torch.Tensor],
217 workspace13: torch.Tensor,
218 workspace2: torch.Tensor,
219 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
220 apply_router_weight_on_input: bool,
221 ):
235 - assert self.block_shape is not None
222 assert a1q_scale is not None
237 - assert w1_scale is not None
238 - assert w2_scale is not None
223 + assert self.a2_scale is None
224 + assert self.block_shape is not None
225 + assert self.w1_scale is not None
226 + assert self.w2_scale is not None
227
228 a1q = hidden_states
229 _, N, K = w1.size()
@@ -270,7 +258,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
258 aq_out=a1q_perm)
259 assert a1q.size(0) == M_sum
260
273 - m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, w1_scale),
261 + m_grouped_fp8_gemm_nt_contiguous((a1q, a1q_scale), (w1, self.w1_scale),
262 mm1_out, expert_ids)
263
264 self.activation(activation, act_out, mm1_out.view(-1, N))
@@ -281,7 +269,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
269 column_major_scales=True,
270 out_q=quant_out)
271
284 - m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, w2_scale),
272 + m_grouped_fp8_gemm_nt_contiguous((a2q, a2q_scale), (w2, self.w2_scale),
273 mm2_out, expert_ids)
274
275 if apply_router_weight_on_input:
@@ -348,9 +336,16 @@ def deep_gemm_moe_fp8(
336 Returns:
337 - torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
338 """
339 + quant_config = fp8_w8a8_moe_quant_config(
340 + w1_scale=w1_scale,
341 + w2_scale=w2_scale,
342 + a1_scale=a1_scale,
343 + a2_scale=a2_scale,
344 + block_shape=deep_gemm_block_shape())
345 +
346 fn = mk.FusedMoEModularKernel(
347 MoEPrepareAndFinalizeNoEP(),
353 - DeepGemmExperts(),
348 + DeepGemmExperts(quant_config),
349 )
350 return fn(
351 hidden_states,
@@ -358,13 +353,9 @@ def deep_gemm_moe_fp8(
353 w2,
354 topk_weights,
355 topk_ids,
361 - inplace,
362 - activation,
363 - global_num_experts,
364 - expert_map,
365 - w1_scale=w1_scale,
366 - w2_scale=w2_scale,
367 - a1_scale=a1_scale,
368 - a2_scale=a2_scale,
356 + inplace=inplace,
357 + activation=activation,
358 + global_num_experts=global_num_experts,
359 + expert_map=expert_map,
360 apply_router_weight_on_input=apply_router_weight_on_input,
361 )
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+4 -9
@@ -183,8 +183,6 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
183 def prepare_async(
184 self,
185 a1: torch.Tensor,
186 - a1_scale: Optional[torch.Tensor],
187 - a2_scale: Optional[torch.Tensor],
186 topk_weights: torch.Tensor,
187 topk_ids: torch.Tensor,
188 num_experts: int,
@@ -204,7 +202,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
202 # Quant and Dispatch
203 a1q, a1q_scale = moe_kernel_quantize_input(
204 a1,
207 - a1_scale,
205 + quant_config.a1_scale,
206 quant_dtype=quant_config.quant_dtype,
207 per_act_token_quant=quant_config.per_act_token_quant,
208 block_shape=quant_config.block_shape,
@@ -215,7 +213,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
213 else:
214 a1q = a1
215 a1q_scale = None
218 - a1_post_scale = a1_scale
216 + a1_post_scale = quant_config.a1_scale
217
218 return (lambda *args: None,
219 self._do_dispatch(tokens=a1q,
@@ -229,8 +227,6 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
227 def prepare(
228 self,
229 a1: torch.Tensor,
232 - a1_scale: Optional[torch.Tensor],
233 - a2_scale: Optional[torch.Tensor],
230 topk_weights: torch.Tensor,
231 topk_ids: torch.Tensor,
232 num_experts: int,
@@ -238,9 +234,8 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
234 apply_router_weight_on_input: bool,
235 quant_config: FusedMoEQuantConfig,
236 ) -> mk.PrepareResultType:
241 - (_, receiver) = self.prepare_async(a1, a1_scale, a2_scale,
242 - topk_weights, topk_ids, num_experts,
243 - expert_map,
237 + (_, receiver) = self.prepare_async(a1, topk_weights, topk_ids,
238 + num_experts, expert_map,
239 apply_router_weight_on_input,
240 quant_config)
241 return receiver()
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+20 -25
@@ -77,15 +77,13 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
77 def _do_quant(
78 self,
79 x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
80 - a1_scale: Optional[torch.Tensor],
80 a1_dtype: torch.dtype,
82 - quant_dtype: Union[torch.dtype, str, None],
83 - per_act_token_quant: bool,
84 - block_shape: Optional[list[int]],
81 + quant_config: FusedMoEQuantConfig,
82 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
83
87 - block_k = block_shape[1] if block_shape is not None else None
84 if self.use_fp8_dispatch:
85 + block_k = quant_config.block_shape[
86 + 1] if quant_config.block_shape is not None else None
87 if block_k == DEEPEP_QUANT_BLOCK_SIZE:
88 # DeepEP kernels did the quantization for us.
89 x, x_scales = x
@@ -101,12 +99,12 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
99
100 # TODO (varun): Optimization - Use a batched version of quant
101 x = x.view((-1, hidden_dim))
104 - x, x_scales = moe_kernel_quantize_input(x, a1_scale, quant_dtype,
105 - per_act_token_quant,
106 - block_shape)
102 + x, x_scales = moe_kernel_quantize_input(
103 + x, quant_config.a1_scale, quant_config.quant_dtype,
104 + quant_config.per_act_token_quant, quant_config.block_shape)
105 x = x.view((num_experts, -1, hidden_dim))
106
109 - if quant_dtype is not None:
107 + if quant_config.quant_dtype is not None:
108 assert x_scales is not None
109 x_scales = normalize_batched_scales_shape(x_scales, num_experts)
110
@@ -118,8 +116,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
116 def prepare_async(
117 self,
118 a1: torch.Tensor,
121 - a1_scale: Optional[torch.Tensor],
122 - a2_scale: Optional[torch.Tensor],
119 topk_weights: torch.Tensor,
120 topk_ids: torch.Tensor,
121 num_experts: int,
@@ -139,9 +135,10 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
135 assert hidden_size % 128 == 0, \
136 "DeepEP kernels quantize the inputs in blocks of shape 128"
137
142 - has_per_token_scales = a1_scale.numel(
143 - ) != 1 if a1_scale is not None else (
144 - a2_scale.numel() != 1 if a2_scale is not None else False)
138 + has_per_token_scales = quant_config.a1_scale.numel(
139 + ) != 1 if quant_config.a1_scale is not None else (
140 + quant_config.a2_scale.numel() != 1
141 + if quant_config.a2_scale is not None else False)
142 assert not has_per_token_scales, (
143 "low_latency kernels doesn't support dispatching per-token scales")
144
@@ -163,20 +160,21 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
160 return_recv_hook=True)
161 self.handles[a2a_idx] = handle
162
166 - return (hook, lambda: self._receiver(expert_x, expert_num_tokens,
167 - a1_scale, a1.dtype, quant_config))
163 + return (
164 + hook,
165 + lambda: self._receiver(expert_x, expert_num_tokens, quant_config.
166 + a1_scale, a1.dtype, quant_config))
167
168 def _receiver(
169 self,
170 expert_x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
171 expert_num_tokens: torch.Tensor,
173 - a1_scale,
174 - a1_dtype,
172 + a1_scale: Optional[torch.Tensor],
173 + a1_dtype: torch.dtype,
174 quant_config: FusedMoEQuantConfig,
175 ) -> mk.PrepareResultType:
177 - expert_x, expert_x_scale = self._do_quant(
178 - expert_x, a1_scale, a1_dtype, quant_config.quant_dtype,
179 - quant_config.per_act_token_quant, quant_config.block_shape)
176 + expert_x, expert_x_scale = self._do_quant(expert_x, a1_dtype,
177 + quant_config)
178
179 expert_tokens_meta = mk.ExpertTokensMetadata(
180 expert_num_tokens=expert_num_tokens, expert_num_tokens_cpu=None)
@@ -186,8 +184,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
184 def prepare(
185 self,
186 a1: torch.Tensor,
189 - a1_scale: Optional[torch.Tensor],
190 - a2_scale: Optional[torch.Tensor],
187 topk_weights: torch.Tensor,
188 topk_ids: torch.Tensor,
189 num_experts: int,
@@ -195,8 +191,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
191 apply_router_weight_on_input: bool,
192 quant_config: FusedMoEQuantConfig,
193 ) -> mk.PrepareResultType:
198 - hook, receiver = self.prepare_async(a1, a1_scale, a2_scale,
199 - topk_weights, topk_ids,
194 + hook, receiver = self.prepare_async(a1, topk_weights, topk_ids,
195 num_experts, expert_map,
196 apply_router_weight_on_input,
197 quant_config)
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+10 -40
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import Optional, Union
3 +from typing import Optional
4
5 import torch
6
@@ -44,33 +44,20 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
44
45 def __init__(
46 self,
47 - g1_alphas: torch.Tensor,
48 - g2_alphas: torch.Tensor,
49 - a1_gscale: torch.Tensor,
50 - a2_gscale: torch.Tensor,
47 out_dtype: torch.dtype,
52 - quant_dtype: Union[torch.dtype, str, None],
48 + quant_config: FusedMoEQuantConfig,
49 ep_rank: int = 0,
50 ep_size: int = 1,
51 tp_rank: int = 0,
52 tp_size: int = 1,
53 ):
58 - super().__init__(
59 - FusedMoEQuantConfig(
60 - quant_dtype=quant_dtype,
61 - per_act_token_quant=False,
62 - block_shape=None,
63 - ))
64 - assert quant_dtype in ("nvfp4", torch.float8_e4m3fn), (
54 + super().__init__(quant_config)
55 + assert quant_config.quant_dtype in ("nvfp4", torch.float8_e4m3fn), (
56 "Only nvfp4,fp8 quantization are currently supported.")
57 self.ep_rank = ep_rank
58 self.ep_size = ep_size
59 self.tp_rank = tp_rank
60 self.tp_size = tp_size
70 - self.g1_alphas = g1_alphas
71 - self.g2_alphas = g2_alphas
72 - self.a1_gscale = a1_gscale
73 - self.a2_gscale = a2_gscale
61 self.out_dtype = out_dtype
62
63 @property
@@ -141,12 +128,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
128 activation: str,
129 global_num_experts: int,
130 expert_map: Optional[torch.Tensor],
144 - w1_scale: Optional[torch.Tensor],
145 - w2_scale: Optional[torch.Tensor],
146 - w1_zp: Optional[torch.Tensor],
147 - w2_zp: Optional[torch.Tensor],
131 a1q_scale: Optional[torch.Tensor],
149 - a2_scale: Optional[torch.Tensor], # Not used
132 workspace13: Optional[torch.Tensor],
133 workspace2: Optional[torch.Tensor],
134 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -162,17 +144,17 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
144 fc2_expert_weights = w2
145 else:
146 # Ensure w1_scale and w2_scale are not None before calling view
165 - assert w1_scale is not None and w2_scale is not None, (
147 + assert self.w1_scale is not None and self.w2_scale is not None, (
148 "w1_scale and w2_scale must not "
149 "be None for FlashInferExperts")
150 # Flashinfer CUTLASS kernel takes scalar global scales,
151 # min because inv_scale.
152 quant_scales = [
153 self.a1_gscale,
172 - w1_scale.view(torch.int32),
154 + self.w1_scale.view(torch.int32),
155 self.g1_alphas,
156 self.a2_gscale,
175 - w2_scale.view(torch.int32),
157 + self.w2_scale.view(torch.int32),
158 self.g2_alphas,
159 ]
160 # FlashInfer API requires weight to be long for nvfp4
@@ -202,12 +184,7 @@ def flashinfer_cutlass_moe_fp4(
184 w2: torch.Tensor,
185 topk_weights: torch.Tensor,
186 topk_ids: torch.Tensor,
205 - w1_scale: torch.Tensor,
206 - w2_scale: torch.Tensor,
207 - g1_alphas: torch.Tensor,
208 - g2_alphas: torch.Tensor,
209 - a1_gscale: torch.Tensor,
210 - a2_gscale: torch.Tensor,
187 + quant_config: FusedMoEQuantConfig,
188 inplace: bool = False,
189 activation: str = "silu",
190 global_num_experts: int = -1,
@@ -216,15 +193,10 @@ def flashinfer_cutlass_moe_fp4(
193 ) -> torch.Tensor:
194
195 fused_experts = mk.FusedMoEModularKernel(
219 - FlashInferCutlassMoEPrepareAndFinalize(use_dp=False,
220 - a1_gscale=a1_gscale),
196 + FlashInferCutlassMoEPrepareAndFinalize(use_dp=False),
197 FlashInferExperts(
222 - g1_alphas=g1_alphas,
223 - g2_alphas=g2_alphas,
224 - a1_gscale=a1_gscale,
225 - a2_gscale=a2_gscale,
198 out_dtype=hidden_states.dtype,
227 - quant_dtype="nvfp4",
199 + quant_config=quant_config,
200 ))
201
202 return fused_experts(
@@ -237,7 +209,5 @@ def flashinfer_cutlass_moe_fp4(
209 activation=activation,
210 global_num_experts=global_num_experts,
211 expert_map=expert_map,
240 - w1_scale=w1_scale,
241 - w2_scale=w2_scale,
212 apply_router_weight_on_input=apply_router_weight_on_input,
213 )
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+1 -6
@@ -22,13 +22,11 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
22 def __init__(
23 self,
24 use_dp: bool,
25 - a1_gscale: Optional[torch.Tensor],
25 num_dispatchers: int = 1,
26 ):
27 super().__init__()
28 self.num_dispatchers_ = num_dispatchers
29 self.use_dp = use_dp
31 - self.a1_gscale = a1_gscale
30 self.local_tokens = None
31
32 @property
@@ -47,14 +45,11 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
45 def prepare(
46 self,
47 a1: torch.Tensor,
50 - a1_scale: Optional[torch.Tensor], # Not used
51 - a2_scale: Optional[torch.Tensor], # Not used
48 topk_weights: torch.Tensor,
49 topk_ids: torch.Tensor,
50 num_experts: int,
51 expert_map: Optional[torch.Tensor],
52 apply_router_weight_on_input: bool,
57 - # TODO(bnell): use quant_config + scales instead of ctor args
53 quant_config: FusedMoEQuantConfig,
54 ) -> mk.PrepareResultType:
55
@@ -67,7 +62,7 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
62
63 a1q, a1q_scale = moe_kernel_quantize_input(
64 a1,
70 - self.a1_gscale,
65 + quant_config.a1_gscale,
66 quant_config.quant_dtype,
67 quant_config.per_act_token_quant,
68 quant_config.block_shape,
vllm/model_executor/layers/fused_moe/flashinfer_trtllm_moe.py new
+185
@@ -0,0 +1,185 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import List # noqa: UP035
4 +from typing import Optional
5 +
6 +import torch
7 +
8 +from vllm.model_executor.layers.fused_moe.utils import (
9 + moe_kernel_quantize_input)
10 +from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
11 + calculate_tile_tokens_dim)
12 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13 + per_token_group_quant_fp8)
14 +from vllm.utils import direct_register_custom_op
15 +
16 +
17 +def flashinfer_fused_moe_blockscale_fp8(
18 + routing_logits: torch.Tensor,
19 + routing_bias: torch.Tensor,
20 + x: torch.Tensor,
21 + w13_weight: torch.Tensor,
22 + w13_weight_scale_inv: torch.Tensor,
23 + w2_weight: torch.Tensor,
24 + w2_weight_scale_inv: torch.Tensor,
25 + global_num_experts: int,
26 + top_k: int,
27 + num_expert_group: int,
28 + topk_group: int,
29 + intermediate_size: int,
30 + expert_offset: int,
31 + local_num_experts: int,
32 + block_shape: List[int], #noqa: UP006
33 + routed_scaling: float = 1.0) -> torch.Tensor:
34 + from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
35 + assert top_k <= global_num_experts
36 + assert top_k <= 8
37 + assert topk_group <= 4
38 + assert global_num_experts > num_expert_group
39 + assert global_num_experts % num_expert_group == 0
40 + assert global_num_experts % 4 == 0
41 + assert top_k < (topk_group * global_num_experts / num_expert_group)
42 + assert block_shape == [128, 128]
43 +
44 + a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])
45 + # NOTE: scales of hidden states have to be transposed!
46 + a_sf_t = a_sf.t().contiguous()
47 + return flashinfer_trtllm_fp8_block_scale_moe(
48 + routing_logits=routing_logits,
49 + routing_bias=routing_bias,
50 + hidden_states=a_q,
51 + hidden_states_scale=a_sf_t,
52 + gemm1_weights=w13_weight,
53 + gemm1_weights_scale=w13_weight_scale_inv,
54 + gemm2_weights=w2_weight,
55 + gemm2_weights_scale=w2_weight_scale_inv,
56 + num_experts=global_num_experts,
57 + top_k=top_k,
58 + n_group=num_expert_group,
59 + topk_group=topk_group,
60 + intermediate_size=intermediate_size,
61 + local_expert_offset=expert_offset,
62 + local_num_experts=local_num_experts,
63 + routed_scaling_factor=routed_scaling,
64 + tile_tokens_dim=calculate_tile_tokens_dim(x.shape[0], top_k,
65 + global_num_experts),
66 + routing_method_type=2, # DeepSeek-styled routing method
67 + use_shuffled_weight=False,
68 + )
69 +
70 +
71 +def flashinfer_fused_moe_blockscale_fp8_fake(
72 + routing_logits: torch.Tensor,
73 + routing_bias: torch.Tensor,
74 + x: torch.Tensor,
75 + w13_weight: torch.Tensor,
76 + w13_weight_scale_inv: torch.Tensor,
77 + w2_weight: torch.Tensor,
78 + w2_weight_scale_inv: torch.Tensor,
79 + global_num_experts: int,
80 + top_k: int,
81 + num_expert_group: int,
82 + topk_group: int,
83 + intermediate_size: int,
84 + expert_offset: int,
85 + local_num_experts: int,
86 + block_shape: list[int],
87 + routed_scaling: float = 1.0) -> torch.Tensor:
88 + return torch.empty_like(x)
89 +
90 +
91 +# TODO(bnell): Does this really need to be a torch.op?
92 +direct_register_custom_op(
93 + op_name="flashinfer_fused_moe_blockscale_fp8",
94 + op_func=flashinfer_fused_moe_blockscale_fp8,
95 + mutates_args=[],
96 + fake_impl=flashinfer_fused_moe_blockscale_fp8_fake,
97 + tags=(torch.Tag.needs_fixed_stride_order, ),
98 +)
99 +
100 +
101 +def flashinfer_fused_moe_per_tensor_scale_fp8(
102 + routing_logits: torch.Tensor,
103 + routing_bias: Optional[torch.Tensor],
104 + hidden_states: torch.Tensor,
105 + input_scale: torch.Tensor,
106 + gemm1_weights: torch.Tensor,
107 + gemm2_weights: torch.Tensor,
108 + output1_scales_scalar: torch.Tensor,
109 + output1_scales_gate_scalar: torch.Tensor,
110 + output2_scales_scalar: torch.Tensor,
111 + num_experts: int,
112 + top_k: int,
113 + num_expert_group: Optional[int],
114 + topk_group: Optional[int],
115 + intermediate_size: int,
116 + local_expert_offset: int,
117 + local_num_experts: int,
118 + use_routing_scales_on_input: bool,
119 + routing_method_type: int,
120 + routed_scaling_factor: float = 1.0) -> torch.Tensor:
121 + num_expert_group = num_expert_group if num_expert_group is not None else 0
122 + topk_group = topk_group if topk_group is not None else 0
123 +
124 + quant_hidden_states, _ = moe_kernel_quantize_input(
125 + hidden_states,
126 + input_scale,
127 + quant_dtype=torch.float8_e4m3fn,
128 + per_act_token_quant=False)
129 +
130 + from vllm.utils.flashinfer import (
131 + flashinfer_trtllm_fp8_per_tensor_scale_moe)
132 + return flashinfer_trtllm_fp8_per_tensor_scale_moe(
133 + routing_logits=routing_logits,
134 + routing_bias=routing_bias,
135 + hidden_states=quant_hidden_states,
136 + gemm1_weights=gemm1_weights,
137 + output1_scales_scalar=output1_scales_scalar,
138 + output1_scales_gate_scalar=output1_scales_gate_scalar,
139 + gemm2_weights=gemm2_weights,
140 + output2_scales_scalar=output2_scales_scalar,
141 + num_experts=num_experts,
142 + top_k=top_k,
143 + n_group=num_expert_group,
144 + topk_group=topk_group,
145 + intermediate_size=intermediate_size,
146 + local_expert_offset=local_expert_offset,
147 + local_num_experts=local_num_experts,
148 + routed_scaling_factor=routed_scaling_factor,
149 + use_routing_scales_on_input=use_routing_scales_on_input,
150 + tile_tokens_dim=calculate_tile_tokens_dim(hidden_states.shape[0],
151 + top_k, num_experts),
152 + routing_method_type=routing_method_type)
153 +
154 +
155 +def flashinfer_fused_moe_per_tensor_scale_fp8_fake(
156 + routing_logits: torch.Tensor,
157 + routing_bias: Optional[torch.Tensor],
158 + hidden_states: torch.Tensor,
159 + input_scale: torch.Tensor,
160 + gemm1_weights: torch.Tensor,
161 + gemm2_weights: torch.Tensor,
162 + output1_scales_scalar: torch.Tensor,
163 + output1_scales_gate_scalar: torch.Tensor,
164 + output2_scales_scalar: torch.Tensor,
165 + num_experts: int,
166 + top_k: int,
167 + num_expert_group: Optional[int],
168 + topk_group: Optional[int],
169 + intermediate_size: int,
170 + local_expert_offset: int,
171 + local_num_experts: int,
172 + use_routing_scales_on_input: bool,
173 + routing_method_type: int,
174 + routed_scaling_factor: float = 1.0) -> torch.Tensor:
175 + return torch.empty_like(hidden_states)
176 +
177 +
178 +# TODO(bnell): Does this really need to be a torch.op?
179 +direct_register_custom_op(
180 + op_name="flashinfer_fused_moe_per_tensor_scale_fp8",
181 + op_func=flashinfer_fused_moe_per_tensor_scale_fp8,
182 + mutates_args=["hidden_states"],
183 + fake_impl=flashinfer_fused_moe_per_tensor_scale_fp8_fake,
184 + tags=(torch.Tag.needs_fixed_stride_order, ),
185 +)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+34 -85
@@ -8,7 +8,7 @@ import torch
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10 from vllm.model_executor.layers.fused_moe.fused_moe import (
11 - get_config_dtype_str, try_get_optimal_moe_config)
11 + try_get_optimal_moe_config)
12 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
13 TopKWeightAndReduceDelegate, TopKWeightAndReduceNaiveBatched)
14 from vllm.model_executor.layers.fused_moe.utils import (
@@ -498,8 +498,6 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
498 def prepare(
499 self,
500 a1: torch.Tensor,
501 - a1_scale: Optional[torch.Tensor],
502 - a2_scale: Optional[torch.Tensor],
501 topk_weights: torch.Tensor,
502 topk_ids: torch.Tensor,
503 num_experts: int,
@@ -545,14 +543,13 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
543 dtype=torch.float32,
544 device=a1.device)
545 else:
548 - assert a1_scale is None
546 + assert quant_config.a1_scale is None
547 b_a1_scale = None
548
549 first_expert = num_local_experts * self.rank
550 last_expert = first_expert + num_local_experts
551
554 - a1_scale = normalize_scales_shape(a1_scale)
555 - a2_scale = normalize_scales_shape(a2_scale)
552 + a1_scale = normalize_scales_shape(quant_config.a1_scale)
553
554 for expert_id in range(first_expert, last_expert):
555 topks = torch.any(topk_ids == expert_id, dim=1).flatten()
@@ -623,28 +620,13 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
620 self,
621 max_num_tokens: int,
622 num_dispatchers: int,
626 - use_fp8_w8a8: bool = False,
627 - use_int8_w8a8: bool = False,
628 - use_int8_w8a16: bool = False,
629 - use_int4_w4a16: bool = False,
630 - use_mxfp4_w4a4: bool = False,
631 - block_shape: Optional[list[int]] = None,
632 - per_act_token_quant: bool = False,
623 + quant_config: FusedMoEQuantConfig,
624 ):
634 - super().__init__(
635 - FusedMoEQuantConfig.make(
636 - use_fp8_w8a8=use_fp8_w8a8,
637 - use_int8_w8a8=use_int8_w8a8,
638 - use_int8_w8a16=use_int8_w8a16,
639 - use_int4_w4a16=use_int4_w4a16,
640 - use_mxfp4_w4a4=use_mxfp4_w4a4,
641 - per_act_token_quant=per_act_token_quant,
642 - block_shape=block_shape,
643 - ))
644 - assert not use_int8_w8a8, "NYI"
645 - assert not use_int8_w8a16, "NYI"
646 - assert not use_int4_w4a16, "NYI"
647 - assert not use_mxfp4_w4a4, "NYI"
625 + super().__init__(quant_config)
626 + assert not self.quant_config.use_int8_w8a8, "NYI"
627 + assert not self.quant_config.use_int8_w8a16, "NYI"
628 + assert not self.quant_config.use_int4_w4a16, "NYI"
629 + assert not self.quant_config.use_mxfp4_w4a4, "NYI"
630 self.max_num_tokens = max_num_tokens
631 self.num_dispatchers = num_dispatchers
632
@@ -705,12 +687,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
687 activation: str,
688 global_num_experts: int,
689 expert_map: Optional[torch.Tensor],
708 - w1_scale: Optional[torch.Tensor],
709 - w2_scale: Optional[torch.Tensor],
710 - w1_zp: Optional[torch.Tensor],
711 - w2_zp: Optional[torch.Tensor],
690 a1q_scale: Optional[torch.Tensor],
713 - a2_scale: Optional[torch.Tensor],
691 workspace13: torch.Tensor,
692 workspace2: torch.Tensor,
693 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -740,10 +717,10 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
717 tmp = _resize_cache(workspace2, (num, N))
718
719 if self.quant_config.is_quantized:
743 - assert a1q_scale is not None and w1_scale is not None
720 + assert a1q_scale is not None and self.w1_scale is not None
721 input = self.dequant(hidden_states[expert, :, :],
722 a1q_scale[expert])
746 - w1_dq = self.dequant(w1[expert], w1_scale[expert])
723 + w1_dq = self.dequant(w1[expert], self.w1_scale[expert])
724 input = input[:num] @ w1_dq.transpose(0, 1)
725 else:
726 input = hidden_states[expert, :num, :] @ w1[expert].transpose(
@@ -752,8 +729,8 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
729 self.activation(activation, tmp, input.to(tmp.dtype))
730
731 if self.quant_config.is_quantized:
755 - assert w2_scale is not None
756 - w2_dq = self.dequant(w2[expert], w2_scale[expert])
732 + assert self.w2_scale is not None
733 + w2_dq = self.dequant(w2[expert], self.w2_scale[expert])
734 else:
735 w2_dq = w2[expert]
736
@@ -840,35 +817,15 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
817 self,
818 max_num_tokens: int,
819 num_dispatchers: int,
843 - use_fp8_w8a8: bool = False,
844 - use_int8_w8a8: bool = False,
845 - use_int8_w8a16: bool = False,
846 - use_int4_w4a16: bool = False,
847 - use_mxfp4_w4a4: bool = False,
848 - per_act_token_quant: bool = False,
849 - block_shape: Optional[list[int]] = None,
820 + quant_config: FusedMoEQuantConfig,
821 ):
851 - super().__init__(
852 - FusedMoEQuantConfig.make(
853 - use_fp8_w8a8=use_fp8_w8a8,
854 - use_int8_w8a8=use_int8_w8a8,
855 - use_int8_w8a16=use_int8_w8a16,
856 - use_int4_w4a16=use_int4_w4a16,
857 - use_mxfp4_w4a4=use_mxfp4_w4a4,
858 - per_act_token_quant=per_act_token_quant,
859 - block_shape=block_shape,
860 - ))
861 - assert not use_int8_w8a8, "NYI"
862 - assert not use_int8_w8a16, "NYI"
863 - assert not use_int4_w4a16, "NYI"
864 - assert not use_mxfp4_w4a4, "NYI"
822 + super().__init__(quant_config)
823 + assert not self.quant_config.use_int8_w8a8, "NYI"
824 + assert not self.quant_config.use_int8_w8a16, "NYI"
825 + assert not self.quant_config.use_int4_w4a16, "NYI"
826 + assert not self.quant_config.use_mxfp4_w4a4, "NYI"
827 assert max_num_tokens > 0
828 assert num_dispatchers > 0
867 - self.use_fp8_w8a8 = use_fp8_w8a8
868 - self.use_int8_w8a8 = use_int8_w8a8
869 - self.use_int4_w4a16 = use_int4_w4a16
870 - self.use_int8_w8a16 = use_int8_w8a16
871 - self.use_mxfp4_w4a4 = use_mxfp4_w4a4
829 self.max_num_tokens = max_num_tokens
830 self.num_dispatchers = num_dispatchers
831
@@ -921,19 +878,14 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
878 activation: str,
879 global_num_experts: int,
880 expert_map: Optional[torch.Tensor],
924 - w1_scale: Optional[torch.Tensor],
925 - w2_scale: Optional[torch.Tensor],
926 - w1_zp: Optional[torch.Tensor],
927 - w2_zp: Optional[torch.Tensor],
881 a1q_scale: Optional[torch.Tensor],
929 - a2_scale: Optional[torch.Tensor],
882 workspace13: torch.Tensor,
883 workspace2: torch.Tensor,
884 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
885 apply_router_weight_on_input: bool,
886 ):
887 # Check constraints.
936 - if self.use_int4_w4a16:
888 + if self.quant_config.use_int4_w4a16:
889 assert hidden_states.size(-1) // 2 == w1.size(2), (
890 "Hidden size mismatch")
891 else:
@@ -958,11 +910,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
910 assert w1.size(0) == E
911 assert w2.size(0) == E
912
961 - config_dtype = get_config_dtype_str(use_fp8_w8a8=self.use_fp8_w8a8,
962 - use_int8_w8a16=self.use_int8_w8a16,
963 - use_int4_w4a16=self.use_int4_w4a16,
964 - use_mxfp4_w4a4=self.use_mxfp4_w4a4,
965 - dtype=hidden_states.dtype)
913 + config_dtype = self.quant_config.config_name(hidden_states.dtype)
914
915 config = try_get_optimal_moe_config(
916 w1.size(),
@@ -992,7 +940,8 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
940 intermediate_cache2 = _resize_cache(workspace2,
941 (E, max_num_tokens, N // 2))
942
995 - if self.use_fp8_w8a8:
943 + # TODO(bnell): should this be done for any quantized type?
944 + if self.quant_config.use_fp8_w8a8:
945 intermediate_cache1.fill_(0)
946
947 a1q_scale = normalize_batched_scales_shape(a1q_scale, E)
@@ -1005,11 +954,11 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
954 expert_num_tokens=expert_num_tokens,
955 compute_type=compute_type,
956 A_scale=a1q_scale,
1008 - B_scale=w1_scale,
1009 - B_zp=w1_zp,
1010 - use_fp8_w8a8=self.use_fp8_w8a8,
1011 - use_int8_w8a16=self.use_int8_w8a16,
1012 - use_int4_w4a16=self.use_int4_w4a16,
957 + B_scale=self.w1_scale,
958 + B_zp=self.w1_zp,
959 + use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
960 + use_int8_w8a16=self.quant_config.use_int8_w8a16,
961 + use_int4_w4a16=self.quant_config.use_int4_w4a16,
962 config=config,
963 per_act_token_quant=self.per_act_token_quant,
964 block_shape=self.block_shape)
@@ -1021,7 +970,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
970 intermediate_cache1.view(-1, N))
971
972 qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
1024 - intermediate_cache2, a2_scale, max_num_tokens, E, N,
973 + intermediate_cache2, self.a2_scale, max_num_tokens, E, N,
974 expert_num_tokens, self.quant_dtype, self.per_act_token_quant,
975 self.block_shape)
976
@@ -1032,11 +981,11 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
981 expert_num_tokens=expert_num_tokens,
982 compute_type=compute_type,
983 A_scale=a2q_scale,
1035 - B_scale=w2_scale,
1036 - B_zp=w2_zp,
1037 - use_fp8_w8a8=self.use_fp8_w8a8,
1038 - use_int8_w8a16=self.use_int8_w8a16,
1039 - use_int4_w4a16=self.use_int4_w4a16,
984 + B_scale=self.w2_scale,
985 + B_zp=self.w2_zp,
986 + use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
987 + use_int8_w8a16=self.quant_config.use_int8_w8a16,
988 + use_int4_w4a16=self.quant_config.use_int4_w4a16,
989 config=config,
990 per_act_token_quant=self.per_act_token_quant,
991 block_shape=self.block_shape)
vllm/model_executor/layers/fused_moe/fused_moe.py
+194 -552
@@ -1,13 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Fused MoE kernel."""
3 +"""Fused MoE Triton kernels."""
4 import functools
5 import json
6 import os
7 # torch.compile needs typing.List. It will fail torch.library.infer_schema
8 # otherwise
9 from typing import List # noqa: UP035
10 -from typing import Any, Callable, Optional
10 +from typing import Any, Callable, Optional, Union
11
12 import torch
13 import torch.nn.functional as F
@@ -18,7 +18,7 @@ from vllm import _custom_ops as ops
18 from vllm.logger import init_logger
19 # yapf: disable
20 from vllm.model_executor.layers.fused_moe.config import (
21 - FusedMoEQuantConfig, get_config_quant_dtype)
21 + FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEQuantConfig, _get_config_dtype_str)
22 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
23 _valid_cutlass_block_scaled_grouped_gemm,
24 run_cutlass_block_scaled_fused_experts)
@@ -32,11 +32,7 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
32 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
33 TopKWeightAndReduceNoOP)
34 from vllm.model_executor.layers.fused_moe.utils import (
35 - _resize_cache, moe_kernel_quantize_input)
36 -from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
37 - calculate_tile_tokens_dim)
38 -from vllm.model_executor.layers.quantization.utils.fp8_utils import (
39 - per_token_group_quant_fp8)
35 + _resize_cache, activation_without_mul, moe_kernel_quantize_input)
36 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
37 dequant_mxfp4)
38 from vllm.platforms import current_platform
@@ -1049,87 +1045,66 @@ def fused_grouped_topk(
1045 return topk_values.to(torch.float32), topk_indices.to(torch.int32)
1046
1047
1052 -def get_config_dtype_str(
1053 - dtype: torch.dtype,
1054 - use_int4_w4a16: Optional[bool] = False,
1055 - use_int8_w8a16: Optional[bool] = False,
1056 - use_fp8_w8a8: Optional[bool] = False,
1057 - use_mxfp4_w4a4: Optional[bool] = False) -> Optional[str]:
1058 - if use_fp8_w8a8:
1059 - return "fp8_w8a8"
1060 - elif use_int8_w8a16:
1061 - return "int8_w8a16"
1062 - elif use_int4_w4a16:
1063 - return "int4_w4a16"
1064 - elif use_mxfp4_w4a4:
1065 - return "mxfp4_w4a4"
1066 - elif dtype == torch.float:
1067 - # avoiding cases where kernel fails when float32 MoE
1068 - # use fp16/bfloat16 configs
1069 - return "float32"
1070 - return None
1071 -
1072 -
1048 def inplace_fused_experts(
1074 - hidden_states: torch.Tensor,
1075 - w1: torch.Tensor,
1076 - w2: torch.Tensor,
1077 - topk_weights: torch.Tensor,
1078 - topk_ids: torch.Tensor,
1079 - activation: str = "silu",
1080 - is_act_and_mul: bool = True,
1081 - apply_router_weight_on_input: bool = False,
1082 - use_fp8_w8a8: bool = False,
1083 - use_int8_w8a8: bool = False,
1084 - use_int8_w8a16: bool = False,
1085 - use_int4_w4a16: bool = False,
1086 - use_mxfp4_w4a4: bool = False,
1087 - per_channel_quant: bool = False,
1088 - global_num_experts: int = -1,
1089 - expert_map: Optional[torch.Tensor] = None,
1090 - w1_scale: Optional[torch.Tensor] = None,
1091 - w2_scale: Optional[torch.Tensor] = None,
1092 - w1_zp: Optional[torch.Tensor] = None,
1093 - w2_zp: Optional[torch.Tensor] = None,
1094 - a1_scale: Optional[torch.Tensor] = None,
1095 - a2_scale: Optional[torch.Tensor] = None,
1096 - block_shape: Optional[List[int]] = None, #noqa: UP006
1097 - w1_bias: Optional[torch.Tensor] = None,
1098 - w2_bias: Optional[torch.Tensor] = None) -> None:
1049 + hidden_states: torch.Tensor,
1050 + w1: torch.Tensor,
1051 + w2: torch.Tensor,
1052 + topk_weights: torch.Tensor,
1053 + topk_ids: torch.Tensor,
1054 + activation: str = "silu",
1055 + apply_router_weight_on_input: bool = False,
1056 + use_fp8_w8a8: bool = False,
1057 + use_int8_w8a8: bool = False,
1058 + use_int8_w8a16: bool = False,
1059 + use_int4_w4a16: bool = False,
1060 + use_mxfp4_w4a4: bool = False,
1061 + per_channel_quant: bool = False,
1062 + global_num_experts: int = -1,
1063 + expert_map: Optional[torch.Tensor] = None,
1064 + w1_scale: Optional[torch.Tensor] = None,
1065 + w2_scale: Optional[torch.Tensor] = None,
1066 + w1_zp: Optional[torch.Tensor] = None,
1067 + w2_zp: Optional[torch.Tensor] = None,
1068 + a1_scale: Optional[torch.Tensor] = None,
1069 + a2_scale: Optional[torch.Tensor] = None,
1070 + block_shape: Optional[List[int]] = None, #noqa: UP006
1071 + w1_bias: Optional[torch.Tensor] = None,
1072 + w2_bias: Optional[torch.Tensor] = None,
1073 +) -> None:
1074 fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
1100 - activation, is_act_and_mul,
1101 - apply_router_weight_on_input, use_fp8_w8a8,
1075 + activation, apply_router_weight_on_input, use_fp8_w8a8,
1076 use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,
1077 use_mxfp4_w4a4, per_channel_quant, global_num_experts,
1078 expert_map, w1_scale, w2_scale, w1_zp, w2_zp, a1_scale,
1079 a2_scale, block_shape, w1_bias, w2_bias)
1080
1081
1108 -def inplace_fused_experts_fake(hidden_states: torch.Tensor,
1109 - w1: torch.Tensor,
1110 - w2: torch.Tensor,
1111 - topk_weights: torch.Tensor,
1112 - topk_ids: torch.Tensor,
1113 - activation: str = "silu",
1114 - is_act_and_mul: bool = True,
1115 - apply_router_weight_on_input: bool = False,
1116 - use_fp8_w8a8: bool = False,
1117 - use_int8_w8a8: bool = False,
1118 - use_int8_w8a16: bool = False,
1119 - use_int4_w4a16: bool = False,
1120 - use_mxfp4_w4a4: bool = False,
1121 - per_channel_quant: bool = False,
1122 - global_num_experts: int = -1,
1123 - expert_map: Optional[torch.Tensor] = None,
1124 - w1_scale: Optional[torch.Tensor] = None,
1125 - w2_scale: Optional[torch.Tensor] = None,
1126 - w1_zp: Optional[torch.Tensor] = None,
1127 - w2_zp: Optional[torch.Tensor] = None,
1128 - a1_scale: Optional[torch.Tensor] = None,
1129 - a2_scale: Optional[torch.Tensor] = None,
1130 - block_shape: Optional[list[int]] = None,
1131 - w1_bias: Optional[torch.Tensor] = None,
1132 - w2_bias: Optional[torch.Tensor] = None) -> None:
1082 +def inplace_fused_experts_fake(
1083 + hidden_states: torch.Tensor,
1084 + w1: torch.Tensor,
1085 + w2: torch.Tensor,
1086 + topk_weights: torch.Tensor,
1087 + topk_ids: torch.Tensor,
1088 + activation: str = "silu",
1089 + apply_router_weight_on_input: bool = False,
1090 + use_fp8_w8a8: bool = False,
1091 + use_int8_w8a8: bool = False,
1092 + use_int8_w8a16: bool = False,
1093 + use_int4_w4a16: bool = False,
1094 + use_mxfp4_w4a4: bool = False,
1095 + per_channel_quant: bool = False,
1096 + global_num_experts: int = -1,
1097 + expert_map: Optional[torch.Tensor] = None,
1098 + w1_scale: Optional[torch.Tensor] = None,
1099 + w2_scale: Optional[torch.Tensor] = None,
1100 + w1_zp: Optional[torch.Tensor] = None,
1101 + w2_zp: Optional[torch.Tensor] = None,
1102 + a1_scale: Optional[torch.Tensor] = None,
1103 + a2_scale: Optional[torch.Tensor] = None,
1104 + block_shape: Optional[List[int]] = None, #noqa: UP006
1105 + w1_bias: Optional[torch.Tensor] = None,
1106 + w2_bias: Optional[torch.Tensor] = None,
1107 +) -> None:
1108 pass
1109
1110
@@ -1143,175 +1118,6 @@ direct_register_custom_op(
1118 )
1119
1120
1146 -def flashinfer_fused_moe_blockscale_fp8(
1147 - routing_logits: torch.Tensor,
1148 - routing_bias: torch.Tensor,
1149 - x: torch.Tensor,
1150 - w13_weight: torch.Tensor,
1151 - w13_weight_scale_inv: torch.Tensor,
1152 - w2_weight: torch.Tensor,
1153 - w2_weight_scale_inv: torch.Tensor,
1154 - global_num_experts: int,
1155 - top_k: int,
1156 - num_expert_group: int,
1157 - topk_group: int,
1158 - intermediate_size: int,
1159 - expert_offset: int,
1160 - local_num_experts: int,
1161 - block_shape: List[int], #noqa: UP006
1162 - routed_scaling: float = 1.0) -> torch.Tensor:
1163 - from vllm.utils.flashinfer import flashinfer_trtllm_fp8_block_scale_moe
1164 - assert top_k <= global_num_experts
1165 - assert top_k <= 8
1166 - assert topk_group <= 4
1167 - assert global_num_experts > num_expert_group
1168 - assert global_num_experts % num_expert_group == 0
1169 - assert global_num_experts % 4 == 0
1170 - assert top_k < (topk_group * global_num_experts / num_expert_group)
1171 - assert block_shape == [128, 128]
1172 -
1173 - a_q, a_sf = per_token_group_quant_fp8(x, block_shape[1])
1174 - # NOTE: scales of hidden states have to be transposed!
1175 - a_sf_t = a_sf.t().contiguous()
1176 - return flashinfer_trtllm_fp8_block_scale_moe(
1177 - routing_logits=routing_logits,
1178 - routing_bias=routing_bias,
1179 - hidden_states=a_q,
1180 - hidden_states_scale=a_sf_t,
1181 - gemm1_weights=w13_weight,
1182 - gemm1_weights_scale=w13_weight_scale_inv,
1183 - gemm2_weights=w2_weight,
1184 - gemm2_weights_scale=w2_weight_scale_inv,
1185 - num_experts=global_num_experts,
1186 - top_k=top_k,
1187 - n_group=num_expert_group,
1188 - topk_group=topk_group,
1189 - intermediate_size=intermediate_size,
1190 - local_expert_offset=expert_offset,
1191 - local_num_experts=local_num_experts,
1192 - routed_scaling_factor=routed_scaling,
1193 - tile_tokens_dim=calculate_tile_tokens_dim(x.shape[0], top_k,
1194 - global_num_experts),
1195 - routing_method_type=2, # DeepSeek-styled routing method
1196 - use_shuffled_weight=False,
1197 - )
1198 -
1199 -
1200 -def flashinfer_fused_moe_blockscale_fp8_fake(
1201 - routing_logits: torch.Tensor,
1202 - routing_bias: torch.Tensor,
1203 - x: torch.Tensor,
1204 - w13_weight: torch.Tensor,
1205 - w13_weight_scale_inv: torch.Tensor,
1206 - w2_weight: torch.Tensor,
1207 - w2_weight_scale_inv: torch.Tensor,
1208 - global_num_experts: int,
1209 - top_k: int,
1210 - num_expert_group: int,
1211 - topk_group: int,
1212 - intermediate_size: int,
1213 - expert_offset: int,
1214 - local_num_experts: int,
1215 - block_shape: list[int],
1216 - routed_scaling: float = 1.0) -> torch.Tensor:
1217 - return torch.empty_like(x)
1218 -
1219 -
1220 -direct_register_custom_op(
1221 - op_name="flashinfer_fused_moe_blockscale_fp8",
1222 - op_func=flashinfer_fused_moe_blockscale_fp8,
1223 - mutates_args=[],
1224 - fake_impl=flashinfer_fused_moe_blockscale_fp8_fake,
1225 - tags=(torch.Tag.needs_fixed_stride_order, ),
1226 -)
1227 -
1228 -
1229 -def flashinfer_fused_moe_per_tensor_scale_fp8(
1230 - routing_logits: torch.Tensor,
1231 - routing_bias: Optional[torch.Tensor],
1232 - hidden_states: torch.Tensor,
1233 - input_scale: torch.Tensor,
1234 - gemm1_weights: torch.Tensor,
1235 - gemm2_weights: torch.Tensor,
1236 - output1_scales_scalar: torch.Tensor,
1237 - output1_scales_gate_scalar: torch.Tensor,
1238 - output2_scales_scalar: torch.Tensor,
1239 - num_experts: int,
1240 - top_k: int,
1241 - num_expert_group: Optional[int],
1242 - topk_group: Optional[int],
1243 - intermediate_size: int,
1244 - local_expert_offset: int,
1245 - local_num_experts: int,
1246 - use_routing_scales_on_input: bool,
1247 - routing_method_type: int,
1248 - routed_scaling_factor: float = 1.0) -> torch.Tensor:
1249 - num_expert_group = num_expert_group if num_expert_group is not None else 0
1250 - topk_group = topk_group if topk_group is not None else 0
1251 -
1252 - quant_hidden_states, _ = moe_kernel_quantize_input(
1253 - hidden_states,
1254 - input_scale,
1255 - quant_dtype=torch.float8_e4m3fn,
1256 - per_act_token_quant=False)
1257 -
1258 - from vllm.utils.flashinfer import (
1259 - flashinfer_trtllm_fp8_per_tensor_scale_moe)
1260 - return flashinfer_trtllm_fp8_per_tensor_scale_moe(
1261 - routing_logits=routing_logits,
1262 - routing_bias=routing_bias,
1263 - hidden_states=quant_hidden_states,
1264 - gemm1_weights=gemm1_weights,
1265 - output1_scales_scalar=output1_scales_scalar,
1266 - output1_scales_gate_scalar=output1_scales_gate_scalar,
1267 - gemm2_weights=gemm2_weights,
1268 - output2_scales_scalar=output2_scales_scalar,
1269 - num_experts=num_experts,
1270 - top_k=top_k,
1271 - n_group=num_expert_group,
1272 - topk_group=topk_group,
1273 - intermediate_size=intermediate_size,
1274 - local_expert_offset=local_expert_offset,
1275 - local_num_experts=local_num_experts,
1276 - routed_scaling_factor=routed_scaling_factor,
1277 - use_routing_scales_on_input=use_routing_scales_on_input,
1278 - tile_tokens_dim=calculate_tile_tokens_dim(hidden_states.shape[0],
1279 - top_k, num_experts),
1280 - routing_method_type=routing_method_type)
1281 -
1282 -
1283 -def flashinfer_fused_moe_per_tensor_scale_fp8_fake(
1284 - routing_logits: torch.Tensor,
1285 - routing_bias: Optional[torch.Tensor],
1286 - hidden_states: torch.Tensor,
1287 - input_scale: torch.Tensor,
1288 - gemm1_weights: torch.Tensor,
1289 - gemm2_weights: torch.Tensor,
1290 - output1_scales_scalar: torch.Tensor,
1291 - output1_scales_gate_scalar: torch.Tensor,
1292 - output2_scales_scalar: torch.Tensor,
1293 - num_experts: int,
1294 - top_k: int,
1295 - num_expert_group: Optional[int],
1296 - topk_group: Optional[int],
1297 - intermediate_size: int,
1298 - local_expert_offset: int,
1299 - local_num_experts: int,
1300 - use_routing_scales_on_input: bool,
1301 - routing_method_type: int,
1302 - routed_scaling_factor: float = 1.0) -> torch.Tensor:
1303 - pass
1304 -
1305 -
1306 -direct_register_custom_op(
1307 - op_name="flashinfer_fused_moe_per_tensor_scale_fp8",
1308 - op_func=flashinfer_fused_moe_per_tensor_scale_fp8,
1309 - mutates_args=["hidden_states"],
1310 - fake_impl=flashinfer_fused_moe_per_tensor_scale_fp8_fake,
1311 - tags=(torch.Tag.needs_fixed_stride_order, ),
1312 -)
1313 -
1314 -
1121 def outplace_fused_experts(
1122 hidden_states: torch.Tensor,
1123 w1: torch.Tensor,
@@ -1319,7 +1125,6 @@ def outplace_fused_experts(
1125 topk_weights: torch.Tensor,
1126 topk_ids: torch.Tensor,
1127 activation: str = "silu",
1322 - is_act_and_mul: bool = True,
1128 apply_router_weight_on_input: bool = False,
1129 use_fp8_w8a8: bool = False,
1130 use_int8_w8a8: bool = False,
@@ -1341,37 +1146,37 @@ def outplace_fused_experts(
1146 ) -> torch.Tensor:
1147 return fused_experts_impl(
1148 hidden_states, w1, w2, topk_weights, topk_ids, False, activation,
1344 - is_act_and_mul, apply_router_weight_on_input, use_fp8_w8a8,
1345 - use_int8_w8a8, use_int8_w8a16, use_int4_w4a16, use_mxfp4_w4a4,
1346 - per_channel_quant, global_num_experts, expert_map, w1_scale, w2_scale,
1347 - w1_zp, w2_zp, a1_scale, a2_scale, block_shape, w1_bias, w2_bias)
1149 + apply_router_weight_on_input, use_fp8_w8a8, use_int8_w8a8,
1150 + use_int8_w8a16, use_int4_w4a16, use_mxfp4_w4a4, per_channel_quant,
1151 + global_num_experts, expert_map, w1_scale, w2_scale, w1_zp, w2_zp,
1152 + a1_scale, a2_scale, block_shape, w1_bias, w2_bias)
1153
1154
1155 def outplace_fused_experts_fake(
1351 - hidden_states: torch.Tensor,
1352 - w1: torch.Tensor,
1353 - w2: torch.Tensor,
1354 - topk_weights: torch.Tensor,
1355 - topk_ids: torch.Tensor,
1356 - activation: str = "silu",
1357 - is_act_and_mul: bool = True,
1358 - use_fp8_w8a8: bool = False,
1359 - use_int8_w8a8: bool = False,
1360 - use_int8_w8a16: bool = False,
1361 - use_int4_w4a16: bool = False,
1362 - use_mxfp4_w4a4: bool = False,
1363 - per_channel_quant: bool = False,
1364 - global_num_experts: int = -1,
1365 - expert_map: Optional[torch.Tensor] = None,
1366 - w1_scale: Optional[torch.Tensor] = None,
1367 - w2_scale: Optional[torch.Tensor] = None,
1368 - w1_zp: Optional[torch.Tensor] = None,
1369 - w2_zp: Optional[torch.Tensor] = None,
1370 - a1_scale: Optional[torch.Tensor] = None,
1371 - a2_scale: Optional[torch.Tensor] = None,
1372 - block_shape: Optional[list[int]] = None,
1373 - w1_bias: Optional[torch.Tensor] = None,
1374 - w2_bias: Optional[torch.Tensor] = None) -> torch.Tensor:
1156 + hidden_states: torch.Tensor,
1157 + w1: torch.Tensor,
1158 + w2: torch.Tensor,
1159 + topk_weights: torch.Tensor,
1160 + topk_ids: torch.Tensor,
1161 + activation: str = "silu",
1162 + use_fp8_w8a8: bool = False,
1163 + use_int8_w8a8: bool = False,
1164 + use_int8_w8a16: bool = False,
1165 + use_int4_w4a16: bool = False,
1166 + use_mxfp4_w4a4: bool = False,
1167 + per_channel_quant: bool = False,
1168 + global_num_experts: int = -1,
1169 + expert_map: Optional[torch.Tensor] = None,
1170 + w1_scale: Optional[torch.Tensor] = None,
1171 + w2_scale: Optional[torch.Tensor] = None,
1172 + w1_zp: Optional[torch.Tensor] = None,
1173 + w2_zp: Optional[torch.Tensor] = None,
1174 + a1_scale: Optional[torch.Tensor] = None,
1175 + a2_scale: Optional[torch.Tensor] = None,
1176 + block_shape: Optional[list[int]] = None,
1177 + w1_bias: Optional[torch.Tensor] = None,
1178 + w2_bias: Optional[torch.Tensor] = None,
1179 +) -> torch.Tensor:
1180 return torch.empty_like(hidden_states)
1181
1182
@@ -1403,45 +1208,36 @@ def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:
1208
1209 # TODO (bnell): replace this with modular op. Can get rid of inplace/outplace
1210 # torch ops.
1406 -def fused_experts(hidden_states: torch.Tensor,
1407 - w1: torch.Tensor,
1408 - w2: torch.Tensor,
1409 - topk_weights: torch.Tensor,
1410 - topk_ids: torch.Tensor,
1411 - inplace: bool = False,
1412 - activation: str = "silu",
1413 - is_act_and_mul: bool = True,
1414 - apply_router_weight_on_input: bool = False,
1415 - use_fp8_w8a8: bool = False,
1416 - use_int8_w8a8: bool = False,
1417 - use_int8_w8a16: bool = False,
1418 - use_int4_w4a16: bool = False,
1419 - use_mxfp4_w4a4: bool = False,
1420 - per_channel_quant: bool = False,
1421 - global_num_experts: int = -1,
1422 - expert_map: Optional[torch.Tensor] = None,
1423 - w1_scale: Optional[torch.Tensor] = None,
1424 - w2_scale: Optional[torch.Tensor] = None,
1425 - w1_zp: Optional[torch.Tensor] = None,
1426 - w2_zp: Optional[torch.Tensor] = None,
1427 - a1_scale: Optional[torch.Tensor] = None,
1428 - a2_scale: Optional[torch.Tensor] = None,
1429 - block_shape: Optional[list[int]] = None,
1430 - allow_deep_gemm: bool = False,
1431 - allow_cutlass_block_scaled_grouped_gemm: bool = False,
1432 - w1_bias: Optional[torch.Tensor] = None,
1433 - w2_bias: Optional[torch.Tensor] = None) -> torch.Tensor:
1211 +def fused_experts(
1212 + hidden_states: torch.Tensor,
1213 + w1: torch.Tensor,
1214 + w2: torch.Tensor,
1215 + topk_weights: torch.Tensor,
1216 + topk_ids: torch.Tensor,
1217 + inplace: bool = False,
1218 + activation: str = "silu",
1219 + apply_router_weight_on_input: bool = False,
1220 + global_num_experts: int = -1,
1221 + expert_map: Optional[torch.Tensor] = None,
1222 + quant_config: Optional[FusedMoEQuantConfig] = None,
1223 + allow_deep_gemm: bool = False,
1224 + allow_cutlass_block_scaled_grouped_gemm: bool = False,
1225 +) -> torch.Tensor:
1226 +
1227 + if quant_config is None:
1228 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
1229 + use_fp8_w8a8 = quant_config.use_fp8_w8a8
1230 +
1231 # For now, disable DeepGemm for small N (<= 512) until better
1232 # permute/unpermute ops are available.
1233 # However, on B200, we use DeepGemm for all cases because they only support
1234 # E8M0 scale, which means we requantize the weight and input to the specific
1235 # scale. Fallen back to cutlass or triton for some cases would cause
1236 # accuracy issue.
1440 - if (allow_deep_gemm and use_fp8_w8a8 and
1237 + if (allow_deep_gemm and quant_config.use_fp8_w8a8 and
1238 (is_deep_gemm_e8m0_used() or _valid_deep_gemm(hidden_states, w1, w2))):
1239 + assert quant_config is not None
1240 assert apply_router_weight_on_input is False
1443 - assert is_act_and_mul, (
1444 - "DeepGemm only supports is_act_and_mul=True for now.")
1241 return deep_gemm_moe_fp8(
1242 hidden_states=hidden_states,
1243 w1=w1,
@@ -1452,22 +1248,23 @@ def fused_experts(hidden_states: torch.Tensor,
1248 activation=activation,
1249 global_num_experts=global_num_experts,
1250 expert_map=expert_map,
1455 - w1_scale=w1_scale,
1456 - w2_scale=w2_scale,
1457 - a1_scale=a1_scale,
1458 - a2_scale=a2_scale,
1251 + w1_scale=quant_config.w1_scale,
1252 + w2_scale=quant_config.w2_scale,
1253 + a1_scale=quant_config.a1_scale,
1254 + a2_scale=quant_config.a2_scale,
1255 apply_router_weight_on_input=apply_router_weight_on_input,
1256 )
1257 elif (allow_cutlass_block_scaled_grouped_gemm and use_fp8_w8a8
1258 and _valid_cutlass_block_scaled_grouped_gemm(
1259 w1, w2, inplace, activation, apply_router_weight_on_input,
1260 expert_map)):
1261 + assert quant_config is not None
1262 return run_cutlass_block_scaled_fused_experts(
1263 a=hidden_states,
1264 w1=w1,
1265 w2=w2,
1469 - w1_scale=w1_scale,
1470 - w2_scale=w2_scale,
1266 + w1_scale=quant_config.w1_scale,
1267 + w2_scale=quant_config.w2_scale,
1268 topk_weights=topk_weights,
1269 topk_ids=topk_ids)
1270 else:
@@ -1478,26 +1275,49 @@ def fused_experts(hidden_states: torch.Tensor,
1275 topk_weights=topk_weights,
1276 topk_ids=topk_ids,
1277 activation=activation,
1481 - is_act_and_mul=is_act_and_mul,
1278 apply_router_weight_on_input=apply_router_weight_on_input,
1483 - use_fp8_w8a8=use_fp8_w8a8,
1484 - use_int8_w8a8=use_int8_w8a8,
1485 - use_int8_w8a16=use_int8_w8a16,
1486 - use_int4_w4a16=use_int4_w4a16,
1487 - use_mxfp4_w4a4=use_mxfp4_w4a4,
1488 - per_channel_quant=per_channel_quant,
1279 + use_fp8_w8a8=quant_config.use_fp8_w8a8,
1280 + use_int8_w8a8=quant_config.use_int8_w8a8,
1281 + use_int8_w8a16=quant_config.use_int8_w8a16,
1282 + use_int4_w4a16=quant_config.use_int4_w4a16,
1283 + use_mxfp4_w4a4=quant_config.use_mxfp4_w4a4,
1284 + per_channel_quant=quant_config.per_act_token_quant,
1285 global_num_experts=global_num_experts,
1286 expert_map=expert_map,
1491 - w1_scale=w1_scale,
1492 - w2_scale=w2_scale,
1493 - w1_zp=w1_zp,
1494 - w2_zp=w2_zp,
1495 - a1_scale=a1_scale,
1496 - a2_scale=a2_scale,
1497 - block_shape=block_shape,
1498 - w1_bias=w1_bias,
1499 - w2_bias=w2_bias,
1500 - )
1287 + w1_scale=quant_config.w1_scale,
1288 + w2_scale=quant_config.w2_scale,
1289 + w1_zp=quant_config.w1_zp,
1290 + w2_zp=quant_config.w2_zp,
1291 + a1_scale=quant_config.a1_scale,
1292 + a2_scale=quant_config.a2_scale,
1293 + block_shape=quant_config.block_shape,
1294 + w1_bias=quant_config.w1_bias,
1295 + w2_bias=quant_config.w2_bias)
1296 +
1297 +
1298 +SILU_NO_MUL: str = activation_without_mul("silu")
1299 +GELU_NO_MUL: str = activation_without_mul("gelu")
1300 +
1301 +
1302 +def _get_config_quant_dtype(
1303 + use_fp8_w8a8: bool,
1304 + use_int8_w8a8: bool,
1305 + use_mxfp4_w4a4: bool,
1306 +) -> Union[None, torch.dtype, str]:
1307 + """
1308 + Get the quantization type based on the quantization strategy flags.
1309 + We don't have a quant_config at this point so we need to work backwards.
1310 + A return type of None means no quantization is required because the
1311 + input is unquantized or has been quantized prior to calling
1312 + fused_experts_impl.
1313 + """
1314 + if use_fp8_w8a8:
1315 + return torch.float8_e4m3fn
1316 + elif use_int8_w8a8:
1317 + return torch.int8
1318 + elif use_mxfp4_w4a4:
1319 + return "mxfp4"
1320 + return None
1321
1322
1323 def fused_experts_impl(
@@ -1508,7 +1328,6 @@ def fused_experts_impl(
1328 topk_ids: torch.Tensor,
1329 inplace: bool = False,
1330 activation: str = "silu",
1511 - is_act_and_mul: bool = True,
1331 apply_router_weight_on_input: bool = False,
1332 use_fp8_w8a8: bool = False,
1333 use_int8_w8a8: bool = False,
@@ -1557,17 +1376,18 @@ def fused_experts_impl(
1376 # https://github.com/vllm-project/vllm/issues/5938
1377 CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
1378 M = min(num_tokens, CHUNK_SIZE)
1560 - config_dtype = get_config_dtype_str(use_fp8_w8a8=use_fp8_w8a8,
1561 - use_int8_w8a16=use_int8_w8a16,
1562 - use_int4_w4a16=use_int4_w4a16,
1563 - use_mxfp4_w4a4=use_mxfp4_w4a4,
1564 - dtype=hidden_states.dtype)
1565 -
1566 - qtype = get_config_quant_dtype(use_fp8_w8a8=use_fp8_w8a8,
1567 - use_int8_w8a8=use_int8_w8a8,
1568 - use_int8_w8a16=use_int8_w8a16,
1569 - use_int4_w4a16=use_int4_w4a16,
1570 - use_mxfp4_w4a4=use_mxfp4_w4a4)
1379 +
1380 + config_dtype = _get_config_dtype_str(use_fp8_w8a8=use_fp8_w8a8,
1381 + use_int8_w8a16=use_int8_w8a16,
1382 + use_int4_w4a16=use_int4_w4a16,
1383 + use_mxfp4_w4a4=use_mxfp4_w4a4,
1384 + dtype=hidden_states.dtype)
1385 +
1386 + # Note: for use_int8_w8a16 or use_int4_w4a16, the activations are
1387 + # quantized prior to calling fused_experts.
1388 + quant_dtype = _get_config_quant_dtype(use_fp8_w8a8=use_fp8_w8a8,
1389 + use_int8_w8a8=use_int8_w8a8,
1390 + use_mxfp4_w4a4=use_mxfp4_w4a4)
1391
1392 get_config_func = functools.partial(
1393 try_get_optimal_moe_config,
@@ -1640,7 +1460,7 @@ def fused_experts_impl(
1460 qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
1461 A=curr_hidden_states,
1462 A_scale=a1_scale,
1643 - quant_dtype=qtype,
1463 + quant_dtype=quant_dtype,
1464 per_act_token_quant=per_channel_quant,
1465 block_shape=block_shape)
1466
@@ -1671,30 +1491,29 @@ def fused_experts_impl(
1491 B_bias=w1_bias)
1492
1493 # Activation function with multiplication
1674 - if activation == "silu" and is_act_and_mul:
1494 + if activation == "silu":
1495 torch.ops._C.silu_and_mul(intermediate_cache2,
1496 intermediate_cache1.view(-1, N))
1677 - elif activation == "gelu" and is_act_and_mul:
1497 + elif activation == "gelu":
1498 torch.ops._C.gelu_and_mul(intermediate_cache2,
1499 intermediate_cache1.view(-1, N))
1680 - elif activation == "swigluoai" and is_act_and_mul:
1500 + elif activation == "swigluoai":
1501 # alpha = 1.702, limit = 7.0
1502 torch.ops._C.swigluoai_and_mul(intermediate_cache2,
1503 intermediate_cache1.view(-1, N))
1504 # Activation function without multiplication
1685 - elif activation == "silu":
1505 + elif activation == SILU_NO_MUL:
1506 intermediate_cache2 = F.silu(intermediate_cache1.view(-1, N))
1687 - elif activation == "gelu":
1507 + elif activation == GELU_NO_MUL:
1508 intermediate_cache2 = F.gelu(intermediate_cache1.view(-1, N))
1509
1510 else:
1691 - raise ValueError(f"Unsupported FusedMoe activation: {activation}, "
1692 - f"with is_act_and_mul={is_act_and_mul}.")
1511 + raise ValueError(f"Unsupported FusedMoe activation: {activation}.")
1512
1513 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1514 A=intermediate_cache2,
1515 A_scale=a2_scale,
1697 - quant_dtype=qtype,
1516 + quant_dtype=quant_dtype,
1517 per_act_token_quant=per_channel_quant,
1518 block_shape=block_shape)
1519
@@ -1726,164 +1545,13 @@ def fused_experts_impl(
1545 return out_hidden_states
1546
1547
1729 -def fused_moe(
1730 - hidden_states: torch.Tensor,
1731 - w1: torch.Tensor,
1732 - w2: torch.Tensor,
1733 - gating_output: torch.Tensor,
1734 - topk: int,
1735 - renormalize: bool,
1736 - inplace: bool = False,
1737 - activation: str = "silu",
1738 - is_act_and_mul: bool = True,
1739 - use_grouped_topk: bool = False,
1740 - num_expert_group: Optional[int] = None,
1741 - topk_group: Optional[int] = None,
1742 - custom_routing_function: Optional[Callable] = None,
1743 - use_fp8_w8a8: bool = False,
1744 - use_int8_w8a8: bool = False,
1745 - use_int8_w8a16: bool = False,
1746 - use_int4_w4a16: bool = False,
1747 - use_mxfp4_w4a4: bool = False,
1748 - per_channel_quant: bool = False,
1749 - global_num_experts: int = -1,
1750 - expert_map: Optional[torch.Tensor] = None,
1751 - w1_scale: Optional[torch.Tensor] = None,
1752 - w2_scale: Optional[torch.Tensor] = None,
1753 - w1_zp: Optional[torch.Tensor] = None,
1754 - w2_zp: Optional[torch.Tensor] = None,
1755 - a1_scale: Optional[torch.Tensor] = None,
1756 - a2_scale: Optional[torch.Tensor] = None,
1757 - block_shape: Optional[list[int]] = None,
1758 - w1_bias: Optional[torch.Tensor] = None,
1759 - w2_bias: Optional[torch.Tensor] = None,
1760 -) -> torch.Tensor:
1761 - """
1762 - This function computes a Mixture of Experts (MoE) layer using two sets of
1763 - weights, w1 and w2, and top-k gating mechanism.
1764 -
1765 - Parameters:
1766 - - hidden_states (torch.Tensor): The input tensor to the MoE layer.
1767 - - w1 (torch.Tensor): The first set of expert weights.
1768 - - w2 (torch.Tensor): The second set of expert weights.
1769 - - gating_output (torch.Tensor): The output of the gating operation
1770 - (before softmax).
1771 - - topk (int): The number of top-k experts to select.
1772 - - renormalize (bool): If True, renormalize the top-k weights to sum to 1.
1773 - - inplace (bool): If True, perform the operation in-place.
1774 - Defaults to False.
1775 - - activation (str): The activation function to apply after the first
1776 - MoE layer.
1777 - - is_act_and_mul (bool): If True, use activation-and-mul function for
1778 - activation (self-gated activation), otherwise use activation function
1779 - for activation (ungated activation).
1780 - - num_expert_group: Optional[int]: additional parameter for grouped_topk
1781 - - topk_group: Optional[int]: additional parameter for grouped_topk
1782 - - use_grouped_topk: If True, use grouped_topk instead of fused_topk
1783 - note: Deepseekv2 model uses grouped_topk
1784 - - use_fp8_w8a8 (bool): If True, use fp8 arithmetic to compute the inner
1785 - products for w1 and w2. Defaults to False.
1786 - - use_int8_w8a8 (bool): If True, use int8 arithmetic to compute the inner
1787 - products for w1 and w2. Defaults to False.
1788 - - use_int8_w8a16 (bool): If True, use matmul of int8 weight and bf16/fp16
1789 - activation to compute the inner products for w1 and w2.
1790 - Defaults to False.
1791 - - use_int4_w4a16 (bool): If True, use matmul of int4 weight and bf16/fp16
1792 - activation to compute the inner products for w1 and w2.
1793 - Defaults to False.
1794 - - use_mxfp4_w4a4 (bool): If True, use matmul of OCP MXFP4 weight and
1795 - OCP MXFP4 activation to compute the inner products for w1 and w2.
1796 - Defaults to False.
1797 - - global_num_experts (int): The total number of experts in the global
1798 - expert space.
1799 - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
1800 - from the global expert space to the local expert space of the expert
1801 - parallel shard.
1802 - - w1_scale (Optional[torch.Tensor]): Optional scale to be used for
1803 - w1.
1804 - - w2_scale (Optional[torch.Tensor]): Optional scale to be used for
1805 - w2.
1806 - - a1_scale (Optional[torch.Tensor]): Optional scale to be used for
1807 - a1.
1808 - - a2_scale (Optional[torch.Tensor]): Optional scale to be used for
1809 - a2.
1810 - - block_shape: (Optional[list[int]]): Optional block size for block-wise
1811 - quantization.
1812 -
1813 - Returns:
1814 - - torch.Tensor: The output tensor after applying the MoE layer.
1815 - """
1816 - if not is_act_and_mul:
1817 - assert inplace is False, (
1818 - "is_act_and_mul=False is not supported with inplace=True")
1819 -
1820 - if use_grouped_topk:
1821 - assert num_expert_group is not None and topk_group is not None
1822 - topk_weights, topk_ids = grouped_topk(hidden_states, gating_output,
1823 - topk, renormalize,
1824 - num_expert_group, topk_group)
1825 - elif custom_routing_function is None:
1826 - topk_weights, topk_ids, token_expert_indices = fused_topk(
1827 - hidden_states, gating_output, topk, renormalize)
1828 - else:
1829 - topk_weights, topk_ids = custom_routing_function(
1830 - hidden_states, gating_output, topk, renormalize)
1831 -
1832 - return fused_experts(hidden_states,
1833 - w1,
1834 - w2,
1835 - topk_weights,
1836 - topk_ids,
1837 - inplace=inplace,
1838 - activation=activation,
1839 - is_act_and_mul=is_act_and_mul,
1840 - use_fp8_w8a8=use_fp8_w8a8,
1841 - use_int8_w8a8=use_int8_w8a8,
1842 - use_int8_w8a16=use_int8_w8a16,
1843 - use_int4_w4a16=use_int4_w4a16,
1844 - use_mxfp4_w4a4=use_mxfp4_w4a4,
1845 - per_channel_quant=per_channel_quant,
1846 - global_num_experts=global_num_experts,
1847 - expert_map=expert_map,
1848 - w1_scale=w1_scale,
1849 - w2_scale=w2_scale,
1850 - w1_zp=w1_zp,
1851 - w2_zp=w2_zp,
1852 - a1_scale=a1_scale,
1853 - a2_scale=a2_scale,
1854 - block_shape=block_shape,
1855 - w1_bias=w1_bias,
1856 - w2_bias=w2_bias)
1857 -
1858 -
1548 class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1549
1550 def __init__(
1551 self,
1863 - use_fp8_w8a8: bool = False,
1864 - use_int8_w8a8: bool = False,
1865 - use_int8_w8a16: bool = False,
1866 - use_int4_w4a16: bool = False,
1867 - use_mxfp4_w4a4: bool = False,
1868 - per_act_token_quant: bool = False,
1869 - block_shape: Optional[list[int]] = None,
1552 + quant_config: FusedMoEQuantConfig,
1553 ):
1871 - super().__init__(
1872 - FusedMoEQuantConfig.make(
1873 - use_fp8_w8a8=use_fp8_w8a8,
1874 - use_int8_w8a8=use_int8_w8a8,
1875 - use_int8_w8a16=use_int8_w8a16,
1876 - use_int4_w4a16=use_int4_w4a16,
1877 - use_mxfp4_w4a4=use_mxfp4_w4a4,
1878 - per_act_token_quant=per_act_token_quant,
1879 - block_shape=block_shape,
1880 - ))
1881 -
1882 - self.use_fp8_w8a8 = use_fp8_w8a8
1883 - self.use_int4_w4a16 = use_int4_w4a16
1884 - self.use_int8_w8a8 = use_int8_w8a8
1885 - self.use_int8_w8a16 = use_int8_w8a16
1886 - self.use_mxfp4_w4a4 = use_mxfp4_w4a4
1554 + super().__init__(quant_config)
1555
1556 @property
1557 def activation_formats(
@@ -1929,19 +1597,14 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1597 activation: str,
1598 global_num_experts: int,
1599 expert_map: Optional[torch.Tensor],
1932 - w1_scale: Optional[torch.Tensor],
1933 - w2_scale: Optional[torch.Tensor],
1934 - w1_zp: Optional[torch.Tensor],
1935 - w2_zp: Optional[torch.Tensor],
1600 a1q_scale: Optional[torch.Tensor],
1937 - a2_scale: Optional[torch.Tensor],
1601 workspace13: torch.Tensor,
1602 workspace2: torch.Tensor,
1603 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
1604 apply_router_weight_on_input: bool,
1605 ):
1606 # Check constraints.
1944 - if self.use_int4_w4a16:
1607 + if self.quant_config.use_int4_w4a16:
1608 assert hidden_states.size(-1) // 2 == w1.size(2), (
1609 "Hidden size mismatch")
1610 else:
@@ -1964,17 +1627,11 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1627 if global_num_experts == -1:
1628 global_num_experts = E
1629
1967 - config_dtype = get_config_dtype_str(use_fp8_w8a8=self.use_fp8_w8a8,
1968 - use_int8_w8a16=self.use_int8_w8a16,
1969 - use_int4_w4a16=self.use_int4_w4a16,
1970 - use_mxfp4_w4a4=self.use_mxfp4_w4a4,
1971 - dtype=hidden_states.dtype)
1972 -
1630 config = try_get_optimal_moe_config(
1631 w1.size(),
1632 w2.size(),
1633 top_k_num,
1977 - config_dtype,
1634 + self.quant_config.config_name(hidden_states.dtype),
1635 num_tokens,
1636 block_shape=self.block_shape,
1637 )
@@ -2008,8 +1665,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1665 w1,
1666 intermediate_cache1,
1667 a1q_scale,
2011 - w1_scale,
2012 - w1_zp,
1668 + self.w1_scale,
1669 + self.w1_zp,
1670 None, # topk_weights
1671 sorted_token_ids,
1672 expert_ids,
@@ -2018,13 +1675,13 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1675 top_k_num,
1676 config,
1677 compute_type=compute_type,
2021 - use_fp8_w8a8=self.use_fp8_w8a8,
2022 - use_int8_w8a8=self.use_int8_w8a8,
2023 - use_int8_w8a16=self.use_int8_w8a16,
2024 - use_int4_w4a16=self.use_int4_w4a16,
1678 + use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
1679 + use_int8_w8a8=self.quant_config.use_int8_w8a8,
1680 + use_int8_w8a16=self.quant_config.use_int8_w8a16,
1681 + use_int4_w4a16=self.quant_config.use_int4_w4a16,
1682 per_channel_quant=self.per_act_token_quant,
1683 block_shape=self.block_shape,
2027 - B_bias=None # TODO support B_bias
1684 + B_bias=self.w1_bias,
1685 )
1686
1687 self.activation(activation, intermediate_cache2,
@@ -2033,7 +1690,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1690 a2q_scale: Optional[torch.Tensor] = None
1691
1692 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
2036 - intermediate_cache2, a2_scale, self.quant_dtype,
1693 + intermediate_cache2, self.a2_scale, self.quant_dtype,
1694 self.per_act_token_quant, self.block_shape)
1695
1696 invoke_fused_moe_kernel(
@@ -2041,8 +1698,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1698 w2,
1699 intermediate_cache3,
1700 a2q_scale,
2044 - w2_scale,
2045 - w2_zp,
1701 + self.w2_scale,
1702 + self.w2_zp,
1703 topk_weights,
1704 sorted_token_ids,
1705 expert_ids,
@@ -2051,36 +1708,21 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1708 1,
1709 config,
1710 compute_type=compute_type,
2054 - use_fp8_w8a8=self.use_fp8_w8a8,
2055 - use_int8_w8a8=self.use_int8_w8a8,
2056 - use_int8_w8a16=self.use_int8_w8a16,
2057 - use_int4_w4a16=self.use_int4_w4a16,
1711 + use_fp8_w8a8=self.quant_config.use_fp8_w8a8,
1712 + use_int8_w8a8=self.quant_config.use_int8_w8a8,
1713 + use_int8_w8a16=self.quant_config.use_int8_w8a16,
1714 + use_int4_w4a16=self.quant_config.use_int4_w4a16,
1715 per_channel_quant=self.per_act_token_quant,
1716 block_shape=self.block_shape,
2060 - B_bias=None # TODO support B_bias
1717 + B_bias=self.w2_bias,
1718 )
1719
1720 ops.moe_sum(intermediate_cache3, output)
1721
1722
1723 def modular_triton_fused_moe(
2067 - use_fp8_w8a8: bool,
2068 - use_int8_w8a8: bool,
2069 - use_int8_w8a16: bool,
2070 - use_int4_w4a16: bool,
2071 - use_mxfp4_w4a4: bool,
2072 - per_act_token_quant: bool,
2073 - block_shape: Optional[list[int]] = None,
2074 -) -> mk.FusedMoEModularKernel:
1724 + quant_config: FusedMoEQuantConfig) -> mk.FusedMoEModularKernel:
1725 return mk.FusedMoEModularKernel(
1726 MoEPrepareAndFinalizeNoEP(),
2077 - TritonExperts(
2078 - use_fp8_w8a8=use_fp8_w8a8,
2079 - use_int8_w8a8=use_int8_w8a8,
2080 - use_int8_w8a16=use_int8_w8a16,
2081 - use_int4_w4a16=use_int4_w4a16,
2082 - use_mxfp4_w4a4=use_mxfp4_w4a4,
2083 - per_act_token_quant=per_act_token_quant,
2084 - block_shape=block_shape,
2085 - ),
1727 + TritonExperts(quant_config),
1728 )
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+24 -71
@@ -1,11 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from typing import TYPE_CHECKING, Optional
3 +from typing import Optional
4
5 import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm.logger import init_logger
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEQuantConfig)
11 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12 TopKWeightAndReduceDelegate)
13 from vllm.utils import has_triton_kernels
@@ -23,9 +25,6 @@ if has_triton_kernels():
25 "Failed to import Triton kernels. Please make sure your triton "
26 "version is compatible.")
27
26 -if TYPE_CHECKING:
27 - from triton_kernels.matmul_ogs import PrecisionConfig
28 -
28
29 def triton_kernel_moe_forward(
30 hidden_states: torch.Tensor,
@@ -35,20 +34,10 @@ def triton_kernel_moe_forward(
34 topk: int,
35 renormalize: bool,
36 activation: str = "silu",
37 + quant_config: Optional[FusedMoEQuantConfig] = None,
38 apply_router_weight_on_input: bool = False,
39 - use_fp8_w8a8: bool = False,
40 - per_channel_quant: bool = False,
39 global_num_experts: int = -1,
40 expert_map: Optional[torch.Tensor] = None,
43 - w1_scale: Optional[torch.Tensor] = None,
44 - w2_scale: Optional[torch.Tensor] = None,
45 - w1_bias: Optional[torch.Tensor] = None,
46 - w2_bias: Optional[torch.Tensor] = None,
47 - w1_precision: Optional["PrecisionConfig"] = None,
48 - w2_precision: Optional["PrecisionConfig"] = None,
49 - a1_scale: Optional[torch.Tensor] = None,
50 - a2_scale: Optional[torch.Tensor] = None,
51 - block_shape: Optional[list[int]] = None,
41 ) -> torch.Tensor:
42
43 routing_data, gather_idx, scatter_idx = routing(gating_output,
@@ -64,20 +53,10 @@ def triton_kernel_moe_forward(
53 gather_idx,
54 scatter_idx,
55 activation=activation,
56 + quant_config=quant_config,
57 apply_router_weight_on_input=apply_router_weight_on_input,
68 - use_fp8_w8a8=use_fp8_w8a8,
69 - per_channel_quant=per_channel_quant,
58 global_num_experts=global_num_experts,
71 - expert_map=expert_map,
72 - w1_scale=w1_scale,
73 - w2_scale=w2_scale,
74 - w1_bias=w1_bias,
75 - w2_bias=w2_bias,
76 - w1_precision=w1_precision,
77 - w2_precision=w2_precision,
78 - a1_scale=a1_scale,
79 - a2_scale=a2_scale,
80 - block_shape=block_shape)
59 + expert_map=expert_map)
60
61
62 # This is a triton implementation of the fused_experts function
@@ -90,28 +69,23 @@ def triton_kernel_fused_experts(
69 gather_indx, # GatherIndx
70 scatter_indx, # ScatterIndx
71 activation: str = "silu",
72 + quant_config: Optional[FusedMoEQuantConfig] = None,
73 swiglu_alpha: float = 1.702,
74 swiglu_limit: float = 7.0,
75 apply_router_weight_on_input: bool = False,
96 - use_fp8_w8a8: bool = False,
97 - per_channel_quant: bool = False,
76 global_num_experts: int = -1,
77 expert_map: Optional[torch.Tensor] = None,
100 - w1_scale: Optional[torch.Tensor] = None,
101 - w2_scale: Optional[torch.Tensor] = None,
102 - w1_bias: Optional[torch.Tensor] = None,
103 - w2_bias: Optional[torch.Tensor] = None,
104 - w1_precision: Optional["PrecisionConfig"] = None,
105 - w2_precision: Optional["PrecisionConfig"] = None,
106 - a1_scale: Optional[torch.Tensor] = None,
107 - a2_scale: Optional[torch.Tensor] = None,
108 - block_shape: Optional[list[int]] = None,
78 + a1q_scale: Optional[torch.Tensor] = None,
79 ) -> torch.Tensor:
80 + if quant_config is None:
81 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
82
83 # type check, uint8 means mxfp4
84 assert hidden_states.dtype == torch.bfloat16
113 - assert w1_bias is None or w1_bias.dtype == torch.float32
114 - assert w2_bias is None or w2_bias.dtype == torch.float32
85 + assert (quant_config.w1_bias is None
86 + or quant_config.w1_bias.dtype == torch.float32)
87 + assert (quant_config.w2_bias is None
88 + or quant_config.w2_bias.dtype == torch.float32)
89
90 # Shape check, only check non-mxfp4
91 assert hidden_states.shape[-1] == w1.shape[-2]
@@ -130,20 +104,20 @@ def triton_kernel_fused_experts(
104 intermediate_cache1 = matmul_ogs(
105 hidden_states,
106 w1,
133 - w1_bias,
107 + quant_config.w1_bias,
108 routing_data,
109 gather_indx=gather_indx,
136 - precision_config=w1_precision,
110 + precision_config=quant_config.w1_precision,
111 gammas=gammas if apply_router_weight_on_input else None,
112 fused_activation=act)
113
114 intermediate_cache3 = matmul_ogs(
115 intermediate_cache1,
116 w2,
143 - w2_bias,
117 + quant_config.w2_bias,
118 routing_data,
119 scatter_indx=scatter_indx,
146 - precision_config=w2_precision,
120 + precision_config=quant_config.w2_precision,
121 gammas=None if apply_router_weight_on_input else gammas,
122 y=output_tensor,
123 )
@@ -154,21 +128,13 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
128
129 def __init__(
130 self,
157 - quant_config,
131 max_num_tokens: int,
132 num_dispatchers: int,
160 - w1_precision: "PrecisionConfig",
161 - w2_precision: "PrecisionConfig",
162 - w1_bias: Optional[torch.Tensor],
163 - w2_bias: Optional[torch.Tensor],
133 + quant_config: FusedMoEQuantConfig,
134 ):
135 super().__init__(quant_config)
136 self.max_num_tokens = max_num_tokens
137 self.num_dispatchers = num_dispatchers
168 - self.w1_precision = w1_precision
169 - self.w2_precision = w2_precision
170 - self.w1_bias = w1_bias
171 - self.w2_bias = w2_bias
138
139 @property
140 def activation_formats(
@@ -212,12 +178,7 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
178 activation: str,
179 global_num_experts: int,
180 expert_map: Optional[torch.Tensor],
215 - w1_scale: Optional[torch.Tensor],
216 - w2_scale: Optional[torch.Tensor],
217 - w1_zp: Optional[torch.Tensor],
218 - w2_zp: Optional[torch.Tensor],
181 a1q_scale: Optional[torch.Tensor],
220 - a2_scale: Optional[torch.Tensor],
182 workspace13: torch.Tensor,
183 workspace2: torch.Tensor,
184 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -228,20 +189,12 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
189 hidden_states,
190 w1,
191 w2,
231 - None,
232 - None,
233 - None,
192 + routing_data=None,
193 + gather_indx=None,
194 + scatter_indx=None,
195 activation=activation,
196 + quant_config=self.quant_config,
197 apply_router_weight_on_input=False,
236 - use_fp8_w8a8=False,
237 - per_channel_quant=False,
198 global_num_experts=global_num_experts,
199 expert_map=expert_map,
240 - w1_scale=w1_scale,
241 - w2_scale=w2_scale,
242 - w1_bias=self.w1_bias,
243 - w2_bias=self.w2_bias,
244 - w1_precision=self.w1_precision,
245 - w2_precision=self.w2_precision,
246 - a1_scale=a1q_scale,
247 - a2_scale=a2_scale)
200 + a1q_scale=a1q_scale)
vllm/model_executor/layers/fused_moe/layer.py
+92 -45
@@ -22,7 +22,8 @@ from vllm.logger import init_logger
22 from vllm.model_executor.custom_op import CustomOp
23 # yapf: disable
24 from vllm.model_executor.layers.fused_moe.config import (
25 - FusedMoEConfig, FusedMoEParallelConfig)
25 + FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEConfig, FusedMoEParallelConfig,
26 + FusedMoEQuantConfig, biased_moe_quant_config)
27 # yapf: enable
28 from vllm.model_executor.layers.fused_moe.modular_kernel import (
29 FusedMoEActivationFormat, FusedMoEModularKernel,
@@ -78,11 +79,11 @@ class FusedMoeWeightScaleSupported(Enum):
79
80 class FusedMoEMethodBase(QuantizeMethodBase):
81
81 - # TODO(bnell): also pass quant_config?
82 def __init__(self, moe: FusedMoEConfig):
83 super().__init__()
84 self.moe = moe
85 - self.fused_experts: Optional[Callable] = None
85 + self.moe_quant_config: Optional[FusedMoEQuantConfig] = None
86 + self.fused_experts: Optional[FusedMoEModularKernel] = None
87 self.topk_indices_dtype = None
88
89 @abstractmethod
@@ -103,23 +104,28 @@ class FusedMoEMethodBase(QuantizeMethodBase):
104
105 @staticmethod
106 def _maybe_make_prepare_finalize(
106 - moe: FusedMoEConfig, ) -> Optional[FusedMoEPrepareAndFinalize]:
107 + moe: FusedMoEConfig,
108 + quant_config: Optional[FusedMoEQuantConfig],
109 + ) -> Optional[FusedMoEPrepareAndFinalize]:
110 all2all_manager = get_ep_group().device_communicator.all2all_manager
111 assert all2all_manager is not None
112
113 prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None
114
115 + # TODO: could allow this now
116 assert not moe.use_flashinfer_cutlass_kernels, \
117 "Must be created in modelopt.py"
118
119 if moe.use_pplx_kernels:
120 + assert quant_config is not None
121 +
122 hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
123 moe.max_num_tokens,
124 moe.hidden_dim,
125 moe.in_dtype,
120 - moe.quant_dtype,
121 - per_act_token_quant=moe.per_act_token_quant,
122 - block_shape=moe.block_shape,
126 + quant_config.quant_dtype,
127 + per_act_token_quant=quant_config.per_act_token_quant,
128 + block_shape=quant_config.block_shape,
129 )
130
131 all_to_all_args = dict(
@@ -165,6 +171,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
171 )
172
173 elif moe.use_deepep_ll_kernels:
174 + assert quant_config is not None
175 all_to_all_args = dict(
176 max_num_tokens_per_dp_rank=moe.max_num_tokens,
177 token_hidden_size=moe.hidden_dim,
@@ -174,13 +181,11 @@ class FusedMoEMethodBase(QuantizeMethodBase):
181 all2all_manager.world_size)
182 handle = all2all_manager.get_handle(all_to_all_args)
183
177 - # Note : We may want to use FP8 dispatch even otherwise just to
178 - # reduce datamovement
179 - use_fp8_dispatch = (moe.quant_config is not None
180 - and moe.quant_config.quant_dtype
181 - == current_platform.fp8_dtype()
182 - and moe.quant_config.block_shape
183 - == DEEPEP_QUANT_BLOCK_SHAPE)
184 + # Note: We may want to use FP8 dispatch just to reduce
185 + # data movement.
186 + use_fp8_dispatch = (
187 + quant_config.quant_dtype == current_platform.fp8_dtype()
188 + and quant_config.block_shape == DEEPEP_QUANT_BLOCK_SHAPE)
189
190 prepare_finalize = DeepEPLLPrepareAndFinalize(
191 handle,
@@ -192,11 +197,10 @@ class FusedMoEMethodBase(QuantizeMethodBase):
197 return prepare_finalize
198
199 def maybe_make_prepare_finalize(
195 - self,
196 - moe: FusedMoEConfig,
197 - ) -> Optional[FusedMoEPrepareAndFinalize]:
198 - if moe.moe_parallel_config.use_all2all_kernels:
199 - return FusedMoEMethodBase._maybe_make_prepare_finalize(moe)
200 + self) -> Optional[FusedMoEPrepareAndFinalize]:
201 + if self.moe.moe_parallel_config.use_all2all_kernels:
202 + return FusedMoEMethodBase._maybe_make_prepare_finalize(
203 + self.moe, self.moe_quant_config)
204 else:
205 return None
206
@@ -204,7 +208,13 @@ class FusedMoEMethodBase(QuantizeMethodBase):
208 # prepare_communication_buffer_for_model.
209 def init_prepare_finalize(self, layer: torch.nn.Module):
210 assert self.moe is not None
207 - prepare_finalize = self.maybe_make_prepare_finalize(self.moe)
211 +
212 + # We must get the quant config here so that the layer is
213 + # completely initialized, i.e. all weights loaded and post
214 + # processed.
215 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)
216 +
217 + prepare_finalize = self.maybe_make_prepare_finalize()
218
219 if prepare_finalize is not None:
220 logger.debug("%s for %s(%s)", prepare_finalize.__class__.__name__,
@@ -213,7 +223,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
223 assert self.fused_experts is None, \
224 f"Attempt to override experts for {id(self)}!"
225 self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
216 - experts = self.select_gemm_impl(prepare_finalize, self.moe, layer)
226 + experts = self.select_gemm_impl(prepare_finalize, layer)
227 self.fused_experts = FusedMoEModularKernel(
228 prepare_finalize,
229 experts,
@@ -223,7 +233,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
233 def select_gemm_impl(
234 self,
235 prepare_finalize: FusedMoEPrepareAndFinalize,
226 - moe: FusedMoEConfig,
236 layer: torch.nn.Module,
237 ) -> FusedMoEPermuteExpertsUnpermute:
238 # based on the all2all implementation, select the appropriate
@@ -232,6 +241,11 @@ class FusedMoEMethodBase(QuantizeMethodBase):
241 f"{self.__class__.__name__} must select appropriate gemm "
242 "implementation based on the prepare_finalize")
243
244 + @abstractmethod
245 + def get_fused_moe_quant_config(
246 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
247 + raise NotImplementedError
248 +
249 @abstractmethod
250 def apply(
251 self,
@@ -265,7 +279,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
279
280 def __init__(self, moe: FusedMoEConfig):
281 super().__init__(moe)
268 - self.has_bias = self.moe.has_bias
282 self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
283 if self.rocm_aiter_moe_enabled:
284 from .rocm_aiter_fused_moe import rocm_aiter_fused_experts
@@ -273,23 +286,30 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
286 else:
287 self.rocm_aiter_fused_experts = None # type: ignore
288
289 + def maybe_make_prepare_finalize(
290 + self) -> Optional[FusedMoEPrepareAndFinalize]:
291 + if self.rocm_aiter_moe_enabled:
292 + return None
293 + else:
294 + return super().maybe_make_prepare_finalize()
295 +
296 def select_gemm_impl(
297 self,
298 prepare_finalize: FusedMoEPrepareAndFinalize,
279 - # TODO(bnell): Remove. Every layer should have an moe config object.
280 - moe: FusedMoEConfig,
299 layer: torch.nn.Module,
300 ) -> FusedMoEPermuteExpertsUnpermute:
301 + assert self.moe_quant_config is not None
302 if (prepare_finalize.activation_format ==
303 FusedMoEActivationFormat.BatchedExperts):
304 logger.debug("BatchedTritonExperts %s", self.moe)
305 return BatchedTritonExperts(
306 max_num_tokens=self.moe.max_num_tokens,
307 num_dispatchers=prepare_finalize.num_dispatchers(),
308 + quant_config=self.moe_quant_config,
309 )
310 else:
311 logger.debug("TritonExperts %s", self.moe)
292 - return TritonExperts()
312 + return TritonExperts(self.moe_quant_config)
313
314 def create_weights(self, layer: torch.nn.Module, num_experts: int,
315 hidden_size: int, intermediate_size_per_partition: int,
@@ -303,7 +323,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
323 requires_grad=False)
324 layer.register_parameter("w13_weight", w13_weight)
325 set_weight_attrs(w13_weight, extra_weight_attrs)
306 - if self.has_bias:
326 + if self.moe.has_bias:
327 w13_bias = torch.nn.Parameter(torch.zeros(
328 num_experts,
329 2 * intermediate_size_per_partition,
@@ -320,7 +340,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
340 requires_grad=False)
341 layer.register_parameter("w2_weight", w2_weight)
342 set_weight_attrs(w2_weight, extra_weight_attrs)
323 - if self.has_bias:
343 + if self.moe.has_bias:
344 w2_bias = torch.nn.Parameter(torch.zeros(num_experts,
345 hidden_size,
346 dtype=params_dtype),
@@ -442,6 +462,16 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
462 logical_replica_count=logical_replica_count,
463 )
464
465 + def get_fused_moe_quant_config(
466 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
467 + if self.moe.has_bias:
468 + return biased_moe_quant_config(
469 + layer.w13_bias,
470 + layer.w2_bias,
471 + )
472 + else:
473 + return FUSED_MOE_UNQUANTIZED_CONFIG
474 +
475 def forward_cuda(
476 self,
477 layer: torch.nn.Module,
@@ -486,6 +516,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
516 logical_replica_count=logical_replica_count)
517
518 if self.rocm_aiter_moe_enabled:
519 + assert self.fused_experts is None
520 return self.rocm_aiter_fused_experts(
521 hidden_states=x,
522 w1=layer.w13_weight,
@@ -496,7 +527,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
527 activation=activation,
528 apply_router_weight_on_input=apply_router_weight_on_input)
529 elif self.fused_experts is not None:
499 - if self.has_bias:
530 + if self.moe.has_bias:
531 raise ValueError(
532 "FusedMoEModularKernel does not support bias.")
533 return self.fused_experts(
@@ -517,12 +548,11 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
548 hidden_states=x,
549 w1=layer.w13_weight,
550 w2=layer.w2_weight,
520 - w1_bias=layer.w13_bias if self.has_bias else None,
521 - w2_bias=layer.w2_bias if self.has_bias else None,
551 topk_weights=topk_weights,
552 topk_ids=topk_ids,
553 inplace=True,
554 activation=activation,
555 + quant_config=self.moe_quant_config,
556 apply_router_weight_on_input=apply_router_weight_on_input,
557 global_num_experts=global_num_experts,
558 expert_map=expert_map,
@@ -933,16 +963,18 @@ class FusedMoE(CustomOp):
963 # since model_config is not set in the pytest test.
964 model_dtype = params_dtype
965
936 - moe = FusedMoEConfig.make(num_experts=self.global_num_experts,
937 - experts_per_token=top_k,
938 - hidden_dim=hidden_size,
939 - num_local_experts=self.local_num_experts,
940 - moe_parallel_config=self.moe_parallel_config,
941 - in_dtype=model_dtype,
942 - max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
943 - quant_config=quant_config,
944 - has_bias=has_bias)
966 + moe = FusedMoEConfig(
967 + num_experts=self.global_num_experts,
968 + experts_per_token=top_k,
969 + hidden_dim=hidden_size,
970 + num_local_experts=self.local_num_experts,
971 + moe_parallel_config=self.moe_parallel_config,
972 + in_dtype=model_dtype,
973 + max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
974 + has_bias=has_bias,
975 + )
976 self.moe_config = moe
977 + self.moe_quant_config: Optional[FusedMoEQuantConfig] = None
978 self.quant_config = quant_config
979
980 # Note: get_quant_method will look at the layer's local_num_experts
@@ -990,6 +1022,9 @@ class FusedMoE(CustomOp):
1022 # Chunked all2all staging tensor
1023 self.batched_hidden_states: Optional[torch.Tensor] = None
1024 self.batched_router_logits: Optional[torch.Tensor] = None
1025 +
1026 + # TODO(bnell): flashinfer uses non-batched format.
1027 + # Does it really need a batched buffer?
1028 if (self.moe_parallel_config.use_pplx_kernels
1029 or self.moe_parallel_config.use_deepep_ll_kernels
1030 or self.moe_config.use_flashinfer_cutlass_kernels):
@@ -1062,7 +1097,9 @@ class FusedMoE(CustomOp):
1097
1098 @property
1099 def use_flashinfer_cutlass_kernels(self):
1065 - return self.moe_config.use_flashinfer_cutlass_kernels
1100 + return (self.moe_quant_config is not None
1101 + and self.moe_quant_config.quant_dtype == "nvfp4"
1102 + and self.moe_config.use_flashinfer_cutlass_kernels)
1103
1104 def update_expert_map(self):
1105 # ep_size and ep_rank should already be updated
@@ -1492,6 +1529,11 @@ class FusedMoE(CustomOp):
1529 self.logical_to_physical_map = logical_to_physical_map[moe_layer_idx]
1530 self.logical_replica_count = logical_replica_count[moe_layer_idx]
1531
1532 + def ensure_moe_quant_config(self):
1533 + if self.quant_method.moe_quant_config is None:
1534 + self.quant_method.moe_quant_config = (
1535 + self.quant_method.get_fused_moe_quant_config(self))
1536 +
1537 @staticmethod
1538 def select_experts(
1539 hidden_states: torch.Tensor,
@@ -1711,6 +1753,8 @@ class FusedMoE(CustomOp):
1753 assert (
1754 self.batched_router_logits.size(-1) == full_router_logits.size(-1))
1755
1756 + self.ensure_moe_quant_config()
1757 +
1758 full_fused_final_hidden_states = torch.empty_like(full_hidden_states)
1759 if self.shared_experts is not None:
1760 full_shared_final_hidden_states = torch.empty_like(
@@ -1825,14 +1869,17 @@ class FusedMoE(CustomOp):
1869 router_logits: torch.Tensor,
1870 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1871 assert self.quant_method is not None
1872 +
1873 + self.ensure_moe_quant_config()
1874 +
1875 # Route to the chunked forward path using the FlashInfer Cutlass kernel
1876 # only when data parallelism (DP) is enabled.
1830 - use_flashinfer_cutlass_kernels = (
1831 - self.dp_size > 1
1832 - and self.moe_config.use_flashinfer_cutlass_kernels)
1877 + _use_flashinfer_cutlass_kernels = (self.dp_size > 1 and
1878 + self.use_flashinfer_cutlass_kernels)
1879 +
1880 if (self.moe_parallel_config.use_pplx_kernels
1881 or self.moe_parallel_config.use_deepep_ll_kernels
1835 - or use_flashinfer_cutlass_kernels):
1882 + or _use_flashinfer_cutlass_kernels):
1883 return self.forward_impl_chunked(hidden_states, router_logits)
1884
1885 do_naive_dispatch_combine: bool = (
vllm/model_executor/layers/fused_moe/modular_kernel.py
+71 -78
@@ -177,8 +177,6 @@ class FusedMoEPrepareAndFinalize(ABC):
177 def prepare(
178 self,
179 a1: torch.Tensor,
180 - a1_scale: Optional[torch.Tensor],
181 - a2_scale: Optional[torch.Tensor],
180 topk_weights: torch.Tensor,
181 topk_ids: torch.Tensor,
182 num_experts: int,
@@ -189,9 +187,6 @@ class FusedMoEPrepareAndFinalize(ABC):
187 """
188 Perform any quantization (and/or) dispatching needed for this kernel.
189 - a1: The (unquantized) input to the MoE layer.
192 - - a1_scale: Optional scales for a1
193 - - a2_scale: Optional scales for the second MoE gemm. Required to make
194 - sure the quantization is consistent for both gemms.
190 - topk_ids: The topk ids.
191 - topk_weights: The topk weights.
192 - num_experts: The total number of experts in the global expert space.
@@ -199,10 +194,11 @@ class FusedMoEPrepareAndFinalize(ABC):
194 space to the local expert space of the expert parallel shard.
195 - apply_router_weight_on_input: When True, apply the weights to the
196 activations, before quantization + dispatching.
197 + - quant_config: Quantization info provided by the fused experts.
198
199 Returns a tuple of:
200 - quantized + dispatched a.
205 - - quantized + dispatched a1_scales.
201 + - Optional quantized + dispatched a1_scales.
202 - Optional ExpertTokensMetadata containing gpu/cpu tensors
203 as big as the number of local experts with the information about the
204 number of tokens assigned to each local expert.
@@ -220,8 +216,6 @@ class FusedMoEPrepareAndFinalize(ABC):
216 def prepare_async(
217 self,
218 a1: torch.Tensor,
223 - a1_scale: Optional[torch.Tensor],
224 - a2_scale: Optional[torch.Tensor],
219 topk_weights: torch.Tensor,
220 topk_ids: torch.Tensor,
221 num_experts: int,
@@ -316,6 +310,7 @@ class FusedMoEPrepareAndFinalize(ABC):
310 raise NotImplementedError
311
312
313 +# TODO: add supported activations method (return string)
314 class FusedMoEPermuteExpertsUnpermute(ABC):
315 """
316 An abstract base class for the [Permute-Experts-Unpermute] step described
@@ -324,12 +319,12 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
319
320 def __init__(
321 self,
327 - quant_config: Optional[FusedMoEQuantConfig],
322 + quant_config: FusedMoEQuantConfig,
323 ):
329 - if quant_config is not None:
330 - self.quant_config = quant_config
331 - else:
332 - self.quant_config = FusedMoEQuantConfig()
324 + """
325 + quant_config: Quantization parameters for this experts instance.
326 + """
327 + self.quant_config = quant_config
328
329 @property
330 @abstractmethod
@@ -341,6 +336,11 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
336 """
337 raise NotImplementedError
338
339 + #
340 + # Various helpers for accessing quantization parameters from the
341 + # quant_config.
342 + #
343 +
344 @property
345 def quant_dtype(self) -> Optional[torch.dtype]:
346 return self.quant_config.quant_dtype
@@ -357,6 +357,54 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
357 def per_out_ch_quant(self) -> bool:
358 return self.quant_config.per_out_ch_quant
359
360 + @property
361 + def a1_scale(self) -> Optional[torch.Tensor]:
362 + return self.quant_config.a1_scale
363 +
364 + @property
365 + def a2_scale(self) -> Optional[torch.Tensor]:
366 + return self.quant_config.a2_scale
367 +
368 + @property
369 + def a1_gscale(self) -> Optional[torch.Tensor]:
370 + return self.quant_config.a1_gscale
371 +
372 + @property
373 + def a2_gscale(self) -> Optional[torch.Tensor]:
374 + return self.quant_config.a2_gscale
375 +
376 + @property
377 + def w1_scale(self) -> Optional[torch.Tensor]:
378 + return self.quant_config.w1_scale
379 +
380 + @property
381 + def w2_scale(self) -> Optional[torch.Tensor]:
382 + return self.quant_config.w2_scale
383 +
384 + @property
385 + def w1_zp(self) -> Optional[torch.Tensor]:
386 + return self.quant_config.w1_zp
387 +
388 + @property
389 + def w2_zp(self) -> Optional[torch.Tensor]:
390 + return self.quant_config.w2_zp
391 +
392 + @property
393 + def w1_bias(self) -> Optional[torch.Tensor]:
394 + return self.quant_config.w1_bias
395 +
396 + @property
397 + def w2_bias(self) -> Optional[torch.Tensor]:
398 + return self.quant_config.w2_bias
399 +
400 + @property
401 + def g1_alphas(self) -> Optional[torch.Tensor]:
402 + return self.quant_config.g1_alphas
403 +
404 + @property
405 + def g2_alphas(self) -> Optional[torch.Tensor]:
406 + return self.quant_config.g2_alphas
407 +
408 # TODO (bnell): make this return a CHUNK_SIZE or None instead?
409 @abstractmethod
410 def supports_chunking(self) -> bool:
@@ -433,12 +481,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
481 activation: str,
482 global_num_experts: int,
483 expert_map: Optional[torch.Tensor],
436 - w1_scale: Optional[torch.Tensor],
437 - w2_scale: Optional[torch.Tensor],
438 - w1_zp: Optional[torch.Tensor],
439 - w2_zp: Optional[torch.Tensor],
484 a1q_scale: Optional[torch.Tensor],
441 - a2_scale: Optional[torch.Tensor],
485 workspace13: torch.Tensor,
486 workspace2: torch.Tensor,
487 expert_tokens_meta: Optional[ExpertTokensMetadata],
@@ -455,7 +498,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
498 - w1 (torch.Tensor): The first set of expert weights.
499 - w2 (torch.Tensor): The second set of expert weights.
500 - topk_weights: A map of row to expert weights. Some implementations
458 - choose to do weight application.
501 + choose to do weight application.
502 - topk_ids (torch.Tensor): A map of row to expert id.
503 - activation (str): The activation function to apply after the first
504 MoE layer.
@@ -464,15 +507,9 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
507 - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
508 from the global expert space to the local expert space of the expert
509 parallel shard.
467 - - w1_scale (Optional[torch.Tensor]): Optional scale to be used for w1.
468 - - w2_scale (Optional[torch.Tensor]): Optional scale to be used for w2.
469 - - w1_zp (Optional[torch.Tensor]): Optional zero points to be used for
470 - w1.
471 - - w2_zp (Optional[torch.Tensor]): Optional zero points to be used for
472 - w2.
510 - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
474 - used for a1.
475 - - a2_scale (Optional[torch.Tensor]): Optional scale to be used for a2.
511 + used for a1. Result of quantization from prepare/finalize and not
512 + from the FusedMoEQuantConfig.
513 - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
514 must be large enough to hold output of either MoE gemm.
515 - workspace2 (torch.Tensor): A scratch tensor used for the activation
@@ -559,12 +596,7 @@ class FusedMoEModularKernel(torch.nn.Module):
596 global_num_experts: int,
597 local_num_experts: int,
598 expert_map: Optional[torch.Tensor],
562 - w1_scale: Optional[torch.Tensor],
563 - w2_scale: Optional[torch.Tensor],
564 - w1_zp: Optional[torch.Tensor],
565 - w2_zp: Optional[torch.Tensor],
599 a1q_scale: Optional[torch.Tensor],
567 - a2_scale: Optional[torch.Tensor],
600 expert_tokens_meta: Optional[ExpertTokensMetadata],
601 apply_router_weight_on_input: bool,
602 ) -> torch.Tensor:
@@ -601,12 +633,7 @@ class FusedMoEModularKernel(torch.nn.Module):
633 activation=activation,
634 global_num_experts=global_num_experts,
635 expert_map=expert_map,
604 - w1_scale=w1_scale,
605 - w2_scale=w2_scale,
606 - w1_zp=w1_zp,
607 - w2_zp=w2_zp,
636 a1q_scale=a1q_scale,
609 - a2_scale=a2_scale,
637 workspace13=workspace13,
638 workspace2=workspace2,
639 expert_tokens_meta=expert_tokens_meta,
@@ -627,12 +654,7 @@ class FusedMoEModularKernel(torch.nn.Module):
654 global_num_experts: int,
655 local_num_experts: int,
656 expert_map: Optional[torch.Tensor],
630 - w1_scale: Optional[torch.Tensor],
631 - w2_scale: Optional[torch.Tensor],
632 - w1_zp: Optional[torch.Tensor],
633 - w2_zp: Optional[torch.Tensor],
657 a1q_scale: Optional[torch.Tensor],
635 - a2_scale: Optional[torch.Tensor],
658 expert_tokens_meta: Optional[ExpertTokensMetadata],
659 apply_router_weight_on_input: bool,
660 ) -> torch.Tensor:
@@ -658,12 +680,7 @@ class FusedMoEModularKernel(torch.nn.Module):
680 global_num_experts=global_num_experts,
681 local_num_experts=local_num_experts,
682 expert_map=expert_map,
661 - w1_scale=w1_scale,
662 - w2_scale=w2_scale,
663 - w1_zp=w1_zp,
664 - w2_zp=w2_zp,
683 a1q_scale=a1q_scale,
666 - a2_scale=a2_scale,
684 expert_tokens_meta=expert_tokens_meta,
685 apply_router_weight_on_input=apply_router_weight_on_input,
686 )
@@ -685,9 +702,13 @@ class FusedMoEModularKernel(torch.nn.Module):
702 Optional[torch.Tensor], torch.Tensor, torch.Tensor]:
703 s = chunk_idx * CHUNK_SIZE
704 e = min(s + CHUNK_SIZE, M)
688 - return (a1q[s:e], _chunk_scales(a1q_scale, s, e),
689 - _chunk_scales(a2_scale, s,
690 - e), topk_ids[s:e], topk_weights[s:e])
705 + return (
706 + a1q[s:e],
707 + _chunk_scales(a1q_scale, s, e),
708 + _chunk_scales(self.fused_experts.a2_scale, s, e),
709 + topk_ids[s:e],
710 + topk_weights[s:e],
711 + )
712
713 def slice_output_tensor(chunk_idx: int) -> torch.Tensor:
714 assert fused_out.size(0) % M == 0, (
@@ -744,12 +765,7 @@ class FusedMoEModularKernel(torch.nn.Module):
765 global_num_experts=global_num_experts,
766 local_num_experts=local_num_experts,
767 expert_map=expert_map,
747 - w1_scale=w1_scale,
748 - w2_scale=w2_scale,
749 - w1_zp=w1_zp,
750 - w2_zp=w2_zp,
768 a1q_scale=c_a1q_scale,
752 - a2_scale=c_a2_scale,
769 expert_tokens_meta=c_expert_tokens_meta,
770 apply_router_weight_on_input=apply_router_weight_on_input,
771 )
@@ -767,12 +783,6 @@ class FusedMoEModularKernel(torch.nn.Module):
783 activation: str = "silu",
784 global_num_experts: int = -1,
785 expert_map: Optional[torch.Tensor] = None,
770 - w1_scale: Optional[torch.Tensor] = None,
771 - w2_scale: Optional[torch.Tensor] = None,
772 - w1_zp: Optional[torch.Tensor] = None,
773 - w2_zp: Optional[torch.Tensor] = None,
774 - a1_scale: Optional[torch.Tensor] = None,
775 - a2_scale: Optional[torch.Tensor] = None,
786 apply_router_weight_on_input: bool = False,
787 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
788 """
@@ -795,14 +805,6 @@ class FusedMoEModularKernel(torch.nn.Module):
805 - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
806 from the global expert space to the local expert space of the expert
807 parallel shard.
798 - - w1_scale (Optional[torch.Tensor]): Optional scale to be used for w1.
799 - - w2_scale (Optional[torch.Tensor]): Optional scale to be used for w2.
800 - - w1_zp (Optional[torch.Tensor]): Optional zero points to be used for
801 - w1.
802 - - w2_zp (Optional[torch.Tensor]): Optional zero points to be used for
803 - w2.
804 - - a1_scale (Optional[torch.Tensor]): Optional scale to be used for a1.
805 - - a2_scale (Optional[torch.Tensor]): Optional scale to be used for a2.
808 - apply_router_weight_on_input (bool): When true, the topk weights are
809 applied directly on the inputs. This is only applicable when topk is
810 1.
@@ -832,8 +834,6 @@ class FusedMoEModularKernel(torch.nn.Module):
834 (a1q, a1q_scale, expert_tokens_meta, _expert_topk_ids,
835 _expert_topk_weights) = self.prepare_finalize.prepare(
836 a1,
835 - a1_scale,
836 - a2_scale,
837 topk_weights,
838 topk_ids,
839 global_num_experts,
@@ -846,8 +846,6 @@ class FusedMoEModularKernel(torch.nn.Module):
846 dbo_maybe_run_recv_hook()
847 hook, receiver = self.prepare_finalize.prepare_async(
848 a1,
849 - a1_scale,
850 - a2_scale,
849 topk_weights,
850 topk_ids,
851 global_num_experts,
@@ -897,12 +895,7 @@ class FusedMoEModularKernel(torch.nn.Module):
895 global_num_experts=global_num_experts,
896 local_num_experts=local_num_experts,
897 expert_map=expert_map,
900 - w1_scale=w1_scale,
901 - w2_scale=w2_scale,
902 - w1_zp=w1_zp,
903 - w2_zp=w2_zp,
898 a1q_scale=a1q_scale,
905 - a2_scale=a2_scale,
899 expert_tokens_meta=expert_tokens_meta,
900 apply_router_weight_on_input=apply_router_weight_on_input,
901 )
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+3 -7
@@ -95,8 +95,6 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
95 def prepare_async(
96 self,
97 a1: torch.Tensor,
98 - a1_scale: Optional[torch.Tensor],
99 - a2_scale: Optional[torch.Tensor],
98 topk_weights: torch.Tensor,
99 topk_ids: torch.Tensor,
100 num_experts: int,
@@ -130,8 +128,10 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
128
129 repeat_cols = 4
130 repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0)
131 + # TODO(bnell): always pass quant_config.a1_scale?
132 a1q, a1q_scale = moe_kernel_quantize_input(
134 - a1, (None if quant_config.per_act_token_quant else a1_scale),
133 + a1, (None if quant_config.per_act_token_quant else
134 + quant_config.a1_scale),
135 quant_dtype=quant_config.quant_dtype,
136 per_act_token_quant=quant_config.per_act_token_quant,
137 block_shape=quant_config.block_shape)
@@ -253,8 +253,6 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
253 def prepare(
254 self,
255 a1: torch.Tensor,
256 - a1_scale: Optional[torch.Tensor],
257 - a2_scale: Optional[torch.Tensor],
256 topk_weights: torch.Tensor,
257 topk_ids: torch.Tensor,
258 num_experts: int,
@@ -264,8 +262,6 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
262 ) -> mk.PrepareResultType:
263 hook, receiver = self.prepare_async(
264 a1,
267 - a1_scale,
268 - a2_scale,
265 topk_weights,
266 topk_ids,
267 num_experts,
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+1 -3
@@ -30,8 +30,6 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
30 def prepare(
31 self,
32 a1: torch.Tensor,
33 - a1_scale: Optional[torch.Tensor],
34 - a2_scale: Optional[torch.Tensor],
33 topk_weights: torch.Tensor,
34 topk_ids: torch.Tensor,
35 num_experts: int,
@@ -48,7 +46,7 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
46 a1.mul_(topk_weights.to(a1.dtype))
47
48 a1q, a1q_scale = moe_kernel_quantize_input(
51 - a1, a1_scale, quant_config.quant_dtype,
49 + a1, quant_config.a1_scale, quant_config.quant_dtype,
50 quant_config.per_act_token_quant, quant_config.block_shape)
51
52 return a1q, a1q_scale, None, None, None
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
+26 -26
@@ -7,6 +7,8 @@ from typing import Optional
7 import torch
8
9 from vllm import envs
10 +from vllm.model_executor.layers.fused_moe.config import (
11 + FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEQuantConfig)
12 from vllm.platforms import current_platform
13 from vllm.utils import direct_register_custom_op
14
@@ -305,21 +307,18 @@ def rocm_aiter_grouped_topk(
307
308
309 def rocm_aiter_fused_experts(
308 - hidden_states: torch.Tensor,
309 - w1: torch.Tensor,
310 - w2: torch.Tensor,
311 - topk_weights: torch.Tensor,
312 - topk_ids: torch.Tensor,
313 - activation: str = "silu",
314 - apply_router_weight_on_input: bool = False,
315 - use_fp8_w8a8: bool = False,
316 - per_channel_quant: bool = False,
317 - w1_scale: Optional[torch.Tensor] = None,
318 - w2_scale: Optional[torch.Tensor] = None,
319 - a1_scale: Optional[torch.Tensor] = None,
320 - a2_scale: Optional[torch.Tensor] = None,
321 - block_shape: Optional[list[int]] = None,
322 - expert_map: Optional[torch.Tensor] = None) -> torch.Tensor:
310 + hidden_states: torch.Tensor,
311 + w1: torch.Tensor,
312 + w2: torch.Tensor,
313 + topk_weights: torch.Tensor,
314 + topk_ids: torch.Tensor,
315 + activation: str = "silu",
316 + apply_router_weight_on_input: bool = False,
317 + expert_map: Optional[torch.Tensor] = None,
318 + quant_config: Optional[FusedMoEQuantConfig] = None,
319 +) -> torch.Tensor:
320 + if quant_config is None:
321 + quant_config = FUSED_MOE_UNQUANTIZED_CONFIG
322
323 activation_method = (ActivationMethod.SILU
324 if activation == "silu" else ActivationMethod.GELU)
@@ -333,7 +332,8 @@ def rocm_aiter_fused_experts(
332 expert_mask = None
333
334 # w8a8 per-channel quantization
336 - if per_channel_quant and apply_router_weight_on_input and use_fp8_w8a8:
335 + if (quant_config.per_act_token_quant and apply_router_weight_on_input
336 + and quant_config.use_fp8_w8a8):
337 # AITER tkw1 kernel for FP8 models with `apply_router_weight_on_input`
338 # This applies topk_weights on the GEMM output of the first FC layer
339 # rather than the second FC.
@@ -349,8 +349,8 @@ def rocm_aiter_fused_experts(
349 w2,
350 topk_weights,
351 topk_ids,
352 - fc1_scale=w1_scale,
353 - fc2_scale=w2_scale,
352 + fc1_scale=quant_config.w1_scale,
353 + fc2_scale=quant_config.w2_scale,
354 fc1_smooth_scale=None,
355 fc2_smooth_scale=None,
356 a16=False,
@@ -362,14 +362,14 @@ def rocm_aiter_fused_experts(
362 quant_method = QuantMethod.NO.value
363
364 # w8a8 block-scaled
365 - if block_shape is not None and use_fp8_w8a8:
365 + if quant_config.block_shape is not None and quant_config.use_fp8_w8a8:
366 assert not apply_router_weight_on_input, (
367 "apply_router_weight_on_input is\
368 not supported for block scaled moe")
369 - assert w1_scale is not None
370 - assert w2_scale is not None
369 + assert quant_config.w1_scale is not None
370 + assert quant_config.w2_scale is not None
371 quant_method = QuantMethod.BLOCK_128x128.value
372 - elif use_fp8_w8a8:
372 + elif quant_config.use_fp8_w8a8:
373 # Currently only per tensor quantization method is enabled.
374 quant_method = QuantMethod.PER_TENSOR.value
375
@@ -390,10 +390,10 @@ def rocm_aiter_fused_experts(
390 expert_mask=expert_mask,
391 quant_method=quant_method,
392 activation_method=activation_method,
393 - w1_scale=w1_scale,
394 - w2_scale=w2_scale,
395 - a1_scale=a1_scale,
396 - a2_scale=a2_scale,
393 + w1_scale=quant_config.w1_scale,
394 + w2_scale=quant_config.w2_scale,
395 + a1_scale=quant_config.a1_scale,
396 + a2_scale=quant_config.a2_scale,
397 doweight_stage1=apply_router_weight_on_input)
398
399
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+9 -39
@@ -7,7 +7,8 @@ import torch
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.deep_gemm_moe import (
10 - DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape,
10 + DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape)
11 +from vllm.model_executor.layers.fused_moe.deep_gemm_utils import (
12 deep_gemm_block_shape)
13 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
14 from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
@@ -17,40 +18,19 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
18
19 def __init__(
20 self,
20 - use_fp8_w8a8: bool = False,
21 - use_int8_w8a8: bool = False,
22 - use_int8_w8a16: bool = False,
23 - use_int4_w4a16: bool = False,
24 - use_mxfp4_w4a4: bool = False,
25 - per_act_token_quant: bool = False,
26 - block_shape: Optional[list[int]] = None,
21 + quant_config: FusedMoEQuantConfig,
22 allow_deep_gemm: bool = False,
23 ):
29 - super().__init__(
30 - FusedMoEQuantConfig.make(
31 - use_fp8_w8a8=use_fp8_w8a8,
32 - use_int8_w8a8=use_int8_w8a8,
33 - use_int8_w8a16=use_int8_w8a16,
34 - use_int4_w4a16=use_int4_w4a16,
35 - use_mxfp4_w4a4=use_mxfp4_w4a4,
36 - per_act_token_quant=per_act_token_quant,
37 - block_shape=block_shape,
38 - ))
39 - self.triton_expert = TritonExperts(
40 - use_fp8_w8a8=use_fp8_w8a8,
41 - use_int8_w8a8=use_int8_w8a8,
42 - use_int4_w4a16=use_int4_w4a16,
43 - use_int8_w8a16=use_int8_w8a16,
44 - use_mxfp4_w4a4=use_mxfp4_w4a4,
45 - per_act_token_quant=per_act_token_quant,
46 - block_shape=block_shape,
47 - )
24 + super().__init__(quant_config)
25 +
26 + self.triton_expert = TritonExperts(quant_config)
27
49 - self.allow_deep_gemm = (allow_deep_gemm and use_fp8_w8a8 and
28 + self.allow_deep_gemm = (allow_deep_gemm
29 + and self.quant_config.use_fp8_w8a8 and
30 self.block_shape == deep_gemm_block_shape())
31
32 self.deep_gemm_expert = DeepGemmExperts(
53 - ) if self.allow_deep_gemm else None
33 + self.quant_config) if self.allow_deep_gemm else None
34
35 @property
36 def activation_formats(
@@ -130,12 +110,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
110 activation: str,
111 global_num_experts: int,
112 expert_map: Optional[torch.Tensor],
133 - w1_scale: Optional[torch.Tensor],
134 - w2_scale: Optional[torch.Tensor],
135 - w1_zp: Optional[torch.Tensor],
136 - w2_zp: Optional[torch.Tensor],
113 a1q_scale: Optional[torch.Tensor],
138 - a2_scale: Optional[torch.Tensor],
114 workspace13: torch.Tensor,
115 workspace2: torch.Tensor,
116 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -158,12 +133,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
133 activation,
134 global_num_experts,
135 expert_map,
161 - w1_scale,
162 - w2_scale,
163 - w1_zp,
164 - w2_zp,
136 a1q_scale,
166 - a2_scale,
137 workspace13,
138 workspace2,
139 expert_tokens_meta,
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+9 -16
@@ -5,7 +5,8 @@ from typing import Optional
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 FusedMoEConfig
8 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
9 + FusedMoEQuantConfig)
10 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
11 TopKWeightAndReduceNoOP)
12 from vllm.utils import next_power_of_2
@@ -16,20 +17,17 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
17 def __init__(
18 self,
19 moe: FusedMoEConfig,
20 + quant_config: FusedMoEQuantConfig,
21 gemm1_alpha,
22 gemm1_beta,
23 gemm1_clamp_limit,
22 - w13_bias,
23 - w2_bias,
24 max_capture_size,
25 ):
26 - super().__init__(moe.quant_config)
26 + super().__init__(quant_config)
27 self.moe = moe
28 self.gemm1_alpha = gemm1_alpha
29 self.gemm1_beta = gemm1_beta
30 self.gemm1_clamp_limit = gemm1_clamp_limit
31 - self.w13_bias = w13_bias
32 - self.w2_bias = w2_bias
31 self.max_capture_size = max_capture_size
32
33 @property
@@ -104,12 +102,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
102 activation: str,
103 global_num_experts: int,
104 expert_map: Optional[torch.Tensor],
107 - w1_scale: Optional[torch.Tensor],
108 - w2_scale: Optional[torch.Tensor],
109 - w1_zp: Optional[torch.Tensor],
110 - w2_zp: Optional[torch.Tensor],
105 a1q_scale: Optional[torch.Tensor],
112 - a2_scale: Optional[torch.Tensor],
106 workspace13: torch.Tensor,
107 workspace2: torch.Tensor,
108 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
@@ -129,8 +122,8 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
122 packed_tensor = (topk_ids.to(torch.int32) << 16) | topk_weights.to(
123 torch.bfloat16).view(torch.int16)
124
132 - assert w1_scale is not None
133 - assert w2_scale is not None
125 + assert self.w1_scale is not None
126 + assert self.w2_scale is not None
127 kwargs = {
128 "topk_ids":
129 packed_tensor,
@@ -143,9 +136,9 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
136 "gemm1_weights":
137 w1,
138 "gemm1_weights_scale":
146 - w1_scale,
139 + self.w1_scale,
140 "gemm1_bias":
148 - self.w13_bias,
141 + self.w1_bias,
142 "gemm1_alpha":
143 self.gemm1_alpha,
144 "gemm1_beta":
@@ -155,7 +148,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
148 "gemm2_weights":
149 w2,
150 "gemm2_weights_scale":
158 - w2_scale,
151 + self.w2_scale,
152 "gemm2_bias":
153 self.w2_bias,
154 "output1_scale_scalar":
vllm/model_executor/layers/fused_moe/utils.py
+4
@@ -268,3 +268,7 @@ def _validate_scale_shape(
268 assert block_shape is not None
269 expected = (a.shape[0], cdiv(a.shape[1], block_shape[1]))
270 assert a_scale.shape == expected, f"{a_scale.shape} == {expected}"
271 +
272 +
273 +def activation_without_mul(activation: str) -> str:
274 + return activation + "_no_mul"
vllm/model_executor/layers/quantization/awq_marlin.py
+7 -1
@@ -9,8 +9,10 @@ from torch.nn import Parameter
9 import vllm.model_executor.layers.fused_moe # noqa
10 from vllm import _custom_ops as ops
11 from vllm.logger import init_logger
12 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
13 + FusedMoEQuantConfig)
14 from vllm.model_executor.layers.fused_moe.layer import (
13 - FusedMoE, FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
15 + FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
16 UnquantizedFusedMoEMethod)
17 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
18 UnquantizedLinearMethod,
@@ -483,6 +485,10 @@ class AWQMoEMethod(FusedMoEMethodBase):
485 if hasattr(layer, "w2_bias") and layer.w2_bias is not None:
486 layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)
487
488 + def get_fused_moe_quant_config(
489 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
490 + return None
491 +
492 def apply(
493 self,
494 layer: torch.nn.Module,
vllm/model_executor/layers/quantization/bitsandbytes.py
+7 -1
@@ -6,8 +6,9 @@ from typing import Any, Callable, Optional, Union
6 import torch
7 from packaging import version
8
9 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
10 + FusedMoEQuantConfig)
11 from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
10 - FusedMoEConfig,
12 FusedMoEMethodBase)
13 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
14 UnquantizedLinearMethod,
@@ -452,6 +453,10 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
453 **extra_weight_attrs,
454 )
455
456 + def get_fused_moe_quant_config(
457 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
458 + return None
459 +
460 def apply(
461 self,
462 layer: torch.nn.Module,
@@ -509,6 +514,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
514 apply_router_weight_on_input=apply_router_weight_on_input,
515 global_num_experts=global_num_experts,
516 expert_map=expert_map,
517 + quant_config=self.moe_quant_config,
518 )
519
520 def _create_weights_4bit(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+215 -180
@@ -16,8 +16,11 @@ from vllm import _custom_ops as ops
16 from vllm.logger import init_logger
17 from vllm.model_executor.layers.fused_moe import (
18 FusedMoE, FusedMoEActivationFormat, FusedMoEConfig, FusedMoEMethodBase,
19 - FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize,
20 - FusedMoeWeightScaleSupported)
19 + FusedMoEPermuteExpertsUnpermute, FusedMoeWeightScaleSupported)
20 +from vllm.model_executor.layers.fused_moe.config import (
21 + FusedMoEQuantConfig, fp8_w8a8_moe_quant_config,
22 + int4_w4a16_moe_quant_config, int8_w8a8_moe_quant_config,
23 + int8_w8a16_moe_quant_config, nvfp4_moe_quant_config)
24 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
25 is_valid_flashinfer_cutlass_fused_moe)
26 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
@@ -122,7 +125,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
125 return CompressedTensorsWNA16MarlinMoEMethod(
126 quant_config, layer.moe_config)
127 elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
125 - return CompressedTensorsW4A4MoeMethod(layer.moe_config, layer)
128 + return CompressedTensorsW4A4MoeMethod(layer.moe_config)
129 elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
130 or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
131 or quant_config._is_fp8_w8a8(weight_quant, input_quant)):
@@ -138,7 +141,7 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
141
142 class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
143
141 - def __init__(self, moe: FusedMoEConfig, layer: torch.nn.Module):
144 + def __init__(self, moe: FusedMoEConfig):
145 from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
146 detect_nvfp4_moe_support)
147 super().__init__(moe)
@@ -147,7 +150,6 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
150 self.allow_flashinfer = _nvfp4.allow_flashinfer
151 self.use_marlin = _nvfp4.use_marlin
152 self.group_size = 16
150 - self.layer = layer
153
154 def create_weights(self, layer: torch.nn.Module, num_experts: int,
155 hidden_size: int, intermediate_size_per_partition: int,
@@ -305,37 +307,46 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
307 (layer.w2_input_global_scale), requires_grad=False)
308
309 def maybe_make_prepare_finalize(
308 - self,
309 - moe: FusedMoEConfig,
310 - ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
311 - if not self.allow_flashinfer:
312 - return super().maybe_make_prepare_finalize(moe)
310 + self) -> Optional[mk.FusedMoEPrepareAndFinalize]:
311 + if self.use_marlin:
312 + return None
313 + elif not self.allow_flashinfer:
314 + return super().maybe_make_prepare_finalize()
315
316 prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(
315 - moe,
316 - a1_gscale=self.layer.w13_input_scale_quant,
317 - )
317 + self.moe)
318 logger.debug_once("%s", prepare_finalize.__class__.__name__)
319 return prepare_finalize
320
321 def select_gemm_impl(
322 self,
323 prepare_finalize: mk.FusedMoEPrepareAndFinalize,
324 - moe: FusedMoEConfig,
324 layer: torch.nn.Module,
325 ) -> mk.FusedMoEPermuteExpertsUnpermute:
326 + assert self.moe_quant_config is not None
327 """Return the appropriate GEMM experts implementation."""
328 experts = select_nvfp4_gemm_impl(
329 - moe,
330 - g1_alphas=self.layer.g1_alphas,
331 - g2_alphas=self.layer.g2_alphas,
332 - a1_gscale=self.layer.w13_input_scale_quant,
333 - a2_gscale=self.layer.w2_input_scale_quant,
329 + self.moe,
330 + self.moe_quant_config,
331 allow_flashinfer=self.allow_flashinfer,
332 )
333 logger.debug_once("Using %s", experts.__class__.__name__)
334 return experts
335
336 + def get_fused_moe_quant_config(
337 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
338 + if self.use_marlin:
339 + return None
340 +
341 + return nvfp4_moe_quant_config(
342 + g1_alphas=layer.g1_alphas,
343 + g2_alphas=layer.g2_alphas,
344 + a1_gscale=layer.w13_input_scale_quant,
345 + a2_gscale=layer.w2_input_scale_quant,
346 + w1_scale=layer.w13_weight_scale,
347 + w2_scale=layer.w2_weight_scale,
348 + )
349 +
350 def apply(
351 self,
352 layer: torch.nn.Module,
@@ -359,8 +370,6 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
370 logical_to_physical_map: Optional[torch.Tensor] = None,
371 logical_replica_count: Optional[torch.Tensor] = None,
372 ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
362 - assert self.fused_experts is None
363 -
373 if enable_eplb:
374 raise NotImplementedError("EPLB not supported for "
375 "`CompressedTensorsW4A4MoeMethod` yet.")
@@ -381,7 +390,12 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
390 indices_type=self.topk_indices_dtype,
391 )
392
393 + #
394 + # Note: the order here is important. self.fused_experts can override
395 + # flashinfer cutlass, cutlass fp4 or fused_experts but not marlin.
396 + #
397 if self.use_marlin:
398 + assert self.fused_experts is None
399 return torch.ops.vllm.fused_marlin_moe(
400 x,
401 layer.w13_weight,
@@ -401,8 +415,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
415 expert_map=expert_map,
416 workspace=layer.workspace)
417
404 - # FlashInfer fused experts path
405 - if self.fused_experts is not None:
418 + elif self.fused_experts is not None:
419 assert is_valid_flashinfer_cutlass_fused_moe(
420 x, layer.w13_weight, layer.w2_weight), (
421 "Flashinfer CUTLASS Fused MoE not applicable!")
@@ -417,11 +430,10 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
430 activation=activation,
431 global_num_experts=global_num_experts,
432 expert_map=expert_map,
420 - w1_scale=layer.w13_weight_scale,
421 - w2_scale=layer.w2_weight_scale,
433 apply_router_weight_on_input=apply_router_weight_on_input,
434 )
435
436 + # FlashInfer fused experts path
437 elif self.allow_flashinfer:
438 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
439 flashinfer_cutlass_moe_fp4)
@@ -430,51 +442,46 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
442 x, layer.w13_weight, layer.w2_weight), (
443 "Flashinfer CUTLASS Fused MoE not applicable!")
444
445 + assert self.moe_quant_config is not None
446 +
447 return flashinfer_cutlass_moe_fp4(
448 hidden_states=x,
449 w1=layer.w13_weight,
450 w2=layer.w2_weight,
451 topk_weights=topk_weights,
452 topk_ids=topk_ids,
453 + quant_config=self.moe_quant_config,
454 inplace=False, # TODO(shuw): fix later, now output is high prec
455 activation=activation,
456 global_num_experts=global_num_experts,
457 expert_map=expert_map,
443 - w1_scale=layer.w13_weight_scale,
444 - w2_scale=layer.w2_weight_scale,
445 - g1_alphas=layer.g1_alphas,
446 - g2_alphas=layer.g2_alphas,
447 - a1_gscale=layer.w13_input_scale_quant,
448 - a2_gscale=layer.w2_input_scale_quant,
458 apply_router_weight_on_input=apply_router_weight_on_input,
459 )
451 -
452 - assert expert_map is None, ("Expert Parallelism / expert_map "
453 - "is currently not supported for "
454 - "CompressedTensorsW4A4MoeMethod.")
455 - from vllm.model_executor.layers.fused_moe.cutlass_moe import (
456 - cutlass_moe_fp4)
457 -
458 - # Cutlass moe takes in activations in BF16/Half precision
459 - # and fp4 quantized weights loaded from the checkpoint
460 - return cutlass_moe_fp4(
461 - a=x,
462 - w1_fp4=layer.w13_weight,
463 - w2_fp4=layer.w2_weight,
464 - w1_blockscale=layer.w13_weight_scale,
465 - w2_blockscale=layer.w2_weight_scale,
466 - g1_alphas=layer.g1_alphas,
467 - g2_alphas=layer.g2_alphas,
468 - a1_gscale=layer.w13_input_scale_quant,
469 - a2_gscale=layer.w2_input_scale_quant,
470 - topk_weights=topk_weights,
471 - topk_ids=topk_ids,
472 - m=x.shape[0],
473 - n=layer.w2_weight.shape[2] * 2,
474 - k=x.shape[1],
475 - e=layer.w13_weight.shape[0],
476 - apply_router_weight_on_input=apply_router_weight_on_input).to(
477 - x.dtype)
460 + else:
461 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
462 + cutlass_moe_fp4)
463 +
464 + assert expert_map is None, ("Expert Parallelism / expert_map "
465 + "is currently not supported for "
466 + "CompressedTensorsW4A4MoeMethod.")
467 + assert self.moe_quant_config is not None
468 +
469 + # Cutlass moe takes in activations in BF16/Half precision
470 + # and fp4 quantized weights loaded from the checkpoint
471 + return cutlass_moe_fp4(
472 + a=x,
473 + w1_fp4=layer.w13_weight,
474 + w2_fp4=layer.w2_weight,
475 + topk_weights=topk_weights,
476 + topk_ids=topk_ids,
477 + quant_config=self.moe_quant_config,
478 + apply_router_weight_on_input=apply_router_weight_on_input,
479 + # TODO(bnell): derive these from arguments
480 + m=x.shape[0],
481 + n=layer.w2_weight.shape[2] * 2,
482 + k=x.shape[1],
483 + e=layer.w13_weight.shape[0],
484 + ).to(x.dtype)
485
486
487 class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
@@ -692,16 +699,11 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
699 layer.w2_weight = torch.nn.Parameter(shuffled_w2,
700 requires_grad=False)
701
695 - self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts
702 elif self.use_marlin:
703 prepare_moe_fp8_layer_for_marlin(layer, False)
704 # Activations not quantized for marlin.
705 del layer.w13_input_scale
706 del layer.w2_input_scale
701 - self.fused_experts_func = None
702 - else:
703 - from vllm.model_executor.layers.fused_moe import fused_experts
704 - self.fused_experts_func = fused_experts
707
708 if self.use_cutlass:
709 device = layer.w13_weight.device
@@ -722,11 +724,20 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
724 device=device,
725 dtype=torch.int64)
726
727 + def maybe_make_prepare_finalize(
728 + self) -> Optional[mk.FusedMoEPrepareAndFinalize]:
729 + if self.use_marlin or self.rocm_aiter_moe_enabled:
730 + return None
731 + else:
732 + return super().maybe_make_prepare_finalize()
733 +
734 def select_gemm_impl(
726 - self, prepare_finalize: FusedMoEPrepareAndFinalize,
727 - moe: FusedMoEConfig,
728 - layer: torch.nn.Module) -> FusedMoEPermuteExpertsUnpermute:
735 + self,
736 + prepare_finalize: mk.FusedMoEPrepareAndFinalize,
737 + layer: torch.nn.Module,
738 + ) -> FusedMoEPermuteExpertsUnpermute:
739 # cutlass path
740 + assert self.moe_quant_config is not None
741 if self.use_cutlass:
742 from vllm.model_executor.layers.fused_moe import (
743 CutlassBatchedExpertsFp8, CutlassExpertsFp8)
@@ -740,26 +751,24 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
751 logger.debug("CutlassBatchedExpertsFp8(%s)",
752 self.__class__.__name__)
753 experts = CutlassBatchedExpertsFp8(
743 - moe.num_local_experts,
754 + self.moe.num_local_experts,
755 num_dispatchers,
745 - moe.in_dtype,
746 - self.input_quant.strategy == QuantizationStrategy.TOKEN,
747 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
756 + self.moe.in_dtype,
757 ab_strides1=self.ab_strides1_c_strides2,
758 ab_strides2=self.ab_strides2,
759 c_strides1=self.c_strides1,
760 c_strides2=self.ab_strides1_c_strides2,
761 + quant_config=self.moe_quant_config,
762 )
763 else:
764 logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
765 experts = CutlassExpertsFp8(
756 - moe.in_dtype,
757 - self.input_quant.strategy == QuantizationStrategy.TOKEN,
758 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
766 + self.moe.in_dtype,
767 ab_strides1=self.ab_strides1_c_strides2,
768 ab_strides2=self.ab_strides2,
769 c_strides1=self.c_strides1,
770 c_strides2=self.ab_strides1_c_strides2,
771 + quant_config=self.moe_quant_config,
772 )
773
774 self.disable_expert_map = (num_dispatchers > 1
@@ -774,29 +783,40 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
783
784 assert not self.rocm_aiter_moe_enabled and not self.use_marlin
785
777 - logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)
778 -
786 if (prepare_finalize.activation_format ==
787 FusedMoEActivationFormat.BatchedExperts):
788 max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank(
789 )
790 assert max_num_tokens_per_rank is not None
791
792 + logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)
793 return BatchedTritonExperts(
794 max_num_tokens=max_num_tokens_per_rank,
795 num_dispatchers=prepare_finalize.num_dispatchers(),
788 - use_fp8_w8a8=True,
789 - block_shape=self.quant_config.weight_block_size,
790 - per_act_token_quant=(
791 - self.input_quant.strategy == QuantizationStrategy.TOKEN),
796 + quant_config=self.moe_quant_config,
797 )
798 else:
794 - return TritonExperts(
795 - use_fp8_w8a8=True,
796 - block_shape=self.quant_config.weight_block_size,
797 - per_act_token_quant=(
798 - self.input_quant.strategy == QuantizationStrategy.TOKEN),
799 - )
799 + logger.debug("TritonExperts(%s)", self.__class__.__name__)
800 + return TritonExperts(self.moe_quant_config)
801 +
802 + def get_fused_moe_quant_config(
803 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
804 + if self.use_marlin:
805 + return None
806 +
807 + per_act_token = (
808 + self.input_quant.strategy == QuantizationStrategy.TOKEN)
809 + per_channel_quant = (
810 + self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
811 +
812 + return fp8_w8a8_moe_quant_config(
813 + w1_scale=layer.w13_weight_scale,
814 + w2_scale=layer.w2_weight_scale,
815 + a1_scale=layer.w13_input_scale,
816 + a2_scale=layer.w2_input_scale,
817 + per_act_token_quant=per_act_token,
818 + per_out_ch_quant=per_channel_quant,
819 + )
820
821 def apply(
822 self,
@@ -841,16 +861,74 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
861 indices_type=self.topk_indices_dtype,
862 )
863
864 + per_act_token = (
865 + self.input_quant.strategy == QuantizationStrategy.TOKEN)
866 + per_channel_quant = (
867 + self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
868 +
869 + #
870 + # Note: the order here is important. self.fused_experts can override
871 + # cutlass fp8 or fused_experts but not marlin or rocm.
872 + #
873 + if self.use_marlin:
874 + assert activation == "silu", (
875 + f"{activation} not supported for Marlin MoE.")
876 + assert self.fused_experts is None
877 + return torch.ops.vllm.fused_marlin_moe(
878 + x,
879 + layer.w13_weight,
880 + layer.w2_weight,
881 + None,
882 + None,
883 + layer.w13_weight_scale,
884 + layer.w2_weight_scale,
885 + router_logits,
886 + topk_weights,
887 + topk_ids,
888 + quant_type_id=scalar_types.float8_e4m3fn.id,
889 + apply_router_weight_on_input=apply_router_weight_on_input,
890 + global_num_experts=global_num_experts,
891 + expert_map=expert_map,
892 + workspace=layer.workspace)
893 +
894 + elif self.rocm_aiter_moe_enabled:
895 + from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa E501
896 + rocm_aiter_fused_experts)
897 + assert per_act_token == per_channel_quant
898 + assert self.moe_quant_config is not None
899 + assert self.fused_experts is None
900 + return rocm_aiter_fused_experts(
901 + hidden_states=x,
902 + w1=layer.w13_weight,
903 + w2=layer.w2_weight,
904 + topk_weights=topk_weights,
905 + topk_ids=topk_ids,
906 + activation=activation,
907 + apply_router_weight_on_input=apply_router_weight_on_input,
908 + expert_map=expert_map,
909 + quant_config=self.moe_quant_config,
910 + )
911 +
912 + elif self.fused_experts is not None:
913 + return self.fused_experts(
914 + x,
915 + layer.w13_weight,
916 + layer.w2_weight,
917 + topk_weights,
918 + topk_ids,
919 + activation=activation,
920 + global_num_experts=global_num_experts,
921 + expert_map=None if self.disable_expert_map else expert_map,
922 + )
923 +
924 # cutlass path
845 - if self.use_cutlass:
846 - per_act_token = (
847 - self.input_quant.strategy == QuantizationStrategy.TOKEN)
848 - per_channel_quant = (
849 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
925 + elif self.use_cutlass:
926 + assert self.moe_quant_config is not None
927
928 # small-batch fallback on SM100
929 if self.is_fp8_w8a8_sm100 and topk_ids.shape[0] <= 8:
930 from vllm.model_executor.layers.fused_moe import fused_experts
931 + assert per_act_token == per_channel_quant
932 return fused_experts(
933 hidden_states=x,
934 w1=layer.w13_weight,
@@ -860,110 +938,48 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
938 inplace=True,
939 activation=activation,
940 apply_router_weight_on_input=apply_router_weight_on_input,
863 - use_fp8_w8a8=True,
864 - per_channel_quant=per_channel_quant,
941 global_num_experts=global_num_experts,
942 expert_map=None if self.disable_expert_map else expert_map,
867 - w1_scale=layer.w13_weight_scale,
868 - w2_scale=layer.w2_weight_scale,
869 - a1_scale=layer.w13_input_scale,
870 - a2_scale=layer.w2_input_scale)
871 -
872 - if self.fused_experts is None:
943 + quant_config=self.moe_quant_config,
944 + )
945 + else:
946 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
947 cutlass_moe_fp8)
948 + assert per_act_token == per_channel_quant
949 + assert self.moe_quant_config is not None
950 return cutlass_moe_fp8(
951 x,
952 layer.w13_weight,
953 layer.w2_weight,
954 topk_weights,
955 topk_ids,
881 - per_act_token=per_act_token,
956 + quant_config=self.moe_quant_config,
957 activation=activation,
958 global_num_experts=global_num_experts,
959 expert_map=None if self.disable_expert_map else expert_map,
885 - w1_scale=layer.w13_weight_scale,
886 - w2_scale=layer.w2_weight_scale,
960 ab_strides1=self.ab_strides1_c_strides2,
961 ab_strides2=self.ab_strides2,
962 c_strides1=self.c_strides1,
963 c_strides2=self.ab_strides1_c_strides2,
891 - a1_scale=layer.w13_input_scale,
892 - a2_scale=layer.w2_input_scale,
893 - )
894 - else:
895 - return self.fused_experts(
896 - x,
897 - layer.w13_weight,
898 - layer.w2_weight,
899 - topk_weights,
900 - topk_ids,
901 - activation=activation,
902 - global_num_experts=global_num_experts,
903 - expert_map=None if self.disable_expert_map else expert_map,
904 - w1_scale=layer.w13_weight_scale,
905 - w2_scale=layer.w2_weight_scale,
906 - a1_scale=layer.w13_input_scale,
907 - a2_scale=layer.w2_input_scale,
964 )
965
910 - if self.rocm_aiter_moe_enabled:
911 - return self.rocm_aiter_fused_experts_func(
966 + else:
967 + from vllm.model_executor.layers.fused_moe import fused_experts
968 + assert per_act_token == per_channel_quant
969 + assert self.moe_quant_config is not None
970 + return fused_experts(
971 hidden_states=x,
972 w1=layer.w13_weight,
973 w2=layer.w2_weight,
974 topk_weights=topk_weights,
975 topk_ids=topk_ids,
976 + inplace=True,
977 activation=activation,
978 apply_router_weight_on_input=apply_router_weight_on_input,
919 - use_fp8_w8a8=True,
920 - per_channel_quant=self.weight_quant.strategy ==
921 - QuantizationStrategy.CHANNEL,
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 - expert_map=expert_map)
927 - if self.use_marlin:
928 - assert activation == "silu", (
929 - f"{activation} not supported for Marlin MoE.")
930 - return torch.ops.vllm.fused_marlin_moe(
931 - x,
932 - layer.w13_weight,
933 - layer.w2_weight,
934 - None,
935 - None,
936 - layer.w13_weight_scale,
937 - layer.w2_weight_scale,
938 - router_logits,
939 - topk_weights,
940 - topk_ids,
941 - quant_type_id=scalar_types.float8_e4m3fn.id,
942 - apply_router_weight_on_input=apply_router_weight_on_input,
979 global_num_experts=global_num_experts,
980 expert_map=expert_map,
945 - workspace=layer.workspace)
946 -
947 - assert self.fused_experts_func is not None
948 -
949 - return self.fused_experts_func(
950 - hidden_states=x,
951 - w1=layer.w13_weight,
952 - w2=layer.w2_weight,
953 - topk_weights=topk_weights,
954 - topk_ids=topk_ids,
955 - inplace=True,
956 - activation=activation,
957 - apply_router_weight_on_input=apply_router_weight_on_input,
958 - use_fp8_w8a8=True,
959 - per_channel_quant=self.weight_quant.strategy ==
960 - QuantizationStrategy.CHANNEL,
961 - global_num_experts=global_num_experts,
962 - expert_map=expert_map,
963 - w1_scale=layer.w13_weight_scale,
964 - w2_scale=layer.w2_weight_scale,
965 - a1_scale=layer.w13_input_scale,
966 - a2_scale=layer.w2_input_scale)
981 + quant_config=self.moe_quant_config,
982 + )
983
984
985 class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
@@ -1049,6 +1065,16 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1065 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1066 pass
1067
1068 + def get_fused_moe_quant_config(
1069 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
1070 + return int8_w8a8_moe_quant_config(
1071 + w1_scale=layer.w13_weight_scale,
1072 + w2_scale=layer.w2_weight_scale,
1073 + a1_scale=layer.w13_input_scale,
1074 + a2_scale=layer.w2_input_scale,
1075 + per_act_token_quant=True,
1076 + )
1077 +
1078 def apply(
1079 self,
1080 layer: torch.nn.Module,
@@ -1104,14 +1130,10 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1130 inplace=True,
1131 activation=activation,
1132 apply_router_weight_on_input=apply_router_weight_on_input,
1107 - use_int8_w8a8=True,
1108 - per_channel_quant=True,
1133 global_num_experts=global_num_experts,
1134 expert_map=expert_map,
1111 - w1_scale=layer.w13_weight_scale,
1112 - w2_scale=layer.w2_weight_scale,
1113 - a1_scale=layer.w13_input_scale,
1114 - a2_scale=layer.w2_input_scale)
1135 + quant_config=self.moe_quant_config,
1136 + )
1137
1138
1139 class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
@@ -1355,6 +1377,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1377
1378 layer.workspace = marlin_make_workspace_new(device, 4)
1379
1380 + def get_fused_moe_quant_config(
1381 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
1382 + return None
1383 +
1384 def apply(
1385 self,
1386 layer: torch.nn.Module,
@@ -1588,6 +1614,20 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1614 layer.w2_weight_scale.transpose(1, 2).contiguous(),
1615 requires_grad=False)
1616
1617 + def get_fused_moe_quant_config(
1618 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
1619 + assert self.num_bits == 4 or self.num_bits == 8
1620 + config_builder = (int4_w4a16_moe_quant_config if self.num_bits == 4
1621 + else int8_w8a16_moe_quant_config)
1622 +
1623 + return config_builder(
1624 + w1_scale=layer.w13_weight_scale,
1625 + w2_scale=layer.w2_weight_scale,
1626 + w1_zp=None,
1627 + w2_zp=None,
1628 + block_shape=[0, self.group_size],
1629 + )
1630 +
1631 def apply(
1632 self,
1633 layer: torch.nn.Module,
@@ -1641,13 +1681,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1681 topk_ids=topk_ids,
1682 inplace=True,
1683 activation=activation,
1644 - use_int4_w4a16=self.num_bits == 4,
1645 - use_int8_w8a16=self.num_bits == 8,
1646 - global_num_experts=global_num_experts,
1684 apply_router_weight_on_input=apply_router_weight_on_input,
1685 + global_num_experts=global_num_experts,
1686 expert_map=expert_map,
1649 - w1_scale=layer.w13_weight_scale,
1650 - w2_scale=layer.w2_weight_scale,
1651 - w1_zp=None,
1652 - w2_zp=None,
1653 - block_shape=[0, self.group_size])
1687 + quant_config=self.moe_quant_config,
1688 + )
vllm/model_executor/layers/quantization/experts_int8.py
+12 -4
@@ -8,6 +8,8 @@ import torch
8 from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
10 FusedMoEMethodBase)
11 +from vllm.model_executor.layers.fused_moe.config import (
12 + FusedMoEQuantConfig, int8_w8a16_moe_quant_config)
13 from vllm.model_executor.layers.linear import (LinearBase,
14 UnquantizedLinearMethod)
15 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -106,6 +108,13 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
108 requires_grad=False)
109 layer.register_parameter("w2_scale", w2_scale)
110
111 + def get_fused_moe_quant_config(
112 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
113 + return int8_w8a16_moe_quant_config(w1_scale=layer.w13_scale,
114 + w2_scale=layer.w2_scale,
115 + w1_zp=None,
116 + w2_zp=None)
117 +
118 def apply(
119 self,
120 layer: torch.nn.Module,
@@ -159,12 +168,11 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
168 topk_ids=topk_ids,
169 inplace=True,
170 activation=activation,
162 - use_int8_w8a16=True,
163 - global_num_experts=global_num_experts,
171 apply_router_weight_on_input=apply_router_weight_on_input,
172 + global_num_experts=global_num_experts,
173 expert_map=expert_map,
166 - w1_scale=layer.w13_scale,
167 - w2_scale=layer.w2_scale)
174 + quant_config=self.moe_quant_config,
175 + )
176
177 @staticmethod
178 def quantizing_weight_loader(layer, weight_loader):
vllm/model_executor/layers/quantization/fp8.py
+80 -79
@@ -14,9 +14,11 @@ from vllm import _custom_ops as ops
14 from vllm.distributed import get_tensor_model_parallel_world_size
15 from vllm.logger import init_logger
16 from vllm.model_executor.layers.fused_moe import (
17 - FusedMoE, FusedMoEActivationFormat, FusedMoEConfig, FusedMoEMethodBase,
17 + FusedMoE, FusedMoEActivationFormat, FusedMoEMethodBase,
18 FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize,
19 FusedMoeWeightScaleSupported)
20 +from vllm.model_executor.layers.fused_moe.config import (
21 + FusedMoEQuantConfig, fp8_w8a8_moe_quant_config)
22 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
23 UnquantizedLinearMethod)
24 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -575,20 +577,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
577 "CutlassBlockScaledGroupedGemm not supported on the current "
578 "platform.")
579
578 - def maybe_make_prepare_finalize(
579 - self,
580 - moe: FusedMoEConfig,
581 - ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
582 - if self.flashinfer_moe_backend != FlashinferMoeBackend.CUTLASS:
583 - return super().maybe_make_prepare_finalize(moe)
584 -
585 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
586 - moe,
587 - layer=self.layer,
588 - )
589 - logger.debug_once("%s", prepare_finalize.__class__.__name__)
590 - return prepare_finalize
591 -
580 def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
581 intermediate_size_per_partition: int,
582 params_dtype: torch.dtype, **extra_weight_attrs):
@@ -928,10 +916,23 @@ class Fp8MoEMethod(FusedMoEMethodBase):
916 layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(
917 layer.w2_weight_scale_inv)
918
919 + def maybe_make_prepare_finalize(
920 + self) -> Optional[mk.FusedMoEPrepareAndFinalize]:
921 + if (self.rocm_aiter_moe_enabled or self.use_marlin
922 + or self.flashinfer_moe_backend
923 + == FlashinferMoeBackend.TENSORRT_LLM):
924 + return None
925 + elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
926 + prepare_finalize = (
927 + build_flashinfer_fp8_cutlass_moe_prepare_finalize(self.moe))
928 + logger.debug_once("%s", prepare_finalize.__class__.__name__)
929 + return prepare_finalize
930 + else:
931 + return super().maybe_make_prepare_finalize()
932 +
933 def select_gemm_impl(
934 self,
935 prepare_finalize: FusedMoEPrepareAndFinalize,
934 - moe: FusedMoEConfig,
936 layer: torch.nn.Module,
937 ) -> FusedMoEPermuteExpertsUnpermute:
938 from vllm.model_executor.layers.fused_moe import (
@@ -940,6 +941,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
941 assert not self.use_marlin and not self.rocm_aiter_moe_enabled, (
942 "Marlin and ROCm AITER are not supported with all2all yet.")
943
944 + assert self.moe_quant_config is not None
945 +
946 if (prepare_finalize.activation_format ==
947 FusedMoEActivationFormat.BatchedExperts):
948 max_num_tokens_per_rank = (
@@ -953,15 +956,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
956 return BatchedTritonOrDeepGemmExperts(
957 max_num_tokens=max_num_tokens_per_rank,
958 num_dispatchers=prepare_finalize.num_dispatchers(),
956 - use_fp8_w8a8=True,
957 - block_shape=self.quant_config.weight_block_size,
958 - per_act_token_quant=False,
959 + quant_config=self.moe_quant_config,
960 allow_deep_gemm=self.allow_deep_gemm,
961 )
962 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
963 experts = select_cutlass_fp8_gemm_impl(
963 - moe,
964 - self.layer,
964 + self.moe,
965 + self.moe_quant_config,
966 )
967 logger.debug_once("Using %s", experts.__class__.__name__)
968 return experts
@@ -971,11 +972,25 @@ class Fp8MoEMethod(FusedMoEMethodBase):
972 self.__class__.__name__, self.quant_config.weight_block_size,
973 False)
974 return TritonOrDeepGemmExperts(
974 - use_fp8_w8a8=True,
975 - block_shape=self.quant_config.weight_block_size,
975 + quant_config=self.moe_quant_config,
976 allow_deep_gemm=self.allow_deep_gemm,
977 )
978
979 + def get_fused_moe_quant_config(
980 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
981 + if self.use_marlin:
982 + return None
983 +
984 + return fp8_w8a8_moe_quant_config(
985 + w1_scale=(layer.w13_weight_scale_inv
986 + if self.block_quant else layer.w13_weight_scale),
987 + w2_scale=(layer.w2_weight_scale_inv
988 + if self.block_quant else layer.w2_weight_scale),
989 + a1_scale=layer.w13_input_scale,
990 + a2_scale=layer.w2_input_scale,
991 + block_shape=self.quant_config.weight_block_size,
992 + )
993 +
994 def apply(
995 self,
996 layer: torch.nn.Module,
@@ -1005,12 +1020,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1020 assert logical_replica_count is not None
1021 assert isinstance(layer, FusedMoE)
1022
1008 - if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1023 + if (self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM
1024 + and self.fused_experts is None):
1025 assert activation == 'silu', (
1026 f"Expected 'silu' activation but got {activation}")
1027 assert scoring_func == 'sigmoid', (
1028 f"Expected 'sigmoid' scoring func but got {scoring_func}")
1029 if self.block_quant:
1030 + import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
1031 assert (renormalize and use_grouped_topk
1032 and custom_routing_function is None)
1033
@@ -1066,9 +1083,14 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1083 logical_replica_count=logical_replica_count,
1084 )
1085
1086 + #
1087 + # Note: the order of checks is important since self.fused_experts
1088 + # can override fused_experts or cutlass but not rocm or marlin.
1089 + #
1090 if self.rocm_aiter_moe_enabled:
1091 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
1092 rocm_aiter_fused_experts)
1093 + assert self.fused_experts is None
1094 return rocm_aiter_fused_experts(
1095 x,
1096 layer.w13_weight,
@@ -1076,19 +1098,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1098 topk_weights=topk_weights,
1099 topk_ids=topk_ids,
1100 activation=activation,
1079 - use_fp8_w8a8=True,
1101 apply_router_weight_on_input=apply_router_weight_on_input,
1081 - w1_scale=(layer.w13_weight_scale_inv
1082 - if self.block_quant else layer.w13_weight_scale),
1083 - w2_scale=(layer.w2_weight_scale_inv
1084 - if self.block_quant else layer.w2_weight_scale),
1085 - a1_scale=layer.w13_input_scale,
1086 - a2_scale=layer.w2_input_scale,
1087 - block_shape=self.quant_config.weight_block_size,
1088 - expert_map=expert_map)
1102 + expert_map=expert_map,
1103 + quant_config=self.moe_quant_config)
1104 elif self.use_marlin:
1105 assert activation == "silu", (
1106 f"{activation} not supported for Marlin MoE.")
1107 + assert self.fused_experts is None
1108 return torch.ops.vllm.fused_marlin_moe(
1109 x,
1110 layer.w13_weight,
@@ -1105,6 +1121,19 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1121 global_num_experts=global_num_experts,
1122 expert_map=expert_map,
1123 workspace=layer.workspace)
1124 + elif self.fused_experts:
1125 + return self.fused_experts(
1126 + hidden_states=x,
1127 + w1=layer.w13_weight,
1128 + w2=layer.w2_weight,
1129 + topk_weights=topk_weights,
1130 + topk_ids=topk_ids,
1131 + inplace=True,
1132 + activation=activation,
1133 + global_num_experts=global_num_experts,
1134 + apply_router_weight_on_input=apply_router_weight_on_input,
1135 + expert_map=expert_map,
1136 + )
1137 elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
1138 assert self.block_quant is None
1139 assert (not renormalize and custom_routing_function is not None)
@@ -1112,33 +1141,21 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1141 f"Expected 'silu' activation but got {activation}")
1142 assert scoring_func == 'sigmoid', (
1143 f"Expected 'sigmoid' scoring func but got {scoring_func}")
1115 - if self.fused_experts is not None:
1116 - return self.fused_experts(
1117 - x,
1118 - layer.w13_weight,
1119 - layer.w2_weight,
1120 - topk_weights,
1121 - topk_ids,
1122 - inplace=False,
1123 - activation=activation,
1124 - global_num_experts=global_num_experts,
1125 - expert_map=expert_map,
1126 - apply_router_weight_on_input=apply_router_weight_on_input,
1127 - )
1128 - else:
1129 - return flashinfer_cutlass_moe_fp8(
1130 - x,
1131 - layer,
1132 - topk_weights,
1133 - topk_ids,
1134 - inplace=False,
1135 - activation=activation,
1136 - global_num_experts=global_num_experts,
1137 - expert_map=expert_map,
1138 - apply_router_weight_on_input=apply_router_weight_on_input,
1139 - )
1144 +
1145 + return flashinfer_cutlass_moe_fp8(
1146 + x,
1147 + layer,
1148 + topk_weights,
1149 + topk_ids,
1150 + inplace=False,
1151 + activation=activation,
1152 + global_num_experts=global_num_experts,
1153 + expert_map=expert_map,
1154 + apply_router_weight_on_input=apply_router_weight_on_input,
1155 + )
1156 else:
1141 - common_kwargs = dict(
1157 + from vllm.model_executor.layers.fused_moe import fused_experts
1158 + return fused_experts(
1159 hidden_states=x,
1160 w1=layer.w13_weight,
1161 w2=layer.w2_weight,
@@ -1149,26 +1166,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1166 global_num_experts=global_num_experts,
1167 apply_router_weight_on_input=apply_router_weight_on_input,
1168 expert_map=expert_map,
1152 - w1_scale=(layer.w13_weight_scale_inv
1153 - if self.block_quant else layer.w13_weight_scale),
1154 - w2_scale=(layer.w2_weight_scale_inv
1155 - if self.block_quant else layer.w2_weight_scale),
1156 - a1_scale=layer.w13_input_scale,
1157 - a2_scale=layer.w2_input_scale,
1158 - )
1159 -
1160 - if self.fused_experts is not None:
1161 - return self.fused_experts(**common_kwargs)
1162 - else:
1163 - from vllm.model_executor.layers.fused_moe import fused_experts
1164 - return fused_experts(
1165 - **common_kwargs,
1166 - use_fp8_w8a8=True,
1167 - block_shape=self.quant_config.weight_block_size,
1168 - allow_deep_gemm=self.allow_deep_gemm,
1169 - allow_cutlass_block_scaled_grouped_gemm=(
1170 - self.allow_cutlass_block_scaled_grouped_gemm),
1171 - )
1169 + quant_config=self.moe_quant_config,
1170 + allow_deep_gemm=self.allow_deep_gemm,
1171 + allow_cutlass_block_scaled_grouped_gemm=(
1172 + self.allow_cutlass_block_scaled_grouped_gemm))
1173
1174
1175 class Fp8KVCacheMethod(BaseKVCacheMethod):
vllm/model_executor/layers/quantization/gguf.py
+6 -1
@@ -10,8 +10,9 @@ from torch.nn.parameter import Parameter, UninitializedParameter
10
11 from vllm import _custom_ops as ops
12 from vllm.logger import init_logger
13 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
14 + FusedMoEQuantConfig)
15 from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
14 - FusedMoEConfig,
16 FusedMoEMethodBase)
17 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
18 UnquantizedLinearMethod)
@@ -518,6 +519,10 @@ class GGUFMoEMethod(FusedMoEMethodBase):
519 set_weight_attrs(w2_qweight_type, extra_weight_attrs)
520 layer.register_parameter("w2_qweight_type", w2_qweight_type)
521
522 + def get_fused_moe_quant_config(
523 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
524 + return None
525 +
526 def apply(
527 self,
528 layer: torch.nn.Module,
vllm/model_executor/layers/quantization/gptq_marlin.py
+7 -1
@@ -9,8 +9,10 @@ import torch
9 import vllm.model_executor.layers.fused_moe # noqa
10 from vllm import _custom_ops as ops
11 from vllm.logger import init_logger
12 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
13 + FusedMoEQuantConfig)
14 from vllm.model_executor.layers.fused_moe.layer import (
13 - FusedMoE, FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
15 + FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
16 UnquantizedFusedMoEMethod)
17 from vllm.model_executor.layers.linear import (LinearMethodBase,
18 set_weight_attrs)
@@ -632,6 +634,10 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
634 if hasattr(layer, "w2_bias") and layer.w2_bias is not None:
635 layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)
636
637 + def get_fused_moe_quant_config(
638 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
639 + return None
640 +
641 def apply(
642 self,
643 layer: torch.nn.Module,
vllm/model_executor/layers/quantization/ipex_quant.py
+5
@@ -11,6 +11,7 @@ from torch.nn.parameter import Parameter
11 from vllm._ipex_ops import ipex_ops as ops
12 from vllm.model_executor.layers.fused_moe import (FusedMoEMethodBase,
13 FusedMoeWeightScaleSupported)
14 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
15 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
16 UnquantizedLinearMethod)
17 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -375,6 +376,10 @@ class XPUFp8MoEMethod(FusedMoEMethodBase):
376 use_prepack=True,
377 )
378
379 + def get_fused_moe_quant_config(
380 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
381 + return None
382 +
383 def apply(
384 self,
385 layer: torch.nn.Module,
vllm/model_executor/layers/quantization/modelopt.py
+130 -107
@@ -11,7 +11,9 @@ import vllm.envs as envs
11 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
13 from vllm.logger import init_logger
14 -from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + FusedMoEConfig, FusedMoEQuantConfig, fp8_w8a8_moe_quant_config,
16 + nvfp4_moe_quant_config)
17 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
18 is_valid_flashinfer_cutlass_fused_moe)
19 from vllm.model_executor.layers.fused_moe.layer import (
@@ -294,8 +296,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
296 cutlass_fp8_supported)
297 self.cutlass_fp8_supported = cutlass_fp8_supported()
298 self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None
297 - self.fused_experts: Optional[
298 - mk.FusedMoEModularKernel] = None # type: ignore
299 if envs.VLLM_USE_FLASHINFER_MOE_FP8 and has_flashinfer_moe():
300 self.flashinfer_moe_backend = get_flashinfer_moe_backend()
301 logger.info_once(
@@ -303,29 +303,27 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
303 )
304
305 def maybe_make_prepare_finalize(
306 - self,
307 - moe: FusedMoEConfig,
308 - ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
309 - if self.fused_experts is not None or \
310 - self.flashinfer_moe_backend != FlashinferMoeBackend.CUTLASS:
311 - return super().maybe_make_prepare_finalize(moe)
312 -
313 - prepare_finalize = build_flashinfer_fp8_cutlass_moe_prepare_finalize(
314 - moe,
315 - layer=self.layer,
316 - )
317 - logger.debug_once("%s", prepare_finalize.__class__.__name__)
318 - return prepare_finalize
306 + self, ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
307 + # TRT LLM not supported with all2all yet.
308 + if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
309 + return None
310 + elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
311 + prepare_finalize = (
312 + build_flashinfer_fp8_cutlass_moe_prepare_finalize(self.moe))
313 + logger.debug_once("%s", prepare_finalize.__class__.__name__)
314 + return prepare_finalize
315 + else:
316 + return super().maybe_make_prepare_finalize()
317
318 def select_gemm_impl(
319 self,
320 prepare_finalize: mk.FusedMoEPrepareAndFinalize,
323 - moe: FusedMoEConfig,
321 layer: torch.nn.Module,
322 ) -> mk.FusedMoEPermuteExpertsUnpermute:
323 + assert self.moe_quant_config is not None
324 experts = select_cutlass_fp8_gemm_impl(
327 - moe,
328 - self.layer,
325 + self.moe,
326 + self.moe_quant_config,
327 )
328 logger.debug_once("Using %s", experts.__class__.__name__)
329 return experts
@@ -479,6 +477,19 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
477 rotate_flashinfer_fp8_moe_weights(layer.w13_weight,
478 layer.w2_weight)
479
480 + def get_fused_moe_quant_config(
481 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
482 + if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
483 + return None
484 +
485 + return fp8_w8a8_moe_quant_config(
486 + w1_scale=layer.w13_weight_scale,
487 + w2_scale=layer.w2_weight_scale,
488 + a1_scale=layer.w13_input_scale,
489 + a2_scale=layer.w2_input_scale,
490 + per_act_token_quant=False,
491 + )
492 +
493 def apply(
494 self,
495 layer: torch.nn.Module,
@@ -507,6 +518,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
518 "EPLB not supported for `ModelOptFp8MoEMethod` yet.")
519
520 if self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
521 + assert self.fused_experts is None
522 assert activation == 'silu', (
523 f"Expected 'silu' activation but got {activation}")
524 assert not renormalize
@@ -537,55 +549,56 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
549 indices_type=self.topk_indices_dtype,
550 )
551
540 - if self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
552 + #
553 + # Note: the order here is important. self.fused_experts can override
554 + # cutlass or fused_experts.
555 + #
556 + if self.fused_experts is not None:
557 + return self.fused_experts(
558 + x,
559 + layer.w13_weight,
560 + layer.w2_weight,
561 + topk_weights,
562 + topk_ids,
563 + inplace=False,
564 + activation=activation,
565 + global_num_experts=global_num_experts,
566 + expert_map=expert_map,
567 + apply_router_weight_on_input=apply_router_weight_on_input,
568 + )
569 + elif self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS:
570 assert not renormalize
571 assert activation == 'silu', (
572 f"Expected 'silu' activation but got {activation}")
544 - if self.fused_experts is not None:
545 - return self.fused_experts(
546 - x,
547 - layer.w13_weight,
548 - layer.w2_weight,
549 - topk_weights,
550 - topk_ids,
551 - inplace=False,
552 - activation=activation,
553 - global_num_experts=global_num_experts,
554 - expert_map=expert_map,
555 - apply_router_weight_on_input=apply_router_weight_on_input,
556 - )
557 - else:
558 - return flashinfer_cutlass_moe_fp8(
559 - x,
560 - layer,
561 - topk_weights,
562 - topk_ids,
563 - inplace=False,
564 - activation=activation,
565 - global_num_experts=global_num_experts,
566 - expert_map=expert_map,
567 - apply_router_weight_on_input=apply_router_weight_on_input,
568 - )
569 - from vllm.model_executor.layers.fused_moe.fused_moe import (
570 - fused_experts)
571 - return fused_experts(
572 - x,
573 - layer.w13_weight,
574 - layer.w2_weight,
575 - topk_weights=topk_weights,
576 - topk_ids=topk_ids,
577 - inplace=True,
578 - activation=activation,
579 - use_fp8_w8a8=True,
580 - per_channel_quant=False,
581 - global_num_experts=global_num_experts,
582 - expert_map=expert_map,
583 - w1_scale=layer.w13_weight_scale,
584 - w2_scale=layer.w2_weight_scale,
585 - a1_scale=layer.w13_input_scale,
586 - a2_scale=layer.w2_input_scale,
587 - apply_router_weight_on_input=apply_router_weight_on_input,
588 - )
573 + return flashinfer_cutlass_moe_fp8(
574 + x,
575 + layer,
576 + topk_weights,
577 + topk_ids,
578 + inplace=False,
579 + activation=activation,
580 + global_num_experts=global_num_experts,
581 + expert_map=expert_map,
582 + apply_router_weight_on_input=apply_router_weight_on_input,
583 + )
584 + else:
585 + from vllm.model_executor.layers.fused_moe.fused_moe import (
586 + fused_experts)
587 + assert self.moe_quant_config is not None
588 +
589 + return fused_experts(
590 + x,
591 + layer.w13_weight,
592 + layer.w2_weight,
593 + topk_weights=topk_weights,
594 + topk_ids=topk_ids,
595 + inplace=True,
596 + activation=activation,
597 + quant_config=self.moe_quant_config,
598 + global_num_experts=global_num_experts,
599 + expert_map=expert_map,
600 + apply_router_weight_on_input=apply_router_weight_on_input,
601 + )
602
603
604 class ModelOptNvFp4Config(QuantizationConfig):
@@ -1034,33 +1047,30 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1047 " for ModelOptNvFp4FusedMoE.")
1048
1049 def maybe_make_prepare_finalize(
1037 - self,
1038 - moe: FusedMoEConfig,
1039 - ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
1040 - if (self.allow_flashinfer and self.flashinfer_moe_backend
1041 - == FlashinferMoeBackend.CUTLASS):
1050 + self) -> Optional[mk.FusedMoEPrepareAndFinalize]:
1051 + if (self.use_marlin
1052 + or (self.allow_flashinfer and self.flashinfer_moe_backend
1053 + == FlashinferMoeBackend.TENSORRT_LLM)):
1054 + return None
1055 + elif (self.allow_flashinfer
1056 + and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS):
1057 + # For now, fp4 moe only works with the flashinfer dispatcher.
1058 prepare_finalize = (
1043 - build_flashinfer_fp4_cutlass_moe_prepare_finalize(
1044 - moe,
1045 - a1_gscale=self.layer.w13_input_scale_quant,
1046 - ))
1059 + build_flashinfer_fp4_cutlass_moe_prepare_finalize(self.moe))
1060 logger.debug_once("%s", prepare_finalize.__class__.__name__)
1061 return prepare_finalize
1049 -
1050 - return super().maybe_make_prepare_finalize(moe)
1062 + else:
1063 + return super().maybe_make_prepare_finalize()
1064
1065 def select_gemm_impl(
1066 self,
1067 prepare_finalize: mk.FusedMoEPrepareAndFinalize,
1055 - moe: FusedMoEConfig,
1068 layer: torch.nn.Module,
1069 ) -> mk.FusedMoEPermuteExpertsUnpermute:
1070 + assert self.moe_quant_config is not None
1071 experts = select_nvfp4_gemm_impl(
1059 - moe,
1060 - g1_alphas=self.layer.g1_alphas,
1061 - g2_alphas=self.layer.g2_alphas,
1062 - a1_gscale=self.layer.w13_input_scale_quant,
1063 - a2_gscale=self.layer.w2_input_scale_quant,
1072 + self.moe,
1073 + self.moe_quant_config,
1074 allow_flashinfer=self.allow_flashinfer,
1075 )
1076 logger.debug_once("Using %s", experts.__class__.__name__)
@@ -1360,6 +1370,21 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1370 layer.w2_weight = Parameter(layer.w2_weight.data,
1371 requires_grad=False)
1372
1373 + def get_fused_moe_quant_config(
1374 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
1375 + if (self.use_marlin or self.flashinfer_moe_backend
1376 + == FlashinferMoeBackend.TENSORRT_LLM):
1377 + return None
1378 +
1379 + return nvfp4_moe_quant_config(
1380 + w1_scale=layer.w13_weight_scale,
1381 + w2_scale=layer.w2_weight_scale,
1382 + g1_alphas=layer.g1_alphas,
1383 + g2_alphas=layer.g2_alphas,
1384 + a1_gscale=layer.w13_input_scale_quant,
1385 + a2_gscale=layer.w2_input_scale_quant,
1386 + )
1387 +
1388 def apply(
1389 self,
1390 layer: torch.nn.Module,
@@ -1388,12 +1413,14 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1413 "EPLB not supported for `ModelOptNvFp4FusedMoE` yet.")
1414 assert activation == "silu", "Only SiLU activation is supported."
1415
1391 - if self.allow_flashinfer and \
1392 - self.flashinfer_moe_backend == FlashinferMoeBackend.TENSORRT_LLM:
1416 + if (self.allow_flashinfer and self.flashinfer_moe_backend
1417 + == FlashinferMoeBackend.TENSORRT_LLM):
1418 import flashinfer
1419
1420 from vllm.model_executor.models.llama4 import Llama4MoE
1421
1422 + assert self.fused_experts is None
1423 +
1424 a1_gscale = layer.w13_input_scale_quant
1425 (hidden_states_fp4,
1426 hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(
@@ -1457,7 +1484,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1484 e_score_correction_bias=e_score_correction_bias,
1485 indices_type=self.topk_indices_dtype)
1486
1487 + #
1488 + # Note: the order here is important. self.fused_experts can override
1489 + # flashinfer cutlass, cutlass fp4 or fused_experts but not marlin or
1490 + # trtllm.
1491 + #
1492 if self.use_marlin:
1493 + assert self.fused_experts is None
1494 return torch.ops.vllm.fused_marlin_moe(
1495 x,
1496 layer.w13_weight,
@@ -1477,7 +1510,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1510 expert_map=expert_map,
1511 workspace=layer.workspace)
1512
1480 - if self.fused_experts is not None:
1513 + elif self.fused_experts is not None:
1514 assert self.allow_flashinfer and \
1515 self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1516
@@ -1485,7 +1518,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1518 x, layer.w13_weight, layer.w2_weight), (
1519 "Flashinfer CUTLASS Fused MoE not applicable!")
1520
1488 - out = self.fused_experts(
1521 + return self.fused_experts(
1522 hidden_states=x,
1523 w1=layer.w13_weight,
1524 w2=layer.w2_weight,
@@ -1495,28 +1528,22 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1528 activation=activation,
1529 global_num_experts=global_num_experts,
1530 expert_map=expert_map,
1498 - w1_scale=layer.w13_weight_scale,
1499 - w2_scale=layer.w2_weight_scale,
1531 apply_router_weight_on_input=apply_router_weight_on_input,
1532 )
1533 elif (self.allow_flashinfer
1534 and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS):
1535 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1536 flashinfer_cutlass_moe_fp4)
1537 + assert self.moe_quant_config is not None
1538
1507 - out = flashinfer_cutlass_moe_fp4(
1539 + return flashinfer_cutlass_moe_fp4(
1540 hidden_states=x,
1541 w1=layer.w13_weight,
1542 w2=layer.w2_weight,
1543 topk_weights=topk_weights,
1544 topk_ids=topk_ids,
1513 - w1_scale=layer.w13_weight_scale,
1514 - w2_scale=layer.w2_weight_scale,
1515 - g1_alphas=layer.g1_alphas,
1516 - g2_alphas=layer.g2_alphas,
1517 - a1_gscale=layer.w13_input_scale_quant,
1518 - a2_gscale=layer.w2_input_scale_quant,
1519 - inplace=False, # TODO(shuw): fix later, now output is high prec
1545 + quant_config=self.moe_quant_config,
1546 + inplace=False,
1547 activation=activation,
1548 global_num_experts=global_num_experts,
1549 expert_map=expert_map,
@@ -1527,23 +1554,19 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1554 # only (no EP).
1555 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
1556 cutlass_moe_fp4)
1530 - out = cutlass_moe_fp4(
1557 + assert self.moe_quant_config is not None
1558 + return cutlass_moe_fp4(
1559 a=x,
1560 w1_fp4=layer.w13_weight,
1561 w2_fp4=layer.w2_weight,
1534 - w1_blockscale=layer.w13_weight_scale,
1535 - w2_blockscale=layer.w2_weight_scale,
1536 - g1_alphas=layer.g1_alphas,
1537 - g2_alphas=layer.g2_alphas,
1538 - a1_gscale=layer.w13_input_scale_quant,
1539 - a2_gscale=layer.w2_input_scale_quant,
1562 topk_weights=topk_weights,
1563 topk_ids=topk_ids,
1564 + quant_config=self.moe_quant_config,
1565 + expert_map=expert_map,
1566 + apply_router_weight_on_input=apply_router_weight_on_input,
1567 + # TODO: derive from arguments
1568 m=x.shape[0],
1569 n=layer.w2_weight.shape[2] * 2,
1570 k=x.shape[1],
1571 e=layer.w13_weight.shape[0],
1546 - expert_map=expert_map,
1547 - apply_router_weight_on_input=apply_router_weight_on_input)
1548 -
1549 - return out
1572 + )
vllm/model_executor/layers/quantization/moe_wna16.py
+22 -11
@@ -6,6 +6,9 @@ from typing import Any, Callable, Optional, Union
6 import torch
7
8 from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9 +from vllm.model_executor.layers.fused_moe.config import (
10 + FusedMoEQuantConfig, int4_w4a16_moe_quant_config,
11 + int8_w8a16_moe_quant_config)
12 from vllm.model_executor.layers.fused_moe.layer import (
13 FusedMoE, FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
14 from vllm.model_executor.layers.linear import (LinearBase,
@@ -283,6 +286,22 @@ class MoeWNA16Method(FusedMoEMethodBase):
286 layer.register_parameter(key, param)
287 set_weight_attrs(param, extra_weight_attrs)
288
289 + def get_fused_moe_quant_config(
290 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
291 + weight_bits = self.quant_config.weight_bits
292 + has_zp = self.quant_config.has_zp
293 + assert weight_bits == 4 or weight_bits == 8
294 + config_builder = (int4_w4a16_moe_quant_config
295 + if weight_bits == 4 else int8_w8a16_moe_quant_config)
296 +
297 + return config_builder(
298 + w1_scale=layer.w13_scales,
299 + w2_scale=layer.w2_scales,
300 + w1_zp=layer.w13_qzeros if has_zp else None,
301 + w2_zp=layer.w2_qzeros if has_zp else None,
302 + block_shape=[0, layer.group_size],
303 + )
304 +
305 def apply(
306 self,
307 layer: torch.nn.Module,
@@ -327,9 +346,6 @@ class MoeWNA16Method(FusedMoEMethodBase):
346 e_score_correction_bias=e_score_correction_bias,
347 indices_type=self.topk_indices_dtype)
348
330 - weight_bits = self.quant_config.weight_bits
331 - has_zp = self.quant_config.has_zp
332 -
349 return fused_experts(
350 x,
351 layer.w13_qweight,
@@ -337,16 +353,11 @@ class MoeWNA16Method(FusedMoEMethodBase):
353 topk_weights=topk_weights,
354 topk_ids=topk_ids,
355 inplace=True,
340 - use_int4_w4a16=weight_bits == 4,
341 - use_int8_w8a16=weight_bits == 8,
342 - global_num_experts=global_num_experts,
356 apply_router_weight_on_input=apply_router_weight_on_input,
357 + global_num_experts=global_num_experts,
358 expert_map=expert_map,
345 - w1_scale=layer.w13_scales,
346 - w2_scale=layer.w2_scales,
347 - w1_zp=layer.w13_qzeros if has_zp else None,
348 - w2_zp=layer.w2_qzeros if has_zp else None,
349 - block_shape=[0, layer.group_size])
359 + quant_config=self.moe_quant_config,
360 + )
361
362 @staticmethod
363 def get_weight_loader(layer, weight_loader):
vllm/model_executor/layers/quantization/mxfp4.py
+27 -10
@@ -12,6 +12,8 @@ from vllm.logger import init_logger
12 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
13 FusedMoEMethodBase)
14 from vllm.model_executor.layers.fused_moe import modular_kernel as mk
15 +from vllm.model_executor.layers.fused_moe.config import (
16 + FusedMoEQuantConfig, mxfp4_w4a4_moe_quant_config)
17 from vllm.model_executor.layers.fused_moe.trtllm_moe import TrtLlmGenExperts
18 from vllm.model_executor.layers.linear import (LinearBase,
19 UnquantizedLinearMethod)
@@ -629,10 +631,29 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
631
632 return tile_tokens_dim
633
634 + def get_fused_moe_quant_config(
635 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
636 +
637 + if self.mxfp4_backend == Mxfp4Backend.MARLIN:
638 + return None
639 +
640 + if self.mxfp4_backend == Mxfp4Backend.TRITON:
641 + w1_scale = layer.w13_precision_config
642 + w2_scale = layer.w2_precision_config
643 + else:
644 + w1_scale = layer.w13_weight_scale
645 + w2_scale = layer.w2_weight_scale
646 +
647 + return mxfp4_w4a4_moe_quant_config(
648 + w1_bias=layer.w13_bias,
649 + w2_bias=layer.w2_bias,
650 + w1_scale=w1_scale,
651 + w2_scale=w2_scale,
652 + )
653 +
654 def select_gemm_impl(
655 self,
656 prepare_finalize: mk.FusedMoEPrepareAndFinalize,
635 - moe: FusedMoEConfig,
657 layer: torch.nn.Module,
658 ) -> mk.FusedMoEPermuteExpertsUnpermute:
659 if (prepare_finalize.activation_format ==
@@ -647,11 +668,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
668 "gemm1_alpha": layer.gemm1_alpha,
669 "gemm1_beta": layer.gemm1_beta,
670 "gemm1_clamp_limit": layer.gemm1_clamp_limit,
650 - "w13_bias": layer.w13_bias,
651 - "w2_bias": layer.w2_bias,
671 + # TODO(bnell): part of quant_config
672 "max_capture_size": self.max_capture_size,
673 }
654 - return TrtLlmGenExperts(moe, **kwargs)
674 + assert self.moe_quant_config is not None
675 + return TrtLlmGenExperts(self.moe, self.moe_quant_config,
676 + **kwargs)
677 else:
678 # Use matmul_ogs from triton_kernels here!
679 raise NotImplementedError(
@@ -710,8 +732,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
732 activation=activation,
733 global_num_experts=global_num_experts,
734 expert_map=expert_map,
713 - w1_scale=layer.w13_weight_scale,
714 - w2_scale=layer.w2_weight_scale,
735 apply_router_weight_on_input=apply_router_weight_on_input,
736 )
737
@@ -941,10 +961,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
961 renormalize=renormalize,
962 global_num_experts=global_num_experts,
963 expert_map=expert_map,
944 - w1_bias=layer.w13_bias,
945 - w2_bias=layer.w2_bias,
946 - w1_precision=self.w13_precision_config,
947 - w2_precision=self.w2_precision_config,
964 + quant_config=self.moe_quant_config,
965 apply_router_weight_on_input=apply_router_weight_on_input,
966 )
967 else:
vllm/model_executor/layers/quantization/quark/quark_moe.py
+27 -19
@@ -11,6 +11,9 @@ from vllm.logger import init_logger
11 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
12 FusedMoEMethodBase,
13 FusedMoeWeightScaleSupported)
14 +from vllm.model_executor.layers.fused_moe.config import (
15 + FusedMoEQuantConfig, fp8_w8a8_moe_quant_config,
16 + mxfp4_w4a4_moe_quant_config)
17 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
18 is_rocm_aiter_moe_enabled)
19 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
@@ -287,6 +290,16 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
290 from vllm.model_executor.layers.fused_moe import fused_experts
291 self.fused_experts_func = fused_experts
292
293 + def get_fused_moe_quant_config(
294 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
295 + return fp8_w8a8_moe_quant_config(
296 + w1_scale=layer.w13_weight_scale,
297 + w2_scale=layer.w2_weight_scale,
298 + a1_scale=layer.w13_input_scale,
299 + a2_scale=layer.w2_input_scale,
300 + per_act_token_quant=self.weight_qscheme == "per_channel",
301 + )
302 +
303 def apply(
304 self,
305 layer: torch.nn.Module,
@@ -339,12 +352,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
352 topk_ids=topk_ids,
353 activation=activation,
354 apply_router_weight_on_input=apply_router_weight_on_input,
342 - use_fp8_w8a8=True,
343 - per_channel_quant=self.weight_qscheme == "per_channel",
344 - w1_scale=layer.w13_weight_scale,
345 - w2_scale=layer.w2_weight_scale,
346 - a1_scale=layer.w13_input_scale,
347 - a2_scale=layer.w2_input_scale,
355 + quant_config=self.moe_quant_config,
356 expert_map=expert_map)
357 if self.use_marlin:
358 assert activation == "silu", (
@@ -376,14 +384,9 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
384 inplace=True,
385 activation=activation,
386 apply_router_weight_on_input=apply_router_weight_on_input,
379 - use_fp8_w8a8=True,
380 - per_channel_quant=self.weight_qscheme == "per_channel",
387 global_num_experts=global_num_experts,
388 expert_map=expert_map,
383 - w1_scale=layer.w13_weight_scale,
384 - w2_scale=layer.w2_weight_scale,
385 - a1_scale=layer.w13_input_scale,
386 - a2_scale=layer.w2_input_scale)
389 + quant_config=self.moe_quant_config)
390
391
392 class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
@@ -487,6 +490,16 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
490 layer.register_parameter("w13_weight_scale", w13_weight_scale)
491 layer.register_parameter("w2_weight_scale", w2_weight_scale)
492
493 + def get_fused_moe_quant_config(
494 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
495 + return mxfp4_w4a4_moe_quant_config(
496 + w1_scale=layer.w13_weight_scale,
497 + w2_scale=layer.w2_weight_scale,
498 + a1_scale=None,
499 + a2_scale=None,
500 + block_shape=None,
501 + )
502 +
503 def apply(
504 self,
505 layer: torch.nn.Module,
@@ -539,15 +552,10 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
552 topk_weights=topk_weights,
553 topk_ids=topk_ids,
554 inplace=True,
542 - use_mxfp4_w4a4=True,
555 + activation=activation,
556 global_num_experts=global_num_experts,
557 apply_router_weight_on_input=apply_router_weight_on_input,
558 expert_map=expert_map,
546 - w1_scale=layer.w13_weight_scale,
547 - w2_scale=layer.w2_weight_scale,
548 - a1_scale=None,
549 - a2_scale=None,
550 - block_shape=None,
551 - activation=activation,
559 + quant_config=self.moe_quant_config,
560 )
561 return out
vllm/model_executor/layers/quantization/rtn.py
+22 -12
@@ -12,6 +12,9 @@ from torch.nn.parameter import Parameter
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
14 FusedMoEMethodBase)
15 +from vllm.model_executor.layers.fused_moe.config import (
16 + FusedMoEQuantConfig, int4_w4a16_moe_quant_config,
17 + int8_w8a16_moe_quant_config)
18 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
19 set_weight_attrs)
20 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -269,6 +272,21 @@ class RTNMoEMethod(FusedMoEMethodBase):
272 fix_weights(layer, "w13_weight", weight_bits == 4)
273 fix_weights(layer, "w2_weight", weight_bits == 4)
274
275 + def get_fused_moe_quant_config(
276 + self, layer: torch.nn.Module) -> Optional[FusedMoEQuantConfig]:
277 + weight_bits = self.quant_config.weight_bits
278 + group_size = self.quant_config.group_size
279 + assert weight_bits == 4 or weight_bits == 8
280 + config_builder = (int4_w4a16_moe_quant_config
281 + if weight_bits == 4 else int8_w8a16_moe_quant_config)
282 + return config_builder(
283 + w1_scale=layer.w13_scale,
284 + w2_scale=layer.w2_scale,
285 + w1_zp=None,
286 + w2_zp=None,
287 + block_shape=[0, group_size],
288 + )
289 +
290 def apply(
291 self,
292 layer: torch.nn.Module,
@@ -314,10 +332,7 @@ class RTNMoEMethod(FusedMoEMethodBase):
332 e_score_correction_bias=e_score_correction_bias,
333 indices_type=self.topk_indices_dtype)
334
317 - weight_bits = self.quant_config.weight_bits
318 - group_size = self.quant_config.group_size
319 -
320 - ret = fused_experts(
335 + return fused_experts(
336 x,
337 layer.w13_weight,
338 layer.w2_weight,
@@ -325,16 +340,11 @@ class RTNMoEMethod(FusedMoEMethodBase):
340 topk_ids=topk_ids,
341 inplace=True,
342 activation=activation,
328 - use_int4_w4a16=weight_bits == 4,
329 - use_int8_w8a16=weight_bits == 8,
330 - global_num_experts=global_num_experts,
331 - w1_scale=layer.w13_scale,
332 - w2_scale=layer.w2_scale,
343 apply_router_weight_on_input=apply_router_weight_on_input,
344 + global_num_experts=global_num_experts,
345 expert_map=expert_map,
335 - block_shape=[0, group_size])
336 -
337 - return ret
346 + quant_config=self.moe_quant_config,
347 + )
348
349
350 def rtn_quantize(tensor: torch.Tensor, num_bits: int,
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+6 -14
@@ -7,7 +7,8 @@ import torch
7
8 import vllm.envs as envs
9 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10 -from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
10 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
11 + FusedMoEQuantConfig)
12 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
13 FlashInferExperts)
14 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
@@ -47,32 +48,23 @@ def reorder_w1w3_to_w3w1(weight: torch.Tensor,
48
49
50 def build_flashinfer_fp4_cutlass_moe_prepare_finalize(
50 - moe: FusedMoEConfig,
51 - a1_gscale: torch.Tensor,
52 -) -> mk.FusedMoEPrepareAndFinalize:
51 + moe: FusedMoEConfig) -> mk.FusedMoEPrepareAndFinalize:
52 """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""
53 use_dp = moe.moe_parallel_config.dp_size > 1
55 - return FlashInferCutlassMoEPrepareAndFinalize(use_dp, a1_gscale=a1_gscale)
54 + return FlashInferCutlassMoEPrepareAndFinalize(use_dp)
55
56
57 def select_nvfp4_gemm_impl(
58 moe: FusedMoEConfig,
60 - g1_alphas: torch.Tensor,
61 - g2_alphas: torch.Tensor,
62 - a1_gscale: torch.Tensor,
63 - a2_gscale: torch.Tensor,
59 + moe_quant_config: FusedMoEQuantConfig,
60 allow_flashinfer: bool,
61 ) -> mk.FusedMoEPermuteExpertsUnpermute:
62 """Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""
63
64 if allow_flashinfer:
65 return FlashInferExperts(
70 - g1_alphas=g1_alphas,
71 - g2_alphas=g2_alphas,
72 - a1_gscale=a1_gscale,
73 - a2_gscale=a2_gscale,
66 out_dtype=moe.in_dtype,
75 - quant_dtype="nvfp4",
67 + quant_config=moe_quant_config,
68 ep_rank=moe.moe_parallel_config.ep_rank,
69 ep_size=moe.moe_parallel_config.ep_size,
70 tp_rank=moe.moe_parallel_config.tp_rank,
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+14 -24
@@ -8,7 +8,8 @@ import torch
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 from vllm import envs
10 from vllm.logger import init_logger
11 -from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
11 +from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
12 + FusedMoEQuantConfig)
13 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
14 FlashInferExperts)
15 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
@@ -99,6 +100,8 @@ def apply_flashinfer_per_tensor_scale_fp8(
100 apply_router_weight_on_input: bool,
101 ) -> torch.Tensor:
102 from flashinfer.fused_moe import RoutingMethodType
103 +
104 + import vllm.model_executor.layers.fused_moe.flashinfer_trtllm_moe # noqa: E501, F401
105 assert layer.output1_scales_scalar is not None, (
106 "Expected output1_scales_scalar to be initialized")
107 assert layer.output1_scales_scalar is not None, (
@@ -167,34 +170,23 @@ def register_moe_scaling_factors(layer: torch.nn.Module) -> None:
170
171
172 def build_flashinfer_fp8_cutlass_moe_prepare_finalize(
170 - moe: Optional[FusedMoEConfig],
171 - layer: torch.nn.Module,
172 -) -> mk.FusedMoEPrepareAndFinalize:
173 + moe: Optional[FusedMoEConfig], ) -> mk.FusedMoEPrepareAndFinalize:
174 """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""
175 use_dp = moe.moe_parallel_config.dp_size > 1 if moe is not None else False
175 - return FlashInferCutlassMoEPrepareAndFinalize(
176 - use_dp, a1_gscale=layer.w13_input_scale)
176 + return FlashInferCutlassMoEPrepareAndFinalize(use_dp)
177
178
179 def select_cutlass_fp8_gemm_impl(
180 moe: Optional[FusedMoEConfig],
181 - layer: torch.nn.Module,
181 + quant_config: FusedMoEQuantConfig,
182 out_dtype: Optional[torch.dtype] = None,
183 ) -> mk.FusedMoEPermuteExpertsUnpermute:
184 """Return a GEMM *experts* implementation for fused-MoE layers"""
185
186 - from vllm.model_executor.models.llama4 import Llama4MoE
187 - assert layer.custom_routing_function == Llama4MoE.custom_routing_function, \
188 - "FusedMoE flashinfer kernels are only supported for Llama4"
189 -
186 if moe is not None:
187 return FlashInferExperts(
192 - g1_alphas=layer.output1_scales_gate_scalar,
193 - g2_alphas=layer.output2_scales_scalar,
194 - a1_gscale=layer.w13_input_scale,
195 - a2_gscale=layer.w2_input_scale_inv,
188 out_dtype=moe.in_dtype,
197 - quant_dtype=torch.float8_e4m3fn,
189 + quant_config=quant_config,
190 ep_rank=moe.moe_parallel_config.ep_rank,
191 ep_size=moe.moe_parallel_config.ep_size,
192 tp_rank=moe.moe_parallel_config.tp_rank,
@@ -204,12 +196,8 @@ def select_cutlass_fp8_gemm_impl(
196 assert out_dtype is not None, (
197 "If moe config is None, out_dtype must be passed")
198 return FlashInferExperts(
207 - g1_alphas=layer.output1_scales_gate_scalar,
208 - g2_alphas=layer.output2_scales_scalar,
209 - a1_gscale=layer.w13_input_scale,
210 - a2_gscale=layer.w2_input_scale_inv,
199 out_dtype=out_dtype,
212 - quant_dtype=torch.float8_e4m3fn,
200 + quant_config=quant_config,
201 )
202
203
@@ -224,11 +212,13 @@ def flashinfer_cutlass_moe_fp8(
212 expert_map: Optional[torch.Tensor] = None,
213 apply_router_weight_on_input: bool = False,
214 ) -> torch.Tensor:
215 + quant_config = layer.quant_method.get_fused_moe_quant_config(layer)
216 + assert quant_config is not None
217 +
218 fused_experts = mk.FusedMoEModularKernel(
228 - build_flashinfer_fp8_cutlass_moe_prepare_finalize(moe=None,
229 - layer=layer),
219 + build_flashinfer_fp8_cutlass_moe_prepare_finalize(moe=None),
220 select_cutlass_fp8_gemm_impl(moe=None,
231 - layer=layer,
221 + quant_config=quant_config,
222 out_dtype=hidden_states.dtype))
223
224 return fused_experts(
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1
@@ -411,6 +411,7 @@ def per_token_group_quant_fp8(
411 x_s = torch.empty(shape, device=x.device, dtype=torch.float32)
412
413 # prefer CUDA kernel if available
414 + # TODO(bnell): this causes some fp8 moe test to fail.
415 if current_platform.is_cuda() and x.is_contiguous():
416 torch.ops._C.per_token_group_fp8_quant(x, x_q, x_s, group_size, eps,
417 fp8_min, fp8_max, use_ue8m0)
vllm/model_executor/models/bert_with_rope.py
+5 -5
@@ -15,8 +15,8 @@ from vllm.distributed import (divide, get_tensor_model_parallel_rank,
15 tensor_model_parallel_all_reduce)
16 from vllm.model_executor.layers.activation import (get_act_and_mul_fn,
17 get_act_fn)
18 -from vllm.model_executor.layers.fused_moe.fused_moe import (
19 - fused_topk, torch_vllm_outplace_fused_experts)
18 +from vllm.model_executor.layers.fused_moe import (activation_without_mul,
19 + fused_topk)
20 from vllm.model_executor.layers.linear import (ColumnParallelLinear,
21 MergedColumnParallelLinear,
22 QKVParallelLinear,
@@ -230,7 +230,7 @@ class NomicMoE(nn.Module):
230 self.hidden_size = hidden_size
231 self.total_intermediate_size = intermediate_size
232 self.intermediate_size = divide(intermediate_size, self.tp_size)
233 - self.hidden_act = hidden_act
233 + self.hidden_act = activation_without_mul(hidden_act)
234
235 if params_dtype is None:
236 params_dtype = torch.get_default_dtype()
@@ -297,14 +297,14 @@ class NomicMoE(nn.Module):
297 router_logits,
298 self.top_k,
299 renormalize=False)
300 - final_hidden_states = torch_vllm_outplace_fused_experts(
300 +
301 + final_hidden_states = torch.ops.vllm.outplace_fused_experts(
302 hidden_states=hidden_states,
303 w1=self.w1,
304 w2=self.w2,
305 topk_weights=topk_weights,
306 topk_ids=topk_ids,
307 activation=self.hidden_act,
307 - is_act_and_mul=False,
308 )
309
310 if self.tp_size > 1:
vllm/model_executor/models/deepseek.py
+14 -8
@@ -37,7 +37,7 @@ from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
37 get_tensor_model_parallel_world_size,
38 tensor_model_parallel_all_reduce)
39 from vllm.model_executor.layers.activation import SiluAndMul
40 -from vllm.model_executor.layers.fused_moe import fused_moe
40 +from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
41 from vllm.model_executor.layers.layernorm import RMSNorm
42 from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
43 QKVParallelLinear,
@@ -163,13 +163,19 @@ class DeepseekMoE(nn.Module):
163 shared_output = self.shared_experts(hidden_states)
164 # router_logits: (num_tokens, n_experts)
165 router_logits, _ = self.gate(hidden_states)
166 - final_hidden_states = fused_moe(hidden_states,
167 - self.w1,
168 - self.w2,
169 - router_logits,
170 - self.top_k,
171 - renormalize=self.config.norm_topk_prob,
172 - inplace=True)
166 +
167 + topk_weights, topk_ids, _ = fused_topk(
168 + hidden_states,
169 + router_logits,
170 + self.top_k,
171 + renormalize=self.config.norm_topk_prob)
172 +
173 + final_hidden_states = fused_experts(hidden_states,
174 + self.w1,
175 + self.w2,
176 + topk_weights,
177 + topk_ids,
178 + inplace=True)
179
180 if self.config.n_shared_experts is not None:
181 final_hidden_states = final_hidden_states + shared_output
vllm/model_executor/models/minicpm.py
+13 -8
@@ -39,7 +39,7 @@ from vllm.distributed import (get_pp_group, get_tensor_model_parallel_rank,
39 get_tensor_model_parallel_world_size,
40 tensor_model_parallel_all_reduce)
41 from vllm.model_executor.layers.activation import FatreluAndMul, SiluAndMul
42 -from vllm.model_executor.layers.fused_moe import fused_moe
42 +from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
43 from vllm.model_executor.layers.layernorm import RMSNorm
44 from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
45 QKVParallelLinear,
@@ -136,13 +136,18 @@ class MiniCPMMoE(nn.Module):
136 hidden_states = hidden_states.view(-1, self.hidden_size)
137 # router_logits: (num_tokens, n_experts)
138 router_logits, _ = self.gate(hidden_states)
139 - final_hidden_states = fused_moe(hidden_states,
140 - self.ws,
141 - self.w2s,
142 - router_logits,
143 - self.top_k,
144 - renormalize=True,
145 - inplace=True)
139 +
140 + topk_weights, topk_ids, _ = fused_topk(hidden_states,
141 + router_logits,
142 + self.top_k,
143 + renormalize=True)
144 +
145 + final_hidden_states = fused_experts(hidden_states,
146 + self.ws,
147 + self.w2s,
148 + topk_weights,
149 + topk_ids,
150 + inplace=True)
151
152 if self.tp_size > 1:
153 final_hidden_states = tensor_model_parallel_all_reduce(
vllm/model_executor/models/qwen3_moe.py
+1 -1
@@ -702,4 +702,4 @@ class Qwen3MoeForCausalLM(nn.Module, SupportsPP, SupportsLoRA,
702 return loader.load_weights(weights)
703
704 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
705 - return self.model.get_expert_mapping()
705 + return self.model.get_expert_mapping()
\ No newline at end of file
vllm/model_executor/warmup/deep_gemm_warmup.py
+8 -3
@@ -81,9 +81,14 @@ def _fp8_linear_may_use_deep_gemm(module: torch.nn.Module) -> bool:
81
82
83 def _fused_moe_grouped_gemm_may_use_deep_gemm(module: torch.nn.Module) -> bool:
84 - if not (isinstance(module, FusedMoE)
85 - and module.moe_config.quant_dtype == torch.float8_e4m3fn
86 - and module.moe_config.block_shape == deep_gemm_block_shape()):
84 + if not isinstance(module, FusedMoE):
85 + return False
86 +
87 + moe_quant_config = module.quant_method.get_fused_moe_quant_config(module)
88 +
89 + if (moe_quant_config is None
90 + or moe_quant_config.quant_dtype != torch.float8_e4m3fn
91 + or moe_quant_config.block_shape != deep_gemm_block_shape()):
92 return False
93
94 if not isinstance(module.quant_method.fused_experts,