@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2cdf92228c
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 6, 19:15 UTC

[Feature]: Remove Chunking From FusedMoE (#34086)

Signed-off-by: SouthWest7 <am1ao@qq.com> Signed-off-by: Southwest <1403572259@qq.com> Signed-off-by: southwest <am1ao@qq.com> Signed-off-by: Xinan Miao <1403572259@qq.com> Co-authored-by: SouthWest7 <am1ao@qq.com>

Xinan Miao committed Mar 13, 2026 at 02:24 UTC 2cdf92228cfcaa7a3829b557bb4656ec2aeaa599
28 files changed +153 -524
docs/design/fused_moe_modular_kernel.md
+2 -5
@@ -167,9 +167,6 @@ FusedMoEExpertsModular performs the core of the FusedMoE operations. The various
167
168 `FusedMoEExpertsModular::activation_formats()`: Return the supported Input and Output activation formats. i.e. Contiguous / Batched format.
169
170 -`FusedMoEExpertsModular::supports_chunking()`: Return True if the implementation supports chunking. Typically
171 -implementations that input `FusedMoEActivationFormat.Standard` support chunking and `FusedMoEActivationFormat.BatchedExperts` do not.
172 -
170 `FusedMoEExpertsModular::supports_expert_map()`: Return True if the implementation supports expert map.
171
172 `FusedMoEExpertsModular::workspace_shapes()` /
@@ -220,8 +217,8 @@ If you are adding some `FusedMoEPrepareAndFinalizeModular` / `FusedMoEExpertsMod
217
218 1. Add the implementation type to `MK_ALL_PREPARE_FINALIZE_TYPES` and `MK_FUSED_EXPERT_TYPES` in [mk_objects.py](../../tests/kernels/moe/modular_kernel_tools/mk_objects.py) respectively.
219 2. Update `Config::is_batched_prepare_finalize()`, `Config::is_batched_fused_experts()`, `Config::is_standard_fused_experts()`,
223 -`Config::is_fe_16bit_supported()`, `Config::is_fe_fp8_supported()`, `Config::is_fe_block_fp8_supported()`,
224 -`Config::is_fe_supports_chunking()` methods in [/tests/kernels/moe/modular_kernel_tools/common.py](../../tests/kernels/moe/modular_kernel_tools/common.py)
220 +`Config::is_fe_16bit_supported()`, `Config::is_fe_fp8_supported()`, `Config::is_fe_block_fp8_supported()`
221 +methods in [/tests/kernels/moe/modular_kernel_tools/common.py](../../tests/kernels/moe/modular_kernel_tools/common.py)
222
223 Doing this will add the new implementation to the test suite.
224
tests/kernels/moe/modular_kernel_tools/cli_args.py
-6
@@ -82,11 +82,6 @@ def make_config_arg_parser(description: str):
82 "--num-experts", type=int, default=32, help="Global num experts"
83 )
84 parser.add_argument("--topk", nargs="+", type=int, default=[4, 1], help="num topk")
85 - parser.add_argument(
86 - "--fused-moe-chunk-size",
87 - type=int,
88 - help="Fused moe chunk size used for the non-batched fused experts impl.",
89 - )
85
86 # Quant args
87 parser.add_argument(
@@ -158,7 +153,6 @@ def make_config(args: argparse.Namespace) -> Config:
153 quant_config=quant_config,
154 prepare_finalize_type=args.pf_type,
155 fused_experts_type=args.experts_type,
161 - fused_moe_chunk_size=args.fused_moe_chunk_size,
156 world_size=args.world_size,
157 torch_trace_dir_path=args.torch_trace_dir_path,
158 )
tests/kernels/moe/modular_kernel_tools/common.py
-15
@@ -68,7 +68,6 @@ class Config:
68 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
69 fused_experts_type: mk.FusedMoEExperts
70
71 - fused_moe_chunk_size: int | None
71 world_size: int
72
73 torch_trace_dir_path: str | None = None
@@ -89,7 +88,6 @@ class Config:
88 s += f" K={self.K}\n"
89 s += f" topk={self.topks}\n"
90 s += f" dtype={self.dtype}\n"
92 - s += f" fused_moe_chunk_size={self.fused_moe_chunk_size}\n"
91 s += " Quant:\n"
92 if self.quant_config is not None:
93 s += f" q_dtype={self.quant_dtype}\n"
@@ -152,11 +150,6 @@ class Config:
150
151 vllm_config.parallel_config.all2all_backend = self.all2all_backend()
152
155 - if self.fused_moe_chunk_size is not None:
156 - env_dict.update(
157 - {"VLLM_FUSED_MOE_CHUNK_SIZE": str(self.fused_moe_chunk_size)}
158 - )
159 -
153 return vllm_config, env_dict
154
155 def is_fp8_block_quantized(self):
@@ -189,10 +182,6 @@ class Config:
182 info = expert_info(self.fused_experts_type)
183 return info.blocked_quantization_support
184
192 - def is_fe_supports_chunking(self):
193 - info = expert_info(self.fused_experts_type)
194 - return info.supports_chunking
195 -
185 def supports_expert_map(self):
186 info = expert_info(self.fused_experts_type)
187 return info.supports_expert_map
@@ -233,10 +222,6 @@ class Config:
222 if not self.is_standard_fused_experts():
223 return False, "Mismatched format."
224
236 - use_chunking = self.fused_moe_chunk_size is not None
237 - if use_chunking and not self.is_fe_supports_chunking():
238 - return False, "Chunking not supported."
239 -
225 # Check quantization sanity
226 if (
227 int(self.is_per_act_token_quant)
tests/kernels/moe/modular_kernel_tools/make_feature_matrix.py
-7
@@ -42,12 +42,6 @@ def rank_worker(
42 ):
43 set_random_seed(pgi.rank)
44
45 - # sanity check
46 - from vllm import envs
47 -
48 - if config.fused_moe_chunk_size is not None:
49 - assert config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE
50 -
45 # get weights to this device
46 weights.to_current_device()
47
@@ -135,7 +129,6 @@ def make_feature_matrix(csv_file_path: str):
129 fused_experts_type=experts_type,
130 quant_config=quant_config,
131 world_size=2,
138 - fused_moe_chunk_size=None,
132 )
133
134 success = None
tests/kernels/moe/modular_kernel_tools/mk_objects.py
-14
@@ -64,7 +64,6 @@ class ExpertInfo:
64 activation_format: mk.FusedMoEActivationFormat
65 supported_dtypes: list[torch.dtype | str]
66 blocked_quantization_support: bool
67 - supports_chunking: bool
67 supports_expert_map: bool
68 needs_matching_quant: bool = False
69 needs_deep_gemm: bool = False
@@ -127,7 +126,6 @@ def register_experts(
126 activation_format: mk.FusedMoEActivationFormat,
127 supported_dtypes: list[torch.dtype | str],
128 blocked_quantization_support: bool,
130 - supports_chunking: bool,
129 supports_expert_map: bool,
130 needs_matching_quant: bool = False,
131 needs_deep_gemm: bool = False,
@@ -141,7 +139,6 @@ def register_experts(
139 activation_format,
140 supported_dtypes,
141 blocked_quantization_support,
144 - supports_chunking,
142 supports_expert_map,
143 needs_matching_quant,
144 needs_deep_gemm,
@@ -176,7 +173,6 @@ register_experts(
173 batched_format,
174 common_float_types,
175 blocked_quantization_support=True,
179 - supports_chunking=False,
176 supports_expert_map=False,
177 needs_matching_quant=True,
178 )
@@ -186,7 +182,6 @@ register_experts(
182 standard_format,
183 common_float_and_int_types,
184 blocked_quantization_support=True,
189 - supports_chunking=True,
185 supports_expert_map=True,
186 needs_matching_quant=True,
187 )
@@ -196,7 +191,6 @@ register_experts(
191 batched_format,
192 common_float_and_int_types,
193 blocked_quantization_support=True,
199 - supports_chunking=False,
194 supports_expert_map=True,
195 )
196
@@ -262,7 +256,6 @@ if has_flashinfer_cutlass_fused_moe() and current_platform.has_device_capability
256 standard_format,
257 nvfp4_types + fp8_types,
258 blocked_quantization_support=True,
265 - supports_chunking=True,
259 # Note: this is a hack to get it to run for now
260 supports_expert_map=True,
261 )
@@ -281,7 +274,6 @@ if has_aiter():
274 standard_format,
275 fp8_types,
276 blocked_quantization_support=True,
284 - supports_chunking=True,
277 supports_expert_map=True,
278 needs_aiter=True,
279 )
@@ -294,7 +286,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
286 batched_format,
287 fp8_types,
288 blocked_quantization_support=True,
297 - supports_chunking=False,
289 supports_expert_map=False,
290 needs_matching_quant=False,
291 needs_deep_gemm=True,
@@ -304,7 +295,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
295 standard_format,
296 fp8_types,
297 blocked_quantization_support=True,
307 - supports_chunking=True,
298 supports_expert_map=True,
299 needs_matching_quant=False,
300 needs_deep_gemm=True,
@@ -314,7 +304,6 @@ if has_deep_gemm() and is_deep_gemm_supported():
304 standard_format,
305 common_float_and_int_types,
306 blocked_quantization_support=True,
317 - supports_chunking=True,
307 supports_expert_map=True,
308 needs_matching_quant=True,
309 needs_deep_gemm=True,
@@ -331,7 +320,6 @@ if cutlass_fp8_supported():
320 standard_format,
321 fp8_types,
322 blocked_quantization_support=False,
334 - supports_chunking=True,
323 supports_expert_map=False,
324 )
325 register_experts(
@@ -339,7 +327,6 @@ if cutlass_fp8_supported():
327 batched_format,
328 fp8_types,
329 blocked_quantization_support=False,
342 - supports_chunking=False,
330 supports_expert_map=False,
331 )
332 else:
@@ -354,7 +341,6 @@ if cutlass_fp4_supported():
341 standard_format,
342 nvfp4_types,
343 blocked_quantization_support=True,
357 - supports_chunking=True,
344 supports_expert_map=False,
345 )
346 else:
tests/kernels/moe/modular_kernel_tools/profile_modular_kernel.py
-6
@@ -85,12 +85,6 @@ def rank_worker(
85 ):
86 set_random_seed(pgi.rank)
87
88 - # sanity check
89 - from vllm import envs
90 -
91 - if config.fused_moe_chunk_size is not None:
92 - assert config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE
93 -
88 # get weights to this device
89 weights.to_current_device()
90
tests/kernels/moe/test_block_fp8.py
+1 -8
@@ -158,8 +158,6 @@ def test_w8a8_block_fp8_fused_moe(
158
159 torch.manual_seed(seed)
160
161 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "2048")
162 -
161 a = torch.randn((M, K), dtype=dtype) / 10
162 score = torch.randn((M, E), dtype=dtype)
163
@@ -226,11 +224,8 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)
224 if not _valid_deep_gemm_shape(M, N, K):
225 pytest.skip(f"Skipping test: invalid size m={M}, n={N}, k={K}")
226
229 - chunk_size = 1024
230 -
227 torch.manual_seed(seed)
228
233 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
229 block_size = get_mk_alignment_for_contiguous_layout()
230 dtype = torch.bfloat16
231
@@ -252,9 +247,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)
247 # setup code in case we are able to revisit this later.
248 use_compile = False
249
255 - use_cudagraph = (
256 - chunk_size < M and N >= 1024 and K >= 1024 and current_platform.is_cuda_alike()
257 - )
250 + use_cudagraph = N >= 1024 and K >= 1024 and current_platform.is_cuda_alike()
251
252 topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
253
tests/kernels/moe/test_cutlass_moe.py
-2
@@ -321,7 +321,6 @@ def test_cutlass_moe_8_bit_no_graph(
321 ep_size: int | None = None,
322 ):
323 set_random_seed(7)
324 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
324 with set_current_vllm_config(vllm_config):
325 mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token, per_out_ch)
326
@@ -376,7 +375,6 @@ def test_cutlass_moe_8_bit_cuda_graph(
375 workspace_init,
376 ):
377 set_random_seed(7)
379 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
378 with set_current_vllm_config(vllm_config):
379 dtype = torch.half
380
tests/kernels/moe/test_flashinfer.py
-2
@@ -204,7 +204,6 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph(
204 if not current_platform.has_device_capability(100):
205 pytest.skip("Test is only supported for sm >= 100")
206 set_random_seed(7)
207 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
207 with set_current_vllm_config(vllm_config):
208 td = TestData.make_moe_tensors_8bit(
209 m, k, n, e, is_trtllm=True, activation=activation
@@ -289,7 +288,6 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(
288 workspace_init,
289 ):
290 set_random_seed(7)
292 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "8192")
291 with set_current_vllm_config(vllm_config):
292 td = TestData.make_moe_tensors_8bit(
293 m, k, n, e, is_trtllm=False, activation=activation
tests/kernels/moe/test_modular_kernel_combinations.py
+3 -17
@@ -84,12 +84,6 @@ def rank_worker(
84
85 set_random_seed(pgi.rank)
86
87 - # sanity check
88 - from vllm import envs
89 -
90 - if base_config.fused_moe_chunk_size is not None:
91 - assert base_config.fused_moe_chunk_size == envs.VLLM_FUSED_MOE_CHUNK_SIZE
92 -
87 # get weights to this device
88 weights.to_current_device()
89
@@ -162,7 +156,6 @@ Ns = [1024]
156 TOPKs = [4, 1]
157 Es = [32]
158 DTYPEs = [torch.bfloat16]
165 -FUSED_MOE_CHUNK_SIZES = [None, 16]
159
160
161 def is_nyi_config(config: Config) -> bool:
@@ -185,14 +178,13 @@ def generate_valid_test_cases(
178 cases = []
179 total = 0
180
188 - for k, n, e, dtype, quant_config, combination, chunk_size in product(
181 + for k, n, e, dtype, quant_config, combination in product(
182 Ks,
183 Ns,
184 Es,
185 DTYPEs,
186 MK_QUANT_CONFIGS,
187 product(prepare_finalize_types, MK_FUSED_EXPERT_TYPES),
195 - FUSED_MOE_CHUNK_SIZES,
188 ):
189 total = total + 1
190
@@ -206,7 +198,6 @@ def generate_valid_test_cases(
198 quant_config=quant_config,
199 prepare_finalize_type=combination[0],
200 fused_experts_type=combination[1],
209 - fused_moe_chunk_size=chunk_size,
201 world_size=world_size,
202 )
203
@@ -234,7 +225,6 @@ def generate_valid_test_cases(
225 quant_config,
226 combination[0],
227 combination[1],
237 - chunk_size,
228 world_size,
229 )
230 )
@@ -245,7 +235,7 @@ def generate_valid_test_cases(
235
236
237 @pytest.mark.parametrize(
248 - "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,chunk_size,world_size",
238 + "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
239 generate_valid_test_cases(
240 world_size=2, prepare_finalize_types=MK_MULTI_GPU_PREPARE_FINALIZE_TYPES
241 ),
@@ -259,7 +249,6 @@ def test_modular_kernel_combinations_multigpu(
249 quant_config: TestMoEQuantConfig | None,
250 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
251 fused_experts_type: mk.FusedMoEExperts,
262 - chunk_size: int | None,
252 world_size: int,
253 pytestconfig,
254 ):
@@ -280,7 +269,6 @@ def test_modular_kernel_combinations_multigpu(
269 quant_config=quant_config,
270 prepare_finalize_type=prepare_finalize_type,
271 fused_experts_type=fused_experts_type,
283 - fused_moe_chunk_size=chunk_size,
272 world_size=world_size,
273 )
274 verbosity = pytestconfig.getoption("verbose")
@@ -288,7 +276,7 @@ def test_modular_kernel_combinations_multigpu(
276
277
278 @pytest.mark.parametrize(
291 - "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,chunk_size,world_size",
279 + "k,n,e,dtype,quant_config,prepare_finalize_type,fused_experts_type,world_size",
280 generate_valid_test_cases(
281 world_size=1, prepare_finalize_types=MK_SINGLE_GPU_PREPARE_FINALIZE_TYPES
282 ),
@@ -301,7 +289,6 @@ def test_modular_kernel_combinations_singlegpu(
289 quant_config: TestMoEQuantConfig | None,
290 prepare_finalize_type: mk.FusedMoEPrepareAndFinalize,
291 fused_experts_type: mk.FusedMoEExperts,
304 - chunk_size: int | None,
292 world_size: int,
293 pytestconfig,
294 workspace_init,
@@ -318,7 +305,6 @@ def test_modular_kernel_combinations_singlegpu(
305 quant_config=quant_config,
306 prepare_finalize_type=prepare_finalize_type,
307 fused_experts_type=fused_experts_type,
321 - fused_moe_chunk_size=chunk_size,
308 world_size=world_size,
309 )
310
tests/kernels/moe/test_moe.py
+4 -15
@@ -287,7 +287,6 @@ def run_moe_test(
287 @pytest.mark.parametrize("ep_size", EP_SIZE)
288 @pytest.mark.parametrize("dtype", [torch.bfloat16])
289 @pytest.mark.parametrize("padding", [True, False])
290 -@pytest.mark.parametrize("chunk_size", [8192])
290 def test_fused_moe(
291 m: int,
292 n: int,
@@ -297,14 +296,11 @@ def test_fused_moe(
296 ep_size: int,
297 dtype: torch.dtype,
298 padding: bool,
300 - chunk_size: int,
299 monkeypatch,
300 workspace_init,
301 ):
302 set_random_seed(7)
303
306 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
307 -
304 #
305 # Setup test data
306 #
@@ -398,12 +394,12 @@ def test_fused_moe(
394 )
395
396
401 -def test_fused_moe_int64_overflow(monkeypatch, workspace_init):
397 +def test_fused_moe_int64_overflow(workspace_init):
398 """Regression test for int32 overflow in stride*offset products.
399
404 - When chunking is disabled and M is large, stride_cm * offs_token can
405 - exceed int32 max. Verifies the offs_token int64 cast (fix for #34413)
406 - prevents overflow and produces correct results.
400 + With large M, stride_cm * offs_token can exceed int32 max. Verifies
401 + the offs_token int64 cast (fix for #34413) prevents overflow and
402 + produces correct results.
403
404 Reproduces the scenario from PR #34279.
405 """
@@ -417,9 +413,6 @@ def test_fused_moe_int64_overflow(monkeypatch, workspace_init):
413 m, n, k, e, topk = 100000, 2048, 1024, 8, 6
414 dtype = torch.bfloat16
415
420 - # Disable chunking to expose the overflow-prone code path
421 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "10000000")
422 -
416 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
417 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
418 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
@@ -452,7 +445,6 @@ def test_fused_moe_int64_overflow(monkeypatch, workspace_init):
445 @pytest.mark.parametrize("topk", TOP_KS_SMALL)
446 @pytest.mark.parametrize("dtype", [torch.bfloat16])
447 @pytest.mark.parametrize("padding", [True, False])
455 -@pytest.mark.parametrize("chunk_size", [8192])
448 def test_naive_block_assignment_moe(
449 m: int,
450 n: int,
@@ -461,14 +453,11 @@ def test_naive_block_assignment_moe(
453 topk: int,
454 dtype: torch.dtype,
455 padding: bool,
464 - chunk_size: int,
456 monkeypatch,
457 workspace_init,
458 ):
459 set_random_seed(7)
460
470 - monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
471 -
461 #
462 # Setup test data
463 #
vllm/envs.py
-11
@@ -53,8 +53,6 @@ if TYPE_CHECKING:
53 VLLM_CPU_SGL_KERNEL: bool = False
54 VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
55 VLLM_XLA_CHECK_RECOMPILATION: bool = False
56 - VLLM_FUSED_MOE_CHUNK_SIZE: int = 16 * 1024
57 - VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING: bool = True
56 VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: Literal["auto", "nccl", "shm"] = "auto"
57 VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM: bool = False
58 VLLM_USE_RAY_WRAPPED_PP_COMM: bool = True
@@ -822,15 +820,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
820 ),
821 # Enable SPMD mode for TPU backend.
822 "VLLM_XLA_USE_SPMD": lambda: bool(int(os.getenv("VLLM_XLA_USE_SPMD", "0"))),
825 - "VLLM_FUSED_MOE_CHUNK_SIZE": lambda: int(
826 - os.getenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(16 * 1024))
827 - ),
828 - # Control whether to use fused MoE activation chunking. Current chunking
829 - # logic is incompatible with torch.compile and causes IMA. See issue
830 - # https://github.com/vllm-project/vllm/issues/19631.
831 - "VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING": lambda: bool(
832 - int(os.getenv("VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING", "1"))
833 - ),
823 # If set, the OpenAI API server will stay alive even after the underlying
824 # AsyncLLMEngine errors and stops serving requests
825 "VLLM_KEEP_ALIVE_ON_ENGINE_DEATH": lambda: bool(
vllm/lora/layers/fused_moe.py
+2 -4
@@ -190,9 +190,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
190 use_int8_w8a16=False,
191 use_int4_w4a16=False,
192 )
193 - CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
193 num_tokens = hidden_states.size(0)
195 - M = min(num_tokens, CHUNK_SIZE)
194 + M = num_tokens
195 max_lora_rank = self.w13_lora_a_stacked[0].shape[-2]
196 shrink_config, expand_config = self._get_lora_moe_configs(
197 op_prefix="w13",
@@ -281,9 +280,8 @@ class FusedMoEWithLoRA(BaseLayerWithLoRA):
280 use_int8_w8a16=False,
281 use_int4_w4a16=False,
282 )
284 - CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
283 num_tokens = hidden_states.size(0)
286 - M = min(num_tokens, CHUNK_SIZE)
284 + M = num_tokens
285 max_lora_rank = self.w2_lora_a_stacked[0].shape[-2]
286 shrink_config, expand_config = self._get_lora_moe_configs(
287 op_prefix="w2",
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
-3
@@ -311,9 +311,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEExpertsModular):
311 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
312 return True
313
314 - def supports_chunking(self) -> bool:
315 - return False
316 -
314 def supports_expert_map(self) -> bool:
315 return False
316
vllm/model_executor/layers/fused_moe/cutlass_moe.py
-12
@@ -400,9 +400,6 @@ class CutlassExpertsFp8(CutlassExpertsFp8Base):
400 or moe_parallel_config.use_deepep_ht_kernels
401 )
402
403 - def supports_chunking(self) -> bool:
404 - return True
405 -
403 def supports_expert_map(self) -> bool:
404 return False
405
@@ -445,9 +442,6 @@ class CutlassBatchedExpertsFp8(CutlassExpertsFp8Base):
442 def activation_format() -> mk.FusedMoEActivationFormat:
443 return mk.FusedMoEActivationFormat.BatchedExperts
444
448 - def supports_chunking(self) -> bool:
449 - return False
450 -
445 def supports_expert_map(self) -> bool:
446 return False
447
@@ -713,9 +707,6 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
707 def supports_expert_map(self) -> bool:
708 return False
709
716 - def supports_chunking(self) -> bool:
717 - return True
718 -
710 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
711 return TopKWeightAndReduceNoOP()
712
@@ -998,9 +989,6 @@ class CutlassExpertsW4A8Fp8(mk.FusedMoEExpertsModular):
989 "This method should not be called."
990 )
991
1001 - def supports_chunking(self) -> bool:
1002 - return True
1003 -
992 def supports_expert_map(self) -> bool:
993 return True
994
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
-3
@@ -154,9 +154,6 @@ class DeepGemmExperts(mk.FusedMoEExpertsModular):
154 # NOTE(rob): discovered an IMA with this combination. Needs investigation.
155 return not moe_parallel_config.use_fi_all2allv_kernels
156
157 - def supports_chunking(self) -> bool:
158 - return True
159 -
157 def supports_expert_map(self) -> bool:
158 return True
159
vllm/model_executor/layers/fused_moe/fallback.py
-10
@@ -92,16 +92,6 @@ class FallbackExperts(mk.FusedMoEExpertsModular, ABC):
92 moe_parallel_config
93 ) and fallback_cls._supports_parallel_config(moe_parallel_config)
94
95 - def supports_chunking(self) -> bool:
96 - assert (
97 - self.experts.supports_chunking()
98 - == self.fallback_experts.supports_chunking()
99 - )
100 - return (
101 - self.experts.supports_chunking()
102 - and self.fallback_experts.supports_chunking()
103 - )
104 -
95 def supports_expert_map(self) -> bool:
96 assert (
97 self.experts.supports_expert_map()
vllm/model_executor/layers/fused_moe/flashinfer_cutedsl_moe.py
-6
@@ -83,12 +83,6 @@ class FlashInferCuteDSLExperts(mk.FusedMoEExpertsModular):
83 def supports_expert_map(self) -> bool:
84 return False
85
86 - def supports_chunking(self) -> bool:
87 - # This refers to TP chunking; DP chunking is handled separately.
88 - # TODO(shuw@nvidia.com): Set to False to be consistent with
89 - # batched_deep_gemm_moe
90 - return False
91 -
86 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
87 # Let PrepareAndFinalize::finalize() decide the impl.
88 return TopKWeightAndReduceDelegate()
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
-4
@@ -195,10 +195,6 @@ class FlashInferExperts(mk.FusedMoEExpertsModular):
195 def supports_expert_map(self) -> bool:
196 return False
197
198 - def supports_chunking(self) -> bool:
199 - # This refers to TP chunking; DP chunking is handled separately.
200 - return True
201 -
198 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
199 return TopKWeightAndReduceNoOP()
200
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
-6
@@ -712,9 +712,6 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
712 "This method should not be called."
713 )
714
715 - def supports_chunking(self) -> bool:
716 - return False
717 -
715 def supports_expert_map(self) -> bool:
716 return False
717
@@ -957,9 +954,6 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
954 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
955 return True
956
960 - def supports_chunking(self) -> bool:
961 - return False
962 -
957 def supports_expert_map(self) -> bool:
958 return False
959
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
-6
@@ -658,9 +658,6 @@ class MarlinExperts(MarlinExpertsBase):
658 def activation_format() -> mk.FusedMoEActivationFormat:
659 return mk.FusedMoEActivationFormat.Standard
660
661 - def supports_chunking(self) -> bool:
662 - return True
663 -
661 def workspace_shapes(
662 self,
663 M: int,
@@ -786,9 +783,6 @@ class BatchedMarlinExperts(MarlinExpertsBase):
783 def activation_format() -> mk.FusedMoEActivationFormat:
784 return mk.FusedMoEActivationFormat.BatchedExperts
785
789 - def supports_chunking(self) -> bool:
790 - return False
791 -
786 def workspace_shapes(
787 self,
788 M: int,
vllm/model_executor/layers/fused_moe/fused_moe.py
+101 -131
@@ -1693,10 +1693,8 @@ def fused_experts_impl(
1693 if global_num_experts == -1:
1694 global_num_experts = E
1695 top_k_num = topk_ids.size(1)
1696 - # We execute the fused_moe kernel in chunks to circumvent this issue:
1697 - # https://github.com/vllm-project/vllm/issues/5938
1698 - CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
1699 - M = min(num_tokens, CHUNK_SIZE)
1696 +
1697 + M = num_tokens
1698
1699 config_dtype = _get_config_dtype_str(
1700 use_fp8_w8a8=use_fp8_w8a8,
@@ -1787,139 +1785,114 @@ def fused_experts_impl(
1785 else:
1786 raise NotImplementedError(f"Unsupported ocp_mx_scheme={ocp_mx_scheme}")
1787
1790 - for chunk in range((num_tokens // CHUNK_SIZE) + 1):
1791 - begin_chunk_idx, end_chunk_idx = (
1792 - chunk * CHUNK_SIZE,
1793 - min((chunk + 1) * CHUNK_SIZE, num_tokens),
1794 - )
1795 - curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
1796 - tokens_in_chunk, _ = curr_hidden_states.size()
1797 -
1798 - if tokens_in_chunk == 0:
1799 - break
1800 -
1801 - if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
1802 - # Adjust the intermediate cache size and config for the last
1803 - # chunk. Note that in most cases we only have one chunk
1804 - # so the cache size and config are already set correctly and
1805 - # do not need to be adjusted.
1806 - intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
1807 - intermediate_cache2 = intermediate_cache2[
1808 - : tokens_in_chunk * topk_ids.size(1)
1809 - ]
1810 - intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
1811 - config = get_config_func(tokens_in_chunk)
1812 -
1813 - curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
1814 - curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1815 - qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
1816 - A=curr_hidden_states,
1817 - A_scale=a1_scale,
1818 - quant_dtype=quant_dtype,
1819 - per_act_token_quant=per_channel_quant,
1820 - block_shape=block_shape,
1821 - ocp_mx_scheme=ocp_mx_scheme,
1822 - )
1788 + qhidden_states, a1q_scale = moe_kernel_quantize_input(
1789 + A=hidden_states,
1790 + A_scale=a1_scale,
1791 + quant_dtype=quant_dtype,
1792 + per_act_token_quant=per_channel_quant,
1793 + block_shape=block_shape,
1794 + ocp_mx_scheme=ocp_mx_scheme,
1795 + )
1796
1824 - # SPARSITY_FACTOR is a heuristic margin ensuring tokens_in_chunk * top_k
1825 - # activates only a small fraction of total experts
1826 - SPARSITY_FACTOR = 4
1827 - # block quantized code path is not implemented yet.
1828 - naive_block_assignment = (
1829 - expert_map is None
1830 - and tokens_in_chunk * top_k_num * SPARSITY_FACTOR <= global_num_experts
1831 - and not (
1832 - (use_int8_w8a16 or use_int4_w4a16)
1833 - and block_shape is not None
1834 - and block_shape[1] > 0
1835 - )
1797 + # SPARSITY_FACTOR is a heuristic margin ensuring num_tokens * top_k
1798 + # activates only a small fraction of total experts
1799 + SPARSITY_FACTOR = 4
1800 + # block quantized code path is not implemented yet.
1801 + naive_block_assignment = (
1802 + expert_map is None
1803 + and num_tokens * top_k_num * SPARSITY_FACTOR <= global_num_experts
1804 + and not (
1805 + (use_int8_w8a16 or use_int4_w4a16)
1806 + and block_shape is not None
1807 + and block_shape[1] > 0
1808 )
1809 + )
1810
1838 - if not naive_block_assignment:
1839 - sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
1840 - curr_topk_ids,
1841 - config["BLOCK_SIZE_M"],
1842 - global_num_experts,
1843 - expert_map,
1844 - ignore_invalid_experts=True,
1845 - )
1846 - else:
1847 - max_num_tokens_padded = topk_ids.numel() * config["BLOCK_SIZE_M"]
1848 - expert_ids = curr_topk_ids.view(-1)
1849 - num_tokens_post_padded = torch.empty(
1850 - (1), dtype=torch.int32, device=topk_ids.device
1851 - )
1852 - num_tokens_post_padded.fill_(max_num_tokens_padded)
1853 - sorted_token_ids = None
1854 -
1855 - dispatch_fused_moe_kernel(
1856 - qcurr_hidden_states,
1857 - w1,
1858 - intermediate_cache1,
1859 - a1q_scale,
1860 - w1_scale,
1861 - w1_zp,
1862 - curr_topk_weights,
1863 - sorted_token_ids,
1864 - expert_ids,
1865 - num_tokens_post_padded,
1866 - apply_router_weight_on_input,
1867 - top_k_num,
1868 - config,
1869 - compute_type=compute_type,
1870 - use_fp8_w8a8=use_fp8_w8a8,
1871 - use_int8_w8a8=use_int8_w8a8,
1872 - use_int8_w8a16=use_int8_w8a16,
1873 - use_int4_w4a16=use_int4_w4a16,
1874 - per_channel_quant=per_channel_quant,
1875 - block_shape=block_shape,
1876 - B_bias=w1_bias,
1811 + if not naive_block_assignment:
1812 + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
1813 + topk_ids,
1814 + config["BLOCK_SIZE_M"],
1815 + global_num_experts,
1816 + expert_map,
1817 + ignore_invalid_experts=True,
1818 )
1878 -
1879 - apply_moe_activation(
1880 - activation_enum, intermediate_cache2, intermediate_cache1.view(-1, N)
1819 + else:
1820 + max_num_tokens_padded = topk_ids.numel() * config["BLOCK_SIZE_M"]
1821 + expert_ids = topk_ids.view(-1)
1822 + num_tokens_post_padded = torch.empty(
1823 + (1), dtype=torch.int32, device=topk_ids.device
1824 )
1825 + num_tokens_post_padded.fill_(max_num_tokens_padded)
1826 + sorted_token_ids = None
1827
1883 - qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1884 - A=intermediate_cache2,
1885 - A_scale=a2_scale,
1886 - quant_dtype=quant_dtype,
1887 - per_act_token_quant=per_channel_quant,
1888 - block_shape=block_shape,
1889 - ocp_mx_scheme=ocp_mx_scheme,
1890 - )
1828 + dispatch_fused_moe_kernel(
1829 + qhidden_states,
1830 + w1,
1831 + intermediate_cache1,
1832 + a1q_scale,
1833 + w1_scale,
1834 + w1_zp,
1835 + topk_weights,
1836 + sorted_token_ids,
1837 + expert_ids,
1838 + num_tokens_post_padded,
1839 + apply_router_weight_on_input,
1840 + top_k_num,
1841 + config,
1842 + compute_type=compute_type,
1843 + use_fp8_w8a8=use_fp8_w8a8,
1844 + use_int8_w8a8=use_int8_w8a8,
1845 + use_int8_w8a16=use_int8_w8a16,
1846 + use_int4_w4a16=use_int4_w4a16,
1847 + per_channel_quant=per_channel_quant,
1848 + block_shape=block_shape,
1849 + B_bias=w1_bias,
1850 + )
1851
1892 - if expert_map is not None:
1893 - intermediate_cache3.zero_()
1852 + apply_moe_activation(
1853 + activation_enum, intermediate_cache2, intermediate_cache1.view(-1, N)
1854 + )
1855
1895 - dispatch_fused_moe_kernel(
1896 - qintermediate_cache2,
1897 - w2,
1898 - intermediate_cache3,
1899 - a2q_scale,
1900 - w2_scale,
1901 - w2_zp,
1902 - curr_topk_weights,
1903 - sorted_token_ids,
1904 - expert_ids,
1905 - num_tokens_post_padded,
1906 - not apply_router_weight_on_input,
1907 - 1,
1908 - config,
1909 - compute_type=compute_type,
1910 - use_fp8_w8a8=use_fp8_w8a8,
1911 - use_int8_w8a8=use_int8_w8a8,
1912 - use_int8_w8a16=use_int8_w8a16,
1913 - use_int4_w4a16=use_int4_w4a16,
1914 - per_channel_quant=per_channel_quant,
1915 - block_shape=block_shape,
1916 - B_bias=w2_bias,
1917 - )
1856 + qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1857 + A=intermediate_cache2,
1858 + A_scale=a2_scale,
1859 + quant_dtype=quant_dtype,
1860 + per_act_token_quant=per_channel_quant,
1861 + block_shape=block_shape,
1862 + ocp_mx_scheme=ocp_mx_scheme,
1863 + )
1864
1919 - ops.moe_sum(
1920 - intermediate_cache3.view(*intermediate_cache3.size()),
1921 - out_hidden_states[begin_chunk_idx:end_chunk_idx],
1922 - )
1865 + if expert_map is not None:
1866 + intermediate_cache3.zero_()
1867 +
1868 + dispatch_fused_moe_kernel(
1869 + qintermediate_cache2,
1870 + w2,
1871 + intermediate_cache3,
1872 + a2q_scale,
1873 + w2_scale,
1874 + w2_zp,
1875 + topk_weights,
1876 + sorted_token_ids,
1877 + expert_ids,
1878 + num_tokens_post_padded,
1879 + not apply_router_weight_on_input,
1880 + 1,
1881 + config,
1882 + compute_type=compute_type,
1883 + use_fp8_w8a8=use_fp8_w8a8,
1884 + use_int8_w8a8=use_int8_w8a8,
1885 + use_int8_w8a16=use_int8_w8a16,
1886 + use_int4_w4a16=use_int4_w4a16,
1887 + per_channel_quant=per_channel_quant,
1888 + block_shape=block_shape,
1889 + B_bias=w2_bias,
1890 + )
1891 +
1892 + ops.moe_sum(
1893 + intermediate_cache3.view(*intermediate_cache3.size()),
1894 + out_hidden_states,
1895 + )
1896
1897 return out_hidden_states
1898
@@ -1994,9 +1967,6 @@ class TritonExperts(mk.FusedMoEExpertsModular):
1967 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
1968 return not moe_parallel_config.use_fi_all2allv_kernels
1969
1997 - def supports_chunking(self) -> bool:
1998 - return True
1999 -
1970 def supports_expert_map(self) -> bool:
1971 return True
1972
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
-6
@@ -609,9 +609,6 @@ class OAITritonExperts(BaseOAITritonExperts):
609 def activation_format() -> mk.FusedMoEActivationFormat:
610 return mk.FusedMoEActivationFormat.Standard
611
612 - def supports_chunking(self) -> bool:
613 - return True
614 -
612 def workspace_shapes(
613 self,
614 M: int,
@@ -696,9 +693,6 @@ class UnfusedOAITritonExperts(BaseOAITritonExperts):
693 def activation_format() -> mk.FusedMoEActivationFormat:
694 return mk.FusedMoEActivationFormat.Standard
695
699 - def supports_chunking(self) -> bool:
700 - return True
701 -
696 def workspace_shapes(
697 self,
698 M: int,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+39 -214
@@ -9,8 +9,6 @@ from typing import final
9
10 import torch
11
12 -import vllm.envs as envs
13 -from vllm.forward_context import get_forward_context, is_forward_context_available
12 from vllm.logger import init_logger
13 from vllm.model_executor.layers.fused_moe.activation import (
14 MoEActivation,
@@ -24,14 +22,12 @@ from vllm.model_executor.layers.fused_moe.config import (
22 )
23 from vllm.model_executor.layers.fused_moe.utils import (
24 _resize_cache,
27 - count_expert_num_tokens,
25 disable_inplace,
26 )
27 from vllm.model_executor.layers.quantization.utils.quant_utils import (
28 QuantKey,
29 )
30 from vllm.platforms import current_platform
34 -from vllm.utils.math_utils import cdiv
31 from vllm.v1.worker.ubatching import (
32 dbo_enabled,
33 dbo_maybe_run_recv_hook,
@@ -719,15 +715,6 @@ class FusedMoEExperts(ABC):
715 def g2_alphas(self) -> torch.Tensor | None:
716 return self.quant_config.g2_alphas
717
722 - # TODO (bnell): make this return a CHUNK_SIZE or None instead?
723 - @abstractmethod
724 - def supports_chunking(self) -> bool:
725 - """
726 - A flag indicating whether or not this class supports activation
727 - chunking.
728 - """
729 - raise NotImplementedError
730 -
718 @abstractmethod
719 def supports_expert_map(self) -> bool:
720 """
@@ -742,11 +729,6 @@ class FusedMoEExperts(ABC):
729 """
730 return False
731
745 - def enable_chunking(self):
746 - return (
747 - envs.VLLM_ENABLE_FUSED_MOE_ACTIVATION_CHUNKING and self.supports_chunking()
748 - )
749 -
732
733 class FusedMoEExpertsModular(FusedMoEExperts):
734 """
@@ -995,17 +977,6 @@ class FusedMoEExpertsMonolithic(FusedMoEExperts):
977 raise NotImplementedError
978
979
998 -def _slice_scales(
999 - scales: torch.Tensor | None, start: int, end: int
1000 -) -> torch.Tensor | None:
1001 - if scales is not None:
1002 - if scales.numel() == 1:
1003 - return scales
1004 - else:
1005 - return scales[start:end]
1006 - return None
1007 -
1008 -
980 ################################################################################
981 # Kernel
982 ################################################################################
@@ -1032,26 +1003,6 @@ class FusedMoEKernelModularImpl:
1003 and moe_parallel_config.use_ep
1004 )
1005
1035 - def _chunk_info(self, M: int) -> tuple[int, int]:
1036 - """
1037 - Compute number of chunks and chunk size for given M.
1038 - If chunking is not supported, set the CHUNK_SIZE to M so we
1039 - get num_chunks == 1. Take max(M, 1) to avoid divide by zero.
1040 - If there are no tokens to process, the number of chunks will be zero.
1041 - """
1042 - CHUNK_SIZE = max(
1043 - 1,
1044 - (
1045 - M
1046 - if not self.fused_experts.enable_chunking()
1047 - else min(M, envs.VLLM_FUSED_MOE_CHUNK_SIZE)
1048 - ),
1049 - )
1050 - num_chunks = cdiv(M, CHUNK_SIZE)
1051 - # If there are no tokens, then there should be no loop iterations.
1052 - assert M > 0 or num_chunks == 0
1053 - return num_chunks, CHUNK_SIZE
1054 -
1006 def _allocate_buffers(
1007 self,
1008 out_dtype: torch.dtype,
@@ -1076,40 +1027,8 @@ class FusedMoEKernelModularImpl:
1027 """
1028 assert M_full > 0 and M_chunk > 0
1029
1079 - num_chunks, _ = self._chunk_info(M_full)
1030 workspace_dtype = self.fused_experts.workspace_dtype(out_dtype)
1031
1082 - # Force worst-case allocation in profiling run for
1083 - # "mk.FusedMoEKernel.Standard" formats where this is only bounded
1084 - # by `VLLM_FUSED_MOE_CHUNK_SIZE` and may not be seen during profiling with
1085 - # DP+EP due to the random token routing.
1086 - is_profile_run = (
1087 - is_forward_context_available()
1088 - and get_forward_context().attn_metadata is None
1089 - )
1090 - if is_profile_run and self.fused_experts.enable_chunking() and self.is_dp_ep:
1091 - max_workspace_13, max_workspace_2, max_fused_out_shape = (
1092 - self.fused_experts.workspace_shapes(
1093 - envs.VLLM_FUSED_MOE_CHUNK_SIZE,
1094 - N,
1095 - K,
1096 - top_k,
1097 - global_num_experts,
1098 - local_num_experts,
1099 - # expert_tokens_meta help in allocating optimal/minimal
1100 - # amount of workspace. Mark it None, so we allocate for
1101 - # the worst-case scenario.
1102 - expert_tokens_meta=None,
1103 - activation=activation,
1104 - )
1105 - )
1106 -
1107 - current_workspace_manager().get_simultaneous(
1108 - (max_workspace_13, workspace_dtype),
1109 - (max_workspace_2, workspace_dtype),
1110 - (max_fused_out_shape, out_dtype),
1111 - )
1112 -
1032 # Get intermediate workspace shapes based off the chunked M size.
1033 workspace13_shape, workspace2_shape, _ = self.fused_experts.workspace_shapes(
1034 M_chunk,
@@ -1136,79 +1055,16 @@ class FusedMoEKernelModularImpl:
1055
1056 # We can reuse the memory between cache1 and cache3 because by the
1057 # time we need cache3, we're done with cache1.
1139 - # Construct the entire output that can then be processed in chunks.
1140 - # Reuse workspace13 for the output in the non-chunked case.
1141 - # This will not always be the case for standard
1142 - # format experts and with experts that have empty workspaces.
1143 - if num_chunks == 1:
1144 - max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
1145 - common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
1146 - ((max_shape_size,), workspace_dtype),
1147 - (workspace2_shape, workspace_dtype),
1148 - )
1149 - workspace13 = _resize_cache(common_workspace, workspace13_shape)
1150 - fused_out = _resize_cache(common_workspace, fused_out_shape)
1151 - else:
1152 - workspace13, workspace2, fused_out = (
1153 - current_workspace_manager().get_simultaneous(
1154 - (workspace13_shape, workspace_dtype),
1155 - (workspace2_shape, workspace_dtype),
1156 - (fused_out_shape, out_dtype),
1157 - )
1158 - )
1159 -
1160 - return workspace13, workspace2, fused_out
1161 -
1162 - @staticmethod
1163 - def _slice_output_tensor(
1164 - fused_out: torch.Tensor,
1165 - chunk_idx: int,
1166 - num_chunks: int,
1167 - CHUNK_SIZE: int,
1168 - M: int,
1169 - ) -> torch.Tensor:
1170 - if num_chunks == 1:
1171 - return fused_out
1172 -
1173 - assert fused_out.size(0) % M == 0, f"fused_out shape {fused_out.shape} vs M {M}"
1174 - factor = fused_out.size(0) // M
1175 - out_chunk_size = CHUNK_SIZE * factor
1176 - s = chunk_idx * out_chunk_size
1177 - e = min(s + out_chunk_size, fused_out.size(0))
1178 - return fused_out[s:e]
1179 -
1180 - @staticmethod
1181 - def _slice_expert_tokens_metadata(
1182 - num_chunks: int,
1183 - full_expert_tokens_meta: ExpertTokensMetadata | None,
1184 - chunk_topk_ids: torch.Tensor,
1185 - local_num_experts: int,
1186 - expert_map: torch.Tensor | None,
1187 - ) -> ExpertTokensMetadata | None:
1188 - if num_chunks == 1 or full_expert_tokens_meta is None:
1189 - return full_expert_tokens_meta
1190 -
1191 - # The existing expert_num_tokens is for the entire a1q
1192 - # input. Chunking forces recomputation of the number
1193 - # of tokens assigned to each expert.
1194 - c_expert_num_tokens = count_expert_num_tokens(
1195 - chunk_topk_ids, local_num_experts, expert_map
1196 - )
1197 -
1198 - c_expert_num_tokens_cpu = None
1199 - need_expert_num_tokens_cpu = (
1200 - full_expert_tokens_meta.expert_num_tokens_cpu is not None
1058 + # Reuse workspace13 for the output since there is only one chunk.
1059 + max_shape_size = max(prod(workspace13_shape), prod(fused_out_shape))
1060 + common_workspace, workspace2 = current_workspace_manager().get_simultaneous(
1061 + ((max_shape_size,), workspace_dtype),
1062 + (workspace2_shape, workspace_dtype),
1063 )
1202 - if need_expert_num_tokens_cpu:
1203 - # This is blocking as some implementations need the count
1204 - # on the CPU to determine appropriate input/out fused-moe
1205 - # buffers
1206 - c_expert_num_tokens_cpu = c_expert_num_tokens.to("cpu", non_blocking=False)
1064 + workspace13 = _resize_cache(common_workspace, workspace13_shape)
1065 + fused_out = _resize_cache(common_workspace, fused_out_shape)
1066
1208 - return ExpertTokensMetadata(
1209 - expert_num_tokens=c_expert_num_tokens,
1210 - expert_num_tokens_cpu=c_expert_num_tokens_cpu,
1211 - )
1067 + return workspace13, workspace2, fused_out
1068
1069 def _prepare(
1070 self,
@@ -1318,18 +1174,6 @@ class FusedMoEKernelModularImpl:
1174 a1q, w1, w2, topk_ids
1175 )
1176
1321 - num_chunks, CHUNK_SIZE = self._chunk_info(M_full)
1322 -
1323 - def input_chunk_range(chunk_idx: int) -> tuple[int, int]:
1324 - if num_chunks == 1:
1325 - # Use a1q.size(0) here since batched format does not
1326 - # keep M in the first dimension.
1327 - return 0, a1q.size(0)
1328 - else:
1329 - s = chunk_idx * CHUNK_SIZE
1330 - e = min(s + CHUNK_SIZE, M_full)
1331 - return s, e
1332 -
1177 # This happens when none of the tokens from the all2all reach this
1178 # EP rank. Also, note that this is only relevant for CUDAGraph
1179 # incompatible all2all kernels like the DeepEP high-throughput
@@ -1337,58 +1181,39 @@ class FusedMoEKernelModularImpl:
1181 # low-latency kernels are always batched and can never run into
1182 # the tensor.numel() == 0 case.
1183 if M_full == 0:
1340 - assert num_chunks == 0
1341 - workspace13 = None
1342 - workspace2 = None
1343 - fused_out = torch.empty_like(a1q, dtype=in_dtype)
1344 - else:
1345 - assert num_chunks > 0
1346 - workspace13, workspace2, fused_out = self._allocate_buffers(
1347 - in_dtype,
1348 - a1q.device,
1349 - CHUNK_SIZE,
1350 - M_full,
1351 - N,
1352 - K,
1353 - top_k,
1354 - global_num_experts,
1355 - local_num_experts,
1356 - expert_tokens_meta,
1357 - activation,
1358 - )
1184 + return torch.empty_like(a1q, dtype=in_dtype)
1185
1360 - for chunk_idx in range(num_chunks):
1361 - s, e = input_chunk_range(chunk_idx)
1362 -
1363 - c_expert_tokens_meta = self._slice_expert_tokens_metadata(
1364 - num_chunks,
1365 - expert_tokens_meta,
1366 - topk_ids[s:e],
1367 - local_num_experts,
1368 - expert_map,
1369 - )
1370 -
1371 - c_fused_out = self._slice_output_tensor(
1372 - fused_out, chunk_idx, num_chunks, CHUNK_SIZE, M_full
1373 - )
1186 + workspace13, workspace2, fused_out = self._allocate_buffers(
1187 + in_dtype,
1188 + a1q.device,
1189 + M_full,
1190 + M_full,
1191 + N,
1192 + K,
1193 + top_k,
1194 + global_num_experts,
1195 + local_num_experts,
1196 + expert_tokens_meta,
1197 + activation,
1198 + )
1199
1375 - self.fused_experts.apply(
1376 - output=c_fused_out,
1377 - hidden_states=a1q[s:e],
1378 - w1=w1,
1379 - w2=w2,
1380 - topk_weights=topk_weights[s:e],
1381 - topk_ids=topk_ids[s:e],
1382 - activation=activation,
1383 - global_num_experts=global_num_experts,
1384 - expert_map=expert_map,
1385 - a1q_scale=_slice_scales(a1q_scale, s, e),
1386 - a2_scale=_slice_scales(self.fused_experts.a2_scale, s, e),
1387 - workspace13=workspace13,
1388 - workspace2=workspace2,
1389 - expert_tokens_meta=c_expert_tokens_meta,
1390 - apply_router_weight_on_input=apply_router_weight_on_input,
1391 - )
1200 + self.fused_experts.apply(
1201 + output=fused_out,
1202 + hidden_states=a1q,
1203 + w1=w1,
1204 + w2=w2,
1205 + topk_weights=topk_weights,
1206 + topk_ids=topk_ids,
1207 + activation=activation,
1208 + global_num_experts=global_num_experts,
1209 + expert_map=expert_map,
1210 + a1q_scale=a1q_scale,
1211 + a2_scale=self.fused_experts.a2_scale,
1212 + workspace13=workspace13,
1213 + workspace2=workspace2,
1214 + expert_tokens_meta=expert_tokens_meta,
1215 + apply_router_weight_on_input=apply_router_weight_on_input,
1216 + )
1217
1218 return fused_out
1219
vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py
-3
@@ -337,9 +337,6 @@ class AiterExperts(mk.FusedMoEExpertsModular):
337 def supports_expert_map(self):
338 return True
339
340 - def supports_chunking(self):
341 - return False
342 -
340 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
341 return TopKWeightAndReduceNoOP()
342
vllm/model_executor/layers/fused_moe/trtllm_moe.py
-3
@@ -83,9 +83,6 @@ class TrtLlmGenExperts(mk.FusedMoEExpertsModular):
83 "This method should not be called."
84 )
85
86 - def supports_chunking(self) -> bool:
87 - return True
88 -
86 def supports_expert_map(self) -> bool:
87 return True
88
vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
-3
@@ -79,9 +79,6 @@ class XPUExperts(mk.FusedMoEExpertsModular):
79 ]
80 return (weight_key, activation_key) in SUPPORTED_W_A
81
82 - def supports_chunking(self) -> bool:
83 - return False
84 -
82 def supports_expert_map(self) -> bool:
83 return True
84
vllm/model_executor/warmup/deep_gemm_warmup.py
+1 -2
@@ -244,8 +244,7 @@ def _get_grouped_gemm_params(
244 device = w1.device
245
246 # Assumes all ranks have the same max_num_batched_tokens
247 - max_tokens_across_dp = get_dp_group().world_size * max_tokens
248 - max_tokens = min(max_tokens_across_dp, envs.VLLM_FUSED_MOE_CHUNK_SIZE)
247 + max_tokens = get_dp_group().world_size * max_tokens
248
249 # This is the maximum GroupedGemm M size that we expect to run
250 # the grouped_gemm with.