@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 53fa457391
Fork 0

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

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

Synced Nov 15, 02:15 UTC

[Misc] Add unit tests for MoE ModularKernel combinations + Profiling utility (#20449)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Jul 11, 2025 at 10:51 UTC 53fa45739176e81b9078f7e3516a8b11fda161de
15 files changed +1727 -22
tests/kernels/moe/modular_kernel_tools/__init__.py
tests/kernels/moe/modular_kernel_tools/cli_args.py new
+160
@@ -0,0 +1,160 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import argparse
5 +
6 +import torch
7 +
8 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10 +
11 +from .common import Config
12 +from .mk_objects import (MK_ALL_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES,
13 + MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
14 +
15 +
16 +def make_config_arg_parser(description: str):
17 +
18 + def to_pf_class_type(s: str) -> mk.FusedMoEPrepareAndFinalize:
19 + for pf in MK_ALL_PREPARE_FINALIZE_TYPES:
20 + if pf.__name__ == s:
21 + return pf
22 + raise ValueError(
23 + f"Cannot find a PrepareFinalize type that matches {s}")
24 +
25 + def to_experts_class_type(s: str) -> mk.FusedMoEPermuteExpertsUnpermute:
26 + for fe in MK_FUSED_EXPERT_TYPES:
27 + if fe.__name__ == s:
28 + return fe
29 + raise ValueError(f"Cannot find a FusedExperts type that matches {s}")
30 +
31 + def to_quant_torch_dtype(s: str) -> torch.dtype:
32 + if s == "torch.float8_e4m3fn":
33 + return torch.float8_e4m3fn
34 + raise ValueError(f"Unsupported quant type {s}")
35 +
36 + parser = argparse.ArgumentParser(description=description)
37 +
38 + parser.add_argument(
39 + "--world-size",
40 + type=int,
41 + default=2,
42 + help="Number of ranks that participate in all2all",
43 + )
44 + parser.add_argument(
45 + "--pf-type",
46 + type=to_pf_class_type,
47 + required=True,
48 + help=("Choose a PrepareFinalize Type : "
49 + f"{[x.__name__ for x in MK_ALL_PREPARE_FINALIZE_TYPES]}"),
50 + )
51 + parser.add_argument(
52 + "--experts-type",
53 + type=to_experts_class_type,
54 + required=True,
55 + help=(f"Choose a FusedExpert type : "
56 + f"{[x.__name__ for x in MK_FUSED_EXPERT_TYPES]}"),
57 + )
58 + parser.add_argument(
59 + "-m",
60 + nargs="+",
61 + type=int,
62 + default=[64],
63 + help="num tokens per rank",
64 + )
65 + parser.add_argument(
66 + "-k",
67 + type=int,
68 + default=7168,
69 + help="hidden-size",
70 + )
71 + parser.add_argument(
72 + "-n",
73 + type=int,
74 + default=1024,
75 + help="N dimension of the first fused-moe matmul",
76 + )
77 + parser.add_argument("--num-experts",
78 + type=int,
79 + default=32,
80 + help="Global num experts")
81 + parser.add_argument("--topk",
82 + nargs="+",
83 + type=int,
84 + default=[4, 1],
85 + help="num topk")
86 + parser.add_argument(
87 + "--fused-moe-chunk-size",
88 + nargs="+",
89 + type=int,
90 + help="Fused moe chunk size used for the non-batched fused experts impl."
91 + )
92 +
93 + # Quant args
94 + parser.add_argument("--quant-dtype",
95 + type=to_quant_torch_dtype,
96 + help="Quant datatype")
97 + parser.add_argument("--per-token-quantized-activations",
98 + action='store_true',
99 + help=("The input activations must be per-token "
100 + "quantized"))
101 + parser.add_argument("--per-channel-quantized-weights",
102 + action="store_true",
103 + help="The weights must be per-channel quantized.")
104 + parser.add_argument("--block-shape",
105 + nargs="+",
106 + type=int,
107 + help="Quantization block shape")
108 +
109 + # Torch trace profile generation args
110 + parser.add_argument("--torch-trace-dir-path",
111 + type=str,
112 + default=None,
113 + help="Get torch trace for single execution")
114 +
115 + return parser
116 +
117 +
118 +def _validate_args(args: argparse.Namespace):
119 +
120 + if args.quant_dtype is not None:
121 + assert args.quant_dtype == torch.float8_e4m3fn
122 + if args.block_shape is not None:
123 + assert len(args.block_shape) == 2, (
124 + f"block shape must have 2 elements. got {args.block_shape}")
125 +
126 + if args.experts_type in MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES:
127 + assert args.world_size == 1, (
128 + "Single GPU objects need world size set to 1")
129 +
130 + if args.torch_trace_dir_path is not None:
131 + from pathlib import Path
132 + assert Path(args.torch_trace_dir_path).is_dir(), (
133 + f"Please create {args.torch_trace_dir_path}")
134 +
135 +
136 +def make_config(args: argparse.Namespace) -> Config:
137 +
138 + _validate_args(args)
139 +
140 + quant_config = None
141 + if args.quant_dtype is not None:
142 + quant_config = FusedMoEQuantConfig(
143 + quant_dtype=args.quant_dtype,
144 + per_act_token_quant=args.per_token_quantized_activations,
145 + per_out_ch_quant=args.per_channel_quantized_weights,
146 + block_shape=args.block_shape)
147 +
148 + return Config(
149 + Ms=args.m,
150 + K=args.k,
151 + N=args.n,
152 + E=args.num_experts,
153 + topks=args.topk,
154 + dtype=torch.bfloat16, # hard-code
155 + quant_config=quant_config,
156 + prepare_finalize_type=args.pf_type,
157 + fused_experts_type=args.experts_type,
158 + fused_moe_chunk_size=args.fused_moe_chunk_size,
159 + world_size=args.world_size,
160 + torch_trace_dir_path=args.torch_trace_dir_path)
tests/kernels/moe/modular_kernel_tools/common.py new
+641
@@ -0,0 +1,641 @@
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 Any, Optional, Union
5 +
6 +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.utils import torch_experts
11 +from vllm.config import VllmConfig
12 +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)
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)
24 +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)
31 +from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
32 +
33 +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)
45 +
46 +
47 +def _describe_tensor(t: Optional[torch.Tensor], name: str) -> str:
48 + if t is None:
49 + return f"{name} : None"
50 + else:
51 + return f"{name} : {t.shape} {t.dtype} {t.device}"
52 +
53 +
54 +@dataclass
55 +class Config:
56 + Ms: Union[list[int], int]
57 + K: int
58 + N: int
59 + E: int
60 + topks: Union[list[int], int]
61 + dtype: torch.dtype
62 + quant_config: Optional[FusedMoEQuantConfig]
63 +
64 + prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
65 + fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute
66 +
67 + fused_moe_chunk_size: Optional[int]
68 + world_size: int
69 +
70 + torch_trace_dir_path: Optional[str] = None
71 +
72 + def describe(self) -> str:
73 + 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 "
83 + 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"
88 + else:
89 + s += " quant=None \n"
90 + return s
91 +
92 + @property
93 + def M(self) -> int:
94 + assert isinstance(self.Ms, int)
95 + return self.Ms
96 +
97 + @property
98 + def quant_dtype(self) -> Optional[torch.dtype]:
99 + if self.quant_config is None:
100 + return None
101 + return self.quant_config.quant_dtype
102 +
103 + @property
104 + def is_per_act_token_quant(self) -> bool:
105 + if self.quant_config is None:
106 + return False
107 + return self.quant_config.per_act_token_quant
108 +
109 + @property
110 + def is_per_tensor_act_quant(self) -> bool:
111 + if self.quant_config is None:
112 + return False
113 + return (not self.is_per_act_token_quant
114 + and self.quant_block_shape is None)
115 +
116 + @property
117 + def is_per_out_ch_quant(self) -> bool:
118 + if self.quant_config is None:
119 + return False
120 + return self.quant_config.per_out_ch_quant
121 +
122 + @property
123 + def quant_block_shape(self) -> Optional[list[int]]:
124 + if self.quant_config is None:
125 + return None
126 + return self.quant_config.block_shape
127 +
128 + @property
129 + def topk(self) -> int:
130 + assert isinstance(self.topks, int)
131 + return self.topks
132 +
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 +
144 + @property
145 + def num_local_experts(self) -> int:
146 + return self.E // self.world_size
147 +
148 + def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
149 + """
150 + make env data for vllm launch.
151 + """
152 + vllm_config = VllmConfig()
153 + vllm_config.parallel_config.data_parallel_size = self.world_size
154 + vllm_config.parallel_config.enable_expert_parallel = True
155 +
156 + env_dict = {
157 + "VLLM_ALL2ALL_BACKEND": self.all2all_backend(),
158 + "VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())),
159 + }
160 + if self.fused_moe_chunk_size is not None:
161 + env_dict.update(
162 + {"VLLM_FUSED_MOE_CHUNK_SIZE": str(self.fused_moe_chunk_size)})
163 + return vllm_config, env_dict
164 +
165 + def is_fp8_block_quantized(self):
166 + return (self.quant_dtype == torch.float8_e4m3fn
167 + and self.quant_block_shape is not None)
168 +
169 + def is_batched_prepare_finalize(self):
170 + return self.prepare_finalize_type in [
171 + PplxPrepareAndFinalize, DeepEPLLPrepareAndFinalize
172 + ]
173 +
174 + def is_batched_fused_experts(self):
175 + return self.fused_experts_type in [
176 + CutlassExpertsFp8, BatchedDeepGemmExperts, BatchedTritonExperts,
177 + NaiveBatchedExperts, BatchedTritonOrDeepGemmExperts
178 + ]
179 +
180 + 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 + ]
214 +
215 + def is_fe_supports_chunking(self):
216 + return self.fused_experts_type in [
217 + CutlassExpertsFp8, DeepGemmExperts, TritonOrDeepGemmExperts,
218 + TritonExperts
219 + ]
220 +
221 + def needs_deep_gemm(self):
222 + return self.fused_experts_type in [
223 + BatchedDeepGemmExperts,
224 + DeepGemmExperts,
225 + ]
226 +
227 + def needs_pplx(self):
228 + return self.prepare_finalize_type in [PplxPrepareAndFinalize]
229 +
230 + def needs_deep_ep(self):
231 + return self.prepare_finalize_type in [
232 + DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize
233 + ]
234 +
235 + 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 + ]
249 +
250 + def is_valid(self):
251 + # Check prepare-finalize and fused-experts compatibility
252 + if self.is_batched_prepare_finalize():
253 + if not self.is_batched_fused_experts():
254 + return False
255 + else:
256 + if not self.is_standard_fused_experts():
257 + return False
258 +
259 + use_chunking = self.fused_moe_chunk_size is not None
260 + if use_chunking and not self.is_fe_supports_chunking():
261 + return False
262 +
263 + # Check quantization sanity
264 + if (int(self.is_per_act_token_quant) +
265 + int(self.is_per_tensor_act_quant) +
266 + int(self.quant_block_shape is not None)) > 1:
267 + # invalid quant config
268 + return False
269 +
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
279 +
280 + # Check fp8 block quanization support
281 + is_block_quatized = self.quant_block_shape is not None
282 + if is_block_quatized and not is_fp8:
283 + return False
284 + if is_block_quatized and not self.is_fe_block_fp8_supported():
285 + return False
286 +
287 + # deep_gemm only works with block-quantized
288 + if self.needs_deep_gemm() and not is_block_quatized:
289 + return False
290 +
291 + # Check dependencies
292 + if self.needs_deep_ep() and not has_deep_ep():
293 + return False
294 + if self.needs_deep_gemm() and not has_deep_gemm():
295 + return False
296 + if self.needs_pplx() and not has_pplx(): # noqa: SIM103
297 + return False
298 +
299 + return True
300 +
301 +
302 +@dataclass
303 +class WeightTensors:
304 + w1: torch.Tensor
305 + w2: torch.Tensor
306 + w1_scale: Optional[torch.Tensor]
307 + w2_scale: Optional[torch.Tensor]
308 +
309 + def describe(self):
310 + s = ""
311 + s += "== Weight Tensors: \n"
312 + s += f' - {_describe_tensor(self.w1, "w1")} \n'
313 + s += f' - {_describe_tensor(self.w2, "w2")} \n'
314 + s += f' - {_describe_tensor(self.w1_scale, "w1_scale")} \n'
315 + s += f' - {_describe_tensor(self.w2_scale, "w2_scale")} \n'
316 + return s
317 +
318 + def to_current_device(self):
319 + self.w1 = self.w1.to(device=torch.cuda.current_device())
320 + self.w2 = self.w2.to(device=torch.cuda.current_device())
321 + is_quantized = self.w1.dtype == torch.float8_e4m3fn
322 + if is_quantized:
323 + assert self.w1_scale is not None
324 + assert self.w2_scale is not None
325 + self.w1_scale = self.w1_scale.to(
326 + device=torch.cuda.current_device())
327 + self.w2_scale = self.w2_scale.to(
328 + device=torch.cuda.current_device())
329 +
330 + def slice_weights(self, rank: int,
331 + num_local_experts: int) -> "WeightTensors":
332 + s = rank * num_local_experts
333 + e = s + num_local_experts
334 + w1 = self.w1[s:e, :, :]
335 + w2 = self.w2[s:e, :, :]
336 + is_quantized = self.w1.dtype == torch.float8_e4m3fn
337 + w1_scale, w2_scale = (None, None)
338 + if is_quantized:
339 + assert self.w1_scale is not None
340 + assert self.w2_scale is not None
341 + w1_scale = self.w1_scale[s:e, :, :]
342 + w2_scale = self.w2_scale[s:e, :, :]
343 + return WeightTensors(w1, w2, w1_scale, w2_scale)
344 +
345 + @staticmethod
346 + def make(config: Config) -> "WeightTensors":
347 +
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)
368 +
369 + assert config.quant_block_shape is not None
370 + w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights(
371 + e=config.E,
372 + n=config.N,
373 + k=config.K,
374 + block_size=config.quant_block_shape,
375 + )
376 + return WeightTensors(w1=w1,
377 + w2=w2,
378 + w1_scale=w1_scale,
379 + w2_scale=w2_scale)
380 +
381 +
382 +@dataclass
383 +class RankTensors:
384 + hidden_states: torch.Tensor
385 + hidden_states_scale: Optional[torch.Tensor]
386 +
387 + topk_weights: torch.Tensor
388 + topk_ids: torch.Tensor
389 + expert_map: Optional[torch.Tensor]
390 +
391 + quant_config: Optional[FusedMoEQuantConfig]
392 +
393 + def describe(self):
394 + s = ""
395 + s += "== Rank Tensors: \n"
396 + s += f' - {_describe_tensor(self.hidden_states, "HS")} \n'
397 + s += f' - {_describe_tensor(self.hidden_states_scale, "HS_scale")} \n'
398 + s += f' - {_describe_tensor(self.topk_weights, "topk_weights")} \n'
399 + s += f' - {_describe_tensor(self.topk_ids, "topk_ids")} \n'
400 + s += f' - {_describe_tensor(self.expert_map, "expert_map")} \n'
401 + return s
402 +
403 + @staticmethod
404 + def make_hidden_states(
405 + config: Config) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
406 + """
407 + Return hidden_states
408 + """
409 + m, k, dtype = (config.M, config.K, config.dtype)
410 + a = (torch.randn(
411 + (m, k), device=torch.cuda.current_device(), dtype=dtype) / 15.0)
412 +
413 + if config.quant_dtype is None:
414 + return a, None
415 +
416 + # We dequant and use that as hidden_states so the tests are stable.
417 + # quantizing and dequantizing yield slightly different results
418 + # depending on the hardware. Here we, quantize and dequantize
419 + # first - so further quantize and dequantize will yeild the same
420 + # values.
421 + if config.is_per_tensor_act_quant:
422 + a_q, a_scales = ops.scaled_fp8_quant(
423 + a, use_per_token_if_dynamic=False)
424 + return a_q.float().mul(a_scales).to(dtype), a_scales
425 +
426 + if config.is_per_act_token_quant:
427 + a_q, a_scales = ops.scaled_fp8_quant(a,
428 + use_per_token_if_dynamic=True)
429 + return a_q.float().mul(a_scales).to(dtype), None
430 +
431 + assert config.quant_block_shape is not None
432 + block_k = config.quant_block_shape[1]
433 + a_q, a_scales = per_token_cast_to_fp8(a, block_size=block_k)
434 + return a_q.float().view(
435 + (-1, block_k)).mul(a_scales.view(-1, 1)).view(m, k).to(dtype), None
436 +
437 + @staticmethod
438 + def make(config: Config, pgi: ProcessGroupInfo):
439 +
440 + dtype = config.dtype
441 + topk, m, _ = (config.topk, config.M, config.K)
442 + hidden_states, hidden_states_scale = RankTensors.make_hidden_states(
443 + config)
444 +
445 + num_local_experts, global_num_experts = (config.num_local_experts,
446 + config.E)
447 + score = torch.randn((m, global_num_experts),
448 + device="cuda",
449 + dtype=dtype)
450 + topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk,
451 + False)
452 + topk_ids = topk_ids.to(config.topk_ids_dtype)
453 +
454 + # distribute topk_ids evenly
455 + for mi in range(m):
456 + topk_ids[mi] = torch.randperm(config.E)[:topk]
457 + topk_ids = topk_ids.to(device=torch.cuda.current_device())
458 +
459 + expert_map = None
460 + if config.world_size > 1:
461 + expert_map = torch.full((global_num_experts, ),
462 + fill_value=-1,
463 + dtype=torch.int32)
464 + s = pgi.rank * num_local_experts
465 + e = s + num_local_experts
466 + expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
467 + expert_map = expert_map.to(device=torch.cuda.current_device(),
468 + dtype=torch.int32)
469 +
470 + return RankTensors(
471 + hidden_states=hidden_states,
472 + hidden_states_scale=hidden_states_scale,
473 + topk_weights=topk_weights,
474 + topk_ids=topk_ids,
475 + expert_map=expert_map,
476 + quant_config=config.quant_config,
477 + )
478 +
479 +
480 +def reference_moe_impl(config: Config, weights: WeightTensors,
481 + rank_tensors: RankTensors) -> torch.Tensor:
482 +
483 + return torch_experts(a=rank_tensors.hidden_states,
484 + w1=weights.w1,
485 + w2=weights.w2,
486 + topk_weight=rank_tensors.topk_weights,
487 + topk_ids=rank_tensors.topk_ids,
488 + global_num_experts=config.E,
489 + 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)
563 +
564 + return experts
565 +
566 +
567 +def make_modular_kernel(config: Config,
568 + vllm_config: VllmConfig) -> mk.FusedMoEModularKernel:
569 +
570 + def next_power_of_2(x):
571 + import math
572 + if x == 0:
573 + return 1
574 + return 2**math.ceil(math.log2(x))
575 +
576 + # make moe config
577 + moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
578 + tp_size_=get_tensor_model_parallel_world_size(),
579 + dp_size_=get_dp_group().world_size,
580 + vllm_parallel_config=vllm_config.parallel_config,
581 + )
582 + moe = FusedMoEConfig(
583 + num_experts=config.E,
584 + experts_per_token=config.topk,
585 + hidden_dim=config.K,
586 + num_local_experts=config.num_local_experts,
587 + moe_parallel_config=moe_parallel_config,
588 + in_dtype=config.dtype,
589 + quant_config=config.quant_config,
590 + max_num_tokens=next_power_of_2(config.M),
591 + )
592 +
593 + # 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())
603 +
604 + modular_kernel = mk.FusedMoEModularKernel(
605 + prepare_finalize=prepare_finalize, fused_experts=fused_experts)
606 +
607 + return modular_kernel
608 +
609 +
610 +def run_modular_kernel(
611 + pgi: ProcessGroupInfo,
612 + vllm_config: VllmConfig,
613 + config: Config,
614 + weights: WeightTensors,
615 + rank_tensors: RankTensors,
616 +) -> torch.Tensor:
617 + assert isinstance(config.Ms, int)
618 + assert isinstance(config.topks, int)
619 +
620 + # weights for rank
621 + rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)
622 +
623 + mk = make_modular_kernel(config, vllm_config)
624 +
625 + mk_kwargs = {
626 + "hidden_states": rank_tensors.hidden_states.clone(
627 + ), # 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,
638 + }
639 + out = mk.forward(**mk_kwargs)
640 +
641 + return out
tests/kernels/moe/modular_kernel_tools/make_feature_matrix.py new
+173
@@ -0,0 +1,173 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import copy
5 +from enum import Enum
6 +from itertools import product
7 +from typing import Optional
8 +
9 +import torch
10 +from tqdm import tqdm
11 +
12 +from vllm.config import VllmConfig, set_current_vllm_config
13 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
14 +from vllm.platforms import current_platform
15 +
16 +from .common import (Config, RankTensors, WeightTensors, reference_moe_impl,
17 + run_modular_kernel)
18 +from .mk_objects import (MK_FUSED_EXPERT_TYPES,
19 + MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_QUANT_CONFIGS)
20 +from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
21 +
22 +
23 +class Result(Enum):
24 + PASS = 1
25 + FAIL = 2
26 + SKIP = 3
27 +
28 +
29 +def rank_worker(
30 + pgi: ProcessGroupInfo,
31 + vllm_config: VllmConfig,
32 + cpu_group,
33 + config: Config,
34 + weights: WeightTensors,
35 +):
36 + current_platform.seed_everything(pgi.rank)
37 +
38 + # sanity check
39 + from vllm import envs
40 + if config.fused_moe_chunk_size is not None:
41 + assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
42 +
43 + # get weights to this device
44 + weights.to_current_device()
45 +
46 + Ms = config.Ms
47 + assert isinstance(Ms, list)
48 + TOPKs = config.topks
49 + assert isinstance(TOPKs, list)
50 +
51 + for m, topk in product(Ms, TOPKs):
52 + print(f"Running m={m}, topk={topk} ...")
53 + # override m and topk
54 + cfgx = copy.deepcopy(config)
55 + cfgx.Ms = m
56 + cfgx.topks = topk
57 +
58 + # inputs for rank
59 + rank_tensors = RankTensors.make(cfgx, pgi)
60 +
61 + # modular kernel out
62 + mk_out = run_modular_kernel(pgi, vllm_config, cfgx, weights,
63 + rank_tensors)
64 +
65 + with set_current_vllm_config(vllm_config):
66 + ref_out = reference_moe_impl(cfgx, weights, rank_tensors)
67 +
68 + torch.testing.assert_close(ref_out, mk_out, atol=3e-2, rtol=3e-2)
69 +
70 +
71 +def make_feature_matrix(csv_file_path: str):
72 +
73 + from dataclasses import asdict
74 +
75 + import pandas as pd
76 +
77 + def add_to_results(config: Config,
78 + success: Result,
79 + results_df: Optional[pd.DataFrame] = None):
80 + config_dict = asdict(config)
81 + config_dict['prepare_finalize_type'] = config_dict[
82 + 'prepare_finalize_type'].__name__
83 + config_dict['fused_experts_type'] = config_dict[
84 + 'fused_experts_type'].__name__
85 + config_dict['per_tensor_act_quant'] = config.is_per_tensor_act_quant
86 + quant_config_dict = config_dict['quant_config']
87 + del config_dict['quant_config']
88 + if quant_config_dict is None:
89 + quant_config = FusedMoEQuantConfig(None)
90 + quant_config_dict = asdict(quant_config)
91 +
92 + config_dict |= quant_config_dict
93 + result_dict = config_dict | {'success': success.name}
94 +
95 + result_df = pd.DataFrame([result_dict])
96 + if results_df is None:
97 + results_df = result_df
98 + else:
99 + results_df = pd.concat([results_df, result_df], ignore_index=True)
100 +
101 + return results_df
102 +
103 + Ms = [64]
104 + Ks = [7168] # hidden sizes
105 + Ns = [2048]
106 + TOPKs = [[4, 1]]
107 + Es = [32]
108 + DTYPEs = [torch.bfloat16]
109 + PF_TYPES = MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
110 + FE_TYPES = MK_FUSED_EXPERT_TYPES
111 + Q_TYPES = MK_QUANT_CONFIGS
112 +
113 + combinations = list(
114 + product(Ms, Ks, Ns, Es, TOPKs, DTYPEs, PF_TYPES, FE_TYPES, Q_TYPES))
115 +
116 + results_df: Optional[pd.DataFrame] = None
117 + for m, k, n, e, topks, dtype, pf_type, experts_type, quant_config in tqdm(
118 + combinations): #noqa: E501
119 + config = Config(Ms=[m],
120 + K=k,
121 + N=n,
122 + E=e,
123 + topks=topks,
124 + dtype=dtype,
125 + prepare_finalize_type=pf_type,
126 + fused_experts_type=experts_type,
127 + quant_config=quant_config,
128 + world_size=2,
129 + fused_moe_chunk_size=None)
130 +
131 + success = None
132 + if config.is_valid():
133 + print(f"Running config : {config.describe()} ...")
134 + try:
135 + weights: WeightTensors = WeightTensors.make(config)
136 + vllm_config, env_dict = config.make_env_data()
137 + parallel_launch_with_config(config.world_size, rank_worker,
138 + vllm_config, env_dict, config,
139 + weights)
140 + success = Result.PASS
141 + except Exception as _:
142 + success = Result.FAIL
143 + else:
144 + success = Result.SKIP
145 +
146 + results_df = add_to_results(config, success, results_df)
147 +
148 + if results_df is not None:
149 + results_df.to_csv(f"{csv_file_path}")
150 +
151 +
152 +if __name__ == '__main__':
153 + import argparse
154 + from pathlib import Path
155 + parser = argparse.ArgumentParser(description=(
156 + "Make ModularKernel feature matrix \n"
157 + "Example : python3 -m tests.kernels.moe.modular_kernel_tools.make_feature_matrix " #noqa: E501
158 + "-f ./feature_matrices/feature_matrix.csv"))
159 +
160 + parser.add_argument("-f",
161 + "--feature-matrix-csv-file-path",
162 + type=str,
163 + required=True,
164 + help="File name to Generate a .csv file")
165 + args = parser.parse_args()
166 +
167 + csv_path = args.feature_matrix_csv_file_path
168 + assert csv_path.endswith(
169 + 'csv'), f"Need a file path ending with .csv, got {csv_path}"
170 + assert Path(csv_path).parent.is_dir(
171 + ), f"Cannot find parent directory for {Path(csv_path).parent}"
172 +
173 + make_feature_matrix(args.feature_matrix_csv_file_path)
tests/kernels/moe/modular_kernel_tools/mk_objects.py new
+87
@@ -0,0 +1,87 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +# Fused experts and PrepareFinalize imports
7 +from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8 + BatchedDeepGemmExperts)
9 +from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
10 + BatchedTritonOrDeepGemmExperts)
11 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12 +from vllm.model_executor.layers.fused_moe.cutlass_moe import CutlassExpertsFp8
13 +from vllm.model_executor.layers.fused_moe.deep_gemm_moe import DeepGemmExperts
14 +from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
15 + BatchedTritonExperts, NaiveBatchedExperts)
16 +from vllm.model_executor.layers.fused_moe.layer import TritonExperts
17 +from vllm.model_executor.layers.fused_moe.prepare_finalize import (
18 + MoEPrepareAndFinalizeNoEP)
19 +from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
20 + TritonOrDeepGemmExperts)
21 +from vllm.utils import has_deep_ep, has_pplx
22 +
23 +if has_deep_ep():
24 + from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
25 + DeepEPHTPrepareAndFinalize)
26 + from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
27 + DeepEPLLPrepareAndFinalize)
28 +
29 +if has_pplx():
30 + from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
31 + PplxPrepareAndFinalize)
32 +
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 + ]
40 +
41 +MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES = [MoEPrepareAndFinalizeNoEP]
42 +
43 +MK_ALL_PREPARE_FINALIZE_TYPES = (MK_MULTI_GPU_PREPARE_FINALIZE_TYPES +
44 + MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
45 +
46 +MK_FUSED_EXPERT_TYPES = [
47 + BatchedDeepGemmExperts,
48 + BatchedTritonExperts,
49 + NaiveBatchedExperts,
50 + BatchedTritonOrDeepGemmExperts,
51 + CutlassExpertsFp8,
52 + DeepGemmExperts,
53 + TritonOrDeepGemmExperts,
54 + TritonExperts,
55 +]
56 +
57 +MK_QUANT_CONFIGS = [
58 + None,
59 + # per-channel / per-column weights and per-tensor activations
60 + FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
61 + per_out_ch_quant=True,
62 + per_act_token_quant=False,
63 + block_shape=None),
64 + # per-channel / per-column weights and per-token activations
65 + FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
66 + per_out_ch_quant=True,
67 + per_act_token_quant=True,
68 + block_shape=None),
69 + # per-tensor weights and per-tensor activations
70 + FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
71 + per_out_ch_quant=False,
72 + per_act_token_quant=False,
73 + block_shape=None),
74 + # per-tensor weights and per-token activations
75 + FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
76 + per_out_ch_quant=False,
77 + per_act_token_quant=True,
78 + block_shape=None),
79 + # block-quantized weights and 128 block per-token activations
80 + FusedMoEQuantConfig(quant_dtype=torch.float8_e4m3fn,
81 + per_out_ch_quant=False,
82 + per_act_token_quant=False,
83 + block_shape=[128, 128]),
84 + # TODO (varun) : Should we test the following combinations ?
85 + # block-quantized weights and per-token activations
86 + # block-quantized weights and per-tensor activations
87 +]
tests/kernels/moe/modular_kernel_tools/parallel_utils.py new
+138
@@ -0,0 +1,138 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import dataclasses
4 +import os
5 +import traceback
6 +from typing import Any, Callable, Optional
7 +
8 +import torch
9 +from torch.multiprocessing import (
10 + spawn) # pyright: ignore[reportPrivateImportUsage]
11 +from typing_extensions import Concatenate, ParamSpec
12 +
13 +from vllm.config import VllmConfig, set_current_vllm_config
14 +from vllm.distributed import (init_distributed_environment,
15 + initialize_model_parallel)
16 +from vllm.utils import get_open_port
17 +
18 +## Parallel Processes Utils
19 +
20 +P = ParamSpec("P")
21 +
22 +
23 +@dataclasses.dataclass
24 +class ProcessGroupInfo:
25 + world_size: int
26 + world_local_size: int
27 + rank: int
28 + node_rank: int
29 + local_rank: int
30 + device: torch.device
31 +
32 +
33 +def _set_vllm_config(vllm_config: VllmConfig, world_size: int, rank: int,
34 + local_rank: int):
35 +
36 + import tempfile
37 + temp_file = tempfile.mkstemp()[1]
38 +
39 + set_current_vllm_config(vllm_config)
40 + with set_current_vllm_config(vllm_config):
41 + init_distributed_environment(
42 + world_size=world_size,
43 + rank=rank,
44 + distributed_init_method=f"file://{temp_file}",
45 + local_rank=local_rank,
46 + backend="nccl",
47 + )
48 +
49 + initialize_model_parallel(
50 + tensor_model_parallel_size=vllm_config.parallel_config.
51 + tensor_parallel_size,
52 + pipeline_model_parallel_size=vllm_config.parallel_config.
53 + pipeline_parallel_size,
54 + )
55 + cpu_group = torch.distributed.new_group(list(range(world_size)),
56 + backend="gloo")
57 + return cpu_group
58 +
59 +
60 +def _worker_parallel_launch(
61 + local_rank: int,
62 + world_size: int,
63 + world_local_size: int,
64 + node_rank: int,
65 + init_method: str,
66 + worker: Callable[Concatenate[ProcessGroupInfo, Optional[VllmConfig], Any,
67 + P], None],
68 + vllm_config: Optional[VllmConfig],
69 + env_dict: Optional[dict],
70 + *args: P.args,
71 + **kwargs: P.kwargs,
72 +) -> None:
73 + rank = node_rank * world_local_size + local_rank
74 + torch.cuda.set_device(local_rank)
75 + device = torch.device("cuda", local_rank)
76 + torch.distributed.init_process_group(
77 + backend="cpu:gloo,cuda:nccl",
78 + init_method=init_method,
79 + rank=rank,
80 + world_size=world_size,
81 + device_id=device,
82 + )
83 + barrier = torch.tensor([rank], device=device)
84 + torch.distributed.all_reduce(barrier)
85 +
86 + if env_dict is not None:
87 + os.environ.update(env_dict)
88 +
89 + cpu_group = None
90 + if vllm_config is not None:
91 + cpu_group = _set_vllm_config(vllm_config, world_size, rank, local_rank)
92 +
93 + try:
94 + worker(
95 + ProcessGroupInfo(
96 + world_size=world_size,
97 + world_local_size=world_local_size,
98 + rank=rank,
99 + node_rank=node_rank,
100 + local_rank=local_rank,
101 + device=device,
102 + ),
103 + vllm_config,
104 + cpu_group,
105 + *args,
106 + **kwargs,
107 + )
108 + except Exception as ex:
109 + print(ex)
110 + traceback.print_exc()
111 + raise
112 + finally:
113 + torch.distributed.destroy_process_group()
114 +
115 +
116 +def parallel_launch_with_config(
117 + world_size: int,
118 + worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig, Any, P], None],
119 + vllm_config: VllmConfig,
120 + env_dict: dict[Any, Any],
121 + *args: P.args,
122 + **kwargs: P.kwargs,
123 +) -> None:
124 + assert not kwargs
125 + spawn(
126 + _worker_parallel_launch,
127 + args=(
128 + world_size,
129 + world_size,
130 + 0,
131 + f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
132 + worker,
133 + vllm_config,
134 + env_dict,
135 + ) + args,
136 + nprocs=world_size,
137 + join=True,
138 + )
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py new
+127
@@ -0,0 +1,127 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import copy
5 +from itertools import product
6 +from typing import Any, Callable
7 +
8 +import torch
9 +
10 +from vllm.config import VllmConfig
11 +from vllm.platforms import current_platform
12 +
13 +from .common import Config, RankTensors, WeightTensors, make_modular_kernel
14 +from .parallel_utils import ProcessGroupInfo, parallel_launch_with_config
15 +
16 +
17 +def do_profile(fn: Callable,
18 + fn_kwargs: dict[Any, Any],
19 + pgi: ProcessGroupInfo,
20 + config: Config,
21 + num_warmups: int = 5):
22 + for _ in range(num_warmups):
23 + fn(**fn_kwargs)
24 +
25 + with torch.profiler.profile(
26 + activities=[
27 + torch.profiler.ProfilerActivity.CPU,
28 + torch.profiler.ProfilerActivity.CUDA,
29 + ],
30 + with_stack=True,
31 + record_shapes=True,
32 + ) as tprof:
33 + fn(**fn_kwargs)
34 + torch.cuda.synchronize(torch.cuda.current_device())
35 +
36 + # TODO (varun): Add a descriptive trace file name
37 + tprof.export_chrome_trace(
38 + f"{config.torch_trace_dir_path}/m{config.M}_{pgi.rank}_trace.json")
39 +
40 +
41 +def profile_modular_kernel(
42 + pgi: ProcessGroupInfo,
43 + vllm_config: VllmConfig,
44 + config: Config,
45 + weights: WeightTensors,
46 + rank_tensors: RankTensors,
47 +) -> None:
48 + assert isinstance(config.Ms, int)
49 + assert isinstance(config.topks, int)
50 +
51 + # weights for rank
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)
56 +
57 + mk_kwargs = {
58 + "hidden_states": rank_tensors.hidden_states,
59 + "w1": rank_weights.w1,
60 + "w2": rank_weights.w2,
61 + "topk_weights": rank_tensors.topk_weights,
62 + "topk_ids": rank_tensors.topk_ids,
63 + "expert_map": rank_tensors.expert_map,
64 + "w1_scale": rank_weights.w1_scale,
65 + "w2_scale": rank_weights.w2_scale,
66 + "a1_scale": rank_tensors.hidden_states_scale,
67 + "global_num_experts": config.E,
68 + "apply_router_weight_on_input": config.topk == 1,
69 + }
70 +
71 + do_profile(mk.forward, mk_kwargs, pgi, config)
72 +
73 +
74 +def rank_worker(
75 + pgi: ProcessGroupInfo,
76 + vllm_config: VllmConfig,
77 + cpu_group,
78 + config: Config,
79 + weights: WeightTensors,
80 +):
81 + current_platform.seed_everything(pgi.rank)
82 +
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)
87 +
88 + # get weights to this device
89 + weights.to_current_device()
90 +
91 + Ms = config.Ms
92 + assert isinstance(Ms, list)
93 + TOPKs = config.topks
94 + assert isinstance(TOPKs, list)
95 +
96 + for m, topk in product(Ms, TOPKs):
97 + print(f"Running m={m}, topk={topk} ...")
98 + # override m and topk
99 + cfgx = copy.deepcopy(config)
100 + cfgx.Ms = m
101 + cfgx.topks = topk
102 +
103 + # inputs for rank
104 + rank_tensors = RankTensors.make(cfgx, pgi)
105 + profile_modular_kernel(pgi, vllm_config, cfgx, weights, rank_tensors)
106 +
107 +
108 +def run(config: Config):
109 + weights: WeightTensors = WeightTensors.make(config)
110 + vllm_config, env_dict = config.make_env_data()
111 + parallel_launch_with_config(config.world_size, rank_worker, vllm_config,
112 + env_dict, config, weights)
113 +
114 +
115 +if __name__ == '__main__':
116 + from .cli_args import make_config, make_config_arg_parser
117 + parser = make_config_arg_parser(description=(
118 + "Run single prepare-finalize & fused-experts combination test"
119 + "Example : python3 -m tests.kernels.moe.modular_kernel_tools.profile_modular_kernel " #noqa: E501
120 + "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts"
121 + ))
122 + args = parser.parse_args()
123 + assert args.torch_trace_dir_path is not None, (
124 + "Please pass in a directory to store torch traces")
125 + config = make_config(args)
126 +
127 + run(config)
tests/kernels/moe/modular_kernel_tools/utils.py new
+142
@@ -0,0 +1,142 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import math
4 +
5 +import torch
6 +
7 +import vllm._custom_ops as ops
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 per_block_cast_to_fp8(
24 + x: torch.Tensor, block_size_k: int,
25 + block_size_n: int) -> tuple[torch.Tensor, torch.Tensor]:
26 + assert x.dim() == 2
27 + m, n = x.shape
28 + x_padded = torch.zeros(
29 + (
30 + int(math.ceil(m / block_size_k)) * block_size_k,
31 + int(math.ceil(n / block_size_n)) * block_size_n,
32 + ),
33 + dtype=x.dtype,
34 + device=x.device,
35 + )
36 + x_padded[:m, :n] = x
37 + x_view = x_padded.view(-1, block_size_k,
38 + x_padded.size(1) // block_size_k, block_size_n)
39 + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
40 + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
41 + x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
42 + scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
43 + return x_scaled_sub, scales
44 +
45 +
46 +def make_non_quant_weights(
47 + e: int,
48 + n: int,
49 + k: int,
50 + dtype: torch.dtype,
51 +) -> tuple[torch.Tensor, torch.Tensor]:
52 + """
53 + Return weights w1, w2
54 + """
55 + device = torch.cuda.current_device()
56 + w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 15
57 + w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 15
58 + return w1, w2
59 +
60 +
61 +def make_block_quant_fp8_weights(
62 + e: int,
63 + n: int,
64 + k: int,
65 + block_size: list[int],
66 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
67 + """
68 + Return weights w1, w2, w1_scale, w2_scale
69 + """
70 + dtype = torch.bfloat16
71 + device = torch.cuda.current_device()
72 +
73 + fp8_info = torch.finfo(torch.float8_e4m3fn)
74 + fp8_max, fp8_min = fp8_info.max, fp8_info.min
75 +
76 + w1_bf16, w2_bf16 = make_non_quant_weights(e, n, k, dtype)
77 + w1_bf16 = w1_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
78 + w2_bf16 = w2_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
79 +
80 + block_n, block_k = block_size[0], block_size[1]
81 + n_tiles_w1 = ((2 * n) + block_n - 1) // block_n
82 + k_tiles_w1 = (k + block_k - 1) // block_k
83 + n_tiles_w2 = (k + block_n - 1) // block_n
84 + k_tiles_w2 = (n + block_k - 1) // block_k
85 +
86 + w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn, device=device)
87 + w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn, device=device)
88 +
89 + w1_s = torch.empty((e, n_tiles_w1, k_tiles_w1),
90 + device=device,
91 + dtype=torch.float32)
92 + w2_s = torch.empty((e, n_tiles_w2, k_tiles_w2),
93 + device=device,
94 + dtype=torch.float32)
95 +
96 + assert w1_s.shape == (e, (2 * n + (block_n - 1)) // block_n,
97 + (k + (block_k - 1)) // block_k)
98 + assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
99 +
100 + for i in range(e):
101 + w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i],
102 + block_size_k=block_k,
103 + block_size_n=block_n)
104 + w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i],
105 + block_size_k=block_k,
106 + block_size_n=block_n)
107 +
108 + return w1, w2, w1_s, w2_s
109 +
110 +
111 +def make_quant_fp8_weights(
112 + e: int,
113 + n: int,
114 + k: int,
115 + per_out_channel_quant: bool,
116 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
117 + """
118 + Return w1, w2, w1_scale, w2_scale
119 + """
120 + q_dtype = torch.float8_e4m3fn
121 +
122 + w1, w2 = make_non_quant_weights(e, n, k, dtype=torch.bfloat16)
123 +
124 + # w1 -> w1_q, w2 -> w2_q
125 + w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=q_dtype)
126 + w2_q = torch.empty((e, k, n), device="cuda", dtype=q_dtype)
127 +
128 + n_b_scales = 2 * n if per_out_channel_quant else 1
129 + k_b_scales = k if per_out_channel_quant else 1
130 + w1_scale = torch.empty((e, n_b_scales, 1),
131 + device="cuda",
132 + dtype=torch.float32)
133 + w2_scale = torch.empty((e, k_b_scales, 1),
134 + device="cuda",
135 + dtype=torch.float32)
136 +
137 + for expert in range(e):
138 + w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
139 + w1[expert], use_per_token_if_dynamic=per_out_channel_quant)
140 + w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
141 + w2[expert], use_per_token_if_dynamic=per_out_channel_quant)
142 + return w1_q, w2_q, w1_scale, w2_scale
tests/kernels/moe/parallel_utils.py
+2 -4
@@ -4,7 +4,6 @@
4 DeepEP test utilities
5 """
6 import dataclasses
7 -import importlib
7 import os
8 import traceback
9 from typing import Callable, Optional
@@ -15,10 +14,9 @@ from torch.multiprocessing import (
14 spawn) # pyright: ignore[reportPrivateImportUsage]
15 from typing_extensions import Concatenate, ParamSpec
16
18 -from vllm.utils import get_open_port
17 +from vllm.utils import get_open_port, has_deep_ep
18
20 -has_deep_ep = importlib.util.find_spec("deep_ep") is not None
21 -if has_deep_ep:
19 +if has_deep_ep():
20 from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
21 DeepEPHTPrepareAndFinalize)
22 from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
tests/kernels/moe/test_modular_kernel_combinations.py new
+214
@@ -0,0 +1,214 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import copy
5 +from itertools import product
6 +from typing import Optional
7 +
8 +import pytest
9 +import torch
10 +
11 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12 +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)
22 +from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
23 +
24 +from .modular_kernel_tools.common import (Config, RankTensors, WeightTensors,
25 + reference_moe_impl,
26 + run_modular_kernel)
27 +from .modular_kernel_tools.mk_objects import (
28 + MK_FUSED_EXPERT_TYPES, MK_MULTI_GPU_PREPARE_FINALIZE_TYPES,
29 + MK_QUANT_CONFIGS, MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES)
30 +from .modular_kernel_tools.parallel_utils import (ProcessGroupInfo,
31 + parallel_launch_with_config)
32 +
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())
35 +
36 +meets_package_requirements = pytest.mark.skipif(
37 + not has_all_packages,
38 + reason="Requires deep_ep & deep_gemm & pplx packages",
39 +)
40 +
41 +
42 +def rank_worker(
43 + pgi: ProcessGroupInfo,
44 + vllm_config: VllmConfig,
45 + cpu_group,
46 + config: Config,
47 + weights: WeightTensors,
48 +):
49 + current_platform.seed_everything(pgi.rank)
50 +
51 + # sanity check
52 + from vllm import envs
53 + if config.fused_moe_chunk_size is not None:
54 + assert (config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE)
55 +
56 + # get weights to this device
57 + weights.to_current_device()
58 +
59 + Ms = config.Ms
60 + assert isinstance(Ms, list)
61 + TOPKs = config.topks
62 + assert isinstance(TOPKs, list)
63 +
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)
77 +
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):
85 + assert config.is_valid()
86 + print(f"Testing config \n{config.describe()} ...")
87 +
88 + weights: WeightTensors = WeightTensors.make(config)
89 +
90 + vllm_config, env_dict = config.make_env_data()
91 + parallel_launch_with_config(config.world_size, rank_worker, vllm_config,
92 + env_dict, config, weights)
93 +
94 +
95 +Ms = [32, 64]
96 +Ks = [7168] # hidden sizes
97 +Ns = [2048]
98 +TOPKs = [4, 1]
99 +Es = [32]
100 +DTYPEs = [torch.bfloat16]
101 +FUSED_MOE_CHUNK_SIZEs = [None, 16]
102 +
103 +
104 +def is_nyi_config(config: Config) -> bool:
105 + # We know these configs to be legitimate. but still fail.
106 +
107 + if (config.fused_experts_type in [
108 + BatchedTritonExperts, BatchedTritonOrDeepGemmExperts,
109 + TritonExperts, TritonOrDeepGemmExperts
110 + ]):
111 + # The triton kernels expect both per-act-token-quant and
112 + # per-out-ch-quant or neither.
113 + unsupported_quant_config = ((config.is_per_act_token_quant +
114 + config.is_per_out_ch_quant) == 1)
115 + return unsupported_quant_config
116 +
117 + # cutlass kernels dont support expert_maps yet.
118 + return config.fused_experts_type == CutlassExpertsFp8
119 +
120 +
121 +@pytest.mark.parametrize("k", Ks)
122 +@pytest.mark.parametrize("n", Ns)
123 +@pytest.mark.parametrize("e", Es)
124 +@pytest.mark.parametrize("dtype", DTYPEs)
125 +@pytest.mark.parametrize("quant_config", MK_QUANT_CONFIGS)
126 +@pytest.mark.parametrize(
127 + "combination",
128 + product(MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))
129 +@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
130 +@pytest.mark.parametrize("world_size", [2])
131 +@meets_package_requirements
132 +def test_modular_kernel_combinations_multigpu(
133 + k: int, n: int, e: int, dtype: torch.dtype,
134 + quant_config: FusedMoEQuantConfig,
135 + combination: tuple[mk.FusedMoEPrepareAndFinalize,
136 + mk.FusedMoEPermuteExpertsUnpermute],
137 + fused_moe_chunk_size: Optional[int], world_size: int):
138 +
139 + config = Config(
140 + Ms=Ms,
141 + K=k,
142 + N=n,
143 + E=e,
144 + topks=TOPKs,
145 + dtype=dtype,
146 + quant_config=quant_config,
147 + prepare_finalize_type=combination[0],
148 + fused_experts_type=combination[1],
149 + fused_moe_chunk_size=fused_moe_chunk_size,
150 + world_size=world_size,
151 + )
152 + if not config.is_valid():
153 + pytest.skip(f"Tests config {config} is not valid. Skipping ...")
154 +
155 + if is_nyi_config(config):
156 + pytest.skip(f"Tests config {config} is nyi. Skipping ...")
157 +
158 + print(f"{config.describe()}")
159 + run(config)
160 +
161 +
162 +@pytest.mark.parametrize("k", Ks)
163 +@pytest.mark.parametrize("n", Ns)
164 +@pytest.mark.parametrize("e", Es)
165 +@pytest.mark.parametrize("dtype", DTYPEs)
166 +@pytest.mark.parametrize("quant_config", MK_QUANT_CONFIGS)
167 +@pytest.mark.parametrize(
168 + "combination",
169 + product(MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES))
170 +@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
171 +@pytest.mark.parametrize("world_size", [1])
172 +@meets_package_requirements
173 +def test_modular_kernel_combinations_singlegpu(
174 + k: int, n: int, e: int, dtype: torch.dtype,
175 + quant_config: FusedMoEQuantConfig,
176 + combination: tuple[mk.FusedMoEPrepareAndFinalize,
177 + mk.FusedMoEPermuteExpertsUnpermute],
178 + fused_moe_chunk_size: Optional[int], world_size: int):
179 + config = Config(
180 + Ms=Ms,
181 + K=k,
182 + N=n,
183 + E=e,
184 + topks=TOPKs,
185 + dtype=dtype,
186 + quant_config=quant_config,
187 + prepare_finalize_type=combination[0],
188 + fused_experts_type=combination[1],
189 + fused_moe_chunk_size=fused_moe_chunk_size,
190 + world_size=world_size,
191 + )
192 +
193 + if not config.is_valid():
194 + pytest.skip(f"Tests config {config} is not valid. Skipping ...")
195 +
196 + if is_nyi_config(config):
197 + pytest.skip(f"Tests config {config} is nyi. Skipping ...")
198 +
199 + run(config)
200 +
201 +
202 +if __name__ == '__main__':
203 + # Ability to test individual PrepareAndFinalize and FusedExperts combination
204 + from .modular_kernel_tools.cli_args import (make_config,
205 + make_config_arg_parser)
206 + parser = make_config_arg_parser(description=(
207 + "Run single prepare-finalize & fused-experts combination test"
208 + "Example : python3 -m tests.kernels.moe.test_modular_kernel_combinations " #noqa: E501
209 + "--pf-type PplxPrepareAndFinalize --experts-type BatchedTritonExperts"
210 + ))
211 + args = parser.parse_args()
212 + config = make_config(args)
213 +
214 + run(config)
tests/kernels/utils.py
+25 -5
@@ -1072,6 +1072,7 @@ def torch_experts(
1072 quant_dtype: Optional[torch.dtype] = None,
1073 per_act_token_quant=False,
1074 block_shape: Optional[list[int]] = None,
1075 + apply_router_weights_on_input: bool = False,
1076 ) -> torch.Tensor:
1077 assert (global_num_experts == -1
1078 or (global_num_experts == w1.shape[0] and expert_map is None)
@@ -1081,11 +1082,17 @@ def torch_experts(
1082 M, K = a.shape
1083 topk = topk_ids.shape[1]
1084
1085 + if apply_router_weights_on_input:
1086 + assert topk == 1
1087 + a = a * topk_weight.to(a.dtype)
1088 +
1089 a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
1090
1091 out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)
1092
1088 - a, a_scale = moe_kernel_quantize_input(a, None, quant_dtype,
1093 + if a1_scale:
1094 + assert not per_act_token_quant and block_shape is None
1095 + a, a_scale = moe_kernel_quantize_input(a, a1_scale, quant_dtype,
1096 per_act_token_quant, block_shape)
1097
1098 num_experts = w1.shape[0]
@@ -1104,6 +1111,7 @@ def torch_experts(
1111 tmp2 = SiluAndMul()(tmp1)
1112 out[mask] = tmp2 @ w2[i].transpose(0, 1)
1113 elif block_shape is not None:
1114 + # block quantized
1115 assert (a_scale is not None and w1_scale is not None
1116 and w2_scale is not None)
1117 tmp1 = native_w8a8_block_matmul(a[mask], w1[i], a_scale[mask],
@@ -1121,15 +1129,27 @@ def torch_experts(
1129 assert (a_scale is not None and w1_scale is not None
1130 and w2_scale is not None)
1131 scales = a_scale if a_scale.numel() == 1 else a_scale[mask]
1132 +
1133 tmp1 = a[mask].to(f32) * scales
1134 w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1)
1126 - tmp1 = tmp1 @ w1_dq
1127 - tmp2 = SiluAndMul()(tmp1)
1135 + tmp1 = (tmp1 @ w1_dq).to(out.dtype)
1136 +
1137 + tmp2 = SiluAndMul()(tmp1).to(out.dtype)
1138 +
1139 + tmp2, b_scale = moe_kernel_quantize_input(
1140 + tmp2, a2_scale, quant_dtype, per_act_token_quant,
1141 + block_shape)
1142 + assert b_scale is not None
1143 +
1144 + tmp2 = tmp2.to(f32) * b_scale
1145 w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1)
1146 out[mask] = (tmp2 @ w2_dq).to(out.dtype)
1147
1131 - return (out.view(M, -1, w2.shape[1]).to(f32) *
1132 - topk_weight.view(M, -1, 1)).sum(dim=1).to(out.dtype)
1148 + if apply_router_weights_on_input:
1149 + return out
1150 + else:
1151 + return (out.view(M, -1, w2.shape[1]).to(f32) *
1152 + topk_weight.view(M, -1, 1)).sum(dim=1).to(out.dtype)
1153
1154
1155 def torch_moe(a: torch.Tensor,
vllm/distributed/device_communicators/base_device_communicator.py
+1 -2
@@ -240,8 +240,7 @@ class DeviceCommunicatorBase:
240 if module.__class__.__name__ == "FusedMoE"
241 ]
242 for module in moe_modules:
243 - module.quant_method.init_prepare_finalize(module.moe_config,
244 - module.quant_config)
243 + module.quant_method.init_prepare_finalize(module.moe_config)
244
245 def dispatch(
246 self, hidden_states: torch.Tensor,
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
-1
@@ -37,7 +37,6 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
37 block_shape=block_shape,
38 per_act_token_quant=per_act_token_quant,
39 ))
40 - self.allow_deep_gemm = allow_deep_gemm
40
41 self.batched_triton_experts = BatchedTritonExperts(
42 max_num_tokens=max_num_tokens,
vllm/model_executor/layers/fused_moe/layer.py
+11 -7
@@ -81,13 +81,12 @@ class FusedMoEMethodBase(QuantizeMethodBase):
81 params_dtype: torch.dtype, **extra_weight_attrs):
82 raise NotImplementedError
83
84 - def init_prepare_finalize(self, moe: FusedMoEConfig,
85 - quant_config: Optional[QuantizationConfig]):
84 + @staticmethod
85 + def maybe_make_prepare_finalize(
86 + moe: FusedMoEConfig) -> Optional[FusedMoEPrepareAndFinalize]:
87 all2all_manager = get_ep_group().device_communicator.all2all_manager
88 assert all2all_manager is not None
89
89 - self.moe = moe
90 -
90 prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None
91
92 if moe.use_pplx_kernels:
@@ -160,8 +159,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
159 and moe.quant_config.block_shape
160 == DEEPEP_QUANT_BLOCK_SHAPE)
161
163 - # Note (varun): Whether to use FP8 dispatch or not needs some
164 - # profiling. Turning it off for now.
162 prepare_finalize = DeepEPLLPrepareAndFinalize(
163 handle,
164 max_tokens_per_rank=moe.max_num_tokens,
@@ -169,11 +166,18 @@ class FusedMoEMethodBase(QuantizeMethodBase):
166 use_fp8_dispatch=use_fp8_dispatch,
167 )
168
169 + return prepare_finalize
170 +
171 + def init_prepare_finalize(self, moe: FusedMoEConfig):
172 + self.moe = moe
173 + prepare_finalize = FusedMoEMethodBase.maybe_make_prepare_finalize(
174 + self.moe)
175 +
176 self.topk_indices_dtype = None
177 if prepare_finalize is not None:
178 logger.debug("%s", prepare_finalize.__class__.__name__)
179 self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
176 - experts = self.select_gemm_impl(prepare_finalize, moe)
180 + experts = self.select_gemm_impl(prepare_finalize, self.moe)
181 self.fused_experts = FusedMoEModularKernel(
182 prepare_finalize,
183 experts,
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+6 -3
@@ -7,7 +7,8 @@ import torch
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
10 - DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape)
10 + DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape,
11 + deep_gemm_block_shape)
12 from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
13 from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used
14
@@ -44,8 +45,10 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
45 per_act_token_quant=per_act_token_quant,
46 block_shape=block_shape,
47 )
47 - self.allow_deep_gemm = (allow_deep_gemm and not per_act_token_quant
48 - and use_fp8_w8a8)
48 +
49 + self.allow_deep_gemm = (allow_deep_gemm and use_fp8_w8a8 and
50 + self.block_shape == deep_gemm_block_shape())
51 +
52 self.deep_gemm_expert = DeepGemmExperts(
53 ) if self.allow_deep_gemm else None
54