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

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

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

Synced Nov 14, 05:30 UTC

[Kernels] Modular kernel refactor (#24812)

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

bnellnm committed Oct 8, 2025 at 17:51 UTC da364615fcc42306c73401f22f57b0c1f8226efb
22 files changed +597 -505
tests/kernels/moe/modular_kernel_tools/common.py
+23 -14
@@ -209,18 +209,18 @@ class Config:
209 info = prepare_finalize_info(self.prepare_finalize_type)
210 return info.backend
211
212 - def is_valid(self):
212 + def is_valid(self) -> tuple[bool, Optional[str]]:
213 # Check prepare-finalize and fused-experts compatibility
214 if self.is_batched_prepare_finalize():
215 if not self.is_batched_fused_experts():
216 - return False
216 + return False, "Mismatched format."
217 else:
218 if not self.is_standard_fused_experts():
219 - return False
219 + return False, "Mismatched format."
220
221 use_chunking = self.fused_moe_chunk_size is not None
222 if use_chunking and not self.is_fe_supports_chunking():
223 - return False
223 + return False, "Chunking not supported."
224
225 # Check quantization sanity
226 if (
@@ -229,7 +229,7 @@ class Config:
229 + int(self.quant_block_shape is not None)
230 ) > 1:
231 # invalid quant config
232 - return False
232 + return False, f"Bad quant_config {self.quant_config}."
233
234 # check type support
235 if self.quant_dtype is None:
@@ -237,34 +237,43 @@ class Config:
237 self.dtype not in self.pf_supported_types()
238 or self.dtype not in self.fe_supported_types()
239 ):
240 - return False
240 + return False, (
241 + f"Unsupported type {self.dtype} not in "
242 + f"{self.pf_supported_types()} and "
243 + f"{self.fe_supported_types()}."
244 + )
245 else:
246 if (
247 self.quant_dtype not in self.pf_supported_types()
248 or self.quant_dtype not in self.fe_supported_types()
249 ):
246 - return False
250 + return False, (
251 + f"Unsupported quant type {self.quant_dtype} "
252 + f"not in {self.pf_supported_types()} and "
253 + f"{self.fe_supported_types()}."
254 + )
255
256 # Check block quanization support
257 is_block_quatized = self.quant_block_shape is not None
258 if is_block_quatized and self.quant_dtype is None:
251 - return False
259 + return False, "No block quantization support."
260 +
261 if is_block_quatized and not self.is_block_quant_supported():
253 - return False
262 + return False, "Mismatched block quantization support."
263
264 # deep_gemm only works with block-quantized
265 if self.needs_deep_gemm() and not is_block_quatized:
257 - return False
266 + return False, "Needs DeepGEMM but not block quantized."
267
268 # Check dependencies (turn into asserts?)
269 if self.needs_deep_ep() and not has_deep_ep():
261 - return False
270 + return False, "Needs DeepEP, but DeepEP not available."
271 if self.needs_deep_gemm() and not has_deep_gemm():
263 - return False
272 + return False, "Needs DeepGEMM, but DeepGEMM not available."
273 if self.needs_pplx() and not has_pplx(): # noqa: SIM103
265 - return False
274 + return False, "Needs PPLX, but PPLX not available."
275
267 - return True
276 + return True, None
277
278
279 @dataclass
tests/kernels/moe/modular_kernel_tools/make_feature_matrix.py
+1 -1
@@ -140,7 +140,7 @@ def make_feature_matrix(csv_file_path: str):
140 )
141
142 success = None
143 - if config.is_valid():
143 + if config.is_valid()[0]:
144 print(f"Running config : {config.describe()} ...")
145 try:
146 weights: WeightTensors = WeightTensors.make(config)
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+12 -14
@@ -244,7 +244,7 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
244 register_prepare_and_finalize(
245 FlashInferCutlassMoEPrepareAndFinalize,
246 standard_format,
247 - nvfp4_types,
247 + nvfp4_types + fp8_types,
248 blocked_quantization_support=True,
249 backend=None,
250 force_multigpu=True,
@@ -254,7 +254,7 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
254 register_experts(
255 FlashInferExperts,
256 standard_format,
257 - nvfp4_types,
257 + nvfp4_types + fp8_types,
258 blocked_quantization_support=True,
259 supports_chunking=True,
260 # Note: this is a hack to get it to run for now
@@ -274,17 +274,15 @@ if has_deep_gemm() and is_deep_gemm_supported():
274 needs_matching_quant=False,
275 needs_deep_gemm=True,
276 )
277 - (
278 - register_experts(
279 - DeepGemmExperts,
280 - standard_format,
281 - fp8_types,
282 - blocked_quantization_support=True,
283 - supports_chunking=True,
284 - supports_expert_map=True,
285 - needs_matching_quant=False,
286 - needs_deep_gemm=True,
287 - ),
277 + register_experts(
278 + DeepGemmExperts,
279 + standard_format,
280 + fp8_types,
281 + blocked_quantization_support=True,
282 + supports_chunking=True,
283 + supports_expert_map=True,
284 + needs_matching_quant=False,
285 + needs_deep_gemm=True,
286 )
287 register_experts(
288 BatchedTritonOrDeepGemmExperts,
@@ -464,7 +462,7 @@ def make_fused_experts(
462 print(f"Making BatchedTritonOrDeepGemmExperts {kwargs} ...")
463 experts = BatchedTritonOrDeepGemmExperts(**kwargs)
464 elif fused_experts_type == DeepGemmExperts:
467 - print("Making DeepGemmExperts {quant_config} ...")
465 + print(f"Making DeepGemmExperts {quant_config} ...")
466 experts = DeepGemmExperts(quant_config)
467 elif fused_experts_type == TritonExperts:
468 kwargs = quant_kwargs
tests/kernels/moe/test_modular_kernel_combinations.py
+96 -48
@@ -5,7 +5,7 @@ import copy
5 import textwrap
6 import traceback
7 from itertools import product
8 -from typing import Optional
8 +from typing import Any, Optional
9
10 import pytest
11 import torch
@@ -13,10 +13,9 @@ import torch
13 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 from vllm.config import VllmConfig, set_current_vllm_config
15 from vllm.platforms import current_platform
16 -from vllm.utils import has_deep_ep, has_deep_gemm, has_pplx
16 +from vllm.utils import cuda_device_count_stateless, has_deep_ep, has_deep_gemm, has_pplx
17 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
18
19 -from ...utils import multi_gpu_test
19 from .modular_kernel_tools.common import (
20 Config,
21 RankTensors,
@@ -132,7 +131,8 @@ def rank_worker(
131
132
133 def run(config: Config, verbose: bool):
135 - assert config.is_valid()
134 + assert config.is_valid()[0]
135 + assert not is_nyi_config(config)
136
137 weights: WeightTensors = WeightTensors.make(config)
138
@@ -168,17 +168,77 @@ def is_nyi_config(config: Config) -> bool:
168 return not info.supports_expert_map
169
170
171 -@pytest.mark.parametrize("k", Ks)
172 -@pytest.mark.parametrize("n", Ns)
173 -@pytest.mark.parametrize("e", Es)
174 -@pytest.mark.parametrize("dtype", DTYPEs)
175 -@pytest.mark.parametrize("quant_config", MK_QUANT_CONFIGS)
171 +def generate_valid_test_cases(
172 + world_size: int, prepare_finalize_types
173 +) -> list[tuple[Any, ...]]:
174 + cases = []
175 + total = 0
176 +
177 + for k, n, e, dtype, quant_config, combination, chunk_size in product(
178 + Ks,
179 + Ns,
180 + Es,
181 + DTYPEs,
182 + MK_QUANT_CONFIGS,
183 + product(prepare_finalize_types, MK_FUSED_EXPERT_TYPES),
184 + FUSED_MOE_CHUNK_SIZEs,
185 + ):
186 + total = total + 1
187 +
188 + config = Config(
189 + Ms=Ms,
190 + K=k,
191 + N=n,
192 + E=e,
193 + topks=TOPKs,
194 + dtype=dtype,
195 + quant_config=quant_config,
196 + prepare_finalize_type=combination[0],
197 + fused_experts_type=combination[1],
198 + fused_moe_chunk_size=chunk_size,
199 + world_size=world_size,
200 + )
201 +
202 + # TODO(bnell): figure out how to get verbose flag here.
203 + verbose = False # pytestconfig.getoption('verbose') > 0
204 +
205 + valid, reason = config.is_valid()
206 +
207 + if not valid:
208 + if verbose:
209 + print(f"Test config {config} is not valid: {reason}")
210 + continue
211 +
212 + if is_nyi_config(config):
213 + if verbose:
214 + print(f"Test config {config} is nyi.")
215 + continue
216 +
217 + cases.append(
218 + (
219 + k,
220 + n,
221 + e,
222 + dtype,
223 + quant_config,
224 + combination[0],
225 + combination[1],
226 + chunk_size,
227 + world_size,
228 + )
229 + )
230 +
231 + print(f"{len(cases)} of {total} valid configs generated.")
232 +
233 + return cases
234 +
235 +
236 @pytest.mark.parametrize(
177 - "combination", product(MK_MULTI_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES)
237 + "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,chunk_size,world_size",
238 + generate_valid_test_cases(
239 + world_size=2, prepare_finalize_types=MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
240 + ),
241 )
179 -@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
180 -@pytest.mark.parametrize("world_size", [2])
181 -@multi_gpu_test(num_gpus=2)
242 @meets_multi_gpu_requirements
243 def test_modular_kernel_combinations_multigpu(
244 k: int,
@@ -186,13 +246,19 @@ def test_modular_kernel_combinations_multigpu(
246 e: int,
247 dtype: torch.dtype,
248 quant_config: Optional[TestMoEQuantConfig],
189 - combination: tuple[
190 - mk.FusedMoEPrepareAndFinalize, mk.FusedMoEPermuteExpertsUnpermute
191 - ],
192 - fused_moe_chunk_size: Optional[int],
249 + prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
250 + fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
251 + chunk_size: Optional[int],
252 world_size: int,
253 pytestconfig,
254 ):
255 + if cuda_device_count_stateless() < world_size:
256 + pytest.skip(
257 + f"Not enough GPUs available to run, got "
258 + f"{cuda_device_count_stateless()} exepected "
259 + f"{world_size}."
260 + )
261 +
262 config = Config(
263 Ms=Ms,
264 K=k,
@@ -201,42 +267,30 @@ def test_modular_kernel_combinations_multigpu(
267 topks=TOPKs,
268 dtype=dtype,
269 quant_config=quant_config,
204 - prepare_finalize_type=combination[0],
205 - fused_experts_type=combination[1],
206 - fused_moe_chunk_size=fused_moe_chunk_size,
270 + prepare_finalize_type=prepare_finalize_type,
271 + fused_experts_type=fused_experts_type,
272 + fused_moe_chunk_size=chunk_size,
273 world_size=world_size,
274 )
209 -
210 - if not config.is_valid():
211 - pytest.skip(f"Tests config {config} is not valid. Skipping ...")
212 -
213 - if is_nyi_config(config):
214 - pytest.skip(f"Tests config {config} is nyi. Skipping ...")
215 -
275 verbosity = pytestconfig.getoption("verbose")
276 run(config, verbosity > 0)
277
278
220 -@pytest.mark.parametrize("k", Ks)
221 -@pytest.mark.parametrize("n", Ns)
222 -@pytest.mark.parametrize("e", Es)
223 -@pytest.mark.parametrize("dtype", DTYPEs)
224 -@pytest.mark.parametrize("quant_config", MK_QUANT_CONFIGS)
279 @pytest.mark.parametrize(
226 - "combination", product(MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES, MK_FUSED_EXPERT_TYPES)
280 + "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,chunk_size,world_size",
281 + generate_valid_test_cases(
282 + world_size=1, prepare_finalize_types=MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
283 + ),
284 )
228 -@pytest.mark.parametrize("fused_moe_chunk_size", FUSED_MOE_CHUNK_SIZEs)
229 -@pytest.mark.parametrize("world_size", [1])
285 def test_modular_kernel_combinations_singlegpu(
286 k: int,
287 n: int,
288 e: int,
289 dtype: torch.dtype,
290 quant_config: Optional[TestMoEQuantConfig],
236 - combination: tuple[
237 - mk.FusedMoEPrepareAndFinalize, mk.FusedMoEPermuteExpertsUnpermute
238 - ],
239 - fused_moe_chunk_size: Optional[int],
291 + prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
292 + fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute,
293 + chunk_size: Optional[int],
294 world_size: int,
295 pytestconfig,
296 ):
@@ -248,18 +302,12 @@ def test_modular_kernel_combinations_singlegpu(
302 topks=TOPKs,
303 dtype=dtype,
304 quant_config=quant_config,
251 - prepare_finalize_type=combination[0],
252 - fused_experts_type=combination[1],
253 - fused_moe_chunk_size=fused_moe_chunk_size,
305 + prepare_finalize_type=prepare_finalize_type,
306 + fused_experts_type=fused_experts_type,
307 + fused_moe_chunk_size=chunk_size,
308 world_size=world_size,
309 )
310
257 - if not config.is_valid():
258 - pytest.skip(f"Tests config {config} is not valid. Skipping ...")
259 -
260 - if is_nyi_config(config):
261 - pytest.skip(f"Tests config {config} is nyi. Skipping ...")
262 -
311 verbosity = pytestconfig.getoption("verbose")
312 run(config, verbosity > 0)
313
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+5 -10
@@ -247,29 +247,24 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
247
248 def workspace_shapes(
249 self,
250 - a: torch.Tensor,
251 - aq: torch.Tensor,
250 M: int,
251 N: int,
252 K: int,
253 topk: int,
254 global_num_experts: int,
255 local_num_experts: int,
258 - expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
259 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
260 - assert a.dim() == 2
256 + expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
257 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
258 # FIXME (varun): We should be able to dispatch only from the leader
259 # DP ranks in the case of TP > 1. At the moment, all the Ranks
260 # end up sending their tokens. This needs to be fixed.
261 num_dispatchers = self.num_dispatchers
262 num_experts = local_num_experts
266 - max_num_tokens = (
267 - a.size(0) if self.max_num_tokens is None else self.max_num_tokens
268 - )
263 + max_num_tokens = M if self.max_num_tokens is None else self.max_num_tokens
264 workspace13 = (num_experts, max_num_tokens * num_dispatchers, max(K, N))
265 workspace2 = (num_experts, max_num_tokens * num_dispatchers, (N // 2))
266 output = (num_experts, max_num_tokens * num_dispatchers, K)
272 - return (workspace13, workspace2, output, a.dtype)
267 + return (workspace13, workspace2, output)
268
269 def apply(
270 self,
@@ -300,7 +295,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
295
296 assert w2.size(1) == K
297
303 - E, max_num_tokens, N, K, top_k_num = self.moe_problem_size(
298 + E, max_num_tokens, N, K, _ = self.moe_problem_size(
299 hidden_states, w1, w2, topk_ids
300 )
301
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+4 -7
@@ -99,10 +99,11 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
99 assert bte_war is not None
100 return bte_war
101
102 + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
103 + return act_dtype
104 +
105 def workspace_shapes(
106 self,
104 - a: torch.Tensor,
105 - aq: torch.Tensor,
107 M: int,
108 N: int,
109 K: int,
@@ -110,15 +111,13 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
111 global_num_experts: int,
112 local_num_experts: int,
113 expert_tokens_metadata: Optional[mk.ExpertTokensMetadata],
113 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
114 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
115 # Note: the deep gemm workspaces are strictly larger than the triton
116 # workspaces so we can be pessimistic here and allocate for DeepGemm
117 # even if we fall back to triton later, e.g. if expert maps are set.
118 if self.allow_deep_gemm:
119 assert self.batched_deep_gemm_experts is not None
120 return self.batched_deep_gemm_experts.workspace_shapes(
120 - a,
121 - aq,
121 M,
122 N,
123 K,
@@ -130,8 +129,6 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
129 else:
130 assert self.batched_triton_experts is not None
131 return self.batched_triton_experts.workspace_shapes(
133 - a,
134 - aq,
132 M,
133 N,
134 K,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+21 -36
@@ -366,10 +366,11 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
366 # topk weights and reduction are fused in moe_unpermute cuda kernel
367 return TopKWeightAndReduceNoOP()
368
369 + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
370 + return self.out_dtype if self.out_dtype is not None else act_dtype
371 +
372 def workspace_shapes(
373 self,
371 - a: torch.Tensor,
372 - aq: torch.Tensor,
374 M: int,
375 N: int,
376 K: int,
@@ -377,16 +378,11 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
378 global_num_experts: int,
379 local_num_experts: int,
380 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
380 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
381 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
382 workspace1 = (M * topk, max(N, K))
383 workspace2 = (M * topk, max(N // 2, K))
384 output = (M, K)
384 - return (
385 - workspace1,
386 - workspace2,
387 - output,
388 - self.out_dtype if self.out_dtype is not None else a.dtype,
389 - )
385 + return (workspace1, workspace2, output)
386
387
388 class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
@@ -428,11 +424,11 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
424 def supports_expert_map(self) -> bool:
425 return False
426
431 - # TODO(bnell): maybe remove need for passing aq to workspace_shapes
427 + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
428 + return self.out_dtype if self.out_dtype is not None else act_dtype
429 +
430 def workspace_shapes(
431 self,
434 - a: torch.Tensor,
435 - aq: torch.Tensor,
432 M: int,
433 N: int,
434 K: int,
@@ -440,19 +436,13 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
436 global_num_experts: int,
437 local_num_experts: int,
438 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
443 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
444 - padded_M = aq.size(1)
439 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
440 num_dp = self.num_dispatchers
441 assert num_dp is not None
447 - workspace1 = (self.max_experts_per_worker, padded_M * num_dp, max(N, K))
448 - workspace2 = (self.max_experts_per_worker, padded_M * num_dp, max(N // 2, K))
449 - output = (self.max_experts_per_worker, padded_M, K)
450 - return (
451 - workspace1,
452 - workspace2,
453 - output,
454 - self.out_dtype if self.out_dtype is not None else a.dtype,
455 - )
442 + workspace1 = (self.max_experts_per_worker, M * num_dp, max(N, K))
443 + workspace2 = (self.max_experts_per_worker, M * num_dp, max(N // 2, K))
444 + output = (self.max_experts_per_worker, M, K)
445 + return (workspace1, workspace2, output)
446
447
448 def cutlass_moe_fp8(
@@ -767,10 +757,11 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
757 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
758 return TopKWeightAndReduceNoOP()
759
760 + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
761 + return self.out_dtype if self.out_dtype is not None else act_dtype
762 +
763 def workspace_shapes(
764 self,
772 - a: torch.Tensor,
773 - aq: torch.Tensor,
765 M: int,
766 N: int,
767 K: int,
@@ -778,25 +769,19 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):
769 global_num_experts: int,
770 local_num_experts: int,
771 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
781 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
772 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
773 workspace1: tuple[int, ...] = ()
774 workspace2: tuple[int, ...] = ()
775 output: tuple[int, ...] = ()
776 if self.use_batched_format:
786 - padded_M = aq.size(1)
787 - workspace1 = (self.max_experts_per_worker, padded_M, max(N, K))
788 - workspace2 = (self.max_experts_per_worker, padded_M, (N // 2))
789 - output = (self.max_experts_per_worker, padded_M, K)
777 + workspace1 = (self.max_experts_per_worker, M, max(N, K))
778 + workspace2 = (self.max_experts_per_worker, M, (N // 2))
779 + output = (self.max_experts_per_worker, M, K)
780 else:
781 workspace1 = (M * topk, max(2 * N, K))
782 workspace2 = (M * topk, N)
783 output = (M, K)
794 - return (
795 - workspace1,
796 - workspace2,
797 - output,
798 - self.out_dtype if self.out_dtype is not None else a.dtype,
799 - )
784 + return (workspace1, workspace2, output)
785
786 def apply(
787 self,
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+2 -4
@@ -198,8 +198,6 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
198
199 def workspace_shapes(
200 self,
201 - a: torch.Tensor,
202 - aq: torch.Tensor,
201 M: int,
202 N: int,
203 K: int,
@@ -207,7 +205,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
205 global_num_experts: int,
206 local_num_experts: int,
207 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
210 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
208 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
209 assert self.block_shape is not None
210 block_m = self.block_shape[0]
211 M_sum = compute_aligned_M(
@@ -218,7 +216,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
216 workspace1 = (M_sum, max(N, K))
217 workspace2 = (M_sum, max(N // 2, K))
218 output = (M, K)
221 - return (workspace1, workspace2, output, a.dtype)
219 + return (workspace1, workspace2, output)
220
221 def apply(
222 self,
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+3
@@ -70,6 +70,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
70 def num_dispatchers(self) -> int:
71 return self.num_dispatchers_
72
73 + def output_is_reduced(self) -> bool:
74 + return True
75 +
76 @property
77 def activation_format(self) -> mk.FusedMoEActivationFormat:
78 return mk.FusedMoEActivationFormat.Standard
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+3
@@ -73,6 +73,9 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
73 def num_dispatchers(self) -> int:
74 return self.num_dispatchers_
75
76 + def output_is_reduced(self) -> bool:
77 + return True
78 +
79 @property
80 def activation_format(self) -> mk.FusedMoEActivationFormat:
81 return mk.FusedMoEActivationFormat.BatchedExperts
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+4 -8
@@ -90,8 +90,6 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
90
91 def workspace_shapes(
92 self,
93 - a: torch.Tensor,
94 - aq: torch.Tensor,
93 M: int,
94 N: int,
95 K: int,
@@ -99,7 +97,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
97 global_num_experts: int,
98 local_num_experts: int,
99 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
102 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
100 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
101 # We use global_num_experts due to how moe_align_block_size handles
102 # expert_maps.
103 """
@@ -118,14 +116,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
116 - Note: in order for activation chunking to work, the first dimension
117 of each tuple must be the number of tokens.
118 """
121 - aq_m, aq_n = aq.shape
119 + workspace1 = (M, K)
120 workspace2 = (0,)
123 - output_shape = (aq_m, aq_n * 2) if self.quant_dtype == "nvfp4" else (aq_m, aq_n)
124 - workspace_dtype = a.dtype
125 - workspace1 = output_shape
121 + output_shape = (M, K * 2 if self.quant_dtype == "nvfp4" else K)
122 # The workspace is determined by `aq`, since it comes after any
123 # potential communication op and is involved in the expert computation.
128 - return (workspace1, workspace2, output_shape, workspace_dtype)
124 + return (workspace1, workspace2, output_shape)
125
126 def apply(
127 self,
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+8
@@ -11,6 +11,9 @@ from vllm.distributed.device_communicators.base_device_communicator import (
11 )
12 from vllm.forward_context import get_forward_context
13 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
14 +from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15 + TopKWeightAndReduceNoOP,
16 +)
17 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
18 from vllm.utils.flashinfer import nvfp4_block_scale_interleave
19
@@ -45,6 +48,9 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
48 def num_dispatchers(self) -> int:
49 return self.num_dispatchers_
50
51 + def output_is_reduced(self) -> bool:
52 + return False
53 +
54 def _apply_router_weight_on_input(
55 self,
56 a1: torch.Tensor,
@@ -194,6 +200,8 @@ class FlashInferAllGatherMoEPrepareAndFinalize(FlashInferCutlassMoEPrepareAndFin
200 apply_router_weight_on_input: bool,
201 weight_and_reduce_impl: mk.TopKWeightAndReduce,
202 ) -> None:
203 + assert isinstance(weight_and_reduce_impl, TopKWeightAndReduceNoOP)
204 +
205 if self.use_dp:
206 fused_expert_output = get_dp_group().reduce_scatterv(
207 fused_expert_output, dim=0, sizes=get_local_sizes()
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+7 -10
@@ -509,6 +509,9 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
509 def num_dispatchers(self) -> int:
510 return self.num_dispatchers_
511
512 + def output_is_reduced(self) -> bool:
513 + return False
514 +
515 def prepare(
516 self,
517 a1: torch.Tensor,
@@ -665,8 +668,6 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
668
669 def workspace_shapes(
670 self,
668 - a: torch.Tensor,
669 - aq: torch.Tensor,
671 M: int,
672 N: int,
673 K: int,
@@ -674,14 +675,13 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
675 global_num_experts: int,
676 local_num_experts: int,
677 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
677 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
678 - assert a.dim() == 2
678 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
679 num_dp = self.num_dispatchers
680 num_experts = local_num_experts
681 workspace13 = (num_experts, self.max_num_tokens * num_dp, K)
682 workspace2 = (self.max_num_tokens * num_dp, N)
683 output = workspace13
684 - return (workspace13, workspace2, output, a.dtype)
684 + return (workspace13, workspace2, output)
685
686 def dequant(self, t: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
687 assert self.quant_config.is_quantized
@@ -862,8 +862,6 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
862
863 def workspace_shapes(
864 self,
865 - a: torch.Tensor,
866 - aq: torch.Tensor,
865 M: int,
866 N: int,
867 K: int,
@@ -871,15 +869,14 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
869 global_num_experts: int,
870 local_num_experts: int,
871 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
874 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
875 - assert a.dim() == 2
872 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
873 num_dp = self.num_dispatchers
874 num_experts = local_num_experts
875 max_num_tokens = self.max_num_tokens
876 workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
877 workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
878 output = (num_experts, max_num_tokens * num_dp, K)
882 - return (workspace13, workspace2, output, a.dtype)
879 + return (workspace13, workspace2, output)
880
881 def apply(
882 self,
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+2 -4
@@ -331,8 +331,6 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
331
332 def workspace_shapes(
333 self,
334 - a: torch.Tensor,
335 - aq: torch.Tensor,
334 M: int,
335 N: int,
336 K: int,
@@ -340,7 +338,7 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
338 global_num_experts: int,
339 local_num_experts: int,
340 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
343 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
341 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
342 # Modular Kernel provisions output buffer from workspace1. However in
343 # the fused_marlin_moe() function, the final torch.sum(), is defined
344 # essentially as,
@@ -360,7 +358,7 @@ class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):
358 workspace2 = (M * topk * max(2 * N, K),)
359 output = (M, K)
360
363 - return (workspace1, workspace2, output, a.dtype)
361 + return (workspace1, workspace2, output)
362
363 def apply(
364 self,
vllm/model_executor/layers/fused_moe/fused_moe.py
+2 -4
@@ -1954,8 +1954,6 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1954
1955 def workspace_shapes(
1956 self,
1957 - a: torch.Tensor,
1958 - aq: torch.Tensor,
1957 M: int,
1958 N: int,
1959 K: int,
@@ -1963,11 +1961,11 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1961 global_num_experts: int,
1962 local_num_experts: int,
1963 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
1966 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
1964 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
1965 workspace1 = (M, topk, max(N // 2, K))
1966 workspace2 = (M, topk, max(N, K))
1967 output = (M, K)
1970 - return (workspace1, workspace2, output, a.dtype)
1968 + return (workspace1, workspace2, output)
1969
1970 def apply(
1971 self,
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+2 -4
@@ -255,8 +255,6 @@ class OAITritonExperts(BaseOAITritonExperts):
255
256 def workspace_shapes(
257 self,
258 - a: torch.Tensor,
259 - aq: torch.Tensor,
258 M: int,
259 N: int,
260 K: int,
@@ -264,12 +262,12 @@ class OAITritonExperts(BaseOAITritonExperts):
262 global_num_experts: int,
263 local_num_experts: int,
264 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
267 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
265 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
266 # workspace are allocated inside the kernel
267 workspace1 = (M, K)
268 workspace2 = (0, 0)
269 output = (M, K)
272 - return (workspace1, workspace2, output, a.dtype)
270 + return (workspace1, workspace2, output)
271
272 def apply(
273 self,
vllm/model_executor/layers/fused_moe/layer.py
+37 -53
@@ -283,6 +283,10 @@ class FusedMoEMethodBase(QuantizeMethodBase):
283 ) -> Optional[FusedMoEQuantConfig]:
284 raise NotImplementedError
285
286 + @property
287 + def using_modular_kernel(self) -> bool:
288 + return self.fused_experts is not None
289 +
290 @abstractmethod
291 def apply(
292 self,
@@ -1237,39 +1241,25 @@ class FusedMoE(CustomOp):
1241 self.batched_hidden_states: Optional[torch.Tensor] = None
1242 self.batched_router_logits: Optional[torch.Tensor] = None
1243
1240 - # TODO(bnell): flashinfer uses non-batched format.
1241 - # Does it really need a batched buffer?
1242 - if (
1243 - self.moe_parallel_config.use_pplx_kernels
1244 - or self.moe_parallel_config.use_deepep_ll_kernels
1245 - or self.moe_config.use_flashinfer_cutlass_kernels
1246 - ):
1247 - if vllm_config.parallel_config.enable_dbo:
1248 - self.batched_hidden_states = torch.zeros(
1249 - (2, moe.max_num_tokens, self.hidden_size),
1250 - dtype=moe.in_dtype,
1251 - device=torch.cuda.current_device(),
1252 - )
1244 + if self.use_dp_chunking:
1245 + states_shape: tuple[int, ...]
1246 + logits_shape: tuple[int, ...]
1247
1254 - # Note here we use `num_experts` which is logical expert count
1255 - self.batched_router_logits = torch.zeros(
1256 - (2, moe.max_num_tokens, num_experts),
1257 - dtype=moe.in_dtype,
1258 - device=torch.cuda.current_device(),
1259 - )
1248 + # Note here we use `num_experts` which is logical expert count
1249 + if vllm_config.parallel_config.enable_dbo:
1250 + states_shape = (2, moe.max_num_tokens, self.hidden_size)
1251 + logits_shape = (2, moe.max_num_tokens, num_experts)
1252 else:
1261 - self.batched_hidden_states = torch.zeros(
1262 - (moe.max_num_tokens, self.hidden_size),
1263 - dtype=moe.in_dtype,
1264 - device=torch.cuda.current_device(),
1265 - )
1253 + states_shape = (moe.max_num_tokens, self.hidden_size)
1254 + logits_shape = (moe.max_num_tokens, num_experts)
1255
1267 - # Note here we use `num_experts` which is logical expert count
1268 - self.batched_router_logits = torch.zeros(
1269 - (moe.max_num_tokens, num_experts),
1270 - dtype=moe.in_dtype,
1271 - device=torch.cuda.current_device(),
1272 - )
1256 + self.batched_hidden_states = torch.zeros(
1257 + states_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1258 + )
1259 +
1260 + self.batched_router_logits = torch.zeros(
1261 + logits_shape, dtype=moe.in_dtype, device=torch.cuda.current_device()
1262 + )
1263
1264 @property
1265 def shared_experts(self) -> Optional[torch.nn.Module]:
@@ -1323,6 +1313,16 @@ class FusedMoE(CustomOp):
1313 and self.moe_config.use_flashinfer_cutlass_kernels
1314 )
1315
1316 + @property
1317 + def use_dp_chunking(self) -> bool:
1318 + # Route to the chunked forward path using the FlashInfer Cutlass kernel
1319 + # only when data parallelism (DP) is enabled.
1320 + return (
1321 + self.moe_parallel_config.use_pplx_kernels
1322 + or self.moe_parallel_config.use_deepep_ll_kernels
1323 + or (self.dp_size > 1 and self.use_flashinfer_cutlass_kernels)
1324 + )
1325 +
1326 def update_expert_map(self):
1327 # ep_size and ep_rank should already be updated
1328 assert self.expert_map is not None
@@ -1987,21 +1987,17 @@ class FusedMoE(CustomOp):
1987 Therefore it is required that we reduce the shared_experts output
1988 early.
1989 """
1990 + assert self.quant_method is not None
1991 return (
1991 - self.use_pplx_kernels
1992 - or self.use_deepep_ht_kernels
1993 - or self.use_deepep_ll_kernels
1992 + self.quant_method.fused_experts is not None
1993 + and self.quant_method.fused_experts.output_is_reduced()
1994 )
1995
1996 def maybe_all_reduce_tensor_model_parallel(self, final_hidden_states: torch.Tensor):
1997 """
1998 - The pplx combine kernel reduces across GPU ranks by default.
1998 + Some combine kernels reduce across GPU ranks by default.
1999 """
2000 - if (
2001 - self.use_pplx_kernels
2002 - or self.use_deepep_ht_kernels
2003 - or self.use_deepep_ll_kernels
2004 - ):
2000 + if self.must_reduce_shared_expert_outputs():
2001 return final_hidden_states
2002 else:
2003 return tensor_model_parallel_all_reduce(final_hidden_states)
@@ -2209,23 +2205,11 @@ class FusedMoE(CustomOp):
2205
2206 self.ensure_moe_quant_config()
2207
2212 - # Route to the chunked forward path using the FlashInfer Cutlass kernel
2213 - # only when data parallelism (DP) is enabled.
2214 - _use_flashinfer_cutlass_kernels = (
2215 - self.dp_size > 1 and self.use_flashinfer_cutlass_kernels
2216 - )
2217 -
2218 - if (
2219 - self.moe_parallel_config.use_pplx_kernels
2220 - or self.moe_parallel_config.use_deepep_ll_kernels
2221 - or _use_flashinfer_cutlass_kernels
2222 - ):
2208 + if self.use_dp_chunking:
2209 return self.forward_impl_chunked(hidden_states, router_logits)
2210
2211 do_naive_dispatch_combine: bool = (
2226 - self.dp_size > 1
2227 - and not self.moe_parallel_config.use_deepep_ht_kernels
2228 - and not self.moe_config.use_flashinfer_cutlass_kernels
2212 + self.dp_size > 1 and not self.quant_method.using_modular_kernel
2213 )
2214
2215 # If there are shared experts but we are not using a modular kernel, the
vllm/model_executor/layers/fused_moe/modular_kernel.py
+354 -273
@@ -337,6 +337,14 @@ class FusedMoEPrepareAndFinalize(ABC):
337 def num_dispatchers(self) -> int:
338 raise NotImplementedError
339
340 + @abstractmethod
341 + def output_is_reduced(self) -> bool:
342 + """
343 + Indicates whether or not the output of finalize is reduced across all
344 + ranks.
345 + """
346 + raise NotImplementedError
347 +
348
349 # TODO: add supported activations method (return string)
350 class FusedMoEPermuteExpertsUnpermute(ABC):
@@ -493,11 +501,15 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
501 """
502 raise NotImplementedError
503
504 + def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
505 + """
506 + Workspace type: The dtype to use for the workspace tensors.
507 + """
508 + return act_dtype
509 +
510 @abstractmethod
511 def workspace_shapes(
512 self,
499 - a: torch.Tensor,
500 - aq: torch.Tensor,
513 M: int,
514 N: int,
515 K: int,
@@ -505,22 +517,33 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
517 global_num_experts: int,
518 local_num_experts: int,
519 expert_tokens_meta: Optional[ExpertTokensMetadata],
508 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
520 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
521 """
522 Compute the shapes for the temporary and final outputs of the two gemms
523 and activation in the fused expert function. Since the gemms are
524 independent, the workspace for the first gemm can be shared with the
525 workspace for the last gemm.
526
527 + Inputs:
528 + - M: number of tokens.
529 + - N: Row (or column) dimension of expert weights.
530 + - K: hidden dimension
531 + - topk: The number of top-k experts to select.
532 + - global_num_experts: global number of experts.
533 + - local_num_experts: local number of experts due to DP/EP.
534 + - expert_tokens_meta: number of tokens per expert metadata for batched
535 + format.
536 +
537 Returns a tuple of:
538 - workspace13 shape tuple: must be large enough to hold the
539 result of either expert gemm.
540 - workspace2 shape tuple: must be large enough to hold the
541 result of the activation function.
542 - output shape tuple: must be exact size of the final gemm output.
521 - - Workspace type: The dtype to use for the workspace tensors.
522 - - Note: in order for activation chunking to work, the first dimension
523 - of each tuple must be the number of tokens.
543 + - Note: workspace shapes can be 0 if the workspace is not needed.
544 + But in order for activation chunking to work, the first dimension
545 + of each tuple must be the number of tokens when the shape is
546 + not 0.
547 """
548 raise NotImplementedError
549
@@ -561,7 +584,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
584 workspace2: torch.Tensor,
585 expert_tokens_meta: Optional[ExpertTokensMetadata],
586 apply_router_weight_on_input: bool,
564 - ):
587 + ) -> None:
588 """
589 This function computes the intermediate result of a Mixture of Experts
590 (MoE) layer using two sets of weights, w1 and w2.
@@ -600,7 +623,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
623 raise NotImplementedError
624
625
603 -def _chunk_scales(
626 +def _slice_scales(
627 scales: Optional[torch.Tensor], start: int, end: int
628 ) -> Optional[torch.Tensor]:
629 if scales is not None:
@@ -615,9 +638,10 @@ class SharedResizableBuffer:
638 def __init__(self):
639 self.buffer = None
640
618 - def get(self, shape: tuple[int, ...], device: torch.device, dtype: torch.dtype):
619 - if shape == () or shape is None:
620 - return None
641 + def get(
642 + self, shape: tuple[int, ...], device: torch.device, dtype: torch.dtype
643 + ) -> torch.Tensor:
644 + assert shape != ()
645 shape_numel = prod(shape)
646 if (
647 self.buffer is None
@@ -678,131 +702,74 @@ class FusedMoEModularKernel(torch.nn.Module):
702 f"{fused_experts.activation_formats[0]}"
703 )
704
681 - def _do_fused_experts(
705 + def output_is_reduced(self) -> bool:
706 + """
707 + Indicates whether or not the output of fused MoE kernel
708 + is reduced across all ranks.
709 + """
710 + return self.prepare_finalize.output_is_reduced()
711 +
712 + def _chunk_info(self, M: int) -> tuple[int, int]:
713 + """
714 + Compute number of chunks and chunk size for given M.
715 + If chunking is not supported, set the CHUNK_SIZE to M so we
716 + get num_chunks == 1. Take max(M, 1) to avoid divide by zero.
717 + If there are no tokens to process, the number of chunks will be zero.
718 + """
719 + CHUNK_SIZE = (
720 + max(M, 1)
721 + if not self.fused_experts.supports_chunking()
722 + else min(M, envs.VLLM_FUSED_MOE_CHUNK_SIZE)
723 + )
724 + num_chunks = cdiv(M, CHUNK_SIZE)
725 + # If there are no tokens, then there should be no loop iterations.
726 + assert M > 0 or num_chunks == 0
727 + return num_chunks, CHUNK_SIZE
728 +
729 + def _allocate_buffers(
730 self,
683 - fused_out: Optional[torch.Tensor],
684 - a1: torch.Tensor,
685 - a1q: torch.Tensor,
686 - w1: torch.Tensor,
687 - w2: torch.Tensor,
688 - topk_weights: torch.Tensor,
689 - topk_ids: torch.Tensor,
690 - activation: str,
731 + out_dtype: torch.dtype,
732 + device: torch.device,
733 + M_chunk: int,
734 + M_full: int,
735 + N: int,
736 + K: int,
737 + top_k: int,
738 global_num_experts: int,
739 local_num_experts: int,
693 - expert_map: Optional[torch.Tensor],
694 - a1q_scale: Optional[torch.Tensor],
695 - a2_scale: Optional[torch.Tensor],
740 expert_tokens_meta: Optional[ExpertTokensMetadata],
697 - apply_router_weight_on_input: bool,
698 - ) -> torch.Tensor:
699 - _, M, N, K, top_k = self.fused_experts.moe_problem_size(a1q, w1, w2, topk_ids)
741 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
742 + """
743 + Allocate temporary and output buffers for the fused experts op.
744 + Inputs:
745 + - out_dtype: output type of workspace and output tensors.
746 + - device: the device of the workspace and output tensors.
747 + See `workspace_shapes` for a description of the remainder of arguments.
748 + Returns a tuple of (workspace13, workspace2, output) tensors.
749 + """
750 + assert M_full > 0 and M_chunk > 0
751
701 - (workspace13_shape, workspace2_shape, fused_out_shape, workspace_dtype) = (
702 - self.fused_experts.workspace_shapes(
703 - a1,
704 - a1q,
705 - M,
706 - N,
707 - K,
708 - top_k,
709 - global_num_experts,
710 - local_num_experts,
711 - expert_tokens_meta,
712 - )
713 - )
752 + num_chunks, _ = self._chunk_info(M_full)
753
754 # select per-ubatch buffers to avoid cross-ubatch reuse under DBO
755 ubatch_idx = dbo_current_ubatch_id()
756 buffers = self.shared_buffers[ubatch_idx]
757 + workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)
758
719 - # We can reuse the memory between cache1 and cache3 because by the
720 - # time we need cache3, we're done with cache1.
721 - workspace13 = buffers.workspace13.get(
722 - workspace13_shape, device=a1.device, dtype=workspace_dtype
723 - )
724 - workspace2 = buffers.workspace2.get(
725 - workspace2_shape, device=a1.device, dtype=workspace_dtype
726 - )
727 -
728 - assert fused_out is None or fused_out.shape == fused_out_shape, (
729 - f"fused_out {fused_out.shape} but expected {fused_out_shape}"
730 - )
731 - if fused_out is None:
732 - # reuse workspace13 for the output
733 - fused_out = _resize_cache(workspace13, fused_out_shape)
734 -
735 - self.fused_experts.apply(
736 - fused_out,
737 - a1q,
738 - w1,
739 - w2,
740 - topk_weights=topk_weights,
741 - topk_ids=topk_ids,
742 - activation=activation,
743 - global_num_experts=global_num_experts,
744 - expert_map=expert_map,
745 - a1q_scale=a1q_scale,
746 - a2_scale=a2_scale,
747 - workspace13=workspace13,
748 - workspace2=workspace2,
749 - expert_tokens_meta=expert_tokens_meta,
750 - apply_router_weight_on_input=apply_router_weight_on_input,
759 + # Get intermediate workspace shapes based off the chunked M size.
760 + workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
761 + M_chunk,
762 + N,
763 + K,
764 + top_k,
765 + global_num_experts,
766 + local_num_experts,
767 + expert_tokens_meta,
768 )
769
753 - return fused_out
754 -
755 - def _maybe_chunk_fused_experts(
756 - self,
757 - a1: torch.Tensor,
758 - a1q: torch.Tensor,
759 - w1: torch.Tensor,
760 - w2: torch.Tensor,
761 - topk_weights: torch.Tensor,
762 - topk_ids: torch.Tensor,
763 - activation: str,
764 - global_num_experts: int,
765 - local_num_experts: int,
766 - expert_map: Optional[torch.Tensor],
767 - a1q_scale: Optional[torch.Tensor],
768 - expert_tokens_meta: Optional[ExpertTokensMetadata],
769 - apply_router_weight_on_input: bool,
770 - ) -> torch.Tensor:
771 - _, M, N, K, top_k = self.fused_experts.moe_problem_size(a1q, w1, w2, topk_ids)
772 -
773 - CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
774 - num_chunks = cdiv(M, CHUNK_SIZE)
775 -
776 - # TODO(bnell): get rid of one level here, update slice functions
777 - # to nops on num_chunks==1
778 -
779 - if not self.fused_experts.supports_chunking() or num_chunks == 1:
780 - return self._do_fused_experts(
781 - fused_out=None,
782 - a1=a1,
783 - a1q=a1q,
784 - w1=w1,
785 - w2=w2,
786 - topk_weights=topk_weights,
787 - topk_ids=topk_ids,
788 - activation=activation,
789 - global_num_experts=global_num_experts,
790 - local_num_experts=local_num_experts,
791 - expert_map=expert_map,
792 - a1q_scale=a1q_scale,
793 - a2_scale=self.fused_experts.a2_scale,
794 - expert_tokens_meta=expert_tokens_meta,
795 - apply_router_weight_on_input=apply_router_weight_on_input,
796 - )
797 -
798 - # Chunking required case
799 - assert num_chunks > 1
800 -
801 - # Construct the entire output that can then be processed in chunks.
802 - (_, _, fused_out_shape, _) = self.fused_experts.workspace_shapes(
803 - a1,
804 - a1q,
805 - M,
770 + # Get final output shape based on the full M size.
771 + _, _, fused_out_shape = self.fused_experts.workspace_shapes(
772 + M_full,
773 N,
774 K,
775 top_k,
@@ -810,150 +777,99 @@ class FusedMoEModularKernel(torch.nn.Module):
777 local_num_experts,
778 expert_tokens_meta,
779 )
813 - ubatch_idx = dbo_current_ubatch_id()
814 - buffers = self.shared_buffers[ubatch_idx]
815 - fused_out = buffers.fused_out.get(
816 - fused_out_shape, device=a1q.device, dtype=a1.dtype
780 +
781 + # We can reuse the memory between cache1 and cache3 because by the
782 + # time we need cache3, we're done with cache1.
783 + workspace13 = buffers.workspace13.get(
784 + workspace13_shape, device=device, dtype=workspace_dtype
785 + )
786 + workspace2 = buffers.workspace2.get(
787 + workspace2_shape, device=device, dtype=workspace_dtype
788 )
789
819 - def slice_input_tensors(
820 - chunk_idx: int,
821 - ) -> tuple[
822 - torch.Tensor,
823 - Optional[torch.Tensor],
824 - Optional[torch.Tensor],
825 - torch.Tensor,
826 - torch.Tensor,
827 - ]:
828 - s = chunk_idx * CHUNK_SIZE
829 - e = min(s + CHUNK_SIZE, M)
830 - return (
831 - a1q[s:e],
832 - _chunk_scales(a1q_scale, s, e),
833 - _chunk_scales(self.fused_experts.a2_scale, s, e),
834 - topk_ids[s:e],
835 - topk_weights[s:e],
790 + # Construct the entire output that can then be processed in chunks.
791 + # Reuse workspace13 for the output in the non-chunked case as long
792 + # as it is large enough. This will not always be the case for standard
793 + # format experts and with experts that have empty workspaces.
794 + if num_chunks == 1 and prod(workspace13_shape) >= prod(fused_out_shape):
795 + fused_out = _resize_cache(workspace13, fused_out_shape)
796 + else:
797 + fused_out = buffers.fused_out.get(
798 + fused_out_shape, device=device, dtype=out_dtype
799 )
800
838 - def slice_output_tensor(chunk_idx: int) -> torch.Tensor:
839 - assert fused_out.size(0) % M == 0, (
840 - f"fused_out shape {fused_out.shape} vs M {M}"
841 - )
842 - factor = fused_out.size(0) // M
843 - out_chunk_size = CHUNK_SIZE * factor
844 - s = chunk_idx * out_chunk_size
845 - e = min(s + out_chunk_size, fused_out.size(0))
846 - return fused_out[s:e]
847 -
848 - def slice_expert_tokens_metadata(
849 - full_expert_tokens_meta: ExpertTokensMetadata,
850 - chunk_topk_ids: torch.Tensor,
851 - local_num_experts: int,
852 - expert_map: Optional[torch.Tensor],
853 - ) -> ExpertTokensMetadata:
854 - # The existing expert_num_tokens is for the entire a1q
855 - # input. Chunking forces recomputation of the number
856 - # of tokens assigned to each expert.
857 - c_expert_num_tokens = count_expert_num_tokens(
858 - chunk_topk_ids, local_num_experts, expert_map
859 - )
801 + return workspace13, workspace2, fused_out
802
861 - c_expert_num_tokens_cpu = None
862 - need_expert_num_tokens_cpu = (
863 - full_expert_tokens_meta.expert_num_tokens_cpu is not None
864 - )
865 - if need_expert_num_tokens_cpu:
866 - # This is blocking as some implementations need the count
867 - # on the CPU to determine appropriate input/out fused-moe
868 - # buffers
869 - c_expert_num_tokens_cpu = c_expert_num_tokens.to(
870 - "cpu", non_blocking=False
871 - )
872 -
873 - return ExpertTokensMetadata(
874 - expert_num_tokens=c_expert_num_tokens,
875 - expert_num_tokens_cpu=c_expert_num_tokens_cpu,
876 - )
803 + @staticmethod
804 + def _slice_output_tensor(
805 + fused_out: torch.Tensor,
806 + chunk_idx: int,
807 + num_chunks: int,
808 + CHUNK_SIZE: int,
809 + M: int,
810 + ) -> torch.Tensor:
811 + if num_chunks == 1:
812 + return fused_out
813
878 - for chunk_idx in range(num_chunks):
879 - c_a1q, c_a1q_scale, c_a2_scale, c_topk_ids, c_topk_weights = (
880 - slice_input_tensors(chunk_idx)
881 - )
814 + assert fused_out.size(0) % M == 0, f"fused_out shape {fused_out.shape} vs M {M}"
815 + factor = fused_out.size(0) // M
816 + out_chunk_size = CHUNK_SIZE * factor
817 + s = chunk_idx * out_chunk_size
818 + e = min(s + out_chunk_size, fused_out.size(0))
819 + return fused_out[s:e]
820
883 - c_expert_tokens_meta = None
884 - if expert_tokens_meta is not None:
885 - c_expert_tokens_meta = slice_expert_tokens_metadata(
886 - expert_tokens_meta, c_topk_ids, local_num_experts, expert_map
887 - )
821 + @staticmethod
822 + def _slice_expert_tokens_metadata(
823 + num_chunks: int,
824 + full_expert_tokens_meta: Optional[ExpertTokensMetadata],
825 + chunk_topk_ids: torch.Tensor,
826 + local_num_experts: int,
827 + expert_map: Optional[torch.Tensor],
828 + ) -> Optional[ExpertTokensMetadata]:
829 + if num_chunks == 1 or full_expert_tokens_meta is None:
830 + return full_expert_tokens_meta
831 +
832 + # The existing expert_num_tokens is for the entire a1q
833 + # input. Chunking forces recomputation of the number
834 + # of tokens assigned to each expert.
835 + c_expert_num_tokens = count_expert_num_tokens(
836 + chunk_topk_ids, local_num_experts, expert_map
837 + )
838
889 - self._do_fused_experts(
890 - fused_out=slice_output_tensor(chunk_idx),
891 - a1=a1,
892 - a1q=c_a1q,
893 - w1=w1,
894 - w2=w2,
895 - topk_weights=c_topk_weights,
896 - topk_ids=c_topk_ids,
897 - activation=activation,
898 - global_num_experts=global_num_experts,
899 - local_num_experts=local_num_experts,
900 - expert_map=expert_map,
901 - a1q_scale=c_a1q_scale,
902 - a2_scale=c_a2_scale,
903 - expert_tokens_meta=c_expert_tokens_meta,
904 - apply_router_weight_on_input=apply_router_weight_on_input,
905 - )
839 + c_expert_num_tokens_cpu = None
840 + need_expert_num_tokens_cpu = (
841 + full_expert_tokens_meta.expert_num_tokens_cpu is not None
842 + )
843 + if need_expert_num_tokens_cpu:
844 + # This is blocking as some implementations need the count
845 + # on the CPU to determine appropriate input/out fused-moe
846 + # buffers
847 + c_expert_num_tokens_cpu = c_expert_num_tokens.to("cpu", non_blocking=False)
848
907 - return fused_out
849 + return ExpertTokensMetadata(
850 + expert_num_tokens=c_expert_num_tokens,
851 + expert_num_tokens_cpu=c_expert_num_tokens_cpu,
852 + )
853
909 - def forward(
854 + def _prepare(
855 self,
856 hidden_states: torch.Tensor,
912 - w1: torch.Tensor,
913 - w2: torch.Tensor,
857 topk_weights: torch.Tensor,
858 topk_ids: torch.Tensor,
916 - inplace: bool = False,
917 - activation: str = "silu",
918 - global_num_experts: int = -1,
919 - expert_map: Optional[torch.Tensor] = None,
920 - apply_router_weight_on_input: bool = False,
921 - ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
859 + global_num_experts: int,
860 + expert_map: Optional[torch.Tensor],
861 + apply_router_weight_on_input: bool,
862 + ) -> tuple[
863 + torch.Tensor,
864 + Optional[torch.Tensor],
865 + Optional[ExpertTokensMetadata],
866 + torch.Tensor,
867 + torch.Tensor,
868 + ]:
869 """
923 - This function computes a Mixture of Experts (MoE) layer using two sets
924 - of weights, w1 and w2, and top-k gating mechanism.
925 -
926 - Parameters:
927 - - hidden_states: (torch.Tensor): The input tensor to the MoE layer.
928 - - w1 (torch.Tensor): The first set of expert weights.
929 - - w2 (torch.Tensor): The second set of expert weights.
930 - - topk_weights (torch.Tensor): The topk weights applied at the end of
931 - the layer.
932 - - topk_ids (torch.Tensor): A map of row to expert id.
933 - - inplace (bool): If True, perform the operation in-place.
934 - Defaults to False.
935 - - activation (str): The activation function to apply after the first
936 - MoE layer.
937 - - global_num_experts (int): The total number of experts in the global
938 - expert space.
939 - - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
940 - from the global expert space to the local expert space of the expert
941 - parallel shard.
942 - - apply_router_weight_on_input (bool): When true, the topk weights are
943 - applied directly on the inputs. This is only applicable when topk is
944 - 1.
945 -
946 - Returns:
947 - - torch.Tensor: The output tensor after applying the MoE layer.
870 + The _prepare method is a wrapper around self.prepare_finalize.prepare
871 + that handles DBO and async.
872 """
949 -
950 - a1 = hidden_states
951 - output = a1 if inplace and self.shared_experts is None else torch.zeros_like(a1)
952 -
953 - local_num_experts = w1.size(0)
954 - if global_num_experts == -1:
955 - global_num_experts = local_num_experts
956 -
873 if not self.prepare_finalize.supports_async():
874 # We shouldn't be running an a2a kernel that doesn't
875 # support async prepare/finalize
@@ -967,7 +883,7 @@ class FusedMoEModularKernel(torch.nn.Module):
883 _expert_topk_ids,
884 _expert_topk_weights,
885 ) = self.prepare_finalize.prepare(
970 - a1,
886 + hidden_states,
887 topk_weights,
888 topk_ids,
889 global_num_experts,
@@ -979,7 +895,7 @@ class FusedMoEModularKernel(torch.nn.Module):
895 # Overlap shared expert compute with all2all dispatch.
896 dbo_maybe_run_recv_hook()
897 prepare_ret = self.prepare_finalize.prepare_async(
982 - a1,
898 + hidden_states,
899 topk_weights,
900 topk_ids,
901 global_num_experts,
@@ -1019,33 +935,114 @@ class FusedMoEModularKernel(torch.nn.Module):
935 topk_weights if _expert_topk_weights is None else _expert_topk_weights
936 )
937
1022 - fused_out = None
938 + return a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights
939
1024 - if a1q.numel() == 0:
1025 - # This happens when none of the tokens from the all2all reach this
1026 - # EP rank. Also, note that this is only relevant for CUDAGraph
1027 - # incompatible all2all kernels like the DeepEP high-throughput
1028 - # kernels. CUDAGraph compatible all2all kernels like the pplx
1029 - # kernels and the DeepEP low-latency kernels are always batched
1030 - # and can never run into the tensor.numel() == 0 case.
1031 - fused_out = torch.empty_like(a1q).to(dtype=a1.dtype)
940 + def _fused_experts(
941 + self,
942 + in_dtype: torch.dtype,
943 + a1q: torch.Tensor,
944 + a1q_scale: Optional[torch.Tensor],
945 + w1: torch.Tensor,
946 + w2: torch.Tensor,
947 + topk_weights: torch.Tensor,
948 + topk_ids: torch.Tensor,
949 + activation: str,
950 + global_num_experts: int,
951 + local_num_experts: int,
952 + expert_map: Optional[torch.Tensor],
953 + apply_router_weight_on_input: bool,
954 + expert_tokens_meta: Optional[ExpertTokensMetadata],
955 + ) -> torch.Tensor:
956 + _, M_full, N, K, top_k = self.fused_experts.moe_problem_size(
957 + a1q, w1, w2, topk_ids
958 + )
959 +
960 + num_chunks, CHUNK_SIZE = self._chunk_info(M_full)
961 +
962 + def input_chunk_range(chunk_idx: int) -> tuple[int, int]:
963 + if num_chunks == 1:
964 + # Use a1q.size(0) here since batched format does not
965 + # keep M in the first dimension.
966 + return 0, a1q.size(0)
967 + else:
968 + s = chunk_idx * CHUNK_SIZE
969 + e = min(s + CHUNK_SIZE, M_full)
970 + return s, e
971 +
972 + # This happens when none of the tokens from the all2all reach this
973 + # EP rank. Also, note that this is only relevant for CUDAGraph
974 + # incompatible all2all kernels like the DeepEP high-throughput
975 + # kernels. CUDAGraph compatible all2all kernels like the pplx
976 + # kernels and the DeepEP low-latency kernels are always batched
977 + # and can never run into the tensor.numel() == 0 case.
978 + if M_full == 0:
979 + assert num_chunks == 0
980 + workspace13 = None
981 + workspace2 = None
982 + fused_out = torch.empty_like(a1q)
983 else:
1033 - fused_out = self._maybe_chunk_fused_experts(
1034 - a1=a1,
1035 - a1q=a1q,
984 + assert num_chunks > 0
985 + workspace13, workspace2, fused_out = self._allocate_buffers(
986 + in_dtype,
987 + a1q.device,
988 + CHUNK_SIZE,
989 + M_full,
990 + N,
991 + K,
992 + top_k,
993 + global_num_experts,
994 + local_num_experts,
995 + expert_tokens_meta,
996 + )
997 +
998 + for chunk_idx in range(num_chunks):
999 + s, e = input_chunk_range(chunk_idx)
1000 +
1001 + c_expert_tokens_meta = self._slice_expert_tokens_metadata(
1002 + num_chunks,
1003 + expert_tokens_meta,
1004 + topk_ids[s:e],
1005 + local_num_experts,
1006 + expert_map,
1007 + )
1008 +
1009 + c_fused_out = self._slice_output_tensor(
1010 + fused_out, chunk_idx, num_chunks, CHUNK_SIZE, M_full
1011 + )
1012 +
1013 + self.fused_experts.apply(
1014 + output=c_fused_out,
1015 + hidden_states=a1q[s:e],
1016 w1=w1,
1017 w2=w2,
1038 - topk_weights=topk_weights,
1039 - topk_ids=topk_ids,
1018 + topk_weights=topk_weights[s:e],
1019 + topk_ids=topk_ids[s:e],
1020 activation=activation,
1021 global_num_experts=global_num_experts,
1042 - local_num_experts=local_num_experts,
1022 expert_map=expert_map,
1044 - a1q_scale=a1q_scale,
1045 - expert_tokens_meta=expert_tokens_meta,
1023 + a1q_scale=_slice_scales(a1q_scale, s, e),
1024 + a2_scale=_slice_scales(self.fused_experts.a2_scale, e, e),
1025 + workspace13=workspace13,
1026 + workspace2=workspace2,
1027 + expert_tokens_meta=c_expert_tokens_meta,
1028 apply_router_weight_on_input=apply_router_weight_on_input,
1029 )
1030
1031 + return fused_out
1032 +
1033 + def _finalize(
1034 + self,
1035 + output: torch.Tensor,
1036 + fused_out: torch.Tensor,
1037 + hidden_states: torch.Tensor,
1038 + topk_weights: torch.Tensor,
1039 + topk_ids: torch.Tensor,
1040 + apply_router_weight_on_input: bool,
1041 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1042 + """
1043 + The _finalize method is a wrapper around self.prepare_finalize.finalize
1044 + that handles DBO, async and shared expert overlap.
1045 + """
1046 shared_output: Optional[torch.Tensor] = None
1047
1048 if not self.prepare_finalize.supports_async():
@@ -1060,7 +1057,7 @@ class FusedMoEModularKernel(torch.nn.Module):
1057 self.fused_experts.finalize_weight_and_reduce_impl(),
1058 )
1059 if self.shared_experts is not None:
1063 - shared_output = self.shared_experts(a1)
1060 + shared_output = self.shared_experts(hidden_states)
1061 else:
1062 finalize_ret = self.prepare_finalize.finalize_async(
1063 output,
@@ -1072,7 +1069,7 @@ class FusedMoEModularKernel(torch.nn.Module):
1069 )
1070
1071 if self.shared_experts is not None:
1075 - shared_output = self.shared_experts(a1)
1072 + shared_output = self.shared_experts(hidden_states)
1073
1074 # TODO(lucas): refactor this in the alternative schedules followup
1075 # currently unpack if we have hook + receiver pair or just
@@ -1100,3 +1097,87 @@ class FusedMoEModularKernel(torch.nn.Module):
1097 else:
1098 assert shared_output is not None
1099 return shared_output, output
1100 +
1101 + def forward(
1102 + self,
1103 + hidden_states: torch.Tensor,
1104 + w1: torch.Tensor,
1105 + w2: torch.Tensor,
1106 + topk_weights: torch.Tensor,
1107 + topk_ids: torch.Tensor,
1108 + inplace: bool = False,
1109 + activation: str = "silu",
1110 + global_num_experts: int = -1,
1111 + expert_map: Optional[torch.Tensor] = None,
1112 + apply_router_weight_on_input: bool = False,
1113 + ) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
1114 + """
1115 + This function computes a Mixture of Experts (MoE) layer using two sets
1116 + of weights, w1 and w2, and top-k gating mechanism.
1117 +
1118 + Parameters:
1119 + - hidden_states: (torch.Tensor): The input tensor to the MoE layer.
1120 + - w1 (torch.Tensor): The first set of expert weights.
1121 + - w2 (torch.Tensor): The second set of expert weights.
1122 + - topk_weights (torch.Tensor): The topk weights applied at the end of
1123 + the layer.
1124 + - topk_ids (torch.Tensor): A map of row to expert id.
1125 + - inplace (bool): If True, perform the operation in-place.
1126 + Defaults to False.
1127 + - activation (str): The activation function to apply after the first
1128 + MoE layer.
1129 + - global_num_experts (int): The total number of experts in the global
1130 + expert space.
1131 + - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
1132 + from the global expert space to the local expert space of the expert
1133 + parallel shard.
1134 + - apply_router_weight_on_input (bool): When true, the topk weights are
1135 + applied directly on the inputs. This is only applicable when topk is
1136 + 1.
1137 +
1138 + Returns:
1139 + - torch.Tensor: The output tensor after applying the MoE layer.
1140 + """
1141 +
1142 + if inplace and self.shared_experts is None:
1143 + output = hidden_states
1144 + else:
1145 + output = torch.zeros_like(hidden_states)
1146 +
1147 + local_num_experts = w1.size(0)
1148 + if global_num_experts == -1:
1149 + global_num_experts = local_num_experts
1150 +
1151 + a1q, a1q_scale, expert_tokens_meta, topk_ids, topk_weights = self._prepare(
1152 + hidden_states,
1153 + topk_weights,
1154 + topk_ids,
1155 + global_num_experts,
1156 + expert_map,
1157 + apply_router_weight_on_input,
1158 + )
1159 +
1160 + fused_out = self._fused_experts(
1161 + in_dtype=hidden_states.dtype,
1162 + a1q=a1q,
1163 + a1q_scale=a1q_scale,
1164 + w1=w1,
1165 + w2=w2,
1166 + topk_weights=topk_weights,
1167 + topk_ids=topk_ids,
1168 + activation=activation,
1169 + global_num_experts=global_num_experts,
1170 + local_num_experts=local_num_experts,
1171 + expert_map=expert_map,
1172 + apply_router_weight_on_input=apply_router_weight_on_input,
1173 + expert_tokens_meta=expert_tokens_meta,
1174 + )
1175 +
1176 + return self._finalize(
1177 + output,
1178 + fused_out,
1179 + hidden_states,
1180 + topk_weights,
1181 + topk_ids,
1182 + apply_router_weight_on_input,
1183 + )
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+3
@@ -91,6 +91,9 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
91 def num_dispatchers(self) -> int:
92 return self.num_dispatchers_
93
94 + def output_is_reduced(self) -> bool:
95 + return True
96 +
97 def supports_async(self) -> bool:
98 return True
99
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+3
@@ -27,6 +27,9 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
27 def num_dispatchers(self) -> int:
28 return 1
29
30 + def output_is_reduced(self) -> bool:
31 + return False
32 +
33 def prepare(
34 self,
35 a1: torch.Tensor,
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+1 -7
@@ -83,8 +83,6 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
83
84 def workspace_shapes(
85 self,
86 - a: torch.Tensor,
87 - aq: torch.Tensor,
86 M: int,
87 N: int,
88 K: int,
@@ -92,7 +90,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
90 global_num_experts: int,
91 local_num_experts: int,
92 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
95 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
93 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
94 # Note: the deep gemm workspaces are strictly larger than the triton
95 # workspaces so we can be pessimistic here and allocate for DeepGemm
96 # even if we fall back to triton later, e.g. if expert maps are set.
@@ -101,8 +99,6 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
99 ):
100 assert self.deep_gemm_expert is not None
101 return self.deep_gemm_expert.workspace_shapes(
104 - a,
105 - aq,
102 M,
103 N,
104 K,
@@ -113,8 +109,6 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
109 )
110 else:
111 return self.triton_expert.workspace_shapes(
116 - a,
117 - aq,
112 M,
113 N,
114 K,
vllm/model_executor/layers/fused_moe/trtllm_moe.py
+4 -8
@@ -52,8 +52,6 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
52
53 def workspace_shapes(
54 self,
55 - a: torch.Tensor,
56 - aq: torch.Tensor,
55 M: int,
56 N: int,
57 K: int,
@@ -61,14 +59,12 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):
59 global_num_experts: int,
60 local_num_experts: int,
61 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
64 - ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
62 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
63 # The workspaces for this implementation are managed by flashinfer.
66 - # TODO(varun) : workspace1 is could be used as the output tensor. This
67 - # is error-prone. Allow the `workspace_shapes` to return None workspaces
68 - workspace1 = (M, K)
69 - workspace2 = (0, 0)
64 + workspace1 = (0,)
65 + workspace2 = (0,)
66 output = (M, K)
71 - return (workspace1, workspace2, output, a.dtype)
67 + return (workspace1, workspace2, output)
68
69 def _get_tile_tokens_dim(self, x: torch.Tensor, top_k: int, local_num_experts: int):
70 # Number of tokens in the input tensor.