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
[Kernels] Clean up FusedMoeMethodBase and modular kernel setup. Remove extra arguments from modular kernel methods. (#22035)
Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
bnellnm committed
Aug 15, 2025 at 14:46 UTC
8ad7285ea28ad3bcc898fa99812120bcda8ea7b4
54 files changed
+2022
-1305
.buildkite/test-pipeline.yaml
+1
@@ -399,6 +399,7 @@ steps:
399
- label: Kernels MoE Test %N
400
mirror_hardwares: [amdexperimental]
401
source_file_dependencies:
402
+ - csrc/quantization/cutlass_w8a8/moe/
403
- csrc/moe/
404
- tests/kernels/moe
405
- vllm/model_executor/layers/fused_moe/
docs/design/fused_moe_modular_kernel.md
+9
-1
@@ -175,11 +175,19 @@ implementations that input `FusedMoEActivationFormat.Standard` support chunking
175
176
### FusedMoEModularKernel Initialization
177
178
-`FusedMoEMethodBase` class has 2 methods that are collectively responsible in creating the `FusedMoEModularKernel` object. They are,
178
+`FusedMoEMethodBase` class has 3 methods that are collectively responsible in creating the `FusedMoEModularKernel` object. They are,
179
180
+* maybe_make_prepare_finalize,
181
* select_gemm_impl, and
182
* init_prepare_finalize
183
184
+#### maybe_make_prepare_finalize
185
+
186
+The `maybe_make_prepare_finalize` method is responsbile for constructing an instance of `FusedMoEPrepareAndFinalize` when appropriate based on the current all2all backend, e.g. when EP + DP is enabled. The base class method currently constructs all the `FusedMoEPrepareAndFinalize` objects for the EP+DP case. Derived classes can override this method to construct prepare/finalize objects for different scenarios, e.g. `ModelOptNvFp4FusedMoE` can construct a `FlashInferCutlassMoEPrepareAndFinalize` for the EP+TP case.
187
+Please refer to the implementations in,
188
+
189
+* `ModelOptNvFp4FusedMoE`
190
+
191
#### select_gemm_impl
192
193
The `select_gemm_impl` method is undefined in the base class. It is the responsibility of the derived class to implement a method that constructs a valid/appropriate `FusedMoEPermuteExpertsUnpermute` object.
examples/offline_inference/data_parallel.py
+22
-1
@@ -70,12 +70,27 @@ def parse_args():
70
default=64,
71
help=("Maximum number of sequences to be processed in a single iteration."),
72
)
73
+ parser.add_argument(
74
+ "--max-model-len",
75
+ type=int,
76
+ help=("Maximum number of tokens to be processed in a single iteration."),
77
+ )
78
+ parser.add_argument(
79
+ "--timeout",
80
+ type=int,
81
+ default=300,
82
+ help=("Number of seconds before unresponsive process is killed."),
83
+ )
84
parser.add_argument(
85
"--gpu-memory-utilization",
86
type=float,
87
default=0.8,
88
help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
89
)
90
+ parser.add_argument(
91
+ "--quantization",
92
+ type=str,
93
+ )
94
return parser.parse_args()
95
96
@@ -90,7 +105,9 @@ def main(
105
enforce_eager,
106
trust_remote_code,
107
max_num_seqs,
108
+ max_model_len,
109
gpu_memory_utilization,
110
+ quantization,
111
):
112
os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
113
os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
@@ -142,7 +159,9 @@ def main(
159
enable_expert_parallel=True,
160
trust_remote_code=trust_remote_code,
161
max_num_seqs=max_num_seqs,
162
+ max_model_len=max_model_len,
163
gpu_memory_utilization=gpu_memory_utilization,
164
+ quantization=quantization,
165
)
166
outputs = llm.generate(prompts, sampling_params)
167
# Print the outputs.
@@ -198,14 +217,16 @@ if __name__ == "__main__":
217
args.enforce_eager,
218
args.trust_remote_code,
219
args.max_num_seqs,
220
+ args.max_model_len,
221
args.gpu_memory_utilization,
222
+ args.quantization,
223
),
224
)
225
proc.start()
226
procs.append(proc)
227
exit_code = 0
228
for proc in procs:
208
- proc.join(timeout=300)
229
+ proc.join(timeout=args.timeout)
230
if proc.exitcode is None:
231
print(f"Killing process {proc.pid} that didn't stop within 5 minutes.")
232
proc.kill()
tests/kernels/moe/modular_kernel_tools/common.py
+263
-277
@@ -7,41 +7,22 @@ import torch
7
8
import vllm._custom_ops as ops
9
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
+from tests.kernels.moe.utils import make_test_weights, per_token_cast_to_fp8
11
+from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
12
+ FLOAT8_E4M3_MAX,
13
+ dequantize_nvfp4_to_dtype)
14
from tests.kernels.utils import torch_experts
15
from vllm.config import VllmConfig
16
from vllm.distributed import get_dp_group, get_tensor_model_parallel_world_size
13
-# Fused experts and PrepareFinalize imports
14
-from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
15
- BatchedDeepGemmExperts)
16
-from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
17
- BatchedTritonOrDeepGemmExperts)
17
+from vllm.forward_context import set_forward_context
18
from vllm.model_executor.layers.fused_moe.config import (
19
FusedMoEConfig, FusedMoEParallelConfig, FusedMoEQuantConfig)
20
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
21
-from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
22
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
23
- BatchedTritonExperts, NaiveBatchedExperts)
20
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
25
-from vllm.model_executor.layers.fused_moe.layer import (FusedMoEMethodBase,
26
- TritonExperts)
27
-from vllm.model_executor.layers.fused_moe.prepare_finalize import (
28
- MoEPrepareAndFinalizeNoEP)
29
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
30
- TritonOrDeepGemmExperts)
21
from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
22
23
+from .mk_objects import (expert_info, make_fused_experts,
24
+ make_prepare_finalize, prepare_finalize_info)
25
from .parallel_utils import ProcessGroupInfo
34
-from .utils import (make_block_quant_fp8_weights, make_non_quant_weights,
35
- make_quant_fp8_weights, per_token_cast_to_fp8)
36
-
37
-if has_pplx():
38
- from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
39
- PplxPrepareAndFinalize)
40
-if has_deep_ep():
41
- from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
42
- DeepEPHTPrepareAndFinalize)
43
- from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
44
- DeepEPLLPrepareAndFinalize)
26
27
28
def _describe_tensor(t: Optional[torch.Tensor], name: str) -> str:
@@ -69,24 +50,31 @@ class Config:
50
51
torch_trace_dir_path: Optional[str] = None
52
53
+ def __post_init__(self):
54
+ if self.quant_config is None:
55
+ self.quant_config = FusedMoEQuantConfig()
56
+
57
def describe(self) -> str:
58
s = ""
74
- s += "== Config: \n"
75
- s += f" world_size={self.world_size} \n"
76
- s += f" PF={self.prepare_finalize_type.__name__} \n"
77
- s += f" FE={self.fused_experts_type.__name__} \n"
78
- s += f" topk={self.topks} \n"
79
- s += f" dtype={self.dtype} \n"
80
- s += f" fused_moe_chunk_size={self.fused_moe_chunk_size} \n"
81
- s += " Quant: \n"
82
- s += f" fused_moe_chunk_size={self.fused_moe_chunk_size} \n "
59
+ s += "== Config:\n"
60
+ s += f" world_size={self.world_size}\n"
61
+ s += f" PF={self.prepare_finalize_type.__name__}\n"
62
+ s += f" FE={self.fused_experts_type.__name__}\n"
63
+ s += f" E={self.E}\n"
64
+ s += f" Ms={self.Ms}\n"
65
+ s += f" N={self.N}\n"
66
+ s += f" K={self.K}\n"
67
+ s += f" topk={self.topks}\n"
68
+ s += f" dtype={self.dtype}\n"
69
+ s += f" fused_moe_chunk_size={self.fused_moe_chunk_size}\n"
70
+ s += " Quant:\n"
71
if self.quant_config is not None:
84
- s += f" q_dtype={self.quant_dtype} \n"
85
- s += f" q_block_shape={self.quant_block_shape} \n"
86
- s += f" q_per_out_ch_quant={self.is_per_out_ch_quant} \n"
87
- s += f" q_per_act_token={self.is_per_act_token_quant} \n"
72
+ s += f" q_dtype={self.quant_dtype}\n"
73
+ s += f" q_block_shape={self.quant_block_shape}\n"
74
+ s += f" q_per_out_ch_quant={self.is_per_out_ch_quant}\n"
75
+ s += f" q_per_act_token={self.is_per_act_token_quant}\n"
76
else:
89
- s += " quant=None \n"
77
+ s += " quant=None\n"
78
return s
79
80
@property
@@ -95,34 +83,28 @@ class Config:
83
return self.Ms
84
85
@property
98
- def quant_dtype(self) -> Optional[torch.dtype]:
99
- if self.quant_config is None:
100
- return None
86
+ def quant_dtype(self) -> Union[torch.dtype, str, None]:
87
+ assert self.quant_config is not None
88
return self.quant_config.quant_dtype
89
90
@property
91
def is_per_act_token_quant(self) -> bool:
105
- if self.quant_config is None:
106
- return False
92
+ assert self.quant_config is not None
93
return self.quant_config.per_act_token_quant
94
95
@property
96
def is_per_tensor_act_quant(self) -> bool:
111
- if self.quant_config is None:
112
- return False
97
return (not self.is_per_act_token_quant
98
and self.quant_block_shape is None)
99
100
@property
101
def is_per_out_ch_quant(self) -> bool:
118
- if self.quant_config is None:
119
- return False
102
+ assert self.quant_config is not None
103
return self.quant_config.per_out_ch_quant
104
105
@property
106
def quant_block_shape(self) -> Optional[list[int]]:
124
- if self.quant_config is None:
125
- return None
107
+ assert self.quant_config is not None
108
return self.quant_config.block_shape
109
110
@property
@@ -130,36 +112,30 @@ class Config:
112
assert isinstance(self.topks, int)
113
return self.topks
114
133
- @property
134
- def topk_ids_dtype(self) -> Optional[torch.dtype]:
135
- topk_ids_dtype = None
136
- if self.prepare_finalize_type == PplxPrepareAndFinalize:
137
- topk_ids_dtype = torch.uint32
138
- elif self.prepare_finalize_type in [
139
- DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
140
- ]:
141
- topk_ids_dtype = torch.int64
142
- return topk_ids_dtype
143
-
115
@property
116
def num_local_experts(self) -> int:
117
return self.E // self.world_size
118
119
def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
120
"""
150
- make env data for vllm launch.
121
+ make env data for vllm launch.
122
"""
123
vllm_config = VllmConfig()
124
vllm_config.parallel_config.data_parallel_size = self.world_size
125
vllm_config.parallel_config.enable_expert_parallel = True
126
127
env_dict = {
157
- "VLLM_ALL2ALL_BACKEND": self.all2all_backend(),
128
"VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())),
129
}
130
+
131
+ backend = self.all2all_backend()
132
+ if backend is not None:
133
+ env_dict.update({"VLLM_ALL2ALL_BACKEND": backend})
134
+
135
if self.fused_moe_chunk_size is not None:
136
env_dict.update(
137
{"VLLM_FUSED_MOE_CHUNK_SIZE": str(self.fused_moe_chunk_size)})
138
+
139
return vllm_config, env_dict
140
141
def is_fp8_block_quantized(self):
@@ -167,85 +143,59 @@ class Config:
143
and self.quant_block_shape is not None)
144
145
def is_batched_prepare_finalize(self):
170
- return self.prepare_finalize_type in [
171
- PplxPrepareAndFinalize, DeepEPLLPrepareAndFinalize
172
- ]
146
+ info = prepare_finalize_info(self.prepare_finalize_type)
147
+ return (mk.FusedMoEActivationFormat.BatchedExperts ==
148
+ info.activation_format)
149
150
def is_batched_fused_experts(self):
175
- return self.fused_experts_type in [
176
- CutlassExpertsFp8, BatchedDeepGemmExperts, BatchedTritonExperts,
177
- NaiveBatchedExperts, BatchedTritonOrDeepGemmExperts
178
- ]
151
+ info = expert_info(self.fused_experts_type)
152
+ return (mk.FusedMoEActivationFormat.BatchedExperts ==
153
+ info.activation_format)
154
155
def is_standard_fused_experts(self):
181
- return self.fused_experts_type in [
182
- CutlassExpertsFp8, DeepGemmExperts, TritonOrDeepGemmExperts,
183
- TritonExperts
184
- ]
185
-
186
- def is_fe_16bit_supported(self):
187
- return self.fused_experts_type in [
188
- BatchedTritonExperts, BatchedTritonOrDeepGemmExperts,
189
- NaiveBatchedExperts, TritonExperts
190
- ]
191
-
192
- def is_fe_fp8_supported(self):
193
- return self.fused_experts_type in [
194
- BatchedDeepGemmExperts,
195
- BatchedTritonExperts,
196
- BatchedTritonOrDeepGemmExperts,
197
- CutlassExpertsFp8,
198
- DeepGemmExperts,
199
- TritonExperts,
200
- TritonOrDeepGemmExperts,
201
- NaiveBatchedExperts,
202
- ]
203
-
204
- def is_fe_block_fp8_supported(self):
205
- return self.fused_experts_type in [
206
- BatchedDeepGemmExperts,
207
- BatchedTritonOrDeepGemmExperts,
208
- DeepGemmExperts,
209
- TritonExperts,
210
- TritonOrDeepGemmExperts,
211
- BatchedTritonExperts,
212
- NaiveBatchedExperts,
213
- ]
156
+ info = expert_info(self.fused_experts_type)
157
+ return mk.FusedMoEActivationFormat.Standard == info.activation_format
158
+
159
+ def fe_supported_types(self):
160
+ info = expert_info(self.fused_experts_type)
161
+ return info.supported_dtypes
162
+
163
+ def pf_supported_types(self):
164
+ info = prepare_finalize_info(self.prepare_finalize_type)
165
+ return info.supported_dtypes
166
+
167
+ def is_block_quant_supported(self):
168
+ info = expert_info(self.fused_experts_type)
169
+ return info.blocked_quantization_support
170
171
def is_fe_supports_chunking(self):
216
- return self.fused_experts_type in [
217
- CutlassExpertsFp8, DeepGemmExperts, TritonOrDeepGemmExperts,
218
- TritonExperts
219
- ]
172
+ info = expert_info(self.fused_experts_type)
173
+ return info.supports_chunking
174
+
175
+ def supports_expert_map(self):
176
+ info = expert_info(self.fused_experts_type)
177
+ return info.supports_expert_map
178
+
179
+ def supports_apply_weight_on_input(self):
180
+ info = prepare_finalize_info(self.prepare_finalize_type)
181
+ return info.supports_apply_weight_on_input
182
183
def needs_deep_gemm(self):
222
- return self.fused_experts_type in [
223
- BatchedDeepGemmExperts,
224
- DeepGemmExperts,
225
- ]
184
+ info = expert_info(self.fused_experts_type)
185
+ return info.needs_deep_gemm
186
187
def needs_pplx(self):
228
- return self.prepare_finalize_type in [PplxPrepareAndFinalize]
188
+ info = prepare_finalize_info(self.prepare_finalize_type)
189
+ return info.backend == "pplx"
190
191
def needs_deep_ep(self):
231
- return self.prepare_finalize_type in [
232
- DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
233
- ]
192
+ info = prepare_finalize_info(self.prepare_finalize_type)
193
+ return (info.backend == "deepep_high_throughput"
194
+ or info.backend == "deepep_low_latency")
195
196
def all2all_backend(self):
236
- if self.needs_pplx():
237
- return "pplx"
238
- if self.prepare_finalize_type == DeepEPHTPrepareAndFinalize:
239
- return "deepep_high_throughput"
240
- if self.prepare_finalize_type == DeepEPLLPrepareAndFinalize:
241
- return "deepep_low_latency"
242
- return "naive"
243
-
244
- def needs_all2all(self):
245
- return self.prepare_finalize_type in [
246
- PplxPrepareAndFinalize, DeepEPHTPrepareAndFinalize,
247
- DeepEPLLPrepareAndFinalize
248
- ]
197
+ info = prepare_finalize_info(self.prepare_finalize_type)
198
+ return info.backend
199
200
def is_valid(self):
201
# Check prepare-finalize and fused-experts compatibility
@@ -267,28 +217,28 @@ class Config:
217
# invalid quant config
218
return False
219
270
- # check bf16 / fp16 support
271
- is_16bit = (self.dtype.itemsize == 2 and self.quant_dtype is None)
272
- if is_16bit and not self.is_fe_16bit_supported():
273
- return False
274
-
275
- # Check fp8 support
276
- is_fp8 = self.quant_dtype == torch.float8_e4m3fn
277
- if is_fp8 and not self.is_fe_fp8_supported():
278
- return False
220
+ # check type support
221
+ if self.quant_dtype is None:
222
+ if (self.dtype not in self.pf_supported_types()
223
+ or self.dtype not in self.fe_supported_types()):
224
+ return False
225
+ else:
226
+ if (self.quant_dtype not in self.pf_supported_types()
227
+ or self.quant_dtype not in self.fe_supported_types()):
228
+ return False
229
280
- # Check fp8 block quanization support
230
+ # Check block quanization support
231
is_block_quatized = self.quant_block_shape is not None
282
- if is_block_quatized and not is_fp8:
232
+ if is_block_quatized and self.quant_dtype is None:
233
return False
284
- if is_block_quatized and not self.is_fe_block_fp8_supported():
234
+ if is_block_quatized and not self.is_block_quant_supported():
235
return False
236
237
# deep_gemm only works with block-quantized
238
if self.needs_deep_gemm() and not is_block_quatized:
239
return False
240
291
- # Check dependencies
241
+ # Check dependencies (turn into asserts?)
242
if self.needs_deep_ep() and not has_deep_ep():
243
return False
244
if self.needs_deep_gemm() and not has_deep_gemm():
@@ -305,6 +255,8 @@ class WeightTensors:
255
w2: torch.Tensor
256
w1_scale: Optional[torch.Tensor]
257
w2_scale: Optional[torch.Tensor]
258
+ w1_gs: Optional[torch.Tensor] = None
259
+ w2_gs: Optional[torch.Tensor] = None
260
261
def describe(self):
262
s = ""
@@ -313,13 +265,20 @@ class WeightTensors:
265
s += f' - {_describe_tensor(self.w2, "w2")} \n'
266
s += f' - {_describe_tensor(self.w1_scale, "w1_scale")} \n'
267
s += f' - {_describe_tensor(self.w2_scale, "w2_scale")} \n'
268
+ s += f' - {_describe_tensor(self.w1_gs, "w1_gs")} \n'
269
+ s += f' - {_describe_tensor(self.w2_gs, "w2_gs")} \n'
270
return s
271
272
+ def is_quantized(self) -> bool:
273
+ # or w1_scale is not None?
274
+ return (self.w1.dtype == torch.float8_e4m3fn
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())
321
- is_quantized = self.w1.dtype == torch.float8_e4m3fn
322
- if is_quantized:
280
+
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(
@@ -327,56 +286,51 @@ class WeightTensors:
286
self.w2_scale = self.w2_scale.to(
287
device=torch.cuda.current_device())
288
289
+ 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())
293
+
294
def slice_weights(self, rank: int,
295
num_local_experts: int) -> "WeightTensors":
296
s = rank * num_local_experts
297
e = s + num_local_experts
298
w1 = self.w1[s:e, :, :]
299
w2 = self.w2[s:e, :, :]
336
- is_quantized = self.w1.dtype == torch.float8_e4m3fn
300
+
301
w1_scale, w2_scale = (None, None)
338
- if is_quantized:
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, :, :]
343
- return WeightTensors(w1, w2, w1_scale, w2_scale)
307
345
- @staticmethod
346
- def make(config: Config) -> "WeightTensors":
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]
314
348
- if config.quant_dtype is None:
349
- # just make normal dtype weights
350
- w1, w2 = make_non_quant_weights(e=config.E,
351
- n=config.N,
352
- k=config.K,
353
- dtype=config.dtype)
354
- return WeightTensors(w1=w1, w2=w2, w1_scale=None, w2_scale=None)
355
-
356
- assert config.quant_dtype == torch.float8_e4m3fn
357
- if not config.is_fp8_block_quantized():
358
- w1, w2, w1_scale, w2_scale = make_quant_fp8_weights(
359
- e=config.E,
360
- n=config.N,
361
- k=config.K,
362
- per_out_channel_quant=config.is_per_out_ch_quant,
363
- )
364
- return WeightTensors(w1=w1,
365
- w2=w2,
366
- w1_scale=w1_scale,
367
- w2_scale=w2_scale)
315
+ return WeightTensors(w1, w2, w1_scale, w2_scale, w1_gs, w2_gs)
316
369
- assert config.quant_block_shape is not None
370
- w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights(
317
+ @staticmethod
318
+ def make(config: Config) -> "WeightTensors":
319
+ (_, w1, w1_scale, w1_gs), (_, w2, w2_scale, w2_gs) = make_test_weights(
320
e=config.E,
321
n=config.N,
322
k=config.K,
374
- block_size=config.quant_block_shape,
323
+ in_dtype=config.dtype,
324
+ quant_dtype=config.quant_dtype,
325
+ block_shape=config.quant_block_shape,
326
+ per_act_token_quant=config.is_per_out_ch_quant,
327
)
328
return WeightTensors(w1=w1,
329
w2=w2,
330
w1_scale=w1_scale,
379
- w2_scale=w2_scale)
331
+ w2_scale=w2_scale,
332
+ w1_gs=w1_gs,
333
+ w2_gs=w2_gs)
334
335
336
@dataclass
@@ -449,7 +403,6 @@ class RankTensors:
403
dtype=dtype)
404
topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk,
405
False)
452
- topk_ids = topk_ids.to(config.topk_ids_dtype)
406
407
# distribute topk_ids evenly
408
for mi in range(m):
@@ -457,7 +410,7 @@ class RankTensors:
410
topk_ids = topk_ids.to(device=torch.cuda.current_device())
411
412
expert_map = None
460
- if config.world_size > 1:
413
+ if config.world_size > 1 and config.supports_expert_map():
414
expert_map = torch.full((global_num_experts, ),
415
fill_value=-1,
416
dtype=torch.int32)
@@ -480,92 +433,100 @@ class RankTensors:
433
def reference_moe_impl(config: Config, weights: WeightTensors,
434
rank_tensors: RankTensors) -> torch.Tensor:
435
483
- return torch_experts(a=rank_tensors.hidden_states,
484
- w1=weights.w1,
485
- w2=weights.w2,
436
+ if config.quant_dtype == "nvfp4":
437
+ quant_blocksize = 16
438
+ dtype = config.dtype
439
+
440
+ w1_q = weights.w1
441
+ w1_blockscale = weights.w1_scale
442
+ w1_gs = weights.w1_gs
443
+
444
+ w2_q = weights.w2
445
+ w2_blockscale = weights.w2_scale
446
+ w2_gs = weights.w2_gs
447
+
448
+ a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(
449
+ rank_tensors.hidden_states.flatten(), dim=-1)).to(torch.float32)
450
+
451
+ assert w1_gs is not None
452
+ assert w2_gs is not None
453
+ assert w1_blockscale is not None
454
+ assert w2_blockscale is not None
455
+
456
+ assert w1_blockscale.shape[1] % 128 == 0
457
+ assert w1_blockscale.shape[2] % 4 == 0
458
+ assert w2_blockscale.shape[1] % 128 == 0
459
+ assert w2_blockscale.shape[2] % 4 == 0
460
+
461
+ a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
462
+ rank_tensors.hidden_states, a_global_scale)
463
+
464
+ a = dequantize_nvfp4_to_dtype(a_fp4,
465
+ a_scale_interleaved,
466
+ a_global_scale,
467
+ dtype=dtype,
468
+ device=a_fp4.device,
469
+ block_size=quant_blocksize)
470
+
471
+ e = w1_q.shape[0]
472
+ n = w1_q.shape[1] // 2
473
+ k = w2_q.shape[1]
474
+
475
+ w1 = torch.zeros((e, 2 * n, k), device="cuda", dtype=dtype)
476
+ w2 = torch.zeros((e, k, n), device="cuda", dtype=dtype)
477
+
478
+ for idx in range(0, e):
479
+ w1[idx] = dequantize_nvfp4_to_dtype(w1_q[idx],
480
+ w1_blockscale[idx],
481
+ w1_gs[idx],
482
+ dtype=dtype,
483
+ device=w1_q.device,
484
+ block_size=quant_blocksize)
485
+ w2[idx] = dequantize_nvfp4_to_dtype(w2_q[idx],
486
+ w2_blockscale[idx],
487
+ w2_gs[idx],
488
+ dtype=dtype,
489
+ device=w2_q.device,
490
+ block_size=quant_blocksize)
491
+ a_scale = None
492
+ w1_scale = None
493
+ w2_scale = None
494
+ quant_dtype = None
495
+ per_act_token_quant = False
496
+ block_shape = None
497
+ else:
498
+ a = rank_tensors.hidden_states
499
+ a_scale = rank_tensors.hidden_states_scale
500
+ w1 = weights.w1
501
+ w1_scale = weights.w1_scale
502
+ w2 = weights.w2
503
+ w2_scale = weights.w2_scale
504
+ quant_dtype = config.quant_dtype
505
+ per_act_token_quant = config.is_per_act_token_quant
506
+ block_shape = config.quant_block_shape
507
+
508
+ return torch_experts(a=a,
509
+ w1=w1,
510
+ w2=w2,
511
topk_weight=rank_tensors.topk_weights,
512
topk_ids=rank_tensors.topk_ids,
513
global_num_experts=config.E,
514
expert_map=None,
490
- w1_scale=weights.w1_scale,
491
- w2_scale=weights.w2_scale,
492
- a1_scale=rank_tensors.hidden_states_scale,
493
- quant_dtype=config.quant_dtype,
494
- per_act_token_quant=config.is_per_act_token_quant,
495
- block_shape=config.quant_block_shape,
496
- apply_router_weights_on_input=config.topk == 1)
497
-
498
-
499
-def make_fused_experts(
500
- config: Config, moe: FusedMoEConfig,
501
- num_dispatchers: int) -> mk.FusedMoEPermuteExpertsUnpermute:
502
-
503
- use_fp8 = config.quant_dtype == torch.float8_e4m3fn
504
- batch_kwargs = {
505
- "max_num_tokens": moe.max_num_tokens,
506
- "num_dispatchers": num_dispatchers,
507
- }
508
- quant_kwargs = {
509
- "use_fp8_w8a8": use_fp8,
510
- "use_int8_w8a8": False,
511
- "use_int8_w8a16": False,
512
- "use_int4_w4a16": False,
513
- "block_shape": config.quant_block_shape,
514
- "per_act_token_quant": config.is_per_act_token_quant,
515
- }
516
- deepgemm_kwargs = {"allow_deep_gemm": has_deep_gemm()}
517
-
518
- if config.fused_experts_type == BatchedDeepGemmExperts:
519
- kwargs = batch_kwargs | {
520
- "block_shape": config.quant_block_shape,
521
- "per_act_token_quant": config.is_per_act_token_quant,
522
- }
523
- print(f"Making BatchedDeepGemmExperts {kwargs} ...")
524
- experts = BatchedDeepGemmExperts(**kwargs)
525
- elif config.fused_experts_type == BatchedTritonExperts:
526
- kwargs = batch_kwargs | quant_kwargs
527
- print(f"Making BatchedTritonExperts {kwargs} ...")
528
- experts = BatchedTritonExperts(**kwargs)
529
- elif config.fused_experts_type == BatchedTritonOrDeepGemmExperts:
530
- kwargs = batch_kwargs | quant_kwargs | deepgemm_kwargs
531
- print(f"Making BatchedTritonOrDeepGemmExperts {kwargs} ...")
532
- experts = BatchedTritonOrDeepGemmExperts(**kwargs)
533
- elif config.fused_experts_type == DeepGemmExperts:
534
- print("Making DeepGemmExperts () ...")
535
- experts = DeepGemmExperts()
536
- elif config.fused_experts_type == TritonExperts:
537
- kwargs = quant_kwargs
538
- print(f"Making TritonExperts {kwargs} ...")
539
- experts = TritonExperts(**kwargs)
540
- elif config.fused_experts_type == TritonOrDeepGemmExperts:
541
- kwargs = quant_kwargs | deepgemm_kwargs
542
- print(f"Making TritonOrDeepGemmExperts {kwargs} ...")
543
- experts = TritonOrDeepGemmExperts(**kwargs)
544
- elif config.fused_experts_type == NaiveBatchedExperts:
545
- kwargs = batch_kwargs | quant_kwargs
546
- print(f"Making NaiveBatchedExperts {kwargs} ...")
547
- experts = NaiveBatchedExperts(**kwargs)
548
- elif config.fused_experts_type == CutlassExpertsFp8:
549
- use_batched_format = config.is_batched_prepare_finalize()
550
- num_experts = (moe.num_local_experts
551
- if use_batched_format else moe.num_experts)
552
- kwargs = {
553
- "max_experts_per_worker": num_experts,
554
- "out_dtype": moe.in_dtype,
555
- "per_act_token_quant": config.is_per_act_token_quant,
556
- "per_out_ch_quant": config.is_per_out_ch_quant,
557
- "block_shape": config.quant_block_shape,
558
- "num_dispatchers": num_dispatchers,
559
- "use_batched_format": use_batched_format
560
- }
561
- print(f"Making CutlassExpertsFp8 {kwargs} ...")
562
- experts = CutlassExpertsFp8(**kwargs)
515
+ w1_scale=w1_scale,
516
+ w2_scale=w2_scale,
517
+ a1_scale=a_scale,
518
+ quant_dtype=quant_dtype,
519
+ per_act_token_quant=per_act_token_quant,
520
+ block_shape=block_shape,
521
+ apply_router_weights_on_input=config.topk == 1
522
+ and config.supports_apply_weight_on_input())
523
564
- return experts
524
566
-
567
-def make_modular_kernel(config: Config,
568
- vllm_config: VllmConfig) -> mk.FusedMoEModularKernel:
525
+def make_modular_kernel(
526
+ config: Config,
527
+ vllm_config: VllmConfig,
528
+ weights: WeightTensors,
529
+) -> mk.FusedMoEModularKernel:
530
531
def next_power_of_2(x):
532
import math
@@ -579,6 +540,7 @@ def make_modular_kernel(config: Config,
540
dp_size_=get_dp_group().world_size,
541
vllm_parallel_config=vllm_config.parallel_config,
542
)
543
+
544
moe = FusedMoEConfig(
545
num_experts=config.E,
546
experts_per_token=config.topk,
@@ -591,15 +553,16 @@ def make_modular_kernel(config: Config,
553
)
554
555
# make modular kernel
594
- prepare_finalize = None
595
- if config.needs_all2all():
596
- prepare_finalize = FusedMoEMethodBase.maybe_make_prepare_finalize(moe)
597
- assert prepare_finalize is not None
598
- else:
599
- prepare_finalize = MoEPrepareAndFinalizeNoEP()
600
-
601
- fused_experts = make_fused_experts(config, moe,
602
- prepare_finalize.num_dispatchers())
556
+ prepare_finalize = make_prepare_finalize(config.prepare_finalize_type,
557
+ config.all2all_backend(), moe)
558
+
559
+ fused_experts = make_fused_experts(
560
+ config.fused_experts_type,
561
+ moe,
562
+ prepare_finalize.num_dispatchers(),
563
+ weights.w1_gs,
564
+ weights.w2_gs,
565
+ )
566
567
modular_kernel = mk.FusedMoEModularKernel(
568
prepare_finalize=prepare_finalize, fused_experts=fused_experts)
@@ -620,22 +583,45 @@ def run_modular_kernel(
583
# weights for rank
584
rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
585
623
- mk = make_modular_kernel(config, vllm_config)
586
+ mk = make_modular_kernel(config, vllm_config, weights)
587
588
mk_kwargs = {
626
- "hidden_states": rank_tensors.hidden_states.clone(
589
+ "hidden_states":
590
+ rank_tensors.hidden_states.clone(
591
), # impls might update the tensor in place
628
- "w1": rank_weights.w1,
629
- "w2": rank_weights.w2,
630
- "topk_weights": rank_tensors.topk_weights,
631
- "topk_ids": rank_tensors.topk_ids,
632
- "expert_map": rank_tensors.expert_map,
633
- "w1_scale": rank_weights.w1_scale,
634
- "w2_scale": rank_weights.w2_scale,
635
- "a1_scale": rank_tensors.hidden_states_scale,
636
- "global_num_experts": config.E,
637
- "apply_router_weight_on_input": config.topk == 1,
592
+ "w1":
593
+ rank_weights.w1,
594
+ "w2":
595
+ rank_weights.w2,
596
+ "topk_weights":
597
+ rank_tensors.topk_weights,
598
+ "topk_ids":
599
+ rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype()),
600
+ "expert_map":
601
+ 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,
608
+ "global_num_experts":
609
+ config.E,
610
+ "apply_router_weight_on_input":
611
+ config.topk == 1 and config.supports_apply_weight_on_input(),
612
}
639
- out = mk.forward(**mk_kwargs)
613
+
614
+ num_tokens = rank_tensors.hidden_states.shape[0]
615
+ num_tokens_across_dp = torch.tensor([num_tokens] * config.world_size,
616
+ device="cuda",
617
+ dtype=torch.int)
618
+
619
+ with set_forward_context(
620
+ None,
621
+ vllm_config,
622
+ num_tokens=num_tokens,
623
+ num_tokens_across_dp=num_tokens_across_dp,
624
+ ):
625
+ out = mk.forward(**mk_kwargs)
626
627
return out
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+436
-25
@@ -1,58 +1,316 @@
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
5
6
import torch
7
8
# Fused experts and PrepareFinalize imports
9
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
11
BatchedDeepGemmExperts)
12
from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
13
BatchedTritonOrDeepGemmExperts)
11
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
14
+from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
15
+ FusedMoEQuantConfig)
16
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
17
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
18
BatchedTritonExperts, NaiveBatchedExperts)
16
-from vllm.model_executor.layers.fused_moe.layer import TritonExperts
19
+from vllm.model_executor.layers.fused_moe.layer import (FusedMoEMethodBase,
20
+ TritonExperts)
21
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
22
MoEPrepareAndFinalizeNoEP)
23
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
24
TritonOrDeepGemmExperts)
21
-from vllm.utils import has_deep_ep, has_pplx
25
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
26
+ cutlass_fp4_supported)
27
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
28
+ cutlass_fp8_supported)
29
+from vllm.platforms import current_platform
30
+from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
31
+from vllm.utils.deep_gemm import is_deep_gemm_supported
32
+from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
33
23
-if has_deep_ep():
34
+
35
+@dataclass
36
+class PrepareFinalizeInfo:
37
+ activation_format: mk.FusedMoEActivationFormat
38
+ supported_dtypes: list[Union[torch.dtype, str]]
39
+ blocked_quantization_support: bool
40
+ backend: Optional[str]
41
+ supports_apply_weight_on_input: bool = True
42
+
43
+
44
+@dataclass
45
+class ExpertInfo:
46
+ activation_format: mk.FusedMoEActivationFormat
47
+ supported_dtypes: list[Union[torch.dtype, str]]
48
+ blocked_quantization_support: bool
49
+ supports_chunking: bool
50
+ supports_expert_map: bool
51
+ needs_matching_quant: bool = False
52
+ needs_deep_gemm: bool = False
53
+
54
+
55
+PREPARE_FINALIZE_INFO: dict[mk.FusedMoEPrepareAndFinalize,
56
+ PrepareFinalizeInfo] = {}
57
+EXPERT_INFO: dict[mk.FusedMoEPermuteExpertsUnpermute, ExpertInfo] = {}
58
+MK_ALL_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []
59
+MK_MULTI_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []
60
+MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES: list[mk.FusedMoEPrepareAndFinalize] = []
61
+MK_FUSED_EXPERT_TYPES: list[mk.FusedMoEPermuteExpertsUnpermute] = []
62
+
63
+standard_format = mk.FusedMoEActivationFormat.Standard
64
+batched_format = mk.FusedMoEActivationFormat.BatchedExperts
65
+common_float_types: list[Union[torch.dtype, str]] = [
66
+ torch.float8_e4m3fn, torch.bfloat16, torch.float16, torch.float32
67
+]
68
+common_float_and_int_types = common_float_types + [torch.int8]
69
+nv_fp4_types = ["nvfp4"]
70
+fp8_types = [torch.float8_e4m3fn]
71
+
72
+
73
+def register_prepare_and_finalize(
74
+ kind,
75
+ activation_format: mk.FusedMoEActivationFormat,
76
+ supported_dtypes: list[Union[torch.dtype, str]],
77
+ blocked_quantization_support: bool,
78
+ backend: Optional[str],
79
+ force_multigpu: bool = False,
80
+ supports_apply_weight_on_input: bool = True,
81
+):
82
+ global PREPARE_FINALIZE_INFO
83
+ global MK_ALL_PREPARE_FINALIZE_TYPES
84
+ global MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
85
+ global MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
86
+ assert kind not in PREPARE_FINALIZE_INFO
87
+
88
+ PREPARE_FINALIZE_INFO[kind] = PrepareFinalizeInfo(
89
+ activation_format,
90
+ supported_dtypes,
91
+ blocked_quantization_support,
92
+ backend,
93
+ supports_apply_weight_on_input,
94
+ )
95
+ MK_ALL_PREPARE_FINALIZE_TYPES.append(kind)
96
+ if backend is not None or force_multigpu:
97
+ MK_MULTI_GPU_PREPARE_FINALIZE_TYPES.append(kind)
98
+ else:
99
+ MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES.append(kind)
100
+
101
+
102
+def register_experts(
103
+ kind,
104
+ activation_format: mk.FusedMoEActivationFormat,
105
+ supported_dtypes: list[Union[torch.dtype, str]],
106
+ blocked_quantization_support: bool,
107
+ supports_chunking: bool,
108
+ supports_expert_map: bool,
109
+ needs_matching_quant: bool = False,
110
+ needs_deep_gemm: bool = False,
111
+):
112
+ global EXPERT_INFO
113
+ global MK_FUSED_EXPERT_TYPES
114
+ assert kind not in EXPERT_INFO
115
+
116
+ EXPERT_INFO[kind] = ExpertInfo(
117
+ activation_format,
118
+ supported_dtypes,
119
+ blocked_quantization_support,
120
+ supports_chunking,
121
+ supports_expert_map,
122
+ needs_matching_quant,
123
+ needs_deep_gemm,
124
+ )
125
+
126
+ MK_FUSED_EXPERT_TYPES.append(kind)
127
+
128
+
129
+def prepare_finalize_info(kind) -> PrepareFinalizeInfo:
130
+ info = PREPARE_FINALIZE_INFO.get(kind)
131
+ assert info is not None
132
+ return info
133
+
134
+
135
+def expert_info(kind) -> ExpertInfo:
136
+ info = EXPERT_INFO.get(kind)
137
+ assert info is not None
138
+ return info
139
+
140
+
141
+register_prepare_and_finalize(
142
+ MoEPrepareAndFinalizeNoEP,
143
+ standard_format,
144
+ common_float_types,
145
+ blocked_quantization_support=True,
146
+ backend=None,
147
+)
148
+
149
+register_experts(
150
+ BatchedTritonExperts,
151
+ batched_format,
152
+ common_float_types,
153
+ blocked_quantization_support=True,
154
+ supports_chunking=False,
155
+ supports_expert_map=False,
156
+ needs_matching_quant=True,
157
+)
158
+
159
+register_experts(
160
+ TritonExperts,
161
+ standard_format,
162
+ common_float_and_int_types,
163
+ blocked_quantization_support=True,
164
+ supports_chunking=True,
165
+ supports_expert_map=True,
166
+ needs_matching_quant=True,
167
+)
168
+
169
+register_experts(
170
+ NaiveBatchedExperts,
171
+ batched_format,
172
+ common_float_and_int_types,
173
+ blocked_quantization_support=True,
174
+ supports_chunking=False,
175
+ supports_expert_map=True,
176
+)
177
+
178
+# Disable on blackwell for now
179
+if has_deep_ep() and not current_platform.has_device_capability(100):
180
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
181
DeepEPHTPrepareAndFinalize)
182
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
183
DeepEPLLPrepareAndFinalize)
184
185
+ register_prepare_and_finalize(
186
+ DeepEPHTPrepareAndFinalize,
187
+ standard_format,
188
+ common_float_types,
189
+ blocked_quantization_support=True,
190
+ backend="deepep_high_throughput",
191
+ )
192
+
193
+ register_prepare_and_finalize(
194
+ DeepEPLLPrepareAndFinalize,
195
+ batched_format,
196
+ common_float_types,
197
+ blocked_quantization_support=True,
198
+ backend="deepep_low_latency",
199
+ )
200
+
201
if has_pplx():
202
from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
203
PplxPrepareAndFinalize)
204
+ register_prepare_and_finalize(
205
+ PplxPrepareAndFinalize,
206
+ batched_format,
207
+ common_float_and_int_types,
208
+ blocked_quantization_support=True,
209
+ backend="pplx",
210
+ )
211
33
-MK_MULTI_GPU_PREPARE_FINALIZE_TYPES = []
34
-if has_pplx():
35
- MK_MULTI_GPU_PREPARE_FINALIZE_TYPES += [PplxPrepareAndFinalize]
36
-if has_deep_ep():
37
- MK_MULTI_GPU_PREPARE_FINALIZE_TYPES += [
38
- DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
39
- ]
212
+if (has_flashinfer_cutlass_fused_moe()
213
+ and current_platform.has_device_capability(100)):
214
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
215
+ FlashInferExperts)
216
+ from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
217
+ FlashInferCutlassMoEPrepareAndFinalize)
218
41
-MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES = [MoEPrepareAndFinalizeNoEP]
219
+ register_prepare_and_finalize(
220
+ FlashInferCutlassMoEPrepareAndFinalize,
221
+ standard_format,
222
+ nv_fp4_types,
223
+ blocked_quantization_support=True,
224
+ backend=None,
225
+ force_multigpu=True,
226
+ supports_apply_weight_on_input=False,
227
+ )
228
43
-MK_ALL_PREPARE_FINALIZE_TYPES = (MK_MULTI_GPU_PREPARE_FINALIZE_TYPES +
44
- MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
229
+ register_experts(
230
+ FlashInferExperts,
231
+ standard_format,
232
+ nv_fp4_types,
233
+ blocked_quantization_support=True,
234
+ supports_chunking=True,
235
+ # Note: this is a hack to get it to run for now
236
+ supports_expert_map=True,
237
+ )
238
+else:
239
+ FlashInferCutlassMoEPrepareAndFinalize = None
240
46
-MK_FUSED_EXPERT_TYPES = [
47
- BatchedDeepGemmExperts,
48
- BatchedTritonExperts,
49
- NaiveBatchedExperts,
50
- BatchedTritonOrDeepGemmExperts,
51
- CutlassExpertsFp8,
52
- DeepGemmExperts,
53
- TritonOrDeepGemmExperts,
54
- TritonExperts,
55
-]
241
+if has_deep_gemm() and is_deep_gemm_supported():
242
+ register_experts(
243
+ BatchedDeepGemmExperts,
244
+ batched_format,
245
+ fp8_types,
246
+ blocked_quantization_support=True,
247
+ supports_chunking=False,
248
+ supports_expert_map=False,
249
+ needs_matching_quant=False,
250
+ needs_deep_gemm=True,
251
+ )
252
+ register_experts(
253
+ DeepGemmExperts,
254
+ standard_format,
255
+ fp8_types,
256
+ blocked_quantization_support=True,
257
+ supports_chunking=True,
258
+ supports_expert_map=True,
259
+ needs_matching_quant=False,
260
+ needs_deep_gemm=True,
261
+ ),
262
+ register_experts(
263
+ BatchedTritonOrDeepGemmExperts,
264
+ batched_format,
265
+ common_float_and_int_types,
266
+ blocked_quantization_support=True,
267
+ supports_chunking=False,
268
+ supports_expert_map=False,
269
+ needs_matching_quant=True,
270
+ needs_deep_gemm=True,
271
+ )
272
+ register_experts(
273
+ TritonOrDeepGemmExperts,
274
+ standard_format,
275
+ common_float_and_int_types,
276
+ blocked_quantization_support=True,
277
+ supports_chunking=True,
278
+ supports_expert_map=True,
279
+ needs_matching_quant=True,
280
+ needs_deep_gemm=True,
281
+ )
282
+
283
+if cutlass_fp8_supported():
284
+ from vllm.model_executor.layers.fused_moe import (CutlassBatchedExpertsFp8,
285
+ CutlassExpertsFp8)
286
+ register_experts(
287
+ CutlassExpertsFp8,
288
+ standard_format,
289
+ fp8_types,
290
+ blocked_quantization_support=False,
291
+ supports_chunking=True,
292
+ supports_expert_map=False,
293
+ )
294
+ register_experts(
295
+ CutlassBatchedExpertsFp8,
296
+ batched_format,
297
+ fp8_types,
298
+ blocked_quantization_support=False,
299
+ supports_chunking=False,
300
+ supports_expert_map=False,
301
+ )
302
+
303
+if cutlass_fp4_supported():
304
+ from vllm.model_executor.layers.fused_moe.cutlass_moe import (
305
+ CutlassExpertsFp4)
306
+ register_experts(
307
+ CutlassExpertsFp4,
308
+ standard_format,
309
+ nv_fp4_types,
310
+ blocked_quantization_support=True,
311
+ supports_chunking=True,
312
+ supports_expert_map=False,
313
+ )
314
315
MK_QUANT_CONFIGS = [
316
None,
@@ -85,3 +343,156 @@ MK_QUANT_CONFIGS = [
343
# block-quantized weights and per-token activations
344
# block-quantized weights and per-tensor activations
345
]
346
+
347
+if cutlass_fp4_supported() or has_flashinfer_cutlass_fused_moe():
348
+ MK_QUANT_CONFIGS += [
349
+ FusedMoEQuantConfig(quant_dtype="nvfp4",
350
+ per_out_ch_quant=False,
351
+ per_act_token_quant=False,
352
+ block_shape=None),
353
+ ]
354
+
355
+
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
+
362
+def make_prepare_finalize(
363
+ prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
364
+ backend: Optional[str],
365
+ moe: FusedMoEConfig,
366
+) -> mk.FusedMoEPrepareAndFinalize:
367
+ if backend != "naive" and backend is not None:
368
+ prepare_finalize = FusedMoEMethodBase._maybe_make_prepare_finalize(moe)
369
+ assert prepare_finalize is not None
370
+ return prepare_finalize
371
+ elif prepare_finalize_type == FlashInferCutlassMoEPrepareAndFinalize:
372
+ return FlashInferCutlassMoEPrepareAndFinalize(
373
+ use_dp=moe.moe_parallel_config.dp_size > 1,
374
+ a1_gscale=_make_gscale(moe.num_local_experts),
375
+ )
376
+ else:
377
+ return MoEPrepareAndFinalizeNoEP()
378
+
379
+
380
+def _slice(rank: int, num_local_experts: int, t: torch.Tensor) -> torch.Tensor:
381
+ s = rank * num_local_experts
382
+ e = s + num_local_experts
383
+ return t[s:e]
384
+
385
+
386
+def make_fused_experts(
387
+ fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
388
+ moe: FusedMoEConfig,
389
+ num_dispatchers: int,
390
+ w1_gs: Optional[torch.Tensor],
391
+ w2_gs: Optional[torch.Tensor],
392
+) -> mk.FusedMoEPermuteExpertsUnpermute:
393
+
394
+ use_fp8 = moe.quant_dtype == torch.float8_e4m3fn
395
+ batch_kwargs = {
396
+ "max_num_tokens": moe.max_num_tokens,
397
+ "num_dispatchers": num_dispatchers,
398
+ }
399
+ 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,
406
+ }
407
+ deepgemm_kwargs = {"allow_deep_gemm": has_deep_gemm()}
408
+
409
+ 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
+ }
414
+ print(f"Making BatchedDeepGemmExperts {kwargs} ...")
415
+ experts = BatchedDeepGemmExperts(**kwargs)
416
+ elif fused_experts_type == BatchedTritonExperts:
417
+ kwargs = batch_kwargs | quant_kwargs
418
+ print(f"Making BatchedTritonExperts {kwargs} ...")
419
+ experts = BatchedTritonExperts(**kwargs)
420
+ elif fused_experts_type == BatchedTritonOrDeepGemmExperts:
421
+ kwargs = batch_kwargs | quant_kwargs | deepgemm_kwargs
422
+ print(f"Making BatchedTritonOrDeepGemmExperts {kwargs} ...")
423
+ experts = BatchedTritonOrDeepGemmExperts(**kwargs)
424
+ elif fused_experts_type == DeepGemmExperts:
425
+ print("Making DeepGemmExperts () ...")
426
+ experts = DeepGemmExperts()
427
+ elif fused_experts_type == TritonExperts:
428
+ kwargs = quant_kwargs
429
+ print(f"Making TritonExperts {kwargs} ...")
430
+ experts = TritonExperts(**kwargs)
431
+ elif fused_experts_type == TritonOrDeepGemmExperts:
432
+ kwargs = quant_kwargs | deepgemm_kwargs
433
+ print(f"Making TritonOrDeepGemmExperts {kwargs} ...")
434
+ experts = TritonOrDeepGemmExperts(**kwargs)
435
+ elif fused_experts_type == NaiveBatchedExperts:
436
+ kwargs = batch_kwargs | quant_kwargs
437
+ print(f"Making NaiveBatchedExperts {kwargs} ...")
438
+ experts = NaiveBatchedExperts(**kwargs)
439
+ elif fused_experts_type == CutlassExpertsFp8:
440
+ kwargs = {
441
+ "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
+ }
446
+ print(f"Making CutlassExpertsFp8 {kwargs} ...")
447
+ experts = CutlassExpertsFp8(**kwargs)
448
+ elif fused_experts_type == CutlassBatchedExpertsFp8:
449
+ kwargs = {
450
+ "max_experts_per_worker": moe.num_local_experts,
451
+ "num_dispatchers": num_dispatchers,
452
+ "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
+ }
457
+ print(f"Making CutlassBatchedExpertsFp8 {kwargs} ...")
458
+ experts = CutlassBatchedExpertsFp8(**kwargs)
459
+ 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
463
+ 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,
473
+ "num_dispatchers": num_dispatchers,
474
+ }
475
+ print(f"Making CutlassExpertsFp4 {kwargs} ...")
476
+ experts = CutlassExpertsFp4(**kwargs)
477
+ 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
481
+ 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),
486
+ "out_dtype": moe.in_dtype,
487
+ "quant_dtype": "nvfp4",
488
+ "ep_rank": moe.ep_rank,
489
+ "ep_size": moe.ep_size,
490
+ "tp_rank": moe.tp_rank,
491
+ "tp_size": moe.tp_size,
492
+ }
493
+ print(f"Making FlashInferExperts {kwargs} ...")
494
+ experts = FlashInferExperts(**kwargs)
495
+ else:
496
+ raise RuntimeError(f"Unknown fused experts type: {fused_experts_type}")
497
+
498
+ return experts
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py
+2
-2
@@ -52,7 +52,7 @@ def profile_modular_kernel(
52
rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
53
54
# make modular kernel
55
- mk = make_modular_kernel(config, vllm_config)
55
+ mk = make_modular_kernel(config, vllm_config, weights)
56
57
mk_kwargs = {
58
"hidden_states": rank_tensors.hidden_states,
@@ -83,7 +83,7 @@ def rank_worker(
83
# sanity check
84
from vllm import envs
85
if config.fused_moe_chunk_size is not None:
86
- assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
86
+ assert config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE
87
88
# get weights to this device
89
weights.to_current_device()
tests/kernels/moe/modular_kernel_tools/utils.py
deleted
-117
@@ -1,117 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import torch
5
-
6
-import vllm._custom_ops as ops
7
-from vllm.utils.deep_gemm import per_block_cast_to_fp8
8
-
9
-
10
-def per_token_cast_to_fp8(
11
- x: torch.Tensor, block_size: int) -> tuple[torch.Tensor, torch.Tensor]:
12
- assert x.dim() == 2
13
- m, n = x.shape
14
- pad_size = (block_size - (n % block_size)) % block_size
15
- x = torch.nn.functional.pad(x,
16
- (0, pad_size), value=0) if pad_size > 0 else x
17
- x_view = x.view(m, -1, block_size)
18
- x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
19
- fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
20
- return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
21
-
22
-
23
-def make_non_quant_weights(
24
- e: int,
25
- n: int,
26
- k: int,
27
- dtype: torch.dtype,
28
-) -> tuple[torch.Tensor, torch.Tensor]:
29
- """
30
- Return weights w1, w2
31
- """
32
- device = torch.cuda.current_device()
33
- w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 15
34
- w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 15
35
- return w1, w2
36
-
37
-
38
-def make_block_quant_fp8_weights(
39
- e: int,
40
- n: int,
41
- k: int,
42
- block_size: list[int],
43
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
44
- """
45
- Return weights w1, w2, w1_scale, w2_scale
46
- """
47
- dtype = torch.bfloat16
48
- device = torch.cuda.current_device()
49
-
50
- fp8_info = torch.finfo(torch.float8_e4m3fn)
51
- fp8_max, fp8_min = fp8_info.max, fp8_info.min
52
-
53
- w1_bf16, w2_bf16 = make_non_quant_weights(e, n, k, dtype)
54
- w1_bf16 = w1_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
55
- w2_bf16 = w2_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
56
-
57
- block_n, block_k = block_size[0], block_size[1]
58
- n_tiles_w1 = ((2 * n) + block_n - 1) // block_n
59
- k_tiles_w1 = (k + block_k - 1) // block_k
60
- n_tiles_w2 = (k + block_n - 1) // block_n
61
- k_tiles_w2 = (n + block_k - 1) // block_k
62
-
63
- w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn, device=device)
64
- w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn, device=device)
65
-
66
- w1_s = torch.empty((e, n_tiles_w1, k_tiles_w1),
67
- device=device,
68
- dtype=torch.float32)
69
- w2_s = torch.empty((e, n_tiles_w2, k_tiles_w2),
70
- device=device,
71
- dtype=torch.float32)
72
-
73
- assert w1_s.shape == (e, (2 * n + (block_n - 1)) // block_n,
74
- (k + (block_k - 1)) // block_k)
75
- assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
76
-
77
- for i in range(e):
78
- w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i],
79
- block_size=[block_k, block_n])
80
- w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i],
81
- block_size=[block_k, block_n])
82
-
83
- return w1, w2, w1_s, w2_s
84
-
85
-
86
-def make_quant_fp8_weights(
87
- e: int,
88
- n: int,
89
- k: int,
90
- per_out_channel_quant: bool,
91
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
92
- """
93
- Return w1, w2, w1_scale, w2_scale
94
- """
95
- q_dtype = torch.float8_e4m3fn
96
-
97
- w1, w2 = make_non_quant_weights(e, n, k, dtype=torch.bfloat16)
98
-
99
- # w1 -> w1_q, w2 -> w2_q
100
- w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=q_dtype)
101
- w2_q = torch.empty((e, k, n), device="cuda", dtype=q_dtype)
102
-
103
- n_b_scales = 2 * n if per_out_channel_quant else 1
104
- k_b_scales = k if per_out_channel_quant else 1
105
- w1_scale = torch.empty((e, n_b_scales, 1),
106
- device="cuda",
107
- dtype=torch.float32)
108
- w2_scale = torch.empty((e, k_b_scales, 1),
109
- device="cuda",
110
- dtype=torch.float32)
111
-
112
- for expert in range(e):
113
- w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
114
- w1[expert], use_per_token_if_dynamic=per_out_channel_quant)
115
- w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
116
- w2[expert], use_per_token_if_dynamic=per_out_channel_quant)
117
- return w1_q, w2_q, w1_scale, w2_scale
tests/kernels/moe/test_batched_moe.py
+2
-2
@@ -133,7 +133,7 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
133
per_act_token_quant=per_act_token_quant,
134
)
135
136
- B, B_q, B_scale, _, _, _ = make_test_weights(
136
+ (B, B_q, B_scale, _), _ = make_test_weights(
137
num_experts,
138
N // 2,
139
K,
@@ -243,7 +243,7 @@ def test_fused_moe_batched_experts(
243
act_dtype = dtype
244
quant_dtype = None
245
246
- w1_16, w1, w1_s, w2_16, w2, w2_s = make_test_weights(
246
+ (w1_16, w1, w1_s, _), (w2_16, w2, w2_s, _) = make_test_weights(
247
e,
248
n,
249
k,
tests/kernels/moe/test_block_fp8.py
+17
-14
@@ -161,18 +161,20 @@ 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 = make_test_weights(E,
165
- N,
166
- K,
167
- dtype,
168
- torch.float8_e4m3fn,
169
- per_act_token_quant=False,
170
- block_shape=block_size)
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)
180
@@ -247,13 +249,14 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
249
a = torch.randn((M, K), dtype=dtype) / 10
250
score = torch.randn((M, E), dtype=dtype)
251
250
- _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
251
- N,
252
- K,
253
- dtype,
254
- torch.float8_e4m3fn,
255
- per_act_token_quant=False,
256
- block_shape=block_size)
252
+ (_, w1, w1_s, _), (_, w2, w2_s,
253
+ _) = make_test_weights(E,
254
+ N,
255
+ K,
256
+ dtype,
257
+ torch.float8_e4m3fn,
258
+ per_act_token_quant=False,
259
+ block_shape=block_size)
260
261
# Note: for now use_compile will error out if the problem size is
262
# large enough to trigger chunking. I'm leaving the flag and
tests/kernels/moe/test_block_int8.py
+8
-7
@@ -118,13 +118,14 @@ def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
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 = make_test_weights(E,
122
- N,
123
- K,
124
- dtype,
125
- torch.int8,
126
- per_act_token_quant=False,
127
- block_shape=block_size)
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)
129
130
# Set the context to avoid lots of warning spam.
131
with set_current_vllm_config(vllm_config):
tests/kernels/moe/test_cutlass_grouped_gemm.py
+2
-15
@@ -9,6 +9,7 @@ import random
9
import pytest
10
import torch
11
12
+from tests.kernels.moe.utils import per_token_cast_to_fp8
13
from tests.kernels.utils import baseline_scaled_mm
14
from vllm import _custom_ops as ops
15
from vllm.platforms import current_platform
@@ -16,20 +17,6 @@ from vllm.utils import cdiv
17
from vllm.utils.deep_gemm import per_block_cast_to_fp8
18
19
19
-def per_token_cast_to_fp8(
20
- x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
21
- assert x.dim() == 2
22
- m, n = x.shape
23
- pad_size = (128 - (n % 128)) % 128
24
- x = torch.nn.functional.pad(x,
25
- (0, pad_size), value=0) if pad_size > 0 else x
26
- x_view = x.view(m, -1, 128)
27
- x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
28
- fp8_data = (x_view *
29
- (448.0 / x_amax.unsqueeze(2))).to(dtype=torch.float8_e4m3fn)
30
- return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
31
-
32
-
20
@pytest.mark.parametrize("num_groups, expected_m_per_group, k, n", [
21
(4, 8192, 7168, 4096),
22
(4, 8192, 2048, 7168),
@@ -76,7 +63,7 @@ def test_cutlass_grouped_gemm(
63
device=device,
64
dtype=torch.float))
65
for i in range(num_groups):
79
- y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i])
66
+ y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i], [128, 128])
67
68
for i in range(num_groups):
69
a = x_fp8[0][ep_offset[i]:ep_offset[i + 1]]
tests/kernels/moe/test_deepep_deepgemm_moe.py
+4
-2
@@ -70,8 +70,10 @@ def make_block_quant_fp8_weights(
70
"""
71
Return weights w1q, w2q, w1_scale, w2_scale
72
"""
73
- w1, w1q, w1_scale, w2, w2q, w2_scale = make_test_weights(
74
- e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_size)
73
+ (_, w1q, w1_scale, _), (_, w2q, w2_scale,
74
+ _) = make_test_weights(e, n, k, torch.bfloat16,
75
+ torch.float8_e4m3fn,
76
+ block_size)
77
return w1q, w2q, w1_scale, w2_scale
78
79
tests/kernels/moe/test_deepgemm.py
+3
-3
@@ -132,9 +132,9 @@ def run_single_case(m, n, k, topk, num_experts, block_size):
132
# Note: W1 has shape (E, 2N, K), so N = 512
133
# can trigger the deepgemm path.
134
MNKs = [
135
- (1024, 512, 128),
136
- (1024, 512, 512),
137
- (2048, 512, 512),
135
+ (1024, 768, 128),
136
+ (1024, 768, 512),
137
+ (2048, 768, 512),
138
(512, 1024, 1024),
139
(512, 2048, 2048),
140
(4096, 4096, 1024),
tests/kernels/moe/test_flashinfer_moe.py
new
+147
@@ -0,0 +1,147 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+
6
+from tests.kernels.moe.utils import make_test_weights
7
+from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
8
+ FLOAT8_E4M3_MAX,
9
+ dequantize_nvfp4_to_dtype)
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.flashinfer_cutlass_moe import (
14
+ FlashInferExperts, is_valid_flashinfer_cutlass_fused_moe)
15
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
16
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
17
+ FusedMoEModularKernel)
18
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
19
+ MoEPrepareAndFinalizeNoEP)
20
+from vllm.platforms import current_platform
21
+from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
22
+
23
+if not has_flashinfer_cutlass_fused_moe(
24
+) or not current_platform.has_device_capability(100):
25
+ pytest.skip("Requires flashinfer_cutlass_fused_moe and nvfp4 support",
26
+ allow_module_level=True)
27
+
28
+MNK_FACTORS = [
29
+ (2, 1024, 1024),
30
+ (2, 1024, 1536),
31
+ (2, 3072, 1024),
32
+ (2, 3072, 1536),
33
+ (64, 1024, 1024),
34
+ (64, 1024, 1536),
35
+ (64, 3072, 1024),
36
+ (64, 2048, 1536),
37
+ (224, 1024, 1024),
38
+ (224, 1024, 1536),
39
+]
40
+
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])
45
+@pytest.mark.parametrize("topk", [1, 6, 8])
46
+@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
47
+@torch.inference_mode()
48
+def test_flashinfer_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
49
+ dtype: torch.dtype):
50
+ current_platform.seed_everything(7)
51
+ with set_current_vllm_config(
52
+ VllmConfig(parallel_config=ParallelConfig(
53
+ pipeline_parallel_size=1))):
54
+
55
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
56
+
57
+ quant_blocksize = 16
58
+
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
+ )
69
+
70
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
71
+ topk_weights, topk_ids, _ = fused_topk(a,
72
+ score,
73
+ topk,
74
+ renormalize=False)
75
+
76
+ a1_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
77
+ a2_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
78
+
79
+ assert is_valid_flashinfer_cutlass_fused_moe(a, w1_q, w2_q)
80
+
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
+
86
+ flashinfer_experts = FusedMoEModularKernel(
87
+ 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
+ ))
96
+
97
+ flashinfer_output = flashinfer_experts(
98
+ hidden_states=a,
99
+ w1=w1_q,
100
+ w1_scale=w1_blockscale,
101
+ w2=w2_q,
102
+ w2_scale=w2_blockscale,
103
+ a1_scale=a1_gs,
104
+ a2_scale=a2_gs,
105
+ topk_weights=topk_weights,
106
+ topk_ids=topk_ids,
107
+ )
108
+
109
+ # Reference check:
110
+ a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
111
+ torch.amax(a.flatten(), dim=-1)).to(torch.float32)
112
+ a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
113
+ _, m_k = a_fp4.shape
114
+ a_in_dtype = dequantize_nvfp4_to_dtype(a_fp4,
115
+ a_scale_interleaved,
116
+ a_global_scale,
117
+ dtype=a.dtype,
118
+ device=a.device,
119
+ block_size=quant_blocksize)
120
+
121
+ w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)
122
+ w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
123
+
124
+ 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)
137
+
138
+ torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk)
139
+
140
+ torch.testing.assert_close(torch_output,
141
+ flashinfer_output,
142
+ atol=1e-1,
143
+ rtol=1e-1)
144
+
145
+
146
+if __name__ == "__main__":
147
+ test_flashinfer_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half)
tests/kernels/moe/test_modular_kernel_combinations.py
+88
-53
@@ -2,6 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import copy
5
+import textwrap
6
+import traceback
7
from itertools import product
8
from typing import Optional
9
@@ -10,41 +12,51 @@ 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
13
-from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
14
- BatchedTritonOrDeepGemmExperts)
15
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
16
-from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
17
-from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
18
- BatchedTritonExperts)
19
-from vllm.model_executor.layers.fused_moe.layer import TritonExperts
20
-from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
21
- TritonOrDeepGemmExperts)
16
from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
17
+from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
18
19
from .modular_kernel_tools.common import (Config, RankTensors, WeightTensors,
20
reference_moe_impl,
21
run_modular_kernel)
22
from .modular_kernel_tools.mk_objects import (
23
MK_FUSED_EXPERT_TYPES, MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
29
- MK_QUANT_CONFIGS, MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
24
+ MK_QUANT_CONFIGS, MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, expert_info)
25
from .modular_kernel_tools.parallel_utils import (ProcessGroupInfo,
26
parallel_launch_with_config)
27
33
-# TODO (varun): These requirements are very strict and could be relaxed.
34
-has_all_packages = (has_deep_ep() and has_deep_gemm() and has_pplx())
28
+has_any_multi_gpu_package = (has_deep_ep() or has_deep_gemm() or has_pplx()
29
+ or has_flashinfer_cutlass_fused_moe())
30
36
-meets_package_requirements = pytest.mark.skipif(
37
- not has_all_packages,
38
- reason="Requires deep_ep & deep_gemm & pplx packages",
31
+meets_multi_gpu_requirements = pytest.mark.skipif(
32
+ not has_any_multi_gpu_package,
33
+ reason="Requires deep_ep or deep_gemm or pplx or flashinfer packages",
34
)
35
36
37
+def format_result(verbose, msg, ex=None):
38
+ if ex is not None:
39
+ x = str(ex)
40
+ newx = x.strip(" \n\t")[:16]
41
+ if len(newx) < len(x):
42
+ newx = newx + " ..."
43
+
44
+ prefix = "E\t"
45
+ print(f"{textwrap.indent(traceback.format_exc(), prefix)}")
46
+ print(f"FAILED {msg} - {newx}\n")
47
+ elif verbose:
48
+ print(f"PASSED {msg}")
49
+ else:
50
+ print(".", end="")
51
+
52
+
53
def rank_worker(
54
pgi: ProcessGroupInfo,
55
vllm_config: VllmConfig,
56
cpu_group,
57
config: Config,
58
weights: WeightTensors,
59
+ verbose: bool,
60
):
61
current_platform.seed_everything(pgi.rank)
62
@@ -61,39 +73,64 @@ def rank_worker(
73
TOPKs = config.topks
74
assert isinstance(TOPKs, list)
75
64
- for m, topk in product(Ms, TOPKs):
65
- print(f"Running m={m}, topk={topk} ...")
66
- # override m and topk
67
- cfgx = copy.deepcopy(config)
68
- cfgx.Ms = m
69
- cfgx.topks = topk
70
-
71
- # inputs for rank
72
- rank_tensors = RankTensors.make(cfgx, pgi)
73
-
74
- # modular kernel out
75
- mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights,
76
- rank_tensors)
76
+ exceptions = []
77
+ count = 0
78
78
- with set_current_vllm_config(vllm_config):
79
- ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
80
-
81
- torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
82
-
83
-
84
-def run(config: Config):
79
+ for m, topk in product(Ms, TOPKs):
80
+ try:
81
+ print(f"Running[{pgi.rank}]: m={m}, topk={topk} ...")
82
+ count = count + 1
83
+ # override m and topk
84
+ cfgx = copy.deepcopy(config)
85
+ cfgx.Ms = m
86
+ cfgx.topks = topk
87
+
88
+ # inputs for rank
89
+ rank_tensors = RankTensors.make(cfgx, pgi)
90
+
91
+ # modular kernel out
92
+ mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights,
93
+ rank_tensors)
94
+
95
+ with set_current_vllm_config(vllm_config):
96
+ ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
97
+
98
+ if config.quant_dtype == "nvfp4":
99
+ atol = 1e-1
100
+ rtol = 1e-1
101
+ else:
102
+ atol = 3e-2
103
+ rtol = 3e-2
104
+
105
+ torch.testing.assert_close(ref_out, mk_out, atol=atol, rtol=rtol)
106
+ format_result(verbose, config.describe())
107
+ except Exception as ex:
108
+ format_result(verbose, config.describe(), ex)
109
+ exceptions.append(ex)
110
+
111
+ if len(exceptions) > 0:
112
+ raise RuntimeError(
113
+ f"{len(exceptions)} of {count} tests failed in child process, "
114
+ f"rank={pgi.rank}.")
115
+ else:
116
+ print(f"{count} of {count} tests passed in child process, "
117
+ f"rank={pgi.rank}.")
118
+
119
+
120
+def run(config: Config, verbose: bool):
121
assert config.is_valid()
86
- print(f"Testing config \n{config.describe()} ...")
122
123
weights: WeightTensors = WeightTensors.make(config)
124
125
vllm_config, env_dict = config.make_env_data()
126
parallel_launch_with_config(config.world_size, rank_worker, vllm_config,
92
- env_dict, config, weights)
127
+ env_dict, config, weights, verbose)
128
129
130
Ms = [32, 64]
96
-Ks = [7168] # hidden sizes
131
+# hidden sizes, making this too large will cause fp4 tests to fail.
132
+# Also needs to be a multiple of 1024 for deep_gemm.
133
+Ks = [2048]
134
Ns = [2048]
135
TOPKs = [4, 1]
136
Es = [32]
@@ -103,19 +140,16 @@ FUSED_MOE_CHUNK_SIZEs = [None, 16]
140
141
def is_nyi_config(config: Config) -> bool:
142
# We know these configs to be legitimate. but still fail.
143
+ info = expert_info(config.fused_experts_type)
144
107
- if (config.fused_experts_type in [
108
- BatchedTritonExperts, BatchedTritonOrDeepGemmExperts,
109
- TritonExperts, TritonOrDeepGemmExperts
110
- ]):
145
+ if info.needs_matching_quant:
146
# The triton kernels expect both per-act-token-quant and
147
# per-out-ch-quant or neither.
148
unsupported_quant_config = ((config.is_per_act_token_quant +
149
config.is_per_out_ch_quant) == 1)
150
return unsupported_quant_config
151
117
- # cutlass kernels dont support expert_maps yet.
118
- return config.fused_experts_type == CutlassExpertsFp8
152
+ return not info.supports_expert_map
153
154
155
@pytest.mark.parametrize("k", Ks)
@@ -128,13 +162,13 @@ def is_nyi_config(config: Config) -> bool:
162
product(MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))
163
@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
164
@pytest.mark.parametrize("world_size", [2])
131
-@meets_package_requirements
165
+@meets_multi_gpu_requirements
166
def test_modular_kernel_combinations_multigpu(
167
k: int, n: int, e: int, dtype: torch.dtype,
134
- quant_config: FusedMoEQuantConfig,
168
+ quant_config: Optional[FusedMoEQuantConfig],
169
combination: tuple[mk.FusedMoEPrepareAndFinalize,
170
mk.FusedMoEPermuteExpertsUnpermute],
137
- fused_moe_chunk_size: Optional[int], world_size: int):
171
+ fused_moe_chunk_size: Optional[int], world_size: int, pytestconfig):
172
173
config = Config(
174
Ms=Ms,
@@ -149,14 +183,15 @@ def test_modular_kernel_combinations_multigpu(
183
fused_moe_chunk_size=fused_moe_chunk_size,
184
world_size=world_size,
185
)
186
+
187
if not config.is_valid():
188
pytest.skip(f"Tests config {config} is not valid. Skipping ...")
189
190
if is_nyi_config(config):
191
pytest.skip(f"Tests config {config} is nyi. Skipping ...")
192
158
- print(f"{config.describe()}")
159
- run(config)
193
+ verbosity = pytestconfig.getoption('verbose')
194
+ run(config, verbosity > 0)
195
196
197
@pytest.mark.parametrize("k", Ks)
@@ -169,13 +204,12 @@ def test_modular_kernel_combinations_multigpu(
204
product(MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))
205
@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
206
@pytest.mark.parametrize("world_size", [1])
172
-@meets_package_requirements
207
def test_modular_kernel_combinations_singlegpu(
208
k: int, n: int, e: int, dtype: torch.dtype,
175
- quant_config: FusedMoEQuantConfig,
209
+ quant_config: Optional[FusedMoEQuantConfig],
210
combination: tuple[mk.FusedMoEPrepareAndFinalize,
211
mk.FusedMoEPermuteExpertsUnpermute],
178
- fused_moe_chunk_size: Optional[int], world_size: int):
212
+ fused_moe_chunk_size: Optional[int], world_size: int, pytestconfig):
213
config = Config(
214
Ms=Ms,
215
K=k,
@@ -196,7 +230,8 @@ def test_modular_kernel_combinations_singlegpu(
230
if is_nyi_config(config):
231
pytest.skip(f"Tests config {config} is nyi. Skipping ...")
232
199
- run(config)
233
+ verbosity = pytestconfig.getoption('verbose')
234
+ run(config, verbosity > 0)
235
236
237
if __name__ == '__main__':
@@ -211,4 +246,4 @@ if __name__ == '__main__':
246
args = parser.parse_args()
247
config = make_config(args)
248
214
- run(config)
249
+ run(config, True)
tests/kernels/moe/test_nvfp4_moe.py
+24
-40
@@ -3,6 +3,7 @@
3
import pytest
4
import torch
5
6
+from tests.kernels.moe.utils import make_test_weights
7
from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
8
FLOAT8_E4M3_MAX,
9
dequantize_nvfp4_to_dtype)
@@ -43,41 +44,20 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
44
VllmConfig(parallel_config=ParallelConfig(
45
pipeline_parallel_size=1))):
46
46
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
47
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
47
quant_blocksize = 16
49
- round_up = lambda x, y: (x + y - 1) // y * y
50
- sf_w1_2n = round_up(2 * n, 128)
51
- sf_w1_k = round_up(k // quant_blocksize, 4)
52
- w1_blockscale = torch.empty((e, sf_w1_2n, sf_w1_k),
53
- device="cuda",
54
- dtype=torch.float8_e4m3fn)
55
-
56
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
57
- sf_w2_k = round_up(k, 128)
58
- sf_w2_n = round_up(n // quant_blocksize, 4)
59
- w2_blockscale = torch.empty((e, sf_w2_k, sf_w2_n),
60
- device="cuda",
61
- dtype=torch.float8_e4m3fn)
62
-
63
- w1_q = torch.empty((e, 2 * n, k // 2),
64
- device="cuda",
65
- dtype=torch.uint8)
66
- w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8)
67
- w1_gs = torch.empty((e, ), device="cuda", dtype=torch.float32)
68
- w2_gs = torch.empty((e, ), device="cuda", dtype=torch.float32)
69
-
70
- for expert in range(e):
71
- w1_amax = torch.abs(w1).max().to(torch.float32)
72
- w2_amax = torch.abs(w2).max().to(torch.float32)
73
- w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
74
- w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
75
-
76
- w1_q[expert], w1_blockscale[expert] = ops.scaled_fp4_quant(
77
- w1[expert], w1_gs[expert])
78
-
79
- w2_q[expert], w2_blockscale[expert] = ops.scaled_fp4_quant(
80
- w2[expert], w2_gs[expert])
48
+
49
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
50
+
51
+ (_, w1_q, w1_blockscale,
52
+ w1_gs), (_, w2_q, w2_blockscale, w2_gs) = make_test_weights(
53
+ e,
54
+ n,
55
+ k,
56
+ in_dtype=dtype,
57
+ quant_dtype="nvfp4",
58
+ block_shape=None, # use quant_blocksize?
59
+ per_act_token_quant=False,
60
+ )
61
62
score = torch.randn((m, e), device="cuda", dtype=dtype)
63
topk_weights, topk_ids, _ = fused_topk(a,
@@ -88,6 +68,11 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
68
a1_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
69
a2_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
70
71
+ assert w1_gs is not None
72
+ assert w2_gs is not None
73
+ assert w1_blockscale is not None
74
+ assert w2_blockscale is not None
75
+
76
cutlass_output = cutlass_moe_fp4(
77
a=a,
78
a1_gscale=a1_gs,
@@ -104,14 +89,13 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
89
n=n,
90
k=k,
91
e=e,
107
- device=a.device,
92
)
93
94
# Reference check:
95
a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
96
torch.amax(a.flatten(), dim=-1)).to(torch.float32)
97
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
114
- _, m_k = a_fp4.shape
98
+
99
a_in_dtype = dequantize_nvfp4_to_dtype(a_fp4,
100
a_scale_interleaved,
101
a_global_scale,
@@ -126,14 +110,14 @@ def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
110
w1_d[idx] = dequantize_nvfp4_to_dtype(w1_q[idx],
111
w1_blockscale[idx],
112
w1_gs[idx],
129
- dtype=w1.dtype,
130
- device=w1.device,
113
+ dtype=dtype,
114
+ device=w1_q.device,
115
block_size=quant_blocksize)
116
w2_d[idx] = dequantize_nvfp4_to_dtype(w2_q[idx],
117
w2_blockscale[idx],
118
w2_gs[idx],
135
- dtype=w2.dtype,
136
- device=w2.device,
119
+ dtype=dtype,
120
+ device=w2_q.device,
121
block_size=quant_blocksize)
122
123
torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk)
tests/kernels/moe/test_pplx_cutlass_moe.py
+4
-7
@@ -9,7 +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.cutlass_moe import CutlassExpertsFp8
12
+from vllm.model_executor.layers.fused_moe.cutlass_moe import (
13
+ CutlassBatchedExpertsFp8)
14
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
15
from vllm.model_executor.layers.fused_moe.modular_kernel import (
16
FusedMoEModularKernel)
@@ -123,12 +124,8 @@ def pplx_cutlass_moe(
124
num_local_experts=num_local_experts,
125
num_dispatchers=num_dispatchers)
126
126
- experts = CutlassExpertsFp8(num_local_experts,
127
- out_dtype,
128
- per_act_token,
129
- per_out_ch,
130
- num_dispatchers=num_dispatchers,
131
- use_batched_format=True)
127
+ experts = CutlassBatchedExpertsFp8(num_local_experts, num_dispatchers,
128
+ out_dtype, per_act_token, per_out_ch)
129
130
fused_cutlass_experts = FusedMoEModularKernel(
131
prepare_finalize,
tests/kernels/moe/test_pplx_moe.py
+2
-2
@@ -770,7 +770,7 @@ def test_pplx_moe_slow(
770
a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
771
score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
772
773
- _, w1, w1_s, _, w2, w2_s = make_test_weights(
773
+ (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
774
e,
775
n,
776
k,
@@ -836,7 +836,7 @@ def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
836
837
args = dict()
838
if make_weights:
839
- _, w1, w1_s, _, w2, w2_s = make_test_weights(
839
+ (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
840
e,
841
n,
842
k,
tests/kernels/moe/utils.py
+56
-19
@@ -1,11 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Optional
3
+from typing import Optional, Union
4
5
import torch
6
7
import vllm._custom_ops as ops
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.fused_moe import fused_experts
12
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
13
BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
@@ -169,28 +171,41 @@ def make_quantized_test_activations(
171
def moe_quantize_weights(
172
w: torch.Tensor,
173
w_s: Optional[torch.Tensor],
172
- quant_dtype: Optional[torch.dtype],
174
+ quant_dtype: Union[torch.dtype, str, None],
175
per_token_quant: bool,
176
block_shape: Optional[list[int]],
175
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
176
- assert (quant_dtype == torch.float8_e4m3fn
177
- or quant_dtype == torch.int8), "only fp8/int8 supported"
177
+) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
178
+ assert (quant_dtype == torch.float8_e4m3fn or quant_dtype == torch.int8
179
+ or quant_dtype == "nvfp4"), "only fp8/int8/nvfp4 supported"
180
+
181
+ w_gs = None
182
183
if block_shape is not None:
184
assert not per_token_quant
185
if quant_dtype == torch.int8:
186
w, w_s = per_block_cast_to_int8(w, block_shape)
183
- else:
187
+ elif quant_dtype == torch.float8_e4m3fn:
188
w, w_s = per_block_cast_to_fp8(w, block_shape)
189
+ elif quant_dtype == "nvfp4":
190
+ raise RuntimeError("blocked quantization not supported for nvfp4")
191
+ else:
192
+ raise RuntimeError(f"Unsupported quant type {quant_dtype}")
193
else:
194
if quant_dtype == torch.int8:
195
w, w_s = ops.scaled_int8_quant(
196
w, w_s, use_per_token_if_dynamic=per_token_quant)
189
- else:
197
+ elif quant_dtype == torch.float8_e4m3fn:
198
w, w_s = ops.scaled_fp8_quant(
199
w, w_s, use_per_token_if_dynamic=per_token_quant)
200
+ elif quant_dtype == "nvfp4":
201
+ assert not per_token_quant
202
+ w_amax = torch.abs(w).max().to(torch.float32)
203
+ w_gs = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w_amax
204
+ w, w_s = ops.scaled_fp4_quant(w, w_gs)
205
+ else:
206
+ raise RuntimeError(f"Unsupported quant type {quant_dtype}")
207
193
- return w, w_s
208
+ return w, w_s, w_gs
209
210
211
def make_test_weight(
@@ -198,21 +213,26 @@ def make_test_weight(
213
rows: int,
214
cols: int,
215
in_dtype: torch.dtype = torch.bfloat16,
201
- quant_dtype: Optional[torch.dtype] = None,
216
+ quant_dtype: Union[torch.dtype, str, None] = None,
217
block_shape: Optional[list[int]] = None,
218
per_act_token_quant: bool = False,
204
-) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
219
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
220
+ Optional[torch.Tensor]]:
221
w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15
222
+ w_gs = None
223
224
if quant_dtype is not None:
225
w_l = [None] * e
226
w_s_l = [None] * e
227
+ w_gs_l = [None] * e
228
for idx in range(e):
211
- w_l[idx], w_s_l[idx] = moe_quantize_weights(
229
+ w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights(
230
w_16[idx], None, quant_dtype, per_act_token_quant, block_shape)
231
232
w = torch.stack(w_l)
233
w_s = torch.stack(w_s_l)
234
+ if e > 0 and w_gs_l[0] is not None:
235
+ w_gs = torch.stack(w_gs_l)
236
if w_s.ndim == 2:
237
assert w_s.shape[-1] == 1
238
w_s = w_s.view(-1, 1, 1)
@@ -225,8 +245,9 @@ def make_test_weight(
245
else:
246
w = w_16
247
w_s = None
248
+ w_gs = None
249
229
- return w_16, w, w_s
250
+ return w_16, w, w_s, w_gs
251
252
253
def make_test_weights(
@@ -234,14 +255,30 @@ def make_test_weights(
255
n: int,
256
k: int,
257
in_dtype: torch.dtype = torch.bfloat16,
237
- quant_dtype: Optional[torch.dtype] = None,
258
+ quant_dtype: Union[torch.dtype, str, None] = None,
259
block_shape: Optional[list[int]] = None,
260
per_act_token_quant: bool = False,
240
-) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor,
241
- torch.Tensor, Optional[torch.Tensor]]:
261
+) -> tuple[tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
262
+ Optional[torch.Tensor]],
263
+ tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor],
264
+ Optional[torch.Tensor]]]:
265
return (
243
- *make_test_weight(e, 2 * n, k, in_dtype, quant_dtype, block_shape,
244
- per_act_token_quant),
245
- *make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape,
246
- per_act_token_quant),
266
+ make_test_weight(e, 2 * n, k, in_dtype, quant_dtype, block_shape,
267
+ per_act_token_quant),
268
+ make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape,
269
+ per_act_token_quant),
270
)
271
+
272
+
273
+def per_token_cast_to_fp8(
274
+ x: torch.Tensor,
275
+ block_size: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
276
+ assert x.dim() == 2
277
+ m, n = x.shape
278
+ pad_size = (block_size - (n % block_size)) % block_size
279
+ x = torch.nn.functional.pad(x,
280
+ (0, pad_size), value=0) if pad_size > 0 else x
281
+ x_view = x.view(m, -1, block_size)
282
+ x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4)
283
+ fp8_data = (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn)
284
+ return fp8_data.view(m, n + pad_size)[:, :n], (x_amax / 448.0).view(m, -1)
vllm/distributed/device_communicators/base_device_communicator.py
+4
-3
@@ -105,7 +105,8 @@ class DeviceCommunicatorBase:
105
# we initialize the all2all manager used in expert parallel.
106
use_ep = config.parallel_config.data_parallel_size > 1
107
108
- self.use_all2all = "ep" in unique_name and use_ep
108
+ self.is_ep_communicator = "ep" in unique_name
109
+ self.use_all2all = self.is_ep_communicator and use_ep
110
self.all2all_manager: Optional[All2AllManagerBase] = None
111
112
def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
@@ -246,7 +247,7 @@ class DeviceCommunicatorBase:
247
"""
248
Prepare the communication buffer for the model.
249
"""
249
- if not self.use_all2all:
250
+ if not self.is_ep_communicator:
251
return
252
253
moe_modules = [
@@ -254,7 +255,7 @@ class DeviceCommunicatorBase:
255
if module.__class__.__name__ == "FusedMoE"
256
]
257
for module in moe_modules:
257
- module.quant_method.init_prepare_finalize(module.moe_config)
258
+ module.quant_method.init_prepare_finalize()
259
260
def dispatch(
261
self, hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/__init__.py
+3
-1
@@ -49,7 +49,8 @@ if HAS_TRITON:
49
from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
50
BatchedTritonOrDeepGemmExperts)
51
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
52
- CutlassExpertsFp8, cutlass_moe_fp4, cutlass_moe_fp8)
52
+ CutlassBatchedExpertsFp8, CutlassExpertsFp8, cutlass_moe_fp4,
53
+ cutlass_moe_fp8)
54
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
55
DeepGemmExperts)
56
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
@@ -69,6 +70,7 @@ if HAS_TRITON:
70
"cutlass_moe_fp8",
71
"cutlass_moe_fp4",
72
"CutlassExpertsFp8",
73
+ "CutlassBatchedExpertsFp8",
74
"TritonExperts",
75
"BatchedTritonExperts",
76
"DeepGemmExperts",
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+23
-13
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional
4
5
import torch
6
@@ -254,18 +254,28 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
254
output = (num_experts, max_num_tokens * num_dispatchers, K)
255
return (workspace13, workspace2, output, a.dtype)
256
257
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
258
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
259
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
260
- expert_map: Optional[torch.Tensor],
261
- w1_scale: Optional[torch.Tensor],
262
- w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
263
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
264
- a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
265
- workspace2: torch.Tensor,
266
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
267
- apply_router_weight_on_input: bool,
268
- extra_expert_args: Optional[dict[str, Any]]):
257
+ def apply(
258
+ self,
259
+ output: torch.Tensor,
260
+ hidden_states: torch.Tensor,
261
+ w1: torch.Tensor,
262
+ w2: torch.Tensor,
263
+ topk_weights: torch.Tensor,
264
+ topk_ids: torch.Tensor,
265
+ activation: str,
266
+ global_num_experts: int,
267
+ expert_map: Optional[torch.Tensor],
268
+ w1_scale: Optional[torch.Tensor],
269
+ w2_scale: Optional[torch.Tensor],
270
+ w1_zp: Optional[torch.Tensor],
271
+ w2_zp: Optional[torch.Tensor],
272
+ a1q_scale: Optional[torch.Tensor],
273
+ a2_scale: Optional[torch.Tensor],
274
+ workspace13: torch.Tensor,
275
+ workspace2: torch.Tensor,
276
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
277
+ apply_router_weight_on_input: bool,
278
+ ):
279
assert expert_tokens_meta is not None
280
expert_num_tokens = expert_tokens_meta.expert_num_tokens
281
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+24
-14
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional
4
5
import torch
6
@@ -132,18 +132,28 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
132
a, aq, M, N, K, topk, global_num_experts, local_num_experts,
133
expert_tokens_metadata)
134
135
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
136
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
137
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
138
- expert_map: Optional[torch.Tensor],
139
- w1_scale: Optional[torch.Tensor],
140
- w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
141
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
142
- a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
143
- workspace2: torch.Tensor,
144
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
145
- apply_router_weight_on_input: bool,
146
- extra_expert_args: Optional[dict[str, Any]]):
135
+ def apply(
136
+ self,
137
+ output: torch.Tensor,
138
+ hidden_states: torch.Tensor,
139
+ w1: torch.Tensor,
140
+ w2: torch.Tensor,
141
+ topk_weights: torch.Tensor,
142
+ topk_ids: torch.Tensor,
143
+ activation: str,
144
+ global_num_experts: int,
145
+ 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],
150
+ a1q_scale: Optional[torch.Tensor],
151
+ a2_scale: Optional[torch.Tensor],
152
+ workspace13: torch.Tensor,
153
+ workspace2: torch.Tensor,
154
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
155
+ apply_router_weight_on_input: bool,
156
+ ):
157
experts = (self.batched_deep_gemm_experts
158
if self.allow_deep_gemm else self.batched_triton_experts)
159
assert experts is not None
@@ -151,4 +161,4 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
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,
154
- apply_router_weight_on_input, extra_expert_args)
164
+ apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/config.py
+6
-5
@@ -45,7 +45,6 @@ def get_quant_config_weight_quant(
45
return _get_quant_config_quantization_args(quant_config, "weights")
46
47
48
-# TODO (bnell): use scalar_type instead of bools?
48
def get_config_quant_dtype(
49
use_fp8_w8a8: bool,
50
use_int8_w8a8: bool,
@@ -65,7 +64,8 @@ def get_config_quant_dtype(
64
@dataclass
65
class FusedMoEQuantConfig:
66
# The post quantization activation type.
68
- quant_dtype: Optional[torch.dtype] = None
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
@@ -141,6 +141,7 @@ class FusedMoEQuantConfig:
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
@@ -334,7 +335,7 @@ class FusedMoEConfig:
335
assert self.max_num_tokens > 0
336
337
@property
337
- def quant_dtype(self) -> Optional[torch.dtype]:
338
+ def quant_dtype(self) -> Union[torch.dtype, str, None]:
339
if self.quant_config is not None:
340
return self.quant_config.quant_dtype
341
else:
@@ -429,7 +430,7 @@ class FusedMoEConfig:
430
block_shape = None
431
per_act_token_quant = False
432
per_out_ch_quant = False
432
- quant_dtype: Optional[torch.dtype] = None
433
+ quant_dtype: Union[torch.dtype, str, None] = None
434
435
input_quant = get_quant_config_input_quant(quant_config)
436
weight_quant = get_quant_config_weight_quant(quant_config)
@@ -453,7 +454,7 @@ class FusedMoEConfig:
454
ModelOptNvFp4Config)
455
if quant_dtype is None and isinstance(quant_config,
456
ModelOptNvFp4Config):
456
- quant_dtype = torch.uint8
457
+ quant_dtype = "nvfp4"
458
459
if weight_quant is not None:
460
per_out_ch_quant = (
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+193
-135
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
""" CUTLASS based Fused MoE kernels."""
4
-from typing import Any, Callable, Optional
4
+from typing import Callable, Optional
5
6
import torch
7
@@ -12,11 +12,10 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
13
MoEPrepareAndFinalizeNoEP)
14
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
- TopKWeightAndReduceDelegate)
15
+ TopKWeightAndReduceDelegate, TopKWeightAndReduceNoOP)
16
from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,
17
_fp8_quantize,
18
- _resize_cache,
19
- extract_required_args)
18
+ _resize_cache)
19
from vllm.scalar_type import scalar_types
20
21
logger = init_logger(__name__)
@@ -213,19 +212,14 @@ def run_cutlass_moe_fp8(
212
output.copy_(c3[c_map].view(M * topk, K), non_blocking=True)
213
214
216
-# TODO (bnell): split class batched vs. non-batched?
217
-# maybe remove need for passing aq to workspace_shapes
218
-class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
215
+class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):
216
217
def __init__(
218
self,
222
- max_experts_per_worker: int,
219
out_dtype: Optional[torch.dtype],
220
per_act_token_quant: bool,
221
per_out_ch_quant: bool,
222
block_shape: Optional[list[int]] = None,
227
- num_dispatchers: Optional[int] = None,
228
- use_batched_format: bool = False,
223
):
224
super().__init__(
225
FusedMoEQuantConfig(
@@ -234,33 +228,84 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
228
per_out_ch_quant=per_out_ch_quant,
229
block_shape=block_shape,
230
))
237
- assert max_experts_per_worker > 0
238
- assert not use_batched_format or num_dispatchers is not None
239
- self.max_experts_per_worker = max_experts_per_worker
240
- self.num_dispatchers = num_dispatchers
231
self.out_dtype = out_dtype
242
- self.use_batched_format = use_batched_format
232
+
233
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
234
+ # Let PrepareAndFinalize::finalize() decide the impl.
235
+ return TopKWeightAndReduceDelegate()
236
+
237
+ def apply(
238
+ self,
239
+ output: torch.Tensor,
240
+ hidden_states: torch.Tensor,
241
+ w1: torch.Tensor,
242
+ w2: torch.Tensor,
243
+ topk_weights: torch.Tensor,
244
+ topk_ids: torch.Tensor,
245
+ activation: str,
246
+ global_num_experts: int,
247
+ expert_map: Optional[torch.Tensor],
248
+ w1_scale: Optional[torch.Tensor],
249
+ w2_scale: Optional[torch.Tensor],
250
+ w1_zp: Optional[torch.Tensor],
251
+ w2_zp: Optional[torch.Tensor],
252
+ a1q_scale: Optional[torch.Tensor],
253
+ a2_scale: Optional[torch.Tensor],
254
+ workspace13: torch.Tensor,
255
+ workspace2: torch.Tensor,
256
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
257
+ apply_router_weight_on_input: bool,
258
+ ):
259
+ assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
260
+ assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
261
+
262
+ expert_num_tokens = None
263
+ if expert_tokens_meta is not None:
264
+ expert_num_tokens = expert_tokens_meta.expert_num_tokens
265
+
266
+ activation_callable = lambda o, i: self.activation(activation, o, i)
267
+
268
+ use_batched_format = self.activation_formats[
269
+ 0] == mk.FusedMoEActivationFormat.BatchedExperts
270
+
271
+ in_dtype = hidden_states.dtype
272
+ run_cutlass_moe_fp8(
273
+ output, hidden_states, w1, w2, topk_ids, activation_callable,
274
+ global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,
275
+ a2_scale, workspace13, workspace2, expert_num_tokens,
276
+ self.out_dtype if self.out_dtype is not None else in_dtype,
277
+ self.per_act_token_quant, self.per_out_ch_quant,
278
+ use_batched_format)
279
+
280
+
281
+class CutlassExpertsFp8(CutlassExpertsFp8Base):
282
+
283
+ def __init__(
284
+ self,
285
+ out_dtype: Optional[torch.dtype],
286
+ per_act_token_quant: bool,
287
+ per_out_ch_quant: bool,
288
+ block_shape: Optional[list[int]] = None,
289
+ ):
290
+ super().__init__(
291
+ out_dtype,
292
+ per_act_token_quant,
293
+ per_out_ch_quant,
294
+ block_shape,
295
+ )
296
297
@property
298
def activation_formats(
299
self
300
) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
248
- if self.use_batched_format:
249
- return (mk.FusedMoEActivationFormat.BatchedExperts,
250
- mk.FusedMoEActivationFormat.BatchedExperts)
251
- else:
252
- return (mk.FusedMoEActivationFormat.Standard,
253
- mk.FusedMoEActivationFormat.Standard)
301
+ return (mk.FusedMoEActivationFormat.Standard,
302
+ mk.FusedMoEActivationFormat.Standard)
303
304
def supports_chunking(self) -> bool:
256
- return not self.use_batched_format
305
+ return True
306
307
def supports_expert_map(self) -> bool:
259
- return not self.use_batched_format
260
-
261
- def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
262
- # Let PrepareAndFinalize::finalize() decide the impl.
263
- return TopKWeightAndReduceDelegate()
308
+ return True
309
310
def workspace_shapes(
311
self,
@@ -274,54 +319,69 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
319
local_num_experts: int,
320
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
321
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
277
- workspace1: tuple[int, ...] = ()
278
- workspace2: tuple[int, ...] = ()
279
- output: tuple[int, ...] = ()
280
- if self.use_batched_format:
281
- padded_M = aq.size(1)
282
- num_dp = self.num_dispatchers
283
- assert num_dp is not None
284
- workspace1 = (self.max_experts_per_worker, padded_M * num_dp,
285
- max(N, K))
286
- workspace2 = (self.max_experts_per_worker, padded_M * num_dp,
287
- (N // 2))
288
- output = (self.max_experts_per_worker, padded_M, K)
289
- else:
290
- workspace1 = (M * topk, max(N, K))
291
- workspace2 = (M * topk, N // 2)
292
- output = (M * topk, K)
322
+ workspace1 = (M * topk, max(N, K))
323
+ workspace2 = (M * topk, N // 2)
324
+ output = (M * topk, K)
325
return (workspace1, workspace2, output,
326
self.out_dtype if self.out_dtype is not None else a.dtype)
327
296
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
297
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
298
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
299
- expert_map: Optional[torch.Tensor],
300
- w1_scale: Optional[torch.Tensor],
301
- w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
302
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
303
- a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
304
- workspace2: torch.Tensor,
305
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
306
- apply_router_weight_on_input: bool,
307
- extra_expert_args: Optional[dict[str, Any]]):
308
- assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
309
- assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
328
311
- expert_num_tokens = None
312
- if expert_tokens_meta is not None:
313
- expert_num_tokens = expert_tokens_meta.expert_num_tokens
329
+class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
330
315
- activation_callable = lambda o, i: self.activation(activation, o, i)
331
+ def __init__(
332
+ self,
333
+ max_experts_per_worker: int,
334
+ num_dispatchers: int,
335
+ out_dtype: Optional[torch.dtype],
336
+ per_act_token_quant: bool,
337
+ per_out_ch_quant: bool,
338
+ block_shape: Optional[list[int]] = None,
339
+ ):
340
+ super().__init__(
341
+ out_dtype,
342
+ per_act_token_quant,
343
+ per_out_ch_quant,
344
+ block_shape,
345
+ )
346
+ assert max_experts_per_worker > 0
347
+ self.max_experts_per_worker = max_experts_per_worker
348
+ self.num_dispatchers = num_dispatchers
349
317
- in_dtype = hidden_states.dtype
318
- run_cutlass_moe_fp8(
319
- output, hidden_states, w1, w2, topk_ids, activation_callable,
320
- global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,
321
- a2_scale, workspace13, workspace2, expert_num_tokens,
322
- self.out_dtype if self.out_dtype is not None else in_dtype,
323
- self.per_act_token_quant, self.per_out_ch_quant,
324
- self.use_batched_format)
350
+ @property
351
+ def activation_formats(
352
+ self
353
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
354
+ return (mk.FusedMoEActivationFormat.BatchedExperts,
355
+ mk.FusedMoEActivationFormat.BatchedExperts)
356
+
357
+ def supports_chunking(self) -> bool:
358
+ return False
359
+
360
+ def supports_expert_map(self) -> bool:
361
+ return False
362
+
363
+ # TODO(bnell): maybe remove need for passing aq to workspace_shapes
364
+ def workspace_shapes(
365
+ self,
366
+ a: torch.Tensor,
367
+ aq: torch.Tensor,
368
+ M: int,
369
+ N: int,
370
+ K: int,
371
+ topk: int,
372
+ global_num_experts: int,
373
+ local_num_experts: int,
374
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
375
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
376
+ padded_M = aq.size(1)
377
+ num_dp = self.num_dispatchers
378
+ assert num_dp is not None
379
+ workspace1 = (self.max_experts_per_worker, padded_M * num_dp,
380
+ max(N, K))
381
+ workspace2 = (self.max_experts_per_worker, padded_M * num_dp, (N // 2))
382
+ output = (self.max_experts_per_worker, padded_M, K)
383
+ return (workspace1, workspace2, output,
384
+ self.out_dtype if self.out_dtype is not None else a.dtype)
385
386
387
def cutlass_moe_fp8(
@@ -387,11 +447,9 @@ def cutlass_moe_fp8(
447
fn = mk.FusedMoEModularKernel(
448
MoEPrepareAndFinalizeNoEP(),
449
CutlassExpertsFp8(
390
- max_experts_per_worker=num_experts,
450
out_dtype=a.dtype,
451
per_act_token_quant=per_act_token,
452
per_out_ch_quant=per_out_ch,
394
- use_batched_format=False,
453
),
454
)
455
@@ -476,8 +534,9 @@ def run_cutlass_moe_fp4(
534
e_w1, nx2_w1, half_k_w1 = w1_fp4.shape
535
e_w2, k_w2, half_n_w2 = w2_fp4.shape
536
479
- assert (e_w1 == e_w2 and e_w1 == e), ("Number of experts must match",
480
- " between weights.")
537
+ assert (e_w1 == e_w2
538
+ and e_w1 == e), ("Number of experts must match",
539
+ f" between weights. {e_w1}, {e_w2}, {e}")
540
assert (k_a == half_k_w1 * 2
541
and k == k_w2), ("Hidden size mismatch between a, w1 and w2")
542
assert (nx2_w1 == n * 2 and half_n_w2 * 2 == n), ("mismatch in "
@@ -554,6 +613,10 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
613
614
def __init__(
615
self,
616
+ g1_alphas: torch.Tensor,
617
+ g2_alphas: torch.Tensor,
618
+ a1_gscale: torch.Tensor,
619
+ a2_gscale: torch.Tensor,
620
max_experts_per_worker: int,
621
out_dtype: torch.dtype,
622
per_act_token_quant: bool,
@@ -562,8 +625,12 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
625
use_batched_format: bool = False,
626
):
627
super().__init__(
628
+ # NVFP4 requires two levels of quantization, which involves
629
+ # computing some scaling factors dynamically. This makes it
630
+ # incompatible with the typical prepare -> MoE -> finalize
631
+ # pipeline. Move the quantization logic into the MoE body.
632
FusedMoEQuantConfig(
566
- quant_dtype=torch.uint8,
633
+ quant_dtype=None, # skip quantization in prepare/finalize
634
per_act_token_quant=per_act_token_quant,
635
per_out_ch_quant=per_out_ch_quant,
636
block_shape=block_shape,
@@ -572,6 +639,12 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
639
self.out_dtype = out_dtype
640
self.use_batched_format = use_batched_format
641
642
+ # TODO(bnell): put this stuff into quant config?
643
+ self.g1_alphas = g1_alphas
644
+ self.g2_alphas = g2_alphas
645
+ self.a1_gscale = a1_gscale
646
+ self.a2_gscale = a2_gscale
647
+
648
@property
649
def activation_formats(
650
self
@@ -590,8 +663,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
663
return True
664
665
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
593
- # Let PrepareAndFinalize::finalize() decide the impl.
594
- return TopKWeightAndReduceDelegate()
666
+ return TopKWeightAndReduceNoOP()
667
668
def workspace_shapes(
669
self,
@@ -620,34 +692,42 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
692
return (workspace1, workspace2, output,
693
self.out_dtype if self.out_dtype is not None else a.dtype)
694
623
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
624
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
625
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
626
- expert_map: Optional[torch.Tensor], w1_scale: torch.Tensor,
627
- w2_scale: torch.Tensor, w1_zp: Optional[torch.Tensor],
628
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
629
- a2_scale: torch.Tensor, workspace13: Optional[torch.Tensor],
630
- workspace2: Optional[torch.Tensor],
631
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
632
- apply_router_weight_on_input: bool,
633
- extra_expert_args: Optional[dict[str, Any]]):
634
- required_keys = [
635
- "g1_alphas", "g2_alphas", "a1_gscale", "a2_gscale", "m", "n", "k",
636
- "e", "device"
637
- ]
638
- (g1_alphas, g2_alphas, a1_gscale, a2_gscale, m, n, k, e,
639
- device) = extract_required_args(extra_expert_args, required_keys)
695
+ def apply(
696
+ self,
697
+ output: torch.Tensor,
698
+ hidden_states: torch.Tensor,
699
+ w1: torch.Tensor,
700
+ w2: torch.Tensor,
701
+ topk_weights: torch.Tensor,
702
+ topk_ids: torch.Tensor,
703
+ activation: str,
704
+ global_num_experts: int,
705
+ expert_map: Optional[torch.Tensor],
706
+ w1_scale: torch.Tensor,
707
+ w2_scale: torch.Tensor,
708
+ w1_zp: Optional[torch.Tensor],
709
+ w2_zp: Optional[torch.Tensor],
710
+ a1q_scale: Optional[torch.Tensor],
711
+ a2_scale: torch.Tensor,
712
+ workspace13: Optional[torch.Tensor],
713
+ workspace2: Optional[torch.Tensor],
714
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
715
+ apply_router_weight_on_input: bool,
716
+ ):
717
+ e, m, n, k, _ = mk._moe_problem_size(hidden_states, w1, w2, topk_ids)
718
+ n = w2.shape[2] * 2
719
+
720
run_cutlass_moe_fp4(
721
output=output,
722
a=hidden_states,
643
- a1_gscale=a1_gscale,
723
+ a1_gscale=self.a1_gscale,
724
w1_fp4=w1,
725
w1_blockscale=w1_scale,
646
- w1_alphas=g1_alphas,
647
- a2_gscale=a2_gscale,
726
+ w1_alphas=self.g1_alphas,
727
+ a2_gscale=self.a2_gscale,
728
w2_fp4=w2,
729
w2_blockscale=w2_scale,
650
- w2_alphas=g2_alphas,
730
+ w2_alphas=self.g2_alphas,
731
topk_weights=topk_weights,
732
topk_ids=topk_ids,
733
workspace13=workspace13,
@@ -656,7 +736,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
736
n=n,
737
k=k,
738
e=e,
659
- device=device,
739
+ device=hidden_states.device,
740
apply_router_weight_on_input=apply_router_weight_on_input,
741
)
742
@@ -677,7 +757,6 @@ def cutlass_moe_fp4(
757
n: int,
758
k: int,
759
e: int,
680
- device: torch.device,
760
expert_map: Optional[torch.Tensor] = None,
761
apply_router_weight_on_input: bool = False) -> torch.Tensor:
762
assert expert_map is None, ("Expert Parallelism / expert_map "
@@ -686,6 +765,10 @@ def cutlass_moe_fp4(
765
fn = mk.FusedMoEModularKernel(
766
MoEPrepareAndFinalizeNoEP(),
767
CutlassExpertsFp4(
768
+ g1_alphas,
769
+ g2_alphas,
770
+ a1_gscale,
771
+ a2_gscale,
772
max_experts_per_worker=e,
773
out_dtype=a.dtype,
774
per_act_token_quant=False,
@@ -693,29 +776,7 @@ def cutlass_moe_fp4(
776
use_batched_format=False,
777
),
778
)
696
- extra_expert_args = {
697
- 'g1_alphas': g1_alphas,
698
- 'g2_alphas': g2_alphas,
699
- 'a1_gscale': a1_gscale,
700
- 'a2_gscale': a2_gscale,
701
- 'm': m,
702
- 'n': n,
703
- 'k': k,
704
- 'e': e,
705
- 'device': device,
706
- }
707
-
708
- # NVFP4 requires two levels of quantization, which involves computing some
709
- # scaling factors dynamically. This makes it incompatible with the typical
710
- # prepare -> MoE -> finalize pipeline. Move the quantization logic into the
711
- # MoE body.
712
- extra_prepare_args = {
713
- 'skip_quant': True,
714
- }
715
- # Similar reason as above.
716
- extra_finalize_args = {
717
- 'skip_weight_reduce': True,
718
- }
779
+
780
return fn(
781
hidden_states=a,
782
w1=w1_fp4,
@@ -731,9 +792,6 @@ def cutlass_moe_fp4(
792
a1_scale=None,
793
a2_scale=None,
794
apply_router_weight_on_input=apply_router_weight_on_input,
734
- extra_expert_args=extra_expert_args,
735
- extra_prepare_args=extra_prepare_args,
736
- extra_finalize_args=extra_finalize_args,
795
)
796
797
@@ -824,16 +882,6 @@ def run_cutlass_block_scaled_fused_experts(
882
k = w1_q.size(1)
883
n = w2_q.size(1)
884
827
- expert_offsets = torch.empty((num_experts + 1, ),
828
- dtype=torch.int32,
829
- device="cuda")
830
- problem_sizes1 = torch.empty((num_experts, 3),
831
- dtype=torch.int32,
832
- device="cuda")
833
- problem_sizes2 = torch.empty((num_experts, 3),
834
- dtype=torch.int32,
835
- device="cuda")
836
-
885
topk = topk_ids.size(1)
886
887
a_q, a1_scale = _fp8_quantize(a,
@@ -842,6 +890,16 @@ def run_cutlass_block_scaled_fused_experts(
890
block_shape=[128, 128])
891
device = a_q.device
892
893
+ expert_offsets = torch.empty((num_experts + 1, ),
894
+ dtype=torch.int32,
895
+ device=device)
896
+ problem_sizes1 = torch.empty((num_experts, 3),
897
+ dtype=torch.int32,
898
+ device=device)
899
+ problem_sizes2 = torch.empty((num_experts, 3),
900
+ dtype=torch.int32,
901
+ device=device)
902
+
903
a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
904
c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
905
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+1
-2
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import functools
4
-from typing import Any, Optional
4
+from typing import Optional
5
6
import torch
7
from tqdm import tqdm
@@ -230,7 +230,6 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
230
workspace2: torch.Tensor,
231
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
232
apply_router_weight_on_input: bool,
233
- extra_expert_args: Optional[dict[str, Any]],
233
):
234
assert self.block_shape is not None
235
assert a1q_scale is not None
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+19
-11
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional
4
5
import deep_ep
6
import torch
@@ -127,12 +127,16 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
127
expert_topk_weights)
128
129
def prepare(
130
- self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
131
- a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
132
- topk_ids: torch.Tensor, num_experts: int,
133
- expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
130
+ self,
131
+ a1: torch.Tensor,
132
+ a1_scale: Optional[torch.Tensor],
133
+ a2_scale: Optional[torch.Tensor],
134
+ topk_weights: torch.Tensor,
135
+ topk_ids: torch.Tensor,
136
+ num_experts: int,
137
+ expert_map: Optional[torch.Tensor],
138
+ apply_router_weight_on_input: bool,
139
quant_config: FusedMoEQuantConfig,
135
- extra_prepare_args: Optional[dict[str, Any]]
140
) -> tuple[torch.Tensor, Optional[torch.Tensor],
141
Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
142
Optional[torch.Tensor]]:
@@ -187,11 +191,15 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
191
return (expert_x, expert_x_scale, expert_tokens_meta, expert_topk_ids,
192
expert_topk_weights)
193
190
- def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
191
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
192
- apply_router_weight_on_input: bool,
193
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
194
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
194
+ def finalize(
195
+ self,
196
+ output: torch.Tensor,
197
+ fused_expert_output: torch.Tensor,
198
+ topk_weights: torch.Tensor,
199
+ topk_ids: torch.Tensor,
200
+ apply_router_weight_on_input: bool,
201
+ weight_and_reduce_impl: mk.TopKWeightAndReduce,
202
+ ) -> None:
203
204
assert self.handle is not None
205
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+20
-12
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional, Union
3
+from typing import Optional, Union
4
5
import deep_ep
6
import torch
@@ -77,7 +77,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
77
a1_scale: Optional[torch.Tensor],
78
a2_scale: Optional[torch.Tensor],
79
a1_dtype: torch.dtype,
80
- quant_dtype: Optional[torch.dtype],
80
+ quant_dtype: Union[torch.dtype, str, None],
81
per_act_token_quant: bool,
82
block_shape: Optional[list[int]],
83
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
@@ -111,12 +111,16 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
111
return x, x_scales
112
113
def prepare(
114
- self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
115
- a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
116
- topk_ids: torch.Tensor, num_experts: int,
117
- expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
114
+ self,
115
+ a1: torch.Tensor,
116
+ a1_scale: Optional[torch.Tensor],
117
+ a2_scale: Optional[torch.Tensor],
118
+ topk_weights: torch.Tensor,
119
+ topk_ids: torch.Tensor,
120
+ num_experts: int,
121
+ expert_map: Optional[torch.Tensor],
122
+ apply_router_weight_on_input: bool,
123
quant_config: FusedMoEQuantConfig,
119
- extra_prepare_args: Optional[dict[str, Any]]
124
) -> tuple[torch.Tensor, Optional[torch.Tensor],
125
Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
126
Optional[torch.Tensor]]:
@@ -162,11 +166,15 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
166
167
return (expert_x, expert_x_scale, expert_tokens_meta, None, None)
168
165
- def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
166
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
167
- apply_router_weight_on_input: bool,
168
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
169
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
169
+ def finalize(
170
+ self,
171
+ output: torch.Tensor,
172
+ fused_expert_output: torch.Tensor,
173
+ topk_weights: torch.Tensor,
174
+ topk_ids: torch.Tensor,
175
+ apply_router_weight_on_input: bool,
176
+ weight_and_reduce_impl: mk.TopKWeightAndReduce,
177
+ ) -> None:
178
assert isinstance(
179
weight_and_reduce_impl, TopKWeightAndReduceDelegate
180
), ("Weight application and reduction happens in the combine kernel.")
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+23
-36
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional, Union
4
5
import torch
6
@@ -8,8 +8,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 FusedMoEQuantConfig
10
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
11
- TopKWeightAndReduceDelegate)
12
-from vllm.model_executor.layers.fused_moe.utils import extract_required_args
11
+ TopKWeightAndReduceNoOP)
12
from vllm.utils.flashinfer import (flashinfer_cutlass_fused_moe,
13
has_flashinfer_cutlass_fused_moe)
14
@@ -20,7 +19,7 @@ def is_valid_flashinfer_cutlass_fused_moe(hidden_states: torch.Tensor,
19
w1: torch.Tensor,
20
w2: torch.Tensor) -> bool:
21
"""
23
- Check if the given problem size is supported by the FlashInfer CUTLASS MoE
22
+ Check if the given problem size is supported by the FlashInfer CUTLASS MoE
23
kernel.
24
"""
25
if not has_flashinfer_cutlass_fused_moe():
@@ -43,31 +42,34 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
42
43
def __init__(
44
self,
46
- use_nvfp4_w4a4: bool = False,
47
- use_fp8_w8a8: bool = False,
48
- use_dp: bool = False,
45
+ g1_alphas: torch.Tensor,
46
+ g2_alphas: torch.Tensor,
47
+ a1_gscale: torch.Tensor,
48
+ a2_gscale: torch.Tensor,
49
+ out_dtype: torch.dtype,
50
+ quant_dtype: Union[torch.dtype, str, None],
51
ep_rank: int = 0,
52
ep_size: int = 1,
53
tp_rank: int = 0,
54
tp_size: int = 1,
53
- num_dispatchers: Optional[int] = None,
54
- use_batched_format: bool = False,
55
):
56
super().__init__(
57
FusedMoEQuantConfig(
58
- quant_dtype=torch.uint8,
58
+ quant_dtype=quant_dtype,
59
per_act_token_quant=False,
60
block_shape=None,
61
))
62
- self.use_nvfp4_w4a4 = use_nvfp4_w4a4
63
- self.use_fp8_w8a8 = use_fp8_w8a8
62
+ assert quant_dtype == "nvfp4", ("Only nvfp4 quantization is "
63
+ "currently supported.")
64
self.ep_rank = ep_rank
65
self.ep_size = ep_size
66
self.tp_rank = tp_rank
67
self.tp_size = tp_size
68
- self.use_dp = use_dp
69
- assert not use_batched_format or num_dispatchers is not None
70
- self.num_dispatchers = num_dispatchers
68
+ self.g1_alphas = g1_alphas
69
+ self.g2_alphas = g2_alphas
70
+ self.a1_gscale = a1_gscale
71
+ self.a2_gscale = a2_gscale
72
+ self.out_dtype = out_dtype
73
74
@property
75
def activation_formats(
@@ -84,8 +86,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
86
return True
87
88
def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
87
- # Let PrepareAndFinalize::finalize() decide the impl.
88
- return TopKWeightAndReduceDelegate()
89
+ return TopKWeightAndReduceNoOP()
90
91
def workspace_shapes(
92
self,
@@ -117,8 +118,6 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
118
- Note: in order for activation chunking to work, the first dimension
119
of each tuple must be the number of tokens.
120
"""
120
- assert self.use_nvfp4_w4a4 is True, ("Only nvfp4 quantization is "
121
- "currently supported.")
121
aq_m, aq_n = aq.shape
122
workspace2 = ()
123
output_shape = (aq_m, aq_n * 2)
@@ -149,21 +148,9 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
148
workspace2: Optional[torch.Tensor],
149
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
150
apply_router_weight_on_input: Optional[bool],
152
- extra_expert_args: Optional[dict[str, Any]],
151
):
154
- assert extra_expert_args is not None, \
155
- "extra_expert_args must be provided"
156
- required_keys = [
157
- 'g1_alphas', 'g2_alphas', 'a1_gscale', 'a2_gscale', 'out_dtype'
158
- ]
159
-
160
- g1_alphas, g2_alphas, a1_gscale, a2_gscale, out_dtype = (
161
- extract_required_args(extra_expert_args, required_keys))
162
-
152
# Flashinfer CUTLASS kernel takes scalar global scales,
153
# min because inv_scale.
165
- assert self.use_nvfp4_w4a4 is True, ("Only nvfp4 quantization is "
166
- "currently supported.")
154
155
# Ensure w1_scale and w2_scale are not None before calling view
156
assert w1_scale is not None and w2_scale is not None, (
@@ -171,12 +158,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
158
"be None for FlashInferExperts")
159
160
quant_scales = [
174
- a1_gscale,
161
+ self.a1_gscale,
162
w1_scale.view(torch.int32),
176
- g1_alphas,
177
- a2_gscale,
163
+ self.g1_alphas,
164
+ self.a2_gscale,
165
w2_scale.view(torch.int32),
179
- g2_alphas,
166
+ self.g2_alphas,
167
]
168
_ = flashinfer_cutlass_fused_moe(
169
input=hidden_states,
@@ -185,7 +172,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
172
# FlashInfer API requires weight to be long for nvfp4
173
fc1_expert_weights=w1.view(torch.long),
174
fc2_expert_weights=w2.view(torch.long),
188
- output_dtype=out_dtype,
175
+ output_dtype=self.out_dtype,
176
quant_scales=quant_scales,
177
input_sf=a1q_scale,
178
tp_size=self.tp_size,
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+24
-28
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional
4
5
import torch
6
@@ -9,7 +9,7 @@ from vllm.distributed import get_dp_group
9
from vllm.forward_context import get_forward_context
10
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
from vllm.model_executor.layers.fused_moe.utils import (
12
- extract_required_args, moe_kernel_quantize_input)
12
+ moe_kernel_quantize_input)
13
from vllm.utils.flashinfer import nvfp4_block_scale_interleave
14
15
@@ -21,16 +21,15 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
21
22
def __init__(
23
self,
24
- quant_dtype: Optional[torch.dtype] = None,
25
- per_channel_quant: bool = False,
26
- block_shape: Optional[list[int]] = None,
24
+ use_dp: bool,
25
+ a1_gscale: Optional[torch.Tensor],
26
num_dispatchers: int = 1,
27
):
28
super().__init__()
30
- self.per_channel_quant = per_channel_quant
31
- self.block_shape = block_shape
32
- self.quant_dtype = quant_dtype
29
self.num_dispatchers_ = num_dispatchers
30
+ self.use_dp = use_dp
31
+ self.a1_gscale = a1_gscale
32
+ self.local_tokens = None
33
34
@property
35
def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -55,10 +54,11 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
54
num_experts: int,
55
expert_map: Optional[torch.Tensor],
56
apply_router_weight_on_input: bool,
57
+ # TODO(bnell): use quant_config + scales instead of ctor args
58
quant_config: FusedMoEQuantConfig,
59
- extra_prepare_args: Optional[dict[str, Any]]
60
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
61
- Optional[torch.Tensor], Optional[torch.Tensor]]:
59
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor],
60
+ Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
61
+ Optional[torch.Tensor]]:
62
63
if apply_router_weight_on_input:
64
topk = topk_ids.size(1)
@@ -67,22 +67,22 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
67
"apply_router_weight_on_input is only implemented for topk=1"
68
a1.mul_(topk_weights.to(a1.dtype))
69
70
- (a1_gscale, use_dp, local_tokens) = extract_required_args(
71
- extra_prepare_args, ['a1_gscale', 'use_dp', 'local_tokens'])
72
-
70
a1q, a1q_scale = moe_kernel_quantize_input(
71
a1,
75
- a1_gscale,
72
+ self.a1_gscale,
73
quant_config.quant_dtype,
77
- self.per_channel_quant,
78
- self.block_shape,
79
- is_fp4_scale_swizzled=not use_dp, # Swizzling after communication
74
+ quant_config.per_act_token_quant,
75
+ quant_config.block_shape,
76
+ # Swizzling after communication
77
+ is_fp4_scale_swizzled=not self.use_dp,
78
)
81
- if use_dp:
79
+ if self.use_dp:
80
topk_weights, topk_ids, a1q, a1q_scale = \
83
- get_dp_group().all_gatherv([topk_weights, topk_ids, a1q, a1q_scale], # noqa: E501
84
- dim=0,
85
- sizes=get_local_sizes())
81
+ get_dp_group().all_gatherv(
82
+ [topk_weights, topk_ids, a1q, a1q_scale],
83
+ dim=0,
84
+ sizes=get_local_sizes(),
85
+ )
86
a1_m, a1_n = a1q.shape
87
a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
88
@@ -91,13 +91,9 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
91
def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
92
topk_weights: torch.Tensor, topk_ids: torch.Tensor,
93
apply_router_weight_on_input: bool,
94
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
95
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
94
+ weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:
95
97
- (use_dp,
98
- local_tokens) = extract_required_args(extra_finalize_args,
99
- ['use_dp', 'local_tokens'])
100
- if use_dp:
96
+ if self.use_dp:
97
fused_expert_output = get_dp_group().reduce_scatterv(
98
fused_expert_output, dim=0, sizes=get_local_sizes())
99
output.copy_(fused_expert_output)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+63
-35
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Fused batched MoE kernel."""
4
-from typing import Any, Optional
4
+from typing import Optional
5
6
import torch
7
@@ -496,12 +496,16 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
496
return self.num_dispatchers_
497
498
def prepare(
499
- self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
500
- a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
501
- topk_ids: torch.Tensor, num_experts: int,
502
- expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
499
+ self,
500
+ a1: torch.Tensor,
501
+ a1_scale: Optional[torch.Tensor],
502
+ a2_scale: Optional[torch.Tensor],
503
+ topk_weights: torch.Tensor,
504
+ topk_ids: torch.Tensor,
505
+ num_experts: int,
506
+ expert_map: Optional[torch.Tensor],
507
+ apply_router_weight_on_input: bool,
508
quant_config: FusedMoEQuantConfig,
504
- extra_prepare_args: Optional[dict[str, Any]]
509
) -> tuple[torch.Tensor, Optional[torch.Tensor],
510
Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
511
Optional[torch.Tensor]]:
@@ -590,11 +594,15 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
594
595
return b_a1, b_a1_scale, expert_tokens_meta, None, None
596
593
- def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
594
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
595
- apply_router_weight_on_input: bool,
596
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
597
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
597
+ def finalize(
598
+ self,
599
+ output: torch.Tensor,
600
+ fused_expert_output: torch.Tensor,
601
+ topk_weights: torch.Tensor,
602
+ topk_ids: torch.Tensor,
603
+ apply_router_weight_on_input: bool,
604
+ weight_and_reduce_impl: mk.TopKWeightAndReduce,
605
+ ) -> None:
606
if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
607
weight_and_reduce_impl = TopKWeightAndReduceNaiveBatched(self.rank)
608
weight_and_reduce_impl.apply(
@@ -688,18 +696,28 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
696
else:
697
return t.to(f32) * group_broadcast(scale, t.shape)
698
691
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
692
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
693
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
694
- expert_map: Optional[torch.Tensor],
695
- w1_scale: Optional[torch.Tensor],
696
- w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
697
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
698
- a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
699
- workspace2: torch.Tensor,
700
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
701
- apply_router_weight_on_input: bool,
702
- extra_expert_args: Optional[dict[str, Any]]):
699
+ def apply(
700
+ self,
701
+ output: torch.Tensor,
702
+ hidden_states: torch.Tensor,
703
+ w1: torch.Tensor,
704
+ w2: torch.Tensor,
705
+ topk_weights: torch.Tensor,
706
+ topk_ids: torch.Tensor,
707
+ activation: str,
708
+ global_num_experts: int,
709
+ expert_map: Optional[torch.Tensor],
710
+ w1_scale: Optional[torch.Tensor],
711
+ w2_scale: Optional[torch.Tensor],
712
+ w1_zp: Optional[torch.Tensor],
713
+ w2_zp: Optional[torch.Tensor],
714
+ a1q_scale: Optional[torch.Tensor],
715
+ a2_scale: Optional[torch.Tensor],
716
+ workspace13: torch.Tensor,
717
+ workspace2: torch.Tensor,
718
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
719
+ apply_router_weight_on_input: bool,
720
+ ):
721
assert hidden_states.dim() == 3
722
assert expert_tokens_meta is not None
723
expert_num_tokens = expert_tokens_meta.expert_num_tokens
@@ -894,18 +912,28 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
912
output = (num_experts, max_num_tokens * num_dp, K)
913
return (workspace13, workspace2, output, a.dtype)
914
897
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
898
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
899
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
900
- expert_map: Optional[torch.Tensor],
901
- w1_scale: Optional[torch.Tensor],
902
- w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
903
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
904
- a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
905
- workspace2: torch.Tensor,
906
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
907
- apply_router_weight_on_input: bool,
908
- extra_expert_args: Optional[dict[str, Any]]):
915
+ def apply(
916
+ self,
917
+ output: torch.Tensor,
918
+ hidden_states: torch.Tensor,
919
+ w1: torch.Tensor,
920
+ w2: torch.Tensor,
921
+ topk_weights: torch.Tensor,
922
+ topk_ids: torch.Tensor,
923
+ activation: str,
924
+ global_num_experts: int,
925
+ expert_map: Optional[torch.Tensor],
926
+ w1_scale: Optional[torch.Tensor],
927
+ w2_scale: Optional[torch.Tensor],
928
+ w1_zp: Optional[torch.Tensor],
929
+ w2_zp: Optional[torch.Tensor],
930
+ a1q_scale: Optional[torch.Tensor],
931
+ a2_scale: Optional[torch.Tensor],
932
+ workspace13: torch.Tensor,
933
+ workspace2: torch.Tensor,
934
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
935
+ apply_router_weight_on_input: bool,
936
+ ):
937
# Check constraints.
938
if self.use_int4_w4a16:
939
assert hidden_states.size(-1) // 2 == w1.size(2), (
vllm/model_executor/layers/fused_moe/fused_moe.py
+3
-4
@@ -1394,9 +1394,9 @@ def fused_experts(hidden_states: torch.Tensor,
1394
# E8M0 scale, which means we requantize the weight and input to the specific
1395
# scale. Fallen back to cutlass or triton for some cases would cause
1396
# accuracy issue.
1397
- should_use_deep_gemm = is_blackwell_deep_gemm_e8m0_used(
1398
- ) or _valid_deep_gemm(hidden_states, w1, w2)
1399
- if (allow_deep_gemm and use_fp8_w8a8 and should_use_deep_gemm):
1397
+ if (allow_deep_gemm and use_fp8_w8a8
1398
+ and (is_blackwell_deep_gemm_e8m0_used()
1399
+ or _valid_deep_gemm(hidden_states, w1, w2))):
1400
assert apply_router_weight_on_input is False
1401
assert is_act_and_mul, (
1402
"DeepGemm only supports is_act_and_mul=True for now.")
@@ -1905,7 +1905,6 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1905
workspace2: torch.Tensor,
1906
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
1907
apply_router_weight_on_input: bool,
1908
- extra_expert_args: Optional[dict[str, Any]],
1908
):
1909
# Check constraints.
1910
if self.use_int4_w4a16:
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+7
-8
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import TYPE_CHECKING, Any, Optional
3
+from typing import TYPE_CHECKING, Optional
4
5
import torch
6
@@ -8,7 +8,6 @@ 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.topk_weight_and_reduce import (
10
TopKWeightAndReduceDelegate)
11
-from vllm.model_executor.layers.fused_moe.utils import extract_required_args
11
from vllm.utils import has_triton_kernels
12
13
logger = init_logger(__name__)
@@ -160,12 +159,16 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
159
num_dispatchers: int,
160
w1_precision: "PrecisionConfig",
161
w2_precision: "PrecisionConfig",
162
+ w1_bias: Optional[torch.Tensor],
163
+ w2_bias: Optional[torch.Tensor],
164
):
165
super().__init__(quant_config)
166
self.max_num_tokens = max_num_tokens
167
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
172
173
@property
174
def activation_formats(
@@ -219,11 +222,7 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
222
workspace2: torch.Tensor,
223
expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
224
apply_router_weight_on_input: bool,
222
- extra_expert_args: Optional[dict[str, Any]],
225
):
224
- w1_bias, w2_bias = (extract_required_args(extra_expert_args,
225
- ["w1_bias", "w2_bias"]))
226
-
226
return triton_kernel_fused_experts(
227
output,
228
hidden_states,
@@ -240,8 +239,8 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
239
expert_map=expert_map,
240
w1_scale=w1_scale,
241
w2_scale=w2_scale,
243
- w1_bias=w1_bias,
244
- w2_bias=w2_bias,
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,
vllm/model_executor/layers/fused_moe/layer.py
+50
-41
@@ -37,7 +37,6 @@ from vllm.platforms import current_platform
37
from vllm.platforms.interface import CpuArchEnum
38
from vllm.utils import (direct_register_custom_op, has_deep_ep, has_pplx,
39
round_up)
40
-from vllm.utils.flashinfer import has_flashinfer
40
41
if current_platform.is_cuda_alike():
42
from .fused_batched_moe import BatchedTritonExperts
@@ -49,9 +48,6 @@ if current_platform.is_cuda_alike():
48
from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
49
from .deepep_ll_prepare_finalize import (DEEPEP_QUANT_BLOCK_SHAPE,
50
DeepEPLLPrepareAndFinalize)
52
- if has_flashinfer():
53
- from .flashinfer_cutlass_prepare_finalize import (
54
- FlashInferCutlassMoEPrepareAndFinalize)
51
else:
52
fused_experts = None # type: ignore
53
FusedMoEPermuteExpertsUnpermute = None # type: ignore
@@ -80,7 +76,12 @@ class FusedMoeWeightScaleSupported(Enum):
76
77
class FusedMoEMethodBase(QuantizeMethodBase):
78
83
- moe: FusedMoEConfig
79
+ # TODO(bnell): also pass quant_config?
80
+ def __init__(self, moe: FusedMoEConfig):
81
+ super().__init__()
82
+ self.moe = moe
83
+ self.fused_experts: Optional[Callable] = None
84
+ self.topk_indices_dtype = None
85
86
@abstractmethod
87
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -99,16 +100,16 @@ class FusedMoEMethodBase(QuantizeMethodBase):
100
return False
101
102
@staticmethod
102
- def maybe_make_prepare_finalize(
103
- moe: FusedMoEConfig) -> Optional[FusedMoEPrepareAndFinalize]:
103
+ def _maybe_make_prepare_finalize(
104
+ moe: FusedMoEConfig, ) -> Optional[FusedMoEPrepareAndFinalize]:
105
all2all_manager = get_ep_group().device_communicator.all2all_manager
106
assert all2all_manager is not None
107
108
prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None
109
109
- if moe.use_flashinfer_cutlass_kernels:
110
- prepare_finalize = FlashInferCutlassMoEPrepareAndFinalize(
111
- quant_dtype=moe.quant_dtype, )
110
+ assert not moe.use_flashinfer_cutlass_kernels, \
111
+ "Must be created in modelopt.py"
112
+
113
if moe.use_pplx_kernels:
114
hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
115
moe.max_num_tokens,
@@ -188,14 +189,25 @@ class FusedMoEMethodBase(QuantizeMethodBase):
189
190
return prepare_finalize
191
191
- def init_prepare_finalize(self, moe: FusedMoEConfig):
192
- self.moe = moe
193
- prepare_finalize = FusedMoEMethodBase.maybe_make_prepare_finalize(
194
- self.moe)
192
+ def maybe_make_prepare_finalize(
193
+ self,
194
+ moe: FusedMoEConfig,
195
+ ) -> Optional[FusedMoEPrepareAndFinalize]:
196
+ if moe.moe_parallel_config.use_all2all_kernels:
197
+ return FusedMoEMethodBase._maybe_make_prepare_finalize(moe)
198
+ else:
199
+ return None
200
+
201
+ def init_prepare_finalize(self):
202
+ assert self.moe is not None
203
+ prepare_finalize = self.maybe_make_prepare_finalize(self.moe)
204
196
- self.topk_indices_dtype = None
205
if prepare_finalize is not None:
198
- logger.debug("%s", prepare_finalize.__class__.__name__)
206
+ logger.debug("%s for %s(%s)", prepare_finalize.__class__.__name__,
207
+ self, id(self))
208
+ assert self.topk_indices_dtype is None
209
+ assert self.fused_experts is None, \
210
+ f"Attempt to override experts for {id(self)}!"
211
self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
212
experts = self.select_gemm_impl(prepare_finalize, self.moe)
213
self.fused_experts = FusedMoEModularKernel(
@@ -214,12 +226,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
226
f"{self.__class__.__name__} must select appropriate gemm "
227
"implementation based on the prepare_finalize")
228
217
- def maybe_swap_experts_impl(
218
- self,
219
- moe_parallel_config: FusedMoEParallelConfig,
220
- ):
221
- pass
222
-
229
@abstractmethod
230
def apply(
231
self,
@@ -251,10 +257,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
257
"""MoE method without quantization."""
258
259
def __init__(self, moe: FusedMoEConfig):
254
- super().__init__()
255
- self.fused_experts = fused_experts # type: ignore
256
- self.topk_indices_dtype = None
257
- self.moe = moe
260
+ super().__init__(moe)
261
self.has_bias = self.moe.has_bias
262
self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
263
if self.rocm_aiter_moe_enabled:
@@ -266,6 +269,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
269
def select_gemm_impl(
270
self,
271
prepare_finalize: FusedMoEPrepareAndFinalize,
272
+ # TODO(bnell): Remove. Every layer should have an moe config object.
273
moe: FusedMoEConfig,
274
) -> FusedMoEPermuteExpertsUnpermute:
275
if (prepare_finalize.activation_format ==
@@ -474,12 +478,30 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
478
expert_map=expert_map,
479
activation=activation,
480
apply_router_weight_on_input=apply_router_weight_on_input)
481
+ elif self.fused_experts is not None:
482
+ if self.has_bias:
483
+ raise ValueError(
484
+ "FusedMoEModularKernel does not support bias.")
485
+ return self.fused_experts(
486
+ hidden_states=x,
487
+ w1=layer.w13_weight,
488
+ w2=layer.w2_weight,
489
+ topk_weights=topk_weights,
490
+ topk_ids=topk_ids,
491
+ inplace=True,
492
+ activation=activation,
493
+ apply_router_weight_on_input=apply_router_weight_on_input,
494
+ global_num_experts=global_num_experts,
495
+ expert_map=expert_map,
496
+ )
497
else:
478
- # add w1_bias/w2_bias to kwargs if they exist
479
- kwargs = dict(
498
+ assert fused_experts is not None
499
+ return fused_experts(
500
hidden_states=x,
501
w1=layer.w13_weight,
502
w2=layer.w2_weight,
503
+ w1_bias=layer.w13_bias if self.has_bias else None,
504
+ w2_bias=layer.w2_bias if self.has_bias else None,
505
topk_weights=topk_weights,
506
topk_ids=topk_ids,
507
inplace=True,
@@ -488,17 +510,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
510
global_num_experts=global_num_experts,
511
expert_map=expert_map,
512
)
491
- if isinstance(self.fused_experts,
492
- FusedMoEModularKernel) and self.has_bias:
493
- raise ValueError(
494
- "FusedMoEModularKernel does not support bias.")
495
- if self.has_bias:
496
- kwargs.update({
497
- "w1_bias": getattr(layer, "w13_bias", None),
498
- "w2_bias": getattr(layer, "w2_bias", None),
499
- })
500
-
501
- return self.fused_experts(**kwargs)
513
514
def forward_cpu(
515
self,
@@ -868,8 +879,6 @@ class FusedMoE(CustomOp):
879
moe_quant_params["intermediate_size_full"] = intermediate_size
880
881
self.quant_method.create_weights(layer=self, **moe_quant_params)
871
- if isinstance(self.quant_method, FusedMoEMethodBase):
872
- self.quant_method.maybe_swap_experts_impl(self.moe_parallel_config)
882
883
# Chunked all2all staging tensor
884
self.batched_hidden_states: Optional[torch.Tensor] = None
vllm/model_executor/layers/fused_moe/modular_kernel.py
+59
-58
@@ -4,7 +4,7 @@ from abc import ABC, abstractmethod
4
from dataclasses import dataclass
5
from enum import Enum
6
from math import prod
7
-from typing import Any, Optional, final
7
+from typing import Optional, final
8
9
import torch
10
@@ -150,15 +150,23 @@ class FusedMoEPrepareAndFinalize(ABC):
150
151
@abstractmethod
152
def prepare(
153
- self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
154
- a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
155
- topk_ids: torch.Tensor, num_experts: int,
156
- expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
153
+ self,
154
+ a1: torch.Tensor,
155
+ a1_scale: Optional[torch.Tensor],
156
+ a2_scale: Optional[torch.Tensor],
157
+ topk_weights: torch.Tensor,
158
+ topk_ids: torch.Tensor,
159
+ num_experts: int,
160
+ expert_map: Optional[torch.Tensor],
161
+ apply_router_weight_on_input: bool,
162
quant_config: FusedMoEQuantConfig,
158
- extra_prepare_args: Optional[dict[str, Any]]
159
- ) -> tuple[torch.Tensor, Optional[torch.Tensor],
160
- Optional[ExpertTokensMetadata], Optional[torch.Tensor],
161
- Optional[torch.Tensor]]:
163
+ ) -> tuple[
164
+ torch.Tensor,
165
+ Optional[torch.Tensor],
166
+ Optional[ExpertTokensMetadata],
167
+ Optional[torch.Tensor],
168
+ Optional[torch.Tensor],
169
+ ]:
170
"""
171
Perform any quantization (and/or) dispatching needed
172
for this kernel.
@@ -186,11 +194,15 @@ class FusedMoEPrepareAndFinalize(ABC):
194
raise NotImplementedError
195
196
@abstractmethod
189
- def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
190
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
191
- apply_router_weight_on_input: bool,
192
- weight_and_reduce_impl: TopKWeightAndReduce,
193
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
197
+ def finalize(
198
+ self,
199
+ output: torch.Tensor,
200
+ fused_expert_output: torch.Tensor,
201
+ topk_weights: torch.Tensor,
202
+ topk_ids: torch.Tensor,
203
+ apply_router_weight_on_input: bool,
204
+ weight_and_reduce_impl: TopKWeightAndReduce,
205
+ ) -> None:
206
"""
207
Perform any combine plus apply weights and perform a reduction on the
208
fused experts output.
@@ -368,7 +380,6 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
380
workspace2: torch.Tensor,
381
expert_tokens_meta: Optional[ExpertTokensMetadata],
382
apply_router_weight_on_input: bool,
371
- extra_expert_args: Optional[dict[str, Any]],
383
):
384
"""
385
This function computes the intermediate result of a Mixture of Experts
@@ -454,18 +465,27 @@ class FusedMoEModularKernel(torch.nn.Module):
465
f"{fused_experts.activation_formats[0]}")
466
467
def _do_fused_experts(
457
- self, fused_out: Optional[torch.Tensor], a1: torch.Tensor,
458
- a1q: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
459
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
460
- activation: str, global_num_experts: int, local_num_experts: int,
461
- expert_map: Optional[torch.Tensor],
462
- w1_scale: Optional[torch.Tensor], w2_scale: Optional[torch.Tensor],
463
- w1_zp: Optional[torch.Tensor], w2_zp: Optional[torch.Tensor],
464
- a1q_scale: Optional[torch.Tensor],
465
- a2_scale: Optional[torch.Tensor],
466
- expert_tokens_meta: Optional[ExpertTokensMetadata],
467
- apply_router_weight_on_input: bool,
468
- extra_expert_args: Optional[dict[str, Any]]) -> torch.Tensor:
468
+ self,
469
+ fused_out: Optional[torch.Tensor],
470
+ a1: torch.Tensor,
471
+ a1q: torch.Tensor,
472
+ w1: torch.Tensor,
473
+ w2: torch.Tensor,
474
+ topk_weights: torch.Tensor,
475
+ topk_ids: torch.Tensor,
476
+ activation: str,
477
+ global_num_experts: int,
478
+ local_num_experts: int,
479
+ expert_map: Optional[torch.Tensor],
480
+ w1_scale: Optional[torch.Tensor],
481
+ w2_scale: Optional[torch.Tensor],
482
+ w1_zp: Optional[torch.Tensor],
483
+ w2_zp: Optional[torch.Tensor],
484
+ a1q_scale: Optional[torch.Tensor],
485
+ a2_scale: Optional[torch.Tensor],
486
+ expert_tokens_meta: Optional[ExpertTokensMetadata],
487
+ apply_router_weight_on_input: bool,
488
+ ) -> torch.Tensor:
489
490
_, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
491
@@ -509,7 +529,7 @@ class FusedMoEModularKernel(torch.nn.Module):
529
workspace2=workspace2,
530
expert_tokens_meta=expert_tokens_meta,
531
apply_router_weight_on_input=apply_router_weight_on_input,
512
- extra_expert_args=extra_expert_args)
532
+ )
533
534
return fused_out
535
@@ -533,7 +553,6 @@ class FusedMoEModularKernel(torch.nn.Module):
553
a2_scale: Optional[torch.Tensor],
554
expert_tokens_meta: Optional[ExpertTokensMetadata],
555
apply_router_weight_on_input: bool,
536
- extra_expert_args: Optional[dict[str, Any]],
556
) -> torch.Tensor:
557
558
_, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
@@ -541,6 +560,9 @@ class FusedMoEModularKernel(torch.nn.Module):
560
CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
561
num_chunks = cdiv(M, CHUNK_SIZE)
562
563
+ # TODO(bnell): get rid of one level here, update slice functions
564
+ # to nops on num_chunks==1
565
+
566
if not self.fused_experts.supports_chunking() or num_chunks == 1:
567
return self._do_fused_experts(
568
fused_out=None,
@@ -562,7 +584,7 @@ class FusedMoEModularKernel(torch.nn.Module):
584
a2_scale=a2_scale,
585
expert_tokens_meta=expert_tokens_meta,
586
apply_router_weight_on_input=apply_router_weight_on_input,
565
- extra_expert_args=extra_expert_args)
587
+ )
588
589
# Chunking required case
590
assert num_chunks > 1
@@ -618,15 +640,6 @@ class FusedMoEModularKernel(torch.nn.Module):
640
expert_num_tokens=c_expert_num_tokens,
641
expert_num_tokens_cpu=c_expert_num_tokens_cpu)
642
621
- m = None
622
- if extra_expert_args is not None and 'm' in extra_expert_args:
623
- m = extra_expert_args.get('m')
624
-
625
- if extra_expert_args is not None:
626
- chunked_extra_expert_args = extra_expert_args
627
- else:
628
- chunked_extra_expert_args = {}
629
-
643
for chunk_idx in range(num_chunks):
644
c_a1q, c_a1q_scale, c_a2_scale, c_topk_ids, c_topk_weights = (
645
slice_input_tensors(chunk_idx))
@@ -637,11 +650,6 @@ class FusedMoEModularKernel(torch.nn.Module):
650
expert_tokens_meta, c_topk_ids, local_num_experts,
651
expert_map)
652
640
- s = chunk_idx * CHUNK_SIZE
641
- e = min(s + CHUNK_SIZE, M)
642
-
643
- if m is not None:
644
- chunked_extra_expert_args['m'] = e - s
653
self._do_fused_experts(
654
fused_out=slice_output_tensor(chunk_idx),
655
a1=a1,
@@ -662,7 +670,7 @@ class FusedMoEModularKernel(torch.nn.Module):
670
a2_scale=c_a2_scale,
671
expert_tokens_meta=c_expert_tokens_meta,
672
apply_router_weight_on_input=apply_router_weight_on_input,
665
- extra_expert_args=chunked_extra_expert_args)
673
+ )
674
675
return fused_out
676
@@ -684,9 +692,6 @@ class FusedMoEModularKernel(torch.nn.Module):
692
a1_scale: Optional[torch.Tensor] = None,
693
a2_scale: Optional[torch.Tensor] = None,
694
apply_router_weight_on_input: bool = False,
687
- extra_expert_args: Optional[dict] = None,
688
- extra_prepare_args: Optional[dict] = None,
689
- extra_finalize_args: Optional[dict] = None,
695
) -> torch.Tensor:
696
"""
697
This function computes a Mixture of Experts (MoE) layer using two sets
@@ -719,12 +724,6 @@ class FusedMoEModularKernel(torch.nn.Module):
724
- apply_router_weight_on_input (bool): When true, the topk weights are
725
applied directly on the inputs. This is only applicable when topk is
726
1.
722
- - extra_expert_args (Optional[dict]): Extra keyword arguments to pass to
723
- fused_experts.apply.
724
- - extra_prepare_args (Optional[dict]): Extra keyword arguments to pass
725
- to prepare.
726
- - extra_finalize_args (Optional[dict]): Extra keyword arguments to pass
727
- to finalize.
727
728
Returns:
729
- torch.Tensor: The output tensor after applying the MoE layer.
@@ -748,7 +747,6 @@ class FusedMoEModularKernel(torch.nn.Module):
747
expert_map,
748
apply_router_weight_on_input,
749
self.fused_experts.quant_config,
751
- extra_prepare_args,
750
)
751
752
# Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
@@ -786,12 +784,15 @@ class FusedMoEModularKernel(torch.nn.Module):
784
a2_scale=a2_scale,
785
expert_tokens_meta=expert_tokens_meta,
786
apply_router_weight_on_input=apply_router_weight_on_input,
789
- extra_expert_args=extra_expert_args)
787
+ )
788
789
self.prepare_finalize.finalize(
792
- output, fused_out, topk_weights, topk_ids,
790
+ output,
791
+ fused_out,
792
+ topk_weights,
793
+ topk_ids,
794
apply_router_weight_on_input,
795
self.fused_experts.finalize_weight_and_reduce_impl(),
795
- extra_finalize_args)
796
+ )
797
798
return output
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+21
-12
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional, Union
4
5
import pplx_kernels as pplx
6
import torch
@@ -21,7 +21,7 @@ def pplx_hidden_dim_scale_bytes(
21
max_num_tokens: int,
22
hidden_dim: int,
23
in_dtype: torch.dtype,
24
- quant_dtype: Optional[torch.dtype],
24
+ quant_dtype: Union[torch.dtype, str, None],
25
per_act_token_quant: bool,
26
block_shape: Optional[list[int]],
27
):
@@ -32,6 +32,7 @@ def pplx_hidden_dim_scale_bytes(
32
# ceil_div(hidden_dim, block_size) * sizeof(float32)
33
# For per-token: set to 4 * sizeof(float32) (x4 for alignment)
34
if quant_dtype is not None:
35
+ assert isinstance(quant_dtype, torch.dtype)
36
assert quant_dtype.itemsize == 1
37
hidden_dim_bytes = hidden_dim * quant_dtype.itemsize
38
elem_size = torch.float32.itemsize
@@ -89,12 +90,16 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
90
return self.num_dispatchers_
91
92
def prepare(
92
- self, a1: torch.Tensor, a1_scale: Optional[torch.Tensor],
93
- a2_scale: Optional[torch.Tensor], topk_weights: torch.Tensor,
94
- topk_ids: torch.Tensor, num_experts: int,
95
- expert_map: Optional[torch.Tensor], apply_router_weight_on_input: bool,
93
+ self,
94
+ a1: torch.Tensor,
95
+ a1_scale: Optional[torch.Tensor],
96
+ a2_scale: Optional[torch.Tensor],
97
+ topk_weights: torch.Tensor,
98
+ topk_ids: torch.Tensor,
99
+ num_experts: int,
100
+ expert_map: Optional[torch.Tensor],
101
+ apply_router_weight_on_input: bool,
102
quant_config: FusedMoEQuantConfig,
97
- extra_prepare_args: Optional[dict[str, Any]]
103
) -> tuple[torch.Tensor, Optional[torch.Tensor],
104
Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
105
Optional[torch.Tensor]]:
@@ -213,11 +218,15 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
218
219
return expert_x, expert_x_scale, expert_tokens_meta, None, None
220
216
- def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
217
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
218
- apply_router_weight_on_input: bool,
219
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
220
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
221
+ def finalize(
222
+ self,
223
+ output: torch.Tensor,
224
+ fused_expert_output: torch.Tensor,
225
+ topk_weights: torch.Tensor,
226
+ topk_ids: torch.Tensor,
227
+ apply_router_weight_on_input: bool,
228
+ weight_and_reduce_impl: mk.TopKWeightAndReduce,
229
+ ) -> None:
230
assert isinstance(
231
weight_and_reduce_impl, TopKWeightAndReduceDelegate
232
), ("Weight application and reduction happens in the combine kernel.")
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+18
-25
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional
4
5
import torch
6
@@ -38,7 +38,6 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
38
expert_map: Optional[torch.Tensor],
39
apply_router_weight_on_input: bool,
40
quant_config: FusedMoEQuantConfig,
41
- extra_prepare_args: Optional[dict[str, Any]],
41
) -> tuple[torch.Tensor, Optional[torch.Tensor],
42
Optional[mk.ExpertTokensMetadata], Optional[torch.Tensor],
43
Optional[torch.Tensor]]:
@@ -50,32 +49,26 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
49
"apply_router_weight_on_input is only implemented for topk=1"
50
a1.mul_(topk_weights.to(a1.dtype))
51
53
- if (extra_prepare_args is not None
54
- and extra_prepare_args.get("skip_quant", True)):
55
- # Skip quantization if explicitly requested
56
- return a1, None, None, None, None
57
-
52
a1q, a1q_scale = moe_kernel_quantize_input(
53
a1, a1_scale, quant_config.quant_dtype,
54
quant_config.per_act_token_quant, quant_config.block_shape)
55
56
return a1q, a1q_scale, None, None, None
57
64
- def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
65
- topk_weights: torch.Tensor, topk_ids: torch.Tensor,
66
- apply_router_weight_on_input: bool,
67
- weight_and_reduce_impl: mk.TopKWeightAndReduce,
68
- extra_finalize_args: Optional[dict[str, Any]]) -> None:
69
- if (extra_finalize_args is not None
70
- and extra_finalize_args.get("skip_weight_reduce", True)):
71
- assert output.shape == fused_expert_output.shape
72
- output.copy_(fused_expert_output)
73
- else:
74
- if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
75
- weight_and_reduce_impl = TopKWeightAndReduceContiguous()
76
- weight_and_reduce_impl.apply(
77
- output=output,
78
- fused_expert_output=fused_expert_output,
79
- topk_weights=topk_weights,
80
- topk_ids=topk_ids,
81
- apply_router_weight_on_input=apply_router_weight_on_input)
58
+ def finalize(
59
+ self,
60
+ output: torch.Tensor,
61
+ fused_expert_output: torch.Tensor,
62
+ topk_weights: torch.Tensor,
63
+ topk_ids: torch.Tensor,
64
+ apply_router_weight_on_input: bool,
65
+ weight_and_reduce_impl: mk.TopKWeightAndReduce,
66
+ ) -> None:
67
+ if isinstance(weight_and_reduce_impl, TopKWeightAndReduceDelegate):
68
+ weight_and_reduce_impl = TopKWeightAndReduceContiguous()
69
+ weight_and_reduce_impl.apply(
70
+ output=output,
71
+ fused_expert_output=fused_expert_output,
72
+ topk_weights=topk_weights,
73
+ topk_ids=topk_ids,
74
+ apply_router_weight_on_input=apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+23
-14
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Any, Optional
3
+from typing import Optional
4
5
import torch
6
@@ -119,18 +119,28 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
119
local_num_experts,
120
expert_tokens_meta)
121
122
- def apply(self, output: torch.Tensor, hidden_states: torch.Tensor,
123
- w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor,
124
- topk_ids: torch.Tensor, activation: str, global_num_experts: int,
125
- expert_map: Optional[torch.Tensor],
126
- w1_scale: Optional[torch.Tensor],
127
- w2_scale: Optional[torch.Tensor], w1_zp: Optional[torch.Tensor],
128
- w2_zp: Optional[torch.Tensor], a1q_scale: Optional[torch.Tensor],
129
- a2_scale: Optional[torch.Tensor], workspace13: torch.Tensor,
130
- workspace2: torch.Tensor,
131
- expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
132
- apply_router_weight_on_input: bool,
133
- extra_expert_args: Optional[dict[str, Any]]):
122
+ def apply(
123
+ self,
124
+ output: torch.Tensor,
125
+ hidden_states: torch.Tensor,
126
+ w1: torch.Tensor,
127
+ w2: torch.Tensor,
128
+ topk_weights: torch.Tensor,
129
+ topk_ids: torch.Tensor,
130
+ activation: str,
131
+ global_num_experts: int,
132
+ 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],
137
+ a1q_scale: Optional[torch.Tensor],
138
+ a2_scale: Optional[torch.Tensor],
139
+ workspace13: torch.Tensor,
140
+ workspace2: torch.Tensor,
141
+ expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
142
+ apply_router_weight_on_input: bool,
143
+ ):
144
use_deep_gemm = (self.allow_deep_gemm
145
and (_valid_deep_gemm(hidden_states, w1, w2)
146
or is_blackwell_deep_gemm_e8m0_used()))
@@ -158,5 +168,4 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
168
workspace2,
169
expert_tokens_meta,
170
apply_router_weight_on_input,
161
- extra_expert_args,
171
)
vllm/model_executor/layers/fused_moe/utils.py
+2
-16
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from math import prod
4
-from typing import Any, Optional, Union
4
+from typing import Optional, Union
5
6
import torch
7
@@ -189,7 +189,7 @@ def moe_kernel_quantize_input(
189
return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
190
elif quant_dtype == torch.int8:
191
return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
192
- elif quant_dtype == torch.uint8: # nvfp4
192
+ elif quant_dtype == "nvfp4":
193
return _fp4_quantize(A,
194
A_scale,
195
is_sf_swizzled_layout=is_fp4_scale_swizzled)
@@ -252,17 +252,3 @@ def _validate_scale_shape(
252
assert block_shape is not None
253
expected = (a.shape[0], cdiv(a.shape[1], block_shape[1]))
254
assert a_scale.shape == expected, f"{a_scale.shape} == {expected}"
255
-
256
-
257
-def extract_required_args(
258
- extra_args: Optional[dict[str, Any]],
259
- required_keys: list[str],
260
-) -> tuple[Any, ...]:
261
- if extra_args is None:
262
- raise ValueError("`extra_args` must be provided.")
263
-
264
- missing_keys = [k for k in required_keys if k not in extra_args]
265
- if missing_keys:
266
- raise ValueError(f"Missing keys in `extra_args`: {missing_keys}")
267
-
268
- return tuple(extra_args[k] for k in required_keys)
vllm/model_executor/layers/quantization/auto_round.py
+2
-2
@@ -241,7 +241,7 @@ class AutoRoundConfig(QuantizationConfig):
241
242
if isinstance(layer, FusedMoE):
243
if use_marlin:
244
- return AWQMoEMethod(quant_args_marlin)
244
+ return AWQMoEMethod(quant_args_marlin, layer.moe)
245
from vllm.model_executor.layers.quantization.moe_wna16 import (
246
MoeWNA16Config)
247
@@ -339,7 +339,7 @@ class AutoRoundConfig(QuantizationConfig):
339
}
340
return MoeWNA16Config.from_config(config).get_quant_method(
341
layer, prefix)
342
- return GPTQMarlinMoEMethod(quant_args_marlin)
342
+ return GPTQMarlinMoEMethod(quant_args_marlin, layer.moe)
343
344
if isinstance(layer, (LinearBase, ParallelLMHead)):
345
if use_marlin:
vllm/model_executor/layers/quantization/awq.py
+1
-1
@@ -113,7 +113,7 @@ class AWQConfig(QuantizationConfig):
113
}
114
awq_marlin_config = AWQMarlinConfig.from_config(
115
marlin_compatible_config_dict)
116
- return AWQMoEMethod(awq_marlin_config)
116
+ return AWQMoEMethod(awq_marlin_config, layer.moe_config)
117
return None
118
119
vllm/model_executor/layers/quantization/awq_marlin.py
+13
-5
@@ -10,7 +10,7 @@ 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.layer import (
13
- FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
13
+ FusedMoE, FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
14
UnquantizedFusedMoEMethod)
15
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
16
UnquantizedLinearMethod,
@@ -151,7 +151,7 @@ class AWQMarlinConfig(QuantizationConfig):
151
"Falling back to Moe WNA16 kernels.")
152
return MoeWNA16Config.from_config(
153
self.full_config).get_quant_method(layer, prefix)
154
- return AWQMoEMethod(self)
154
+ return AWQMoEMethod(self, layer.moe_config)
155
return None
156
157
@classmethod
@@ -328,7 +328,12 @@ class AWQMarlinLinearMethod(LinearMethodBase):
328
329
class AWQMoEMethod(FusedMoEMethodBase):
330
331
- def __init__(self, quant_config: AWQMarlinConfig):
331
+ def __init__(
332
+ self,
333
+ quant_config: AWQMarlinConfig,
334
+ moe: FusedMoEConfig,
335
+ ):
336
+ super().__init__(moe)
337
self.quant_config = quant_config
338
if self.quant_config.weight_bits != 4:
339
raise ValueError("AWQMoEMethod only supports 4bit now.")
@@ -500,6 +505,8 @@ class AWQMoEMethod(FusedMoEMethodBase):
505
logical_to_physical_map: Optional[torch.Tensor] = None,
506
logical_replica_count: Optional[torch.Tensor] = None,
507
) -> torch.Tensor:
508
+ assert self.fused_experts is None
509
+
510
if enable_eplb:
511
raise NotImplementedError(
512
"EPLB not supported for `AWQMoEMethod` yet.")
@@ -516,7 +523,8 @@ class AWQMoEMethod(FusedMoEMethodBase):
523
num_expert_group=num_expert_group,
524
custom_routing_function=custom_routing_function,
525
scoring_func=scoring_func,
519
- e_score_correction_bias=e_score_correction_bias)
526
+ e_score_correction_bias=e_score_correction_bias,
527
+ indices_type=self.topk_indices_dtype)
528
529
return torch.ops.vllm.fused_marlin_moe(
530
x,
@@ -535,4 +543,4 @@ class AWQMoEMethod(FusedMoEMethodBase):
543
expert_map=expert_map,
544
w1_zeros=layer.w13_qzeros,
545
w2_zeros=layer.w2_qzeros,
538
- workspace=layer.workspace)
\ No newline at end of file
546
+ workspace=layer.workspace)
vllm/model_executor/layers/quantization/bitsandbytes.py
+9
-3
@@ -7,6 +7,7 @@ import torch
7
from packaging import version
8
9
from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
10
+ FusedMoEConfig,
11
FusedMoEMethodBase)
12
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
13
UnquantizedLinearMethod,
@@ -132,7 +133,7 @@ class BitsAndBytesConfig(QuantizationConfig):
133
return UnquantizedLinearMethod()
134
return BitsAndBytesLinearMethod(self)
135
elif isinstance(layer, FusedMoE):
135
- return BitsAndBytesMoEMethod(self)
136
+ return BitsAndBytesMoEMethod(self, layer.moe_config)
137
return None
138
139
@@ -411,7 +412,12 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
412
quant_config: The BitsAndBytes quantization config.
413
"""
414
414
- def __init__(self, quant_config: BitsAndBytesConfig):
415
+ def __init__(
416
+ self,
417
+ quant_config: BitsAndBytesConfig,
418
+ moe: FusedMoEConfig,
419
+ ):
420
+ super().__init__(moe)
421
try:
422
import bitsandbytes
423
if version.parse(
@@ -422,7 +428,6 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
428
raise ImportError("Please install bitsandbytes>=0.46.1 via "
429
"`pip install bitsandbytes>=0.46.1` to use "
430
"bitsandbytes quantizer.") from err
425
- self.topk_indices_dtype = None
431
self.quant_config = quant_config
432
433
def create_weights(
@@ -470,6 +475,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):
475
logical_replica_count: Optional[torch.Tensor] = None,
476
) -> torch.Tensor:
477
from vllm.model_executor.layers.fused_moe import fused_experts
478
+ assert self.fused_experts is None
479
480
if enable_eplb:
481
raise NotImplementedError(
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+112
-56
@@ -11,20 +11,21 @@ from compressed_tensors.quantization import (ActivationOrdering,
11
QuantizationStrategy)
12
13
import vllm.envs as envs
14
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
15
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)
20
-from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa
21
- FlashInferCutlassMoEPrepareAndFinalize)
21
+from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
22
+ is_valid_flashinfer_cutlass_fused_moe)
23
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
24
WNA16_SUPPORTED_BITS, WNA16_SUPPORTED_TYPES_MAP)
25
from vllm.model_executor.layers.quantization.utils import replace_parameter
26
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
26
- build_flashinfer_fp4_cutlass_moe_kernel,
27
- flashinfer_fp4_cutlass_moe_forward, reorder_w1w3_to_w3w1)
27
+ build_flashinfer_fp4_cutlass_moe_prepare_finalize, reorder_w1w3_to_w3w1,
28
+ select_nvfp4_gemm_impl)
29
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
30
check_moe_marlin_supports_layer, marlin_make_workspace_new,
31
marlin_moe_permute_scales)
@@ -58,6 +59,9 @@ __all__ = [
59
60
class CompressedTensorsMoEMethod(FusedMoEMethodBase):
61
62
+ def __init_(self, moe: FusedMoEConfig):
63
+ super().__init__(moe)
64
+
65
@staticmethod
66
def get_moe_method(
67
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
@@ -81,18 +85,22 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
85
"WNA16MoE is not supported with actorder=group/dynamic."
86
)
87
logger.info_once("Using CompressedTensorsWNA16MoEMethod")
84
- return CompressedTensorsWNA16MoEMethod(quant_config)
88
+ return CompressedTensorsWNA16MoEMethod(quant_config,
89
+ layer.moe_config)
90
else:
91
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
87
- return CompressedTensorsWNA16MarlinMoEMethod(quant_config)
92
+ return CompressedTensorsWNA16MarlinMoEMethod(
93
+ quant_config, layer.moe_config)
94
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
89
- return CompressedTensorsW4A4MoeMethod()
95
+ return CompressedTensorsW4A4MoeMethod(layer.moe_config, layer)
96
elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
97
or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
98
or quant_config._is_fp8_w8a8(weight_quant, input_quant)):
93
- return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
99
+ return CompressedTensorsW8A8Fp8MoEMethod(quant_config,
100
+ layer.moe_config)
101
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
95
- return CompressedTensorsW8A8Int8MoEMethod(quant_config)
102
+ return CompressedTensorsW8A8Int8MoEMethod(quant_config,
103
+ layer.moe_config)
104
else:
105
raise RuntimeError(
106
f"Unsupported FusedMoe scheme: {weight_quant}, {input_quant}")
@@ -100,15 +108,16 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
108
109
class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
110
103
- def __init__(self):
111
+ def __init__(self, moe: FusedMoEConfig, layer: torch.nn.Module):
112
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
113
detect_nvfp4_moe_support)
114
+ super().__init__(moe)
115
_nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)
116
self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported
117
self.allow_flashinfer = _nvfp4.allow_flashinfer
118
self.use_marlin = _nvfp4.use_marlin
119
self.group_size = 16
111
- self.fused_experts = None # type: ignore[assignment]
120
+ self.layer = layer
121
122
def create_weights(self, layer: torch.nn.Module, num_experts: int,
123
hidden_size: int, intermediate_size_per_partition: int,
@@ -265,19 +274,36 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
274
layer.w2_input_scale_quant = torch.nn.Parameter(
275
(layer.w2_input_global_scale), requires_grad=False)
276
268
- def maybe_swap_experts_impl(self, moe_parallel_config):
277
+ def maybe_make_prepare_finalize(
278
+ self,
279
+ moe: FusedMoEConfig,
280
+ ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
281
if not self.allow_flashinfer:
270
- return
271
- self.fused_experts = build_flashinfer_fp4_cutlass_moe_kernel(
272
- moe_parallel_config)
282
+ return super().maybe_make_prepare_finalize(moe)
283
274
- def select_gemm_impl(self, prepare_finalize, moe):
275
- """Return the appropriate GEMM experts implementation."""
276
- assert moe is not None and prepare_finalize is not None
277
- from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( # noqa: E501
278
- select_nvfp4_gemm_impl)
284
+ prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(
285
+ moe,
286
+ a1_gscale=self.layer.w13_input_scale_quant,
287
+ )
288
+ logger.debug_once("%s", prepare_finalize.__class__.__name__)
289
+ return prepare_finalize
290
280
- return select_nvfp4_gemm_impl(self.allow_flashinfer, moe, logger)
291
+ def select_gemm_impl(
292
+ self,
293
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
294
+ moe: FusedMoEConfig,
295
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
296
+ """Return the appropriate GEMM experts implementation."""
297
+ experts = select_nvfp4_gemm_impl(
298
+ moe,
299
+ g1_alphas=self.layer.g1_alphas,
300
+ g2_alphas=self.layer.g2_alphas,
301
+ a1_gscale=self.layer.w13_input_scale_quant,
302
+ a2_gscale=self.layer.w2_input_scale_quant,
303
+ allow_flashinfer=self.allow_flashinfer,
304
+ )
305
+ logger.debug_once("Using %s", experts.__class__.__name__)
306
+ return experts
307
308
def apply(
309
self,
@@ -301,6 +327,8 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
327
logical_to_physical_map: Optional[torch.Tensor] = None,
328
logical_replica_count: Optional[torch.Tensor] = None,
329
) -> torch.Tensor:
330
+ assert self.fused_experts is None
331
+
332
if enable_eplb:
333
raise NotImplementedError("EPLB not supported for "
334
"`CompressedTensorsW4A4MoeMethod` yet.")
@@ -317,6 +345,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
345
custom_routing_function=custom_routing_function,
346
scoring_func=scoring_func,
347
e_score_correction_bias=e_score_correction_bias,
348
+ indices_type=self.topk_indices_dtype,
349
)
350
351
if self.use_marlin:
@@ -340,15 +369,22 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
369
370
# FlashInfer fused experts path
371
if self.fused_experts is not None:
343
- return flashinfer_fp4_cutlass_moe_forward(
344
- self.fused_experts,
345
- layer,
346
- x,
347
- topk_weights,
348
- topk_ids,
372
+ assert is_valid_flashinfer_cutlass_fused_moe(
373
+ x, layer.w13_weight, layer.w2_weight), (
374
+ "Flashinfer CUTLASS Fused MoE not applicable!")
375
+
376
+ return self.fused_experts(
377
+ hidden_states=x,
378
+ w1=layer.w13_weight,
379
+ w2=layer.w2_weight,
380
+ topk_weights=topk_weights,
381
+ topk_ids=topk_ids,
382
+ inplace=False, # TODO(shuw): fix later, now output is high prec
383
activation=activation,
384
global_num_experts=global_num_experts,
385
expert_map=expert_map,
386
+ w1_scale=layer.w13_blockscale_swizzled,
387
+ w2_scale=layer.w2_blockscale_swizzled,
388
apply_router_weight_on_input=apply_router_weight_on_input,
389
)
390
@@ -376,7 +412,6 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
412
n=layer.w2_weight.shape[2] * 2,
413
k=x.shape[1],
414
e=layer.w13_weight.shape[0],
379
- device=x.device,
415
apply_router_weight_on_input=apply_router_weight_on_input).to(
416
x.dtype)
417
@@ -384,15 +419,16 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
419
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
420
421
def __init__(
387
- self,
388
- quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
422
+ self,
423
+ quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
424
+ moe: FusedMoEConfig,
425
):
426
+ super().__init__(moe)
427
self.quant_config = quant_config
428
self.weight_quant = self.quant_config.target_scheme_map["Linear"].get(
429
"weights")
430
self.input_quant = self.quant_config.target_scheme_map["Linear"].get(
431
"input_activations")
395
- self.topk_indices_dtype = None
432
433
per_tensor = (self.weight_quant.strategy == QuantizationStrategy.TENSOR
434
and self.input_quant.strategy
@@ -429,7 +465,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
465
self.weight_quant, self.input_quant)
466
self.use_cutlass = (quant_config._is_fp8_w8a8_sm90(
467
self.weight_quant, self.input_quant) or self.is_fp8_w8a8_sm100)
432
- self.fused_experts = None # type: ignore[assignment]
468
self.disable_expert_map = False
469
470
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -614,25 +649,31 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
649
) -> FusedMoEPermuteExpertsUnpermute:
650
# cutlass path
651
if self.use_cutlass:
617
- from vllm.model_executor.layers.fused_moe import CutlassExpertsFp8
652
+ from vllm.model_executor.layers.fused_moe import (
653
+ CutlassBatchedExpertsFp8, CutlassExpertsFp8)
654
619
- use_batched_format = (prepare_finalize.activation_format ==
620
- FusedMoEActivationFormat.BatchedExperts)
655
+ experts: FusedMoEPermuteExpertsUnpermute
656
657
num_dispatchers = prepare_finalize.num_dispatchers()
623
- num_experts = (moe.num_local_experts
624
- if use_batched_format else moe.num_experts)
625
-
626
- logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
658
628
- experts = CutlassExpertsFp8(
629
- num_experts,
630
- moe.in_dtype,
631
- self.input_quant.strategy == QuantizationStrategy.TOKEN,
632
- self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
633
- num_dispatchers=num_dispatchers,
634
- use_batched_format=use_batched_format,
635
- )
659
+ if (prepare_finalize.activation_format ==
660
+ FusedMoEActivationFormat.BatchedExperts):
661
+ logger.debug("CutlassBatchedExpertsFp8(%s)",
662
+ self.__class__.__name__)
663
+ experts = CutlassBatchedExpertsFp8(
664
+ moe.num_local_experts,
665
+ num_dispatchers,
666
+ moe.in_dtype,
667
+ self.input_quant.strategy == QuantizationStrategy.TOKEN,
668
+ self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
669
+ )
670
+ else:
671
+ logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
672
+ experts = CutlassExpertsFp8(
673
+ moe.in_dtype,
674
+ self.input_quant.strategy == QuantizationStrategy.TOKEN,
675
+ self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
676
+ )
677
678
self.disable_expert_map = (num_dispatchers > 1
679
or not experts.supports_expert_map())
@@ -834,9 +875,11 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
875
class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
876
877
def __init__(
837
- self,
838
- quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
878
+ self,
879
+ quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
880
+ moe: FusedMoEConfig,
881
):
882
+ super().__init__(moe)
883
self.quant_config = quant_config
884
self.weight_quant = self.quant_config.target_scheme_map["Linear"].get(
885
"weights")
@@ -934,6 +977,8 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
977
logical_to_physical_map: Optional[torch.Tensor] = None,
978
logical_replica_count: Optional[torch.Tensor] = None,
979
) -> torch.Tensor:
980
+ assert self.fused_experts is None
981
+
982
if enable_eplb:
983
raise NotImplementedError(
984
"EPLB not supported for "
@@ -951,7 +996,8 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
996
num_expert_group=num_expert_group,
997
custom_routing_function=custom_routing_function,
998
scoring_func=scoring_func,
954
- e_score_correction_bias=e_score_correction_bias)
999
+ e_score_correction_bias=e_score_correction_bias,
1000
+ indices_type=self.topk_indices_dtype)
1001
1002
return fused_experts(
1003
hidden_states=x,
@@ -975,9 +1021,11 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1021
class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1022
1023
def __init__(
978
- self,
979
- quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
1024
+ self,
1025
+ quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1026
+ moe: FusedMoEConfig,
1027
):
1028
+ super().__init__(moe)
1029
self.quant_config = quant_config
1030
# TODO: @dsikka: refactor this to use schemes as other kernels
1031
# are supported + check if the layer is being ignored.
@@ -1233,6 +1281,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1281
logical_to_physical_map: Optional[torch.Tensor] = None,
1282
logical_replica_count: Optional[torch.Tensor] = None,
1283
) -> torch.Tensor:
1284
+ assert self.fused_experts is None
1285
+
1286
if enable_eplb:
1287
raise NotImplementedError(
1288
"EPLB not supported for "
@@ -1251,7 +1301,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1301
num_expert_group=num_expert_group,
1302
custom_routing_function=custom_routing_function,
1303
scoring_func=scoring_func,
1254
- e_score_correction_bias=e_score_correction_bias)
1304
+ e_score_correction_bias=e_score_correction_bias,
1305
+ indices_type=self.topk_indices_dtype)
1306
1307
return torch.ops.vllm.fused_marlin_moe(
1308
x,
@@ -1279,9 +1330,11 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1330
class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1331
1332
def __init__(
1282
- self,
1283
- quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
1333
+ self,
1334
+ quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1335
+ moe: FusedMoEConfig,
1336
):
1337
+ super().__init__(moe)
1338
self.quant_config = quant_config
1339
# TODO: @dsikka: refactor this to use schemes as other kernels
1340
# are supported + check if the layer is being ignored.
@@ -1459,6 +1512,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1512
logical_to_physical_map: Optional[torch.Tensor] = None,
1513
logical_replica_count: Optional[torch.Tensor] = None,
1514
) -> torch.Tensor:
1515
+ assert self.fused_experts is None
1516
+
1517
if enable_eplb:
1518
raise NotImplementedError("EPLB not supported for "
1519
"`CompressedTensorsWNA16MoEMethod` yet.")
@@ -1475,7 +1530,8 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1530
num_expert_group=num_expert_group,
1531
custom_routing_function=custom_routing_function,
1532
scoring_func=scoring_func,
1478
- e_score_correction_bias=e_score_correction_bias)
1533
+ e_score_correction_bias=e_score_correction_bias,
1534
+ indices_type=self.topk_indices_dtype)
1535
1536
return fused_experts(
1537
x,
vllm/model_executor/layers/quantization/experts_int8.py
+13
-4
@@ -6,7 +6,8 @@ from typing import Any, Callable, Optional
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 import FusedMoE, FusedMoEMethodBase
9
+from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
10
+ FusedMoEMethodBase)
11
from vllm.model_executor.layers.linear import (LinearBase,
12
UnquantizedLinearMethod)
13
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -46,13 +47,18 @@ class ExpertsInt8Config(QuantizationConfig):
47
if isinstance(layer, LinearBase):
48
return UnquantizedLinearMethod()
49
elif isinstance(layer, FusedMoE):
49
- return ExpertsInt8MoEMethod(self)
50
+ return ExpertsInt8MoEMethod(self, layer.moe_config)
51
return None
52
53
54
class ExpertsInt8MoEMethod(FusedMoEMethodBase):
55
55
- def __init__(self, quant_config: ExpertsInt8Config):
56
+ def __init__(
57
+ self,
58
+ quant_config: ExpertsInt8Config,
59
+ moe: FusedMoEConfig,
60
+ ):
61
+ super().__init__(moe)
62
self.quant_config = quant_config
63
64
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -122,6 +128,8 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
128
logical_to_physical_map: Optional[torch.Tensor] = None,
129
logical_replica_count: Optional[torch.Tensor] = None,
130
) -> torch.Tensor:
131
+ assert self.fused_experts is None
132
+
133
if enable_eplb:
134
raise NotImplementedError(
135
"EPLB not supported for `ExpertsInt8MoEMethod` yet.")
@@ -138,7 +146,8 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):
146
num_expert_group=num_expert_group,
147
custom_routing_function=custom_routing_function,
148
scoring_func=scoring_func,
141
- e_score_correction_bias=e_score_correction_bias)
149
+ e_score_correction_bias=e_score_correction_bias,
150
+ indices_type=self.topk_indices_dtype)
151
152
return fused_experts(
153
x,
vllm/model_executor/layers/quantization/fp8.py
+28
-15
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import functools
4
from typing import TYPE_CHECKING, Any, Callable, Optional
5
6
import torch
@@ -142,7 +141,7 @@ class Fp8Config(QuantizationConfig):
141
return UnquantizedLinearMethod()
142
return Fp8LinearMethod(self)
143
elif isinstance(layer, FusedMoE):
145
- return Fp8MoEMethod(self)
144
+ return Fp8MoEMethod(self, layer.moe_config)
145
elif isinstance(layer, Attention):
146
return Fp8KVCacheMethod(self)
147
return None
@@ -479,9 +478,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
478
quant_config: The quantization config.
479
"""
480
482
- def __init__(self, quant_config: Fp8Config):
483
-
484
- from vllm.model_executor.layers.fused_moe import fused_experts
481
+ def __init__(self, quant_config: Fp8Config, moe: FusedMoEConfig):
482
+ super().__init__(moe)
483
self.quant_config = quant_config
484
self.block_quant = self.quant_config.weight_block_size is not None
485
@@ -529,15 +527,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
527
"CutlassBlockScaledGroupedGemm not supported on the current "
528
"platform.")
529
532
- self.topk_indices_dtype = None
533
- self.fused_experts = functools.partial( # type: ignore
534
- fused_experts,
535
- use_fp8_w8a8=True,
536
- block_shape=self.quant_config.weight_block_size,
537
- allow_deep_gemm=self.allow_deep_gemm,
538
- allow_cutlass_block_scaled_grouped_gemm=(
539
- self.allow_cutlass_block_scaled_grouped_gemm))
540
-
530
def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
531
intermediate_size_per_partition: int,
532
params_dtype: torch.dtype, **extra_weight_attrs):
@@ -1033,7 +1022,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1022
num_expert_group=num_expert_group,
1023
topk_group=topk_group,
1024
apply_router_weight_on_input=apply_router_weight_on_input)
1036
- else:
1025
+ elif self.fused_experts is not None:
1026
return self.fused_experts(
1027
hidden_states=x,
1028
w1=layer.w13_weight,
@@ -1052,6 +1041,30 @@ class Fp8MoEMethod(FusedMoEMethodBase):
1041
a1_scale=layer.w13_input_scale,
1042
a2_scale=layer.w2_input_scale,
1043
)
1044
+ else:
1045
+ from vllm.model_executor.layers.fused_moe import fused_experts
1046
+ return fused_experts(
1047
+ hidden_states=x,
1048
+ w1=layer.w13_weight,
1049
+ w2=layer.w2_weight,
1050
+ topk_weights=topk_weights,
1051
+ topk_ids=topk_ids,
1052
+ inplace=True,
1053
+ activation=activation,
1054
+ global_num_experts=global_num_experts,
1055
+ apply_router_weight_on_input=apply_router_weight_on_input,
1056
+ expert_map=expert_map,
1057
+ w1_scale=(layer.w13_weight_scale_inv
1058
+ if self.block_quant else layer.w13_weight_scale),
1059
+ w2_scale=(layer.w2_weight_scale_inv
1060
+ if self.block_quant else layer.w2_weight_scale),
1061
+ a1_scale=layer.w13_input_scale,
1062
+ a2_scale=layer.w2_input_scale,
1063
+ use_fp8_w8a8=True,
1064
+ block_shape=self.quant_config.weight_block_size,
1065
+ allow_deep_gemm=self.allow_deep_gemm,
1066
+ allow_cutlass_block_scaled_grouped_gemm=(
1067
+ self.allow_cutlass_block_scaled_grouped_gemm))
1068
1069
1070
class Fp8KVCacheMethod(BaseKVCacheMethod):
vllm/model_executor/layers/quantization/gguf.py
+12
-3
@@ -11,6 +11,7 @@ from torch.nn.parameter import Parameter, UninitializedParameter
11
from vllm import _custom_ops as ops
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,
14
+ FusedMoEConfig,
15
FusedMoEMethodBase)
16
from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase
17
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -58,7 +59,7 @@ class GGUFConfig(QuantizationConfig):
59
elif isinstance(layer, VocabParallelEmbedding):
60
return GGUFEmbeddingMethod(self)
61
elif isinstance(layer, FusedMoE):
61
- return GGUFMoEMethod(self)
62
+ return GGUFMoEMethod(self, layer.moe_config)
63
return None
64
65
@@ -445,7 +446,12 @@ class GGUFMoEMethod(FusedMoEMethodBase):
446
quant_config: The GGUF quantization config.
447
"""
448
448
- def __init__(self, quant_config: GGUFConfig):
449
+ def __init__(
450
+ self,
451
+ quant_config: GGUFConfig,
452
+ moe: FusedMoEConfig,
453
+ ):
454
+ super().__init__(moe)
455
self.quant_config = quant_config
456
457
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -525,6 +531,8 @@ class GGUFMoEMethod(FusedMoEMethodBase):
531
logical_to_physical_map: Optional[torch.Tensor] = None,
532
logical_replica_count: Optional[torch.Tensor] = None,
533
):
534
+ assert self.fused_experts is None
535
+
536
if enable_eplb:
537
raise NotImplementedError(
538
"EPLB not supported for `GGUFMoEMethod` yet.")
@@ -545,7 +553,8 @@ class GGUFMoEMethod(FusedMoEMethodBase):
553
num_expert_group=num_expert_group,
554
custom_routing_function=custom_routing_function,
555
scoring_func=scoring_func,
548
- e_score_correction_bias=e_score_correction_bias)
556
+ e_score_correction_bias=e_score_correction_bias,
557
+ indices_type=self.topk_indices_dtype)
558
return fused_moe_gguf(x, layer.w13_qweight, layer.w2_qweight,
559
topk_weights, topk_ids,
560
layer.w13_qweight_type.weight_type,
vllm/model_executor/layers/quantization/gptq_marlin.py
+11
-3
@@ -10,7 +10,7 @@ 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.layer import (
13
- FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
13
+ FusedMoE, FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported,
14
UnquantizedFusedMoEMethod)
15
from vllm.model_executor.layers.linear import (LinearMethodBase,
16
set_weight_attrs)
@@ -375,7 +375,12 @@ class GPTQMarlinLinearMethod(LinearMethodBase):
375
class GPTQMarlinMoEMethod(FusedMoEMethodBase):
376
"""MoE Marlin method with quantization."""
377
378
- def __init__(self, quant_config: GPTQMarlinConfig) -> None:
378
+ def __init__(
379
+ self,
380
+ quant_config: GPTQMarlinConfig,
381
+ moe: FusedMoEConfig,
382
+ ) -> None:
383
+ super().__init__(moe)
384
self.quant_config = quant_config
385
if self.quant_config.quant_type.size_bits == 4:
386
self.quant_type = scalar_types.uint4b8
@@ -646,6 +651,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
651
logical_to_physical_map: Optional[torch.Tensor] = None,
652
logical_replica_count: Optional[torch.Tensor] = None,
653
) -> torch.Tensor:
654
+ assert self.fused_experts is None
655
+
656
if enable_eplb:
657
raise NotImplementedError(
658
"EPLB not supported for `GPTQMarlinMoEMethod` yet.")
@@ -662,7 +669,8 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):
669
num_expert_group=num_expert_group,
670
custom_routing_function=custom_routing_function,
671
scoring_func=scoring_func,
665
- e_score_correction_bias=e_score_correction_bias)
672
+ e_score_correction_bias=e_score_correction_bias,
673
+ indices_type=self.topk_indices_dtype)
674
675
return torch.ops.vllm.fused_marlin_moe(
676
x,
vllm/model_executor/layers/quantization/modelopt.py
+66
-33
@@ -12,7 +12,9 @@ import vllm.envs as envs
12
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
14
from vllm.logger import init_logger
15
-from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
15
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
16
+from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
17
+ is_valid_flashinfer_cutlass_fused_moe)
18
from vllm.model_executor.layers.fused_moe.layer import (
19
FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
20
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
@@ -22,8 +24,8 @@ from vllm.model_executor.layers.quantization.base_config import (
24
QuantizationConfig, QuantizeMethodBase)
25
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
26
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
25
- build_flashinfer_fp4_cutlass_moe_kernel,
26
- flashinfer_fp4_cutlass_moe_forward, reorder_w1w3_to_w3w1)
27
+ build_flashinfer_fp4_cutlass_moe_prepare_finalize, reorder_w1w3_to_w3w1,
28
+ select_nvfp4_gemm_impl)
29
from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
30
apply_flashinfer_per_tensor_scale_fp8, register_moe_scaling_factors,
31
rotate_flashinfer_fp8_moe_weights, swap_w13_to_w31)
@@ -177,7 +179,7 @@ class ModelOptFp8Config(QuantizationConfig):
179
elif isinstance(layer, Attention):
180
return ModelOptFp8KVCacheMethod(self)
181
elif isinstance(layer, FusedMoE):
180
- return ModelOptFp8MoEMethod(self)
182
+ return ModelOptFp8MoEMethod(self, layer.moe_config)
183
return None
184
185
@@ -273,7 +275,12 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
275
quant_config: The ModelOpt quantization config.
276
"""
277
276
- def __init__(self, quant_config: ModelOptFp8Config) -> None:
278
+ def __init__(
279
+ self,
280
+ quant_config: ModelOptFp8Config,
281
+ moe: FusedMoEConfig,
282
+ ) -> None:
283
+ super().__init__(moe)
284
self.quant_config = quant_config
285
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
286
cutlass_fp8_supported)
@@ -454,6 +461,8 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
461
logical_to_physical_map: Optional[torch.Tensor] = None,
462
logical_replica_count: Optional[torch.Tensor] = None,
463
) -> torch.Tensor:
464
+ assert self.fused_experts is None
465
+
466
if enable_eplb:
467
raise NotImplementedError(
468
"EPLB not supported for `ModelOptFp8MoEMethod` yet.")
@@ -484,6 +493,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
493
custom_routing_function=custom_routing_function,
494
scoring_func=scoring_func,
495
e_score_correction_bias=e_score_correction_bias,
496
+ indices_type=self.topk_indices_dtype,
497
)
498
from vllm.model_executor.layers.fused_moe.fused_moe import (
499
fused_experts)
@@ -699,7 +709,7 @@ class ModelOptNvFp4Config(QuantizationConfig):
709
elif isinstance(layer, Attention):
710
return ModelOptFp8KVCacheMethod(self)
711
elif isinstance(layer, FusedMoE):
702
- return ModelOptNvFp4FusedMoE(self)
712
+ return ModelOptNvFp4FusedMoE(self, layer.moe_config, layer)
713
return None
714
715
@@ -923,10 +933,17 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
933
quant_config: NVFP4 Quant Config
934
"""
935
926
- def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
927
- self.quant_config = quant_config
936
+ def __init__(
937
+ self,
938
+ quant_config: ModelOptNvFp4Config,
939
+ moe: FusedMoEConfig,
940
+ layer: torch.nn.Module,
941
+ ) -> None:
942
from vllm.model_executor.layers.quantization.utils.nvfp4_moe_support import ( # noqa: E501
943
detect_nvfp4_moe_support)
944
+ super().__init__(moe)
945
+ self.quant_config = quant_config
946
+ self.layer = layer
947
_nvfp4 = detect_nvfp4_moe_support(self.__class__.__name__)
948
self.cutlass_nvfp4_supported = _nvfp4.cutlass_supported
949
self.allow_flashinfer = _nvfp4.allow_flashinfer
@@ -952,27 +969,35 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
969
self.fused_experts: Optional[
970
mk.FusedMoEModularKernel] = None # type: ignore[assignment]
971
955
- def maybe_swap_experts_impl(
972
+ def maybe_make_prepare_finalize(
973
self,
957
- moe_parallel_config: FusedMoEParallelConfig,
958
- ):
974
+ moe: FusedMoEConfig,
975
+ ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
976
if not self.allow_flashinfer:
960
- return
961
- self.fused_experts = build_flashinfer_fp4_cutlass_moe_kernel(
962
- moe_parallel_config)
977
+ return super().maybe_make_prepare_finalize(moe)
978
964
- # This method update self.fused_experts
965
- # only prepare_finalize is not None call select_gemm_impl
966
- # so when native cutlass fp4, fused_expert is in fuse_moe.py fused_expert
967
- # when it's not called(TP case), we still have 2 kernels to use.
968
- def select_gemm_impl(self, prepare_finalize,
969
- moe) -> mk.FusedMoEPermuteExpertsUnpermute:
970
-
971
- assert moe is not None and prepare_finalize is not None
972
- from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( # noqa: E501
973
- select_nvfp4_gemm_impl)
979
+ prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(
980
+ moe,
981
+ a1_gscale=self.layer.w13_input_scale_quant,
982
+ )
983
+ logger.debug_once("%s", prepare_finalize.__class__.__name__)
984
+ return prepare_finalize
985
975
- return select_nvfp4_gemm_impl(self.allow_flashinfer, moe, logger)
986
+ def select_gemm_impl(
987
+ self,
988
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
989
+ moe: FusedMoEConfig,
990
+ ) -> mk.FusedMoEPermuteExpertsUnpermute:
991
+ experts = select_nvfp4_gemm_impl(
992
+ moe,
993
+ g1_alphas=self.layer.g1_alphas,
994
+ g2_alphas=self.layer.g2_alphas,
995
+ a1_gscale=self.layer.w13_input_scale_quant,
996
+ a2_gscale=self.layer.w2_input_scale_quant,
997
+ allow_flashinfer=self.allow_flashinfer,
998
+ )
999
+ logger.debug_once("Using %s", experts.__class__.__name__)
1000
+ return experts
1001
1002
def uses_weight_scale_2_pattern(self) -> bool:
1003
"""
@@ -1362,7 +1387,8 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1387
num_expert_group=num_expert_group,
1388
custom_routing_function=custom_routing_function,
1389
scoring_func=scoring_func,
1365
- e_score_correction_bias=e_score_correction_bias)
1390
+ e_score_correction_bias=e_score_correction_bias,
1391
+ indices_type=self.topk_indices_dtype)
1392
1393
if self.use_marlin:
1394
return torch.ops.vllm.fused_marlin_moe(
@@ -1404,21 +1430,28 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1430
n=layer.w2_weight.shape[2] * 2,
1431
k=x.shape[1],
1432
e=layer.w13_weight.shape[0],
1407
- device=x.device,
1433
expert_map=expert_map,
1434
apply_router_weight_on_input=apply_router_weight_on_input)
1435
else:
1436
assert self.allow_flashinfer and \
1437
self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1413
- out = flashinfer_fp4_cutlass_moe_forward(
1414
- self.fused_experts,
1415
- layer,
1416
- x,
1417
- topk_weights,
1418
- topk_ids,
1438
+
1439
+ assert is_valid_flashinfer_cutlass_fused_moe(
1440
+ x, layer.w13_weight, layer.w2_weight), (
1441
+ "Flashinfer CUTLASS Fused MoE not applicable!")
1442
+
1443
+ out = self.fused_experts(
1444
+ hidden_states=x,
1445
+ w1=layer.w13_weight,
1446
+ w2=layer.w2_weight,
1447
+ topk_weights=topk_weights,
1448
+ topk_ids=topk_ids,
1449
+ inplace=False, # TODO(shuw): fix later, now output is high prec
1450
activation=activation,
1451
global_num_experts=global_num_experts,
1452
expert_map=expert_map,
1453
+ w1_scale=layer.w13_blockscale_swizzled,
1454
+ w2_scale=layer.w2_blockscale_swizzled,
1455
apply_router_weight_on_input=apply_router_weight_on_input,
1456
)
1457
vllm/model_executor/layers/quantization/moe_wna16.py
+12
-4
@@ -7,7 +7,7 @@ import torch
7
8
from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
9
from vllm.model_executor.layers.fused_moe.layer import (
10
- FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
10
+ FusedMoE, FusedMoEConfig, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
11
from vllm.model_executor.layers.linear import (LinearBase,
12
UnquantizedLinearMethod)
13
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -160,7 +160,7 @@ class MoeWNA16Config(QuantizationConfig):
160
else:
161
raise ValueError("moe_wna16 only support gptq and awq.")
162
elif isinstance(layer, FusedMoE):
163
- return MoeWNA16Method(self)
163
+ return MoeWNA16Method(self, layer.moe_config)
164
return None
165
166
@@ -175,7 +175,12 @@ class MoeWNA16Method(FusedMoEMethodBase):
175
quant_config: The MOE WNA16 (W8A16/W4A16) quantization config.
176
"""
177
178
- def __init__(self, quant_config: MoeWNA16Config):
178
+ def __init__(
179
+ self,
180
+ quant_config: MoeWNA16Config,
181
+ moe: FusedMoEConfig,
182
+ ):
183
+ super().__init__(moe)
184
self.quant_config = quant_config
185
186
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -302,6 +307,8 @@ class MoeWNA16Method(FusedMoEMethodBase):
307
logical_to_physical_map: Optional[torch.Tensor] = None,
308
logical_replica_count: Optional[torch.Tensor] = None,
309
) -> torch.Tensor:
310
+ assert self.fused_experts is None
311
+
312
if enable_eplb:
313
raise NotImplementedError(
314
"EPLB not supported for `MoeWNA16Method` yet.")
@@ -318,7 +325,8 @@ class MoeWNA16Method(FusedMoEMethodBase):
325
num_expert_group=num_expert_group,
326
custom_routing_function=custom_routing_function,
327
scoring_func=scoring_func,
321
- e_score_correction_bias=e_score_correction_bias)
328
+ e_score_correction_bias=e_score_correction_bias,
329
+ indices_type=self.topk_indices_dtype)
330
331
weight_bits = self.quant_config.weight_bits
332
has_zp = self.quant_config.has_zp
vllm/model_executor/layers/quantization/mxfp4.py
+1
-1
@@ -82,7 +82,7 @@ class Mxfp4Config(QuantizationConfig):
82
class Mxfp4MoEMethod(FusedMoEMethodBase):
83
84
def __init__(self, moe: FusedMoEConfig):
85
- super().__init__()
85
+ super().__init__(moe)
86
self.topk_indices_dtype = None
87
self.moe = moe
88
self.use_marlin = self._should_use_marlin()
vllm/model_executor/layers/quantization/quark/quark_moe.py
+30
-9
@@ -7,7 +7,8 @@ import torch
7
8
from vllm import _custom_ops as ops
9
from vllm.logger import init_logger
10
-from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
10
+from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
11
+ FusedMoEMethodBase,
12
FusedMoeWeightScaleSupported)
13
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
14
OCP_MX_BLOCK_SIZE)
@@ -25,6 +26,9 @@ __all__ = [
26
27
class QuarkMoEMethod(FusedMoEMethodBase):
28
29
+ def __init__(self, moe: FusedMoEConfig):
30
+ super().__init__(moe)
31
+
32
@staticmethod
33
def get_moe_method(
34
quant_config: "QuarkConfig", # type: ignore # noqa E501 # noqa F821
@@ -42,17 +46,24 @@ class QuarkMoEMethod(FusedMoEMethodBase):
46
input_config = layer_quant_config.get("input_tensors")
47
48
if quant_config._is_fp8_w8a8(weight_config, input_config):
45
- return QuarkW8A8Fp8MoEMethod(weight_config, input_config)
49
+ return QuarkW8A8Fp8MoEMethod(weight_config, input_config,
50
+ module.moe_config)
51
elif quant_config._is_mx_fp4(weight_config, input_config):
47
- return QuarkW4A4MXFp4MoEMethod(weight_config, input_config)
52
+ return QuarkW4A4MXFp4MoEMethod(weight_config, input_config,
53
+ module.moe_config)
54
else:
55
raise RuntimeError("Unsupported FusedMoe scheme")
56
57
58
class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
59
54
- def __init__(self, weight_config: dict[str, Any], input_config: dict[str,
55
- Any]):
60
+ def __init__(
61
+ self,
62
+ weight_config: dict[str, Any],
63
+ input_config: dict[str, Any],
64
+ moe: FusedMoEConfig,
65
+ ):
66
+ super().__init__(moe)
67
self.weight_quant = weight_config
68
self.input_quant = input_config
69
@@ -215,6 +226,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
226
logical_to_physical_map: Optional[torch.Tensor] = None,
227
logical_replica_count: Optional[torch.Tensor] = None,
228
) -> torch.Tensor:
229
+ assert self.fused_experts is None
230
+
231
if enable_eplb:
232
raise NotImplementedError(
233
"EPLB not supported for `QuarkW8A8Fp8MoEMethod` yet.")
@@ -231,7 +244,8 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
244
num_expert_group=num_expert_group,
245
custom_routing_function=custom_routing_function,
246
scoring_func=scoring_func,
234
- e_score_correction_bias=e_score_correction_bias)
247
+ e_score_correction_bias=e_score_correction_bias,
248
+ indices_type=self.topk_indices_dtype)
249
250
return fused_experts(
251
x,
@@ -253,8 +267,13 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
267
268
class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
269
256
- def __init__(self, weight_config: dict[str, Any], input_config: dict[str,
257
- Any]):
270
+ def __init__(
271
+ self,
272
+ weight_config: dict[str, Any],
273
+ input_config: dict[str, Any],
274
+ moe: FusedMoEConfig,
275
+ ):
276
+ super().__init__(moe)
277
self.weight_quant = weight_config
278
self.input_quant = input_config
279
@@ -369,6 +388,7 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
388
logical_to_physical_map: Optional[torch.Tensor] = None,
389
logical_replica_count: Optional[torch.Tensor] = None,
390
) -> torch.Tensor:
391
+ assert self.fused_experts is None
392
393
if enable_eplb:
394
raise NotImplementedError(
@@ -386,7 +406,8 @@ class QuarkW4A4MXFp4MoEMethod(QuarkMoEMethod):
406
num_expert_group=num_expert_group,
407
custom_routing_function=custom_routing_function,
408
scoring_func=scoring_func,
389
- e_score_correction_bias=e_score_correction_bias)
409
+ e_score_correction_bias=e_score_correction_bias,
410
+ indices_type=self.topk_indices_dtype)
411
412
out = fused_experts(
413
x,
vllm/model_executor/layers/quantization/rtn.py
+9
-4
@@ -10,7 +10,8 @@ import torch.nn.functional as F
10
from torch.nn.parameter import Parameter
11
12
from vllm.logger import init_logger
13
-from vllm.model_executor.layers.fused_moe import FusedMoE, FusedMoEMethodBase
13
+from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
14
+ FusedMoEMethodBase)
15
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
16
set_weight_attrs)
17
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -76,7 +77,7 @@ class RTNConfig(QuantizationConfig):
77
if isinstance(layer, LinearBase):
78
return RTNLinearMethod(self)
79
elif isinstance(layer, FusedMoE):
79
- return RTNMoEMethod(self)
80
+ return RTNMoEMethod(self, layer.moe_config)
81
return None
82
83
@@ -210,7 +211,8 @@ class RTNLinearMethod(LinearMethodBase):
211
212
class RTNMoEMethod(FusedMoEMethodBase):
213
213
- def __init__(self, quant_config: RTNConfig):
214
+ def __init__(self, quant_config: RTNConfig, moe: FusedMoEConfig):
215
+ super().__init__(moe)
216
self.quant_config = quant_config
217
218
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -289,6 +291,8 @@ class RTNMoEMethod(FusedMoEMethodBase):
291
logical_to_physical_map: Optional[torch.Tensor] = None,
292
logical_replica_count: Optional[torch.Tensor] = None,
293
) -> torch.Tensor:
294
+ assert self.fused_experts is None
295
+
296
if enable_eplb:
297
raise NotImplementedError(
298
"EPLB not supported for `RTNMoEMethod` yet.")
@@ -305,7 +309,8 @@ class RTNMoEMethod(FusedMoEMethodBase):
309
num_expert_group=num_expert_group,
310
custom_routing_function=custom_routing_function,
311
scoring_func=scoring_func,
308
- e_score_correction_bias=e_score_correction_bias)
312
+ e_score_correction_bias=e_score_correction_bias,
313
+ indices_type=self.topk_indices_dtype)
314
315
weight_bits = self.quant_config.weight_bits
316
group_size = self.quant_config.group_size
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+27
-102
@@ -3,33 +3,30 @@
3
"""Utility helpers for NVFP4 + FlashInfer fused-MoE path"""
4
from __future__ import annotations
5
6
-from typing import Optional
7
-
6
import torch
7
8
import vllm.envs as envs
9
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
-from vllm.logger import init_logger
13
-from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
10
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
11
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
15
- FlashInferExperts, is_valid_flashinfer_cutlass_fused_moe)
12
+ FlashInferExperts)
13
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
14
FlashInferCutlassMoEPrepareAndFinalize)
15
from vllm.platforms import current_platform
19
-
20
-logger = init_logger(__name__)
16
+from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
17
18
__all__ = [
19
"is_flashinfer_fp4_cutlass_moe_available",
20
"reorder_w1w3_to_w3w1",
25
- "build_flashinfer_fp4_cutlass_moe_kernel",
26
- "flashinfer_fp4_cutlass_moe_forward",
21
+ "build_flashinfer_fp4_cutlass_moe_prepare_finalize",
22
]
23
24
25
def is_flashinfer_fp4_cutlass_moe_available() -> bool:
26
"""Return ``True`` when FlashInfer CUTLASS NV-FP4 kernels can be used."""
32
- return (envs.VLLM_USE_FLASHINFER_MOE_FP4 and current_platform.is_cuda()
27
+ return (envs.VLLM_USE_FLASHINFER_MOE_FP4
28
+ and has_flashinfer_cutlass_fused_moe()
29
+ and current_platform.is_cuda()
30
and current_platform.is_device_capability(100))
31
32
@@ -49,105 +46,33 @@ def reorder_w1w3_to_w3w1(weight: torch.Tensor,
46
dim=dim).contiguous())
47
48
52
-def build_flashinfer_fp4_cutlass_moe_kernel(
53
- moe_parallel_config: FusedMoEParallelConfig, ) -> mk.FusedMoEModularKernel:
54
- """Create *and return* a FlashInfer CUTLASS fused-MoE modular kernel"""
55
- experts = FlashInferExperts(
56
- use_nvfp4_w4a4=True,
57
- use_dp=moe_parallel_config.dp_size > 1,
58
- ep_rank=moe_parallel_config.ep_rank,
59
- ep_size=moe_parallel_config.ep_size,
60
- tp_rank=moe_parallel_config.tp_rank,
61
- tp_size=moe_parallel_config.tp_size,
62
- )
63
- logger.debug_once("FlashInferExperts (util)")
64
- return mk.FusedMoEModularKernel(
65
- FlashInferCutlassMoEPrepareAndFinalize(quant_dtype=torch.uint8),
66
- experts,
67
- )
68
-
69
-
70
-def flashinfer_fp4_cutlass_moe_forward(
71
- fused_experts: mk.FusedMoEModularKernel,
72
- layer: torch.nn.Module,
73
- x: torch.Tensor,
74
- topk_weights: torch.Tensor,
75
- topk_ids: torch.Tensor,
76
- activation: str,
77
- global_num_experts: int,
78
- expert_map: Optional[torch.Tensor],
79
- apply_router_weight_on_input: bool,
80
-) -> torch.Tensor:
81
- """Common forward wrapper for FlashInfer NV-FP4 fused-MoE"""
82
-
83
- assert is_valid_flashinfer_cutlass_fused_moe(
84
- x, layer.w13_weight,
85
- layer.w2_weight), ("FlashInfer CUTLASS fused-MoE not applicable!")
86
-
87
- a1_gscale = layer.w13_input_scale_quant
88
- a2_gscale = layer.w2_input_scale_quant
89
-
90
- extra_expert_args = {
91
- "g1_alphas": layer.g1_alphas,
92
- "g2_alphas": layer.g2_alphas,
93
- # Avoid confusion with a1_scale and a2_scale
94
- # where are batch size related.
95
- "a1_gscale": a1_gscale,
96
- "a2_gscale": a2_gscale,
97
- "out_dtype": x.dtype,
98
- }
99
- extra_prepare_args = {
100
- "use_dp": layer.dp_size > 1,
101
- "local_tokens": x.shape[0],
102
- "a1_gscale": a1_gscale,
103
- }
104
- extra_finalize_args = {
105
- "use_dp": layer.dp_size > 1,
106
- "local_tokens": x.shape[0],
107
- }
108
-
109
- return fused_experts(
110
- hidden_states=x,
111
- w1=layer.w13_weight,
112
- w2=layer.w2_weight,
113
- topk_weights=topk_weights,
114
- topk_ids=topk_ids,
115
- inplace=False, # TODO(shuw): fix later, now output is high prec
116
- activation=activation,
117
- global_num_experts=global_num_experts,
118
- expert_map=expert_map,
119
- w1_scale=layer.w13_blockscale_swizzled,
120
- w2_scale=layer.w2_blockscale_swizzled,
121
- apply_router_weight_on_input=apply_router_weight_on_input,
122
- extra_expert_args=extra_expert_args,
123
- extra_prepare_args=extra_prepare_args,
124
- extra_finalize_args=extra_finalize_args,
125
- )
49
+def build_flashinfer_fp4_cutlass_moe_prepare_finalize(
50
+ moe: FusedMoEConfig,
51
+ a1_gscale: torch.Tensor,
52
+) -> mk.FusedMoEPrepareAndFinalize:
53
+ """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""
54
+ use_dp = moe.moe_parallel_config.dp_size > 1
55
+ return FlashInferCutlassMoEPrepareAndFinalize(use_dp, a1_gscale=a1_gscale)
56
57
58
def select_nvfp4_gemm_impl(
129
- allow_flashinfer: bool,
130
- moe, # FusedMoEConfig
131
- logger):
59
+ moe: FusedMoEConfig,
60
+ g1_alphas: torch.Tensor,
61
+ g2_alphas: torch.Tensor,
62
+ a1_gscale: torch.Tensor,
63
+ a2_gscale: torch.Tensor,
64
+ allow_flashinfer: bool,
65
+) -> mk.FusedMoEPermuteExpertsUnpermute:
66
"""Return a GEMM *experts* implementation for NV-FP4 fused-MoE layers"""
67
134
- # lazy import
135
- from vllm.distributed import get_ep_group
136
-
137
- all2all_manager = get_ep_group().device_communicator.all2all_manager
138
- assert all2all_manager is not None
139
-
68
if allow_flashinfer:
141
- flashinfer_backend = envs.VLLM_FLASHINFER_MOE_BACKEND
142
- if flashinfer_backend != "throughput":
143
- raise ValueError(
144
- f"Only throughput backend is supported for FlashInferExperts, "
145
- f"but got {flashinfer_backend}.")
146
- logger.debug_once(
147
- "Initializing FlashInferExperts with throughput backend.")
69
return FlashInferExperts(
149
- use_nvfp4_w4a4=True,
150
- use_dp=moe.moe_parallel_config.dp_size > 1,
70
+ g1_alphas=g1_alphas,
71
+ g2_alphas=g2_alphas,
72
+ a1_gscale=a1_gscale,
73
+ a2_gscale=a2_gscale,
74
+ out_dtype=moe.in_dtype,
75
+ quant_dtype="nvfp4",
76
ep_rank=moe.moe_parallel_config.ep_rank,
77
ep_size=moe.moe_parallel_config.ep_size,
78
tp_rank=moe.moe_parallel_config.tp_rank,