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

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

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

Synced Nov 15, 06:00 UTC

[Kernels] Add activation chunking logic to FusedMoEModularKernel (#19168)

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

bnellnm committed Jun 11, 2025 at 12:53 UTC 29fa5cac1cd731026f59084d93a822921507573c
15 files changed +458 -396
tests/kernels/moe/test_cutlass_moe.py
+1
@@ -29,6 +29,7 @@ MNK_FACTORS = [
29 (224, 1024, 1536),
30 (224, 3072, 1024),
31 (224, 3072, 1536),
32 + (1024 * 128, 1024, 1024),
33 ]
34
35 vllm_config = VllmConfig(parallel_config=ParallelConfig(
tests/kernels/moe/test_moe.py
+22 -1
@@ -15,7 +15,8 @@ import vllm.model_executor.layers.fused_moe # noqa
15 from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
16 from vllm.config import VllmConfig, set_current_vllm_config
17 from vllm.model_executor.layers.fused_moe import fused_moe
18 -from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
18 +from vllm.model_executor.layers.fused_moe.fused_moe import (
19 + fused_topk, modular_triton_fused_moe)
20 from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
21 fused_moe as iterative_moe)
22 from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
@@ -76,6 +77,13 @@ def test_fused_moe(
77 else:
78 e_map = None
79
80 + m_fused_moe = modular_triton_fused_moe(use_fp8_w8a8=False,
81 + use_int8_w8a8=False,
82 + use_int8_w8a16=False,
83 + use_int4_w4a16=False,
84 + per_channel_quant=False,
85 + block_shape=None)
86 +
87 with set_current_vllm_config(vllm_config):
88 torch_output = torch_moe(a, w1, w2, score, topk, e_map)
89 iterative_output = iterative_moe(a,
@@ -103,7 +111,20 @@ def test_fused_moe(
111 expert_map=e_map,
112 renormalize=False)
113
114 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
115 + m_triton_output = m_fused_moe(a,
116 + w1,
117 + w2,
118 + topk_weights,
119 + topk_ids,
120 + global_num_experts=e,
121 + expert_map=e_map)
122 +
123 torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
124 + torch.testing.assert_close(m_triton_output,
125 + torch_output,
126 + atol=2e-2,
127 + rtol=0)
128 torch.testing.assert_close(iterative_output,
129 torch_output,
130 atol=2e-2,
tests/kernels/moe/test_pplx_cutlass_moe.py
+33 -10
@@ -1,10 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from typing import Optional
5 +
6 import pytest
7 import torch
8
7 -from tests.pplx_utils import ProcessGroupInfo, parallel_launch
9 from vllm import _custom_ops as ops
10 from vllm.config import VllmConfig, set_current_vllm_config
11 from vllm.model_executor.layers.activation import SiluAndMul
@@ -14,6 +15,8 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
15 FusedMoEModularKernel)
16 from vllm.platforms import current_platform
17
18 +from .deepep_utils import ProcessGroupInfo, parallel_launch
19 +
20 try:
21 from pplx_kernels import AllToAll
22 from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
@@ -64,6 +67,7 @@ def pplx_cutlass_moe(
67 out_dtype,
68 per_act_token: bool,
69 per_out_ch: bool,
70 + group_name: Optional[str],
71 ):
72 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
73 PplxPrepareAndFinalize)
@@ -84,7 +88,7 @@ def pplx_cutlass_moe(
88 else:
89 scale_elems = (hidden_dim + block_size - 1) // block_size
90
87 - ata = AllToAll.internode(
91 + args = dict(
92 max_num_tokens=max_num_tokens,
93 num_experts=num_experts,
94 experts_per_token=topk,
@@ -96,6 +100,12 @@ def pplx_cutlass_moe(
100 hidden_dim_scale_bytes=scale_elems * torch.float32.itemsize,
101 )
102
103 + if group_name is None:
104 + ata = AllToAll.internode(**args)
105 + else:
106 + args["group_name"] = group_name
107 + ata = AllToAll.intranode(**args)
108 +
109 w1 = w1.to(device)
110 w2 = w2.to(device)
111 w1_scale = w1_scale.to(device)
@@ -113,7 +123,10 @@ def pplx_cutlass_moe(
123 )
124
125 experts = CutlassExpertsFp8((num_experts + world_size - 1) // world_size,
116 - out_dtype, per_act_token, per_out_ch)
126 + out_dtype,
127 + per_act_token,
128 + per_out_ch,
129 + use_batched_format=True)
130
131 fused_cutlass_experts = FusedMoEModularKernel(
132 prepare_finalize,
@@ -184,11 +197,17 @@ def _pplx_moe(
197 w2_full: torch.Tensor,
198 per_act_token: bool,
199 per_out_ch: bool,
200 + use_internode: bool,
201 ):
188 - uid = nvshmem_get_unique_id(
189 - ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
190 - torch.distributed.broadcast(uid, src=0)
191 - nvshmem_init(uid, pgi.rank, pgi.world_size)
202 + if use_internode:
203 + uid = nvshmem_get_unique_id(
204 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
205 + torch.distributed.broadcast(uid, src=0)
206 + nvshmem_init(uid, pgi.rank, pgi.world_size)
207 + else:
208 + group_ranks = list(range(pgi.world_size))
209 + cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
210 + group_name = cpu_group.group_name
211
212 with set_current_vllm_config(vllm_config):
213 torch_output = torch_moe2(a_full, w1_full, w2_full, topk_weights,
@@ -196,7 +215,7 @@ def _pplx_moe(
215 pplx_output = pplx_cutlass_moe(pgi, dp_size, a, w1, w2, w1_scale,
216 w2_scale, topk_weights, topk_ids,
217 a1_scale, out_dtype, per_act_token,
199 - per_out_ch)
218 + per_out_ch, group_name)
219
220 torch_output = chunk_by_rank(torch_output, pgi.rank,
221 pgi.world_size).to(pplx_output.device)
@@ -207,7 +226,8 @@ def _pplx_moe(
226
227 torch.testing.assert_close(pplx_output, torch_output, atol=0.05, rtol=0)
228
210 - nvshmem_finalize()
229 + if use_internode:
230 + nvshmem_finalize()
231
232
233 @pytest.mark.parametrize("m", [2, 224])
@@ -218,6 +238,7 @@ def _pplx_moe(
238 @pytest.mark.parametrize("per_act_token", [True, False])
239 @pytest.mark.parametrize("per_out_ch", [True, False])
240 @pytest.mark.parametrize("world_dp_size", [[2, 1]]) #, [4, 2]])
241 +@pytest.mark.parametrize("use_internode", [False])
242 @pytest.mark.skipif(
243 (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
244 current_platform.get_device_capability()),
@@ -232,6 +253,7 @@ def test_cutlass_moe_pplx(
253 per_act_token: bool,
254 per_out_ch: bool,
255 world_dp_size: tuple[int, int],
256 + use_internode: bool,
257 ):
258 current_platform.seed_everything(7)
259
@@ -284,4 +306,5 @@ def test_cutlass_moe_pplx(
306
307 parallel_launch(world_size, _pplx_moe, dp_size, a, w1_q, w2_q,
308 w1_scale, w2_scale, topk_weights, topk_ids, a_scale1,
287 - dtype, a, w1_d, w2_d, per_act_token, per_out_ch)
309 + dtype, a, w1_d, w2_d, per_act_token, per_out_ch,
310 + use_internode)
tests/kernels/moe/test_pplx_moe.py
+67 -22
@@ -18,7 +18,6 @@ try:
18 except ImportError:
19 has_pplx = False
20
21 -from tests.pplx_utils import ProcessGroupInfo, parallel_launch
21 from vllm.config import VllmConfig, set_current_vllm_config
22 from vllm.model_executor.layers.activation import SiluAndMul
23 from vllm.model_executor.layers.fused_moe import override_config
@@ -30,6 +29,8 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
29 FusedMoEModularKernel)
30 from vllm.platforms import current_platform
31
32 +from .deepep_utils import ProcessGroupInfo, parallel_launch
33 +
34 requires_pplx = pytest.mark.skipif(
35 not has_pplx,
36 reason="Requires PPLX kernels",
@@ -153,7 +154,10 @@ def batched_moe(
154 num_experts = w1.shape[0]
155
156 fused_experts = FusedMoEModularKernel(
156 - BatchedPrepareAndFinalize(a.shape[0], world_size=1, dp_size=1, rank=0),
157 + BatchedPrepareAndFinalize(max_num_tokens=a.shape[0],
158 + world_size=1,
159 + dp_size=1,
160 + rank=0),
161 BatchedExperts(max_num_tokens=a.shape[0], dp_size=1, world_size=1))
162
163 return fused_experts(a, w1, w2, topk_weight, topk_ids, num_experts)
@@ -229,9 +233,15 @@ def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor:
233 return t[(r * chunk):(r + 1) * chunk]
234
235
232 -def pplx_prepare_finalize(pgi: ProcessGroupInfo, dp_size: int, a: torch.Tensor,
233 - topk_weight: torch.Tensor, topk_ids: torch.Tensor,
234 - num_experts: int) -> torch.Tensor:
236 +def pplx_prepare_finalize(
237 + pgi: ProcessGroupInfo,
238 + dp_size: int,
239 + a: torch.Tensor,
240 + topk_weight: torch.Tensor,
241 + topk_ids: torch.Tensor,
242 + num_experts: int,
243 + group_name: Optional[str],
244 +) -> torch.Tensor:
245 from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
246 PplxPrepareAndFinalize)
247
@@ -245,7 +255,7 @@ def pplx_prepare_finalize(pgi: ProcessGroupInfo, dp_size: int, a: torch.Tensor,
255 world_size = pgi.world_size
256 max_num_tokens = rank_chunk(num_tokens, 0, world_size)
257
248 - ata = AllToAll.internode(
258 + args = dict(
259 max_num_tokens=max_num_tokens,
260 num_experts=num_experts,
261 experts_per_token=topk,
@@ -259,6 +269,12 @@ def pplx_prepare_finalize(pgi: ProcessGroupInfo, dp_size: int, a: torch.Tensor,
269 torch.float32.itemsize)),
270 )
271
272 + if group_name is None:
273 + ata = AllToAll.internode(**args)
274 + else:
275 + args["group_name"] = group_name
276 + ata = AllToAll.intranode(**args)
277 +
278 topk_ids = topk_ids.to(dtype=torch.uint32)
279
280 prepare_finalize = PplxPrepareAndFinalize(
@@ -318,11 +334,19 @@ def _pplx_prepare_finalize(
334 score: torch.Tensor,
335 topk: torch.Tensor,
336 num_experts: int,
337 + use_internode: bool,
338 ):
322 - uid = nvshmem_get_unique_id(
323 - ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
324 - torch.distributed.broadcast(uid, src=0)
325 - nvshmem_init(uid, pgi.rank, pgi.world_size)
339 + if use_internode:
340 + uid = nvshmem_get_unique_id(
341 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
342 + torch.distributed.broadcast(uid, src=0)
343 + nvshmem_init(uid, pgi.rank, pgi.world_size)
344 + group_name = None
345 + else:
346 + group_ranks = list(range(pgi.world_size))
347 + cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
348 + group_name = cpu_group.group_name
349 +
350 device = pgi.device
351
352 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
@@ -335,14 +359,15 @@ def _pplx_prepare_finalize(
359 a.dtype)
360
361 pplx_output = pplx_prepare_finalize(pgi, dp_size, a, topk_weight, topk_ids,
338 - num_experts)
362 + num_experts, group_name)
363
364 torch_output = chunk_by_rank(torch_output, pgi.rank,
365 pgi.world_size).to(pplx_output.device)
366
367 torch.testing.assert_close(pplx_output, torch_output, atol=2e-2, rtol=0)
368
345 - nvshmem_finalize()
369 + if use_internode:
370 + nvshmem_finalize()
371
372
373 # TODO (bnell): this test point does not work for odd M due to how the test is
@@ -353,6 +378,7 @@ def _pplx_prepare_finalize(
378 @pytest.mark.parametrize("topk", TOP_KS)
379 @pytest.mark.parametrize("dtype", [torch.bfloat16])
380 @pytest.mark.parametrize("world_dp_size", [[2, 1]])
381 +@pytest.mark.parametrize("use_internode", [False])
382 @requires_pplx
383 def test_pplx_prepare_finalize(
384 mnk: tuple[int, int, int],
@@ -360,6 +386,7 @@ def test_pplx_prepare_finalize(
386 topk: int,
387 dtype: torch.dtype,
388 world_dp_size: tuple[int, int],
389 + use_internode: bool,
390 ):
391 current_platform.seed_everything(7)
392 m, n, k = mnk
@@ -369,10 +396,11 @@ def test_pplx_prepare_finalize(
396 score = torch.randn((m, e), device=device, dtype=dtype)
397
398 parallel_launch(world_size, _pplx_prepare_finalize, dp_size, a, score,
372 - topk, e)
399 + topk, e, use_internode)
400
401
402 def pplx_moe(
403 + group_name: Optional[str],
404 rank: int,
405 world_size: int,
406 dp_size: int,
@@ -394,7 +422,7 @@ def pplx_moe(
422 topk = topk_ids.shape[1]
423 max_num_tokens = rank_chunk(a.shape[0], 0, world_size)
424
397 - ata = AllToAll.internode(
425 + args = dict(
426 max_num_tokens=max_num_tokens,
427 num_experts=num_experts,
428 experts_per_token=topk,
@@ -408,6 +436,12 @@ def pplx_moe(
436 torch.float32.itemsize)),
437 )
438
439 + if group_name is None:
440 + ata = AllToAll.internode(**args)
441 + else:
442 + args["group_name"] = group_name
443 + ata = AllToAll.intranode(**args)
444 +
445 topk_ids = topk_ids.to(dtype=torch.uint32)
446
447 prepare_finalize = PplxPrepareAndFinalize(
@@ -522,11 +556,18 @@ def _pplx_moe(
556 w2: torch.Tensor,
557 score: torch.Tensor,
558 topk: int,
559 + use_internode: bool,
560 ):
526 - uid = nvshmem_get_unique_id(
527 - ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
528 - torch.distributed.broadcast(uid, src=0)
529 - nvshmem_init(uid, pgi.rank, pgi.world_size)
561 + if use_internode:
562 + uid = nvshmem_get_unique_id(
563 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
564 + torch.distributed.broadcast(uid, src=0)
565 + nvshmem_init(uid, pgi.rank, pgi.world_size)
566 + group_name = None
567 + else:
568 + group_ranks = list(range(pgi.world_size))
569 + cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
570 + group_name = cpu_group.group_name
571
572 m, k = a.shape
573 e, _, n = w2.shape
@@ -536,8 +577,8 @@ def _pplx_moe(
577 with set_current_vllm_config(vllm_config), override_config(moe_config):
578 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
579 torch_output = torch_moe2(a, w1, w2, topk_weight, topk_ids)
539 - pplx_output = pplx_moe(pgi.rank, pgi.world_size, dp_size, a, w1, w2,
540 - topk_weight, topk_ids)
580 + pplx_output = pplx_moe(group_name, pgi.rank, pgi.world_size, dp_size,
581 + a, w1, w2, topk_weight, topk_ids)
582 # TODO (bnell): fix + re-enable
583 #batched_output = _batched_moe(pgi, dp_size, a, w1, w2, topk_weight,
584 # topk_ids)
@@ -548,7 +589,8 @@ def _pplx_moe(
589 torch.testing.assert_close(pplx_output, torch_output, atol=2e-2, rtol=0)
590 #torch.testing.assert_close(batched_output, torch_output, atol=2e-2, rtol=0)
591
551 - nvshmem_finalize()
592 + if use_internode:
593 + nvshmem_finalize()
594
595
596 @pytest.mark.parametrize("mnk", PPLX_MOE_COMBOS)
@@ -556,6 +598,7 @@ def _pplx_moe(
598 @pytest.mark.parametrize("topk", TOP_KS)
599 @pytest.mark.parametrize("dtype", [torch.bfloat16])
600 @pytest.mark.parametrize("world_dp_size", [[2, 1]])
601 +@pytest.mark.parametrize("use_internode", [False])
602 @requires_pplx
603 def test_pplx_moe(
604 mnk: tuple[int, int, int],
@@ -563,6 +606,7 @@ def test_pplx_moe(
606 topk: int,
607 dtype: torch.dtype,
608 world_dp_size: tuple[int, int],
609 + use_internode: bool,
610 ):
611 current_platform.seed_everything(7)
612 m, n, k = mnk
@@ -572,4 +616,5 @@ def test_pplx_moe(
616 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
617 score = torch.randn((m, e), device="cuda", dtype=dtype)
618
575 - parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk)
619 + parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk,
620 + use_internode)
tests/kernels/quantization/test_block_fp8.py
+25 -2
@@ -13,7 +13,8 @@ from vllm.model_executor.layers.activation import SiluAndMul
13 from vllm.model_executor.layers.fused_moe import fused_moe
14 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
15 _valid_deep_gemm_shape, deep_gemm_moe_fp8)
16 -from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
16 +from vllm.model_executor.layers.fused_moe.fused_moe import (
17 + fused_topk, modular_triton_fused_moe)
18 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
19 moe_align_block_size)
20 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
@@ -45,7 +46,7 @@ N = [128, 512, 7168, 7748, 13824]
46 K = [256, 3884, 4096, 13824, 16384]
47 # Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8
48 # and its hidden size is 7168.
48 -M_moe = [1, 2, 7, 83, 128, 2048]
49 +M_moe = [1, 2, 7, 83, 128, 2048, 1024 * 128]
50 M_moe_dg = [128, 192, 1335, 2048]
51 N_moe = [128, 256, 1024, 4608] # [13824]
52 K_moe = [256, 512, 7168] # [13824]
@@ -214,6 +215,13 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
215
216 score = torch.randn((M, E), dtype=dtype)
217
218 + m_fused_moe = modular_triton_fused_moe(use_fp8_w8a8=True,
219 + use_int8_w8a8=False,
220 + use_int8_w8a16=False,
221 + use_int4_w4a16=False,
222 + per_channel_quant=False,
223 + block_shape=block_size)
224 +
225 # Set the context to avoid lots of warning spam.
226 with set_current_vllm_config(vllm_config):
227 out = fused_moe(
@@ -231,6 +239,16 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
239 ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk,
240 block_size)
241
242 + topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
243 + m_out = m_fused_moe(a,
244 + w1,
245 + w2,
246 + topk_weights,
247 + topk_ids,
248 + global_num_experts=E,
249 + w1_scale=w1_s,
250 + w2_scale=w2_s)
251 +
252 #print(f"{out.sum()=}")
253 #print(f"{ref_out.sum()=}")
254
@@ -239,6 +257,11 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
257 torch.mean(torch.abs(ref_out.to(torch.float32))))
258 assert rel_diff < 0.03
259
260 + rel_diff = (torch.mean(
261 + torch.abs(m_out.to(torch.float32) - ref_out.to(torch.float32))) /
262 + torch.mean(torch.abs(ref_out.to(torch.float32))))
263 + assert rel_diff < 0.03
264 +
265
266 def per_block_cast_to_fp8(
267 x: torch.Tensor,
tests/pplx_utils.py deleted
-123
@@ -1,123 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -
4 -import dataclasses
5 -import os
6 -import traceback
7 -from typing import Callable
8 -
9 -import torch
10 -from torch.multiprocessing import (
11 - spawn) # pyright: ignore[reportPrivateImportUsage]
12 -from typing_extensions import Concatenate, ParamSpec
13 -
14 -P = ParamSpec("P")
15 -
16 -
17 -@dataclasses.dataclass
18 -class ProcessGroupInfo:
19 - world_size: int
20 - world_local_size: int
21 - rank: int
22 - node_rank: int
23 - local_rank: int
24 - device: torch.device
25 -
26 -
27 -def _worker_parallel_launch(
28 - local_rank: int,
29 - world_size: int,
30 - world_local_size: int,
31 - node_rank: int,
32 - init_method: str,
33 - worker: Callable[Concatenate[ProcessGroupInfo, P], None],
34 - *args: P.args,
35 - **kwargs: P.kwargs,
36 -) -> None:
37 - rank = node_rank * world_local_size + local_rank
38 - torch.cuda.set_device(local_rank)
39 - device = torch.device("cuda", local_rank)
40 - torch.distributed.init_process_group(
41 - backend="cpu:gloo,cuda:nccl",
42 - init_method=init_method,
43 - rank=rank,
44 - world_size=world_size,
45 - device_id=device,
46 - )
47 - barrier = torch.tensor([rank], device=device)
48 - torch.distributed.all_reduce(barrier)
49 -
50 - try:
51 - worker(
52 - ProcessGroupInfo(
53 - world_size=world_size,
54 - world_local_size=world_local_size,
55 - rank=rank,
56 - node_rank=node_rank,
57 - local_rank=local_rank,
58 - device=device,
59 - ),
60 - *args,
61 - **kwargs,
62 - )
63 - except Exception as ex:
64 - print(ex)
65 - traceback.print_exc()
66 - raise
67 - finally:
68 - torch.distributed.destroy_process_group()
69 -
70 -
71 -def parallel_launch(
72 - world_size: int,
73 - worker: Callable[Concatenate[ProcessGroupInfo, P], None],
74 - *args: P.args,
75 - **kwargs: P.kwargs,
76 -) -> None:
77 - assert not kwargs
78 - spawn(
79 - _worker_parallel_launch,
80 - args=(
81 - world_size,
82 - world_size,
83 - 0,
84 - "tcp://localhost:29500",
85 - worker,
86 - ) + args,
87 - nprocs=world_size,
88 - join=True,
89 - )
90 -
91 -
92 -def parallel_launch_from_env(
93 - worker: Callable[Concatenate[ProcessGroupInfo, P], None],
94 - *args: P.args,
95 - **kwargs: P.kwargs,
96 -) -> None:
97 - """
98 - Launches a worker function in parallel across all processes in the current
99 - environment. The environment must have the following variables set:
100 - - WORLD_SIZE: The total number of processes.
101 - - WORLD_LOCAL_SIZE: The number of processes on the current node.
102 - - NODE_RANK: The rank of the current
103 - - MASTER_ADDR: The address of the master process.
104 - - MASTER_PORT: The port of the master process.
105 - """
106 - assert not kwargs
107 - world_size = int(os.environ["WORLD_SIZE"])
108 - world_local_size = int(os.environ["WORLD_LOCAL_SIZE"])
109 - node_rank = int(os.environ["NODE_RANK"])
110 - assert "MASTER_ADDR" in os.environ
111 - assert "MASTER_PORT" in os.environ
112 - spawn(
113 - _worker_parallel_launch,
114 - args=(
115 - world_size,
116 - world_local_size,
117 - node_rank,
118 - "env://",
119 - worker,
120 - ) + args,
121 - nprocs=world_local_size,
122 - join=True,
123 - )
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+11 -9
@@ -36,6 +36,9 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
36 assert (len(self.block_shape) == 2 and all(
37 [v == self.DEEPGEMM_BLOCK_SHAPE for v in self.block_shape]))
38
39 + def supports_chunking(self) -> bool:
40 + return False
41 +
42 def workspace_shapes(
43 self,
44 a: torch.Tensor,
@@ -45,17 +48,19 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
48 K: int,
49 topk: int,
50 num_experts: int,
48 - ) -> tuple[int, int, torch.dtype]:
51 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
52 assert a.dim() == 2
53 num_dp = self.world_size // self.dp_size
54 max_num_tokens = a.size(
55 0) if self.max_num_tokens is None else self.max_num_tokens
53 - workspace13 = num_experts * max_num_tokens * num_dp * max(K, N)
54 - workspace2 = num_experts * max_num_tokens * num_dp * (N // 2)
55 - return (workspace13, workspace2, a.dtype)
56 + workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
57 + workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
58 + output = (num_experts, max_num_tokens * num_dp, K)
59 + return (workspace13, workspace2, output, a.dtype)
60
61 def apply(
62 self,
63 + output: torch.Tensor,
64 hidden_states: torch.Tensor,
65 w1: torch.Tensor,
66 w2: torch.Tensor,
@@ -72,7 +77,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
77 workspace13: torch.Tensor,
78 workspace2: torch.Tensor,
79 expert_num_tokens: Optional[torch.Tensor],
75 - ) -> torch.Tensor:
80 + ):
81 import deep_gemm as dg
82 assert hidden_states.ndim == 3
83
@@ -89,7 +94,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
94
95 workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))
96 workspace2 = _resize_cache(workspace2, (E, max_num_tokens, N // 2))
92 - workspace3 = _resize_cache(workspace13, (E, max_num_tokens, K))
97
98 # (from deepgemm docs) : A value hint (which is a value on CPU)
99 # for the M expectation of each batch, correctly setting this value
@@ -118,8 +122,6 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
122
123 dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a2q, a2q_scale),
124 (w2, w2_scale),
121 - out=workspace3,
125 + out=output,
126 masked_m=expert_num_tokens,
127 expected_m=expected_m)
124 -
125 - return workspace3
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+16 -6
@@ -64,6 +64,15 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
64 block_shape=self.block_shape, # type: ignore[arg-type]
65 ) if (self.allow_deep_gemm and is_fp8_128_block_quantized) else None
66
67 + assert (self.batched_deep_gemm_experts is not None
68 + or self.batched_triton_experts is not None)
69 +
70 + def supports_chunking(self) -> bool:
71 + bdge = self.batched_deep_gemm_experts
72 + bte = self.batched_triton_experts
73 + return ((bdge is None or bdge.supports_chunking())
74 + and (bte is None or bte.supports_chunking()))
75 +
76 def workspace_shapes(
77 self,
78 a: torch.Tensor,
@@ -73,7 +82,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
82 K: int,
83 topk: int,
84 num_experts: int,
76 - ) -> tuple[int, int, torch.dtype]:
85 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
86 # Note: the deep gemm workspaces are strictly larger than the triton
87 # workspaces so we can be pessimistic here and allocate for DeepGemm
88 # even if we fall back to triton later, e.g. if expert maps are set.
@@ -87,6 +96,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
96
97 def apply(
98 self,
99 + output: torch.Tensor,
100 hidden_states: torch.Tensor,
101 w1: torch.Tensor,
102 w2: torch.Tensor,
@@ -103,7 +113,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
113 workspace13: torch.Tensor,
114 workspace2: torch.Tensor,
115 expert_num_tokens: Optional[torch.Tensor],
106 - ) -> torch.Tensor:
116 + ):
117 use_batched_deep_gemm_experts = (self.allow_deep_gemm
118 and self.batched_deep_gemm_experts
119 is not None)
@@ -111,7 +121,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
121 if use_batched_deep_gemm_experts else
122 self.batched_triton_experts)
123 assert experts is not None
114 - return experts.apply(hidden_states, w1, w2, topk_ids, activation,
115 - global_num_experts, expert_map, w1_scale,
116 - w2_scale, w1_zp, w2_zp, a1q_scale, a2_scale,
117 - workspace13, workspace2, expert_num_tokens)
124 + experts.apply(output, hidden_states, w1, w2, topk_ids, activation,
125 + global_num_experts, expert_map, w1_scale, w2_scale,
126 + w1_zp, w2_zp, a1q_scale, a2_scale, workspace13,
127 + workspace2, expert_num_tokens)
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+49 -28
@@ -14,6 +14,7 @@ from vllm.scalar_type import scalar_types
14
15
16 def run_cutlass_moe_fp8(
17 + output: torch.Tensor,
18 hidden_states: torch.Tensor,
19 w1: torch.Tensor,
20 w2: torch.Tensor,
@@ -31,7 +32,8 @@ def run_cutlass_moe_fp8(
32 out_dtype: torch.dtype,
33 per_act_token: bool,
34 per_out_ch: bool,
34 -) -> torch.Tensor:
35 + use_batched_format: bool,
36 +):
37 a1q = hidden_states
38
39 assert w1_scale is not None
@@ -61,23 +63,20 @@ def run_cutlass_moe_fp8(
63 if expert_map is not None:
64 assert expert_num_tokens is None
65
64 - # We have two modes: PPLX and non-PPLX. We differentiate them by checking
65 - # if expert_num_tokens is None (expert_num_tokens is a tensor which PPLX
66 - # uses to track the number of tokens per expert).
67 - # In the non-PPLX mode, the input tokens are not padded: thus, the shape
66 + # We have two modes: batched experts and non-batched experts.
67 + # In the non-batched mode, the input tokens are not padded: thus, the shape
68 # of the input is [total_num_tokens, hidden_size]. The input and output
69 # require shuffling by a_map and c_map such that the tokens assigned to
70 # each expert are contiguous.
71 - # In the PPLX mode, the input tokens are padded per expert to ensure that
72 - # the PPLX dispatch and combine functions work correctly: thus, the shape
71 + # In the batched mode, the input tokens are padded per expert to ensure that
72 + # the batched dispatch and combine functions work correctly: thus, the shape
73 # of the input is [num_experts, max_num_tokens_per_expert, hidden_size].
74 - # The PPLX input and output require no shuffling by a_map and c_map since
74 + # The batched input and output require no shuffling by a_map and c_map since
75 # their tokens are already contiguous for each expert as a result of
76 # the dispatch function.
77 - is_pplx = expert_num_tokens is not None
77
79 - M = a1q.shape[0] # no pplx
80 - padded_M = a1q.shape[1] # pplx
78 + M = a1q.shape[0] # non batched expert M
79 + padded_M = a1q.shape[1] # batched expert M
80 _, K, N = w2.shape
81 device = a1q.device
82
@@ -95,7 +94,9 @@ def run_cutlass_moe_fp8(
94 topk = local_topk_ids.shape[1]
95 local_E = w1.shape[0]
96
98 - if is_pplx:
97 + if use_batched_format:
98 + assert expert_num_tokens is not None
99 +
100 expert_offsets = torch.empty((local_E),
101 dtype=torch.int32,
102 device=device)
@@ -167,7 +168,7 @@ def run_cutlass_moe_fp8(
168 device=device,
169 dtype=torch.int64)
170
170 - if is_pplx:
171 + if use_batched_format:
172 c1 = _resize_cache(workspace13, (local_E * padded_M, N * 2))
173 c2 = _resize_cache(workspace2, (local_E * padded_M, N))
174 c3 = _resize_cache(workspace13, (local_E * padded_M, K))
@@ -192,12 +193,15 @@ def run_cutlass_moe_fp8(
193 problem_sizes2, ab_strides2, ab_strides2, c_strides2,
194 per_act_token, per_out_ch)
195
195 - if is_pplx:
196 - return c3.reshape(local_E, padded_M, K)
196 + if use_batched_format:
197 + output.copy_(c3.reshape(local_E, padded_M, K), non_blocking=True)
198 else:
198 - return c3[c_map].view(M, topk, K)
199 + # We can't do this inplace because output may point to the same tensor
200 + # as c3.
201 + output.copy_(c3[c_map].view(M * topk, K), non_blocking=True)
202
203
204 +# TODO (bnell): split class batched vs. non-batched?
205 class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
206
207 def __init__(
@@ -206,12 +210,17 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
210 out_dtype: torch.dtype,
211 per_act_token: bool,
212 per_out_ch: bool,
213 + use_batched_format: bool = False,
214 ):
215 super().__init__()
216 self.max_experts_per_worker = max_experts_per_worker
217 self.out_dtype = out_dtype
218 self.per_act_token = per_act_token
219 self.per_out_ch = per_out_ch
220 + self.use_batched_format = use_batched_format
221 +
222 + def supports_chunking(self) -> bool:
223 + return not self.use_batched_format
224
225 def workspace_shapes(
226 self,
@@ -222,14 +231,24 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
231 K: int,
232 topk: int,
233 num_experts: int,
225 - ) -> tuple[int, int, torch.dtype]:
226 - padded_M = aq.shape[1]
227 - workspace1 = self.max_experts_per_worker * padded_M * max(N, K)
228 - workspace2 = self.max_experts_per_worker * padded_M * (N // 2)
229 - return (workspace1, workspace2, self.out_dtype)
234 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
235 + workspace1: tuple[int, ...] = ()
236 + workspace2: tuple[int, ...] = ()
237 + output: tuple[int, ...] = ()
238 + if self.use_batched_format:
239 + padded_M = aq.shape[1]
240 + workspace1 = (self.max_experts_per_worker, padded_M, max(N, K))
241 + workspace2 = (self.max_experts_per_worker, padded_M, (N // 2))
242 + output = (self.max_experts_per_worker, padded_M, K)
243 + else:
244 + workspace1 = (M * topk, max(2 * N, K))
245 + workspace2 = (M * topk, N)
246 + output = (M * topk, K)
247 + return (workspace1, workspace2, output, self.out_dtype)
248
249 def apply(
250 self,
251 + output: torch.Tensor,
252 hidden_states: torch.Tensor,
253 w1: torch.Tensor,
254 w2: torch.Tensor,
@@ -246,16 +265,17 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
265 workspace13: torch.Tensor,
266 workspace2: torch.Tensor,
267 expert_num_tokens: Optional[torch.Tensor],
249 - ) -> torch.Tensor:
268 + ):
269 assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
270 assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
271 activation_callable = lambda i, o: self.activation(activation, i, o)
253 - return run_cutlass_moe_fp8(hidden_states, w1, w2, topk_ids,
254 - activation_callable, global_num_experts,
255 - expert_map, w1_scale, w2_scale, a1q_scale,
256 - a2_scale, workspace13, workspace2,
257 - expert_num_tokens, self.out_dtype,
258 - self.per_act_token, self.per_out_ch)
272 + run_cutlass_moe_fp8(output, hidden_states, w1, w2, topk_ids,
273 + activation_callable, global_num_experts,
274 + expert_map, w1_scale, w2_scale, a1q_scale,
275 + a2_scale, workspace13, workspace2,
276 + expert_num_tokens, self.out_dtype,
277 + self.per_act_token, self.per_out_ch,
278 + self.use_batched_format)
279
280
281 def cutlass_moe_fp8(
@@ -325,6 +345,7 @@ def cutlass_moe_fp8(
345 out_dtype=out_dtype,
346 per_act_token=per_act_token,
347 per_out_ch=per_out_ch,
348 + use_batched_format=False,
349 ),
350 )
351
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+11 -11
@@ -70,6 +70,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
70 super().__init__()
71 self.block_shape = deep_gemm_block_shape()
72
73 + def supports_chunking(self) -> bool:
74 + return True
75 +
76 def workspace_shapes(
77 self,
78 a: torch.Tensor,
@@ -79,18 +82,18 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
82 K: int,
83 topk: int,
84 num_experts: int,
82 - ) -> tuple[int, int, torch.dtype]:
83 -
85 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
86 block_m = self.block_shape[0]
87 M_sum = (M * topk) + num_experts * (block_m - 1)
88 M_sum = round_up(M_sum, block_m)
87 - workspace1 = M_sum * max(N * 2, K)
88 - workspace2 = M_sum * max(N, K)
89 -
90 - return (workspace1, workspace2, a.dtype)
89 + workspace1 = (M_sum, max(N * 2, K))
90 + workspace2 = (M_sum, max(N, K))
91 + output = (M * topk, K)
92 + return (workspace1, workspace2, output, a.dtype)
93
94 def apply(
95 self,
96 + output: torch.Tensor,
97 hidden_states: torch.Tensor,
98 w1: torch.Tensor,
99 w2: torch.Tensor,
@@ -107,7 +110,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
110 workspace13: torch.Tensor,
111 workspace2: torch.Tensor,
112 expert_num_tokens: Optional[torch.Tensor],
110 - ) -> torch.Tensor:
113 + ):
114 import deep_gemm as dg
115
116 a1q = hidden_states
@@ -143,7 +146,6 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
146 quant_out = _resize_cache(workspace13.view(dtype=torch.float8_e4m3fn),
147 (M_sum, N // 2))
148 mm2_out = _resize_cache(workspace2, (M_sum, K))
146 - out = _resize_cache(workspace13, (inv_perm.size(0), K))
149
150 dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
151 (a1q, a1q_scale), (w1, w1_scale), mm1_out, expert_ids)
@@ -159,9 +161,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
161 dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
162 (a2q, a2q_scale), (w2, w2_scale), mm2_out, expert_ids)
163
162 - torch.index_select(mm2_out, 0, inv_perm, out=out)
163 -
164 - return out
164 + torch.index_select(mm2_out, 0, inv_perm, out=output)
165
166
167 def deep_gemm_moe_fp8(
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+29 -44
@@ -335,9 +335,6 @@ def invoke_moe_batched_triton_kernel(
335 BLOCK_M = config['BLOCK_SIZE_M']
336 BLOCK_N = config['BLOCK_SIZE_N']
337 BLOCK_K = config['BLOCK_SIZE_K']
338 - assert (torch.compiler.is_compiling()
339 - or torch.cuda.is_current_stream_capturing()
340 - or max_num_tokens % BLOCK_M == 0)
338
339 grid = (expert_num_tokens.size(0), triton.cdiv(max_num_tokens, BLOCK_M) *
340 triton.cdiv(B.size(1), BLOCK_N))
@@ -390,8 +387,8 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
387 that the PPLX dispatch/combine kernels use.
388 """
389
393 - def __init__(self, max_num_tokens: Optional[int], world_size: int,
394 - dp_size: int, rank: int):
390 + def __init__(self, max_num_tokens: int, world_size: int, dp_size: int,
391 + rank: int):
392 super().__init__()
393 self.world_size = world_size
394 self.dp_size = dp_size
@@ -430,14 +427,9 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
427 num_tokens, hidden_dim = a1.size()
428 topk = topk_ids.size(1)
429
433 - if self.max_num_tokens is None:
434 - tokens_per_expert = torch.bincount(topk_ids.view(-1),
435 - minlength=num_experts)
436 - self.max_num_tokens = int(tokens_per_expert.max().item())
437 - else:
438 - tokens_per_expert = torch.zeros(num_experts,
439 - dtype=torch.int,
440 - device=a1.device)
430 + tokens_per_expert = torch.zeros(num_experts,
431 + dtype=torch.int,
432 + device=a1.device)
433
434 assert num_experts % self.world_size == 0
435
@@ -497,9 +489,9 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
489
490 def __init__(
491 self,
492 + max_num_tokens: int,
493 world_size: int,
494 dp_size: int,
502 - max_num_tokens: Optional[int] = None,
495 use_fp8_w8a8: bool = False,
496 use_int8_w8a8: bool = False,
497 use_int8_w8a16: bool = False,
@@ -518,6 +510,9 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
510 self.world_size = world_size
511 self.dp_size = dp_size
512
513 + def supports_chunking(self) -> bool:
514 + return False
515 +
516 def workspace_shapes(
517 self,
518 a: torch.Tensor,
@@ -527,18 +522,16 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
522 K: int,
523 topk: int,
524 num_experts: int,
530 - ) -> tuple[int, int, torch.dtype]:
525 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
526 assert a.dim() == 2
527 num_dp = self.world_size // self.dp_size
533 - max_num_tokens = a.size(
534 - 0) if self.max_num_tokens is None else self.max_num_tokens
535 - #print(f"WORKSPACE {max_num_tokens} {num_dp}")
536 - workspace13 = num_experts * max_num_tokens * num_dp * K
537 - workspace2 = max_num_tokens * num_dp * N
538 - return (workspace13, workspace2, a.dtype)
528 + workspace13 = (num_experts, self.max_num_tokens * num_dp, K)
529 + workspace2 = (self.max_num_tokens * num_dp, N)
530 + return (workspace13, workspace2, workspace13, a.dtype)
531
532 def apply(
533 self,
534 + output: torch.Tensor,
535 hidden_states: torch.Tensor,
536 w1: torch.Tensor,
537 w2: torch.Tensor,
@@ -555,20 +548,12 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
548 workspace13: torch.Tensor,
549 workspace2: torch.Tensor,
550 expert_num_tokens: Optional[torch.Tensor],
558 - ) -> torch.Tensor:
551 + ):
552 assert hidden_states.dim() == 3
553 assert expert_num_tokens is not None
561 - hidden_dim = hidden_states.size(-1)
562 -
563 - if self.max_num_tokens is None:
564 - max_num_tokens = hidden_states.size(1)
565 - else:
566 - max_num_tokens = self.max_num_tokens
554
555 + max_num_tokens = self.max_num_tokens
556 num_dp = self.world_size // self.dp_size
569 - num_experts = global_num_experts
570 - out = _resize_cache(workspace13,
571 - (num_experts, max_num_tokens * num_dp, hidden_dim))
557 num_local_experts = w1.size(0)
558 assert num_local_experts == w1.size(0), (
559 f"{num_local_experts} == {w1.size(0)}")
@@ -585,15 +570,13 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
570 # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor
571 if (torch.compiler.is_compiling()
572 or torch.cuda.is_current_stream_capturing()):
588 - num = max_num_tokens * num_dp
573 + num = hidden_states.shape[1]
574 else:
575 num = int(expert_num_tokens[expert].item())
576 tmp = _resize_cache(workspace2, (num, N))
577 input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
578 self.activation(activation, tmp, input)
594 - out[expert, :num, :] = tmp @ w2[expert].transpose(0, 1)
595 -
596 - return out
579 + output[expert, :num, :] = tmp @ w2[expert].transpose(0, 1)
580
581
582 class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -630,6 +613,9 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
613 assert not use_int4_w4a16, "NYI"
614 assert self.block_shape is None, "NYI"
615
616 + def supports_chunking(self) -> bool:
617 + return False
618 +
619 def workspace_shapes(
620 self,
621 a: torch.Tensor,
@@ -639,17 +625,19 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
625 K: int,
626 topk: int,
627 num_experts: int,
642 - ) -> tuple[int, int, torch.dtype]:
628 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
629 assert a.dim() == 2
630 num_dp = self.world_size // self.dp_size
631 max_num_tokens = a.size(
632 0) if self.max_num_tokens is None else self.max_num_tokens
647 - workspace13 = num_experts * max_num_tokens * num_dp * max(K, N)
648 - workspace2 = num_experts * max_num_tokens * num_dp * (N // 2)
649 - return (workspace13, workspace2, a.dtype)
633 + workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
634 + workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
635 + output = (num_experts, max_num_tokens * num_dp, K)
636 + return (workspace13, workspace2, output, a.dtype)
637
638 def apply(
639 self,
640 + output: torch.Tensor,
641 hidden_states: torch.Tensor,
642 w1: torch.Tensor,
643 w2: torch.Tensor,
@@ -666,7 +654,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
654 workspace13: torch.Tensor,
655 workspace2: torch.Tensor,
656 expert_num_tokens: Optional[torch.Tensor],
669 - ) -> torch.Tensor:
657 + ):
658 # Check constraints.
659 if self.use_int4_w4a16:
660 assert hidden_states.size(-1) // 2 == w1.size(2), (
@@ -723,8 +711,6 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
711 (E, max_num_tokens, N))
712 intermediate_cache2 = _resize_cache(workspace2,
713 (E, max_num_tokens, N // 2))
726 - intermediate_cache3 = _resize_cache(workspace13,
727 - (E, max_num_tokens, K))
714
715 # MM1
716 invoke_moe_batched_triton_kernel(A=hidden_states,
@@ -761,7 +747,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
747
748 invoke_moe_batched_triton_kernel(A=qintermediate_cache2,
749 B=w2,
764 - C=intermediate_cache3,
750 + C=output,
751 expert_num_tokens=expert_num_tokens,
752 compute_type=compute_type,
753 A_scale=a2q_scale,
@@ -772,4 +758,3 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
758 use_int4_w4a16=self.use_int4_w4a16,
759 config=config,
760 block_shape=self.block_shape)
775 - return intermediate_cache3
vllm/model_executor/layers/fused_moe/fused_moe.py
+11 -11
@@ -1542,6 +1542,9 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1542 use_int4_w4a16=use_int4_w4a16)
1543 self.per_channel_quant = per_channel_quant
1544
1545 + def supports_chunking(self) -> bool:
1546 + return True
1547 +
1548 def workspace_shapes(
1549 self,
1550 a: torch.Tensor,
@@ -1551,14 +1554,15 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1554 K: int,
1555 topk: int,
1556 num_experts: int,
1554 - ) -> tuple[int, int, torch.dtype]:
1555 - factor = num_experts if a.dim() == 3 else 1
1556 - workspace1 = M * topk * max(N * 2, K) * factor
1557 - workspace2 = M * topk * N * factor
1558 - return (workspace1, workspace2, a.dtype)
1557 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
1558 + workspace1 = (M, topk, max(N * 2, K))
1559 + workspace2 = (M, topk, N)
1560 + output = (M, topk, K)
1561 + return (workspace1, workspace2, output, a.dtype)
1562
1563 def apply(
1564 self,
1565 + output: torch.Tensor,
1566 hidden_states: torch.Tensor,
1567 w1: torch.Tensor,
1568 w2: torch.Tensor,
@@ -1575,7 +1579,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1579 workspace13: torch.Tensor,
1580 workspace2: torch.Tensor,
1581 expert_num_tokens: Optional[torch.Tensor],
1578 - ) -> torch.Tensor:
1582 + ):
1583 # Check constraints.
1584 if self.use_int4_w4a16:
1585 assert hidden_states.size(-1) // 2 == w1.size(2), (
@@ -1632,8 +1636,6 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1636 (num_tokens, top_k_num, N))
1637 intermediate_cache2 = _resize_cache(workspace2,
1638 (num_tokens * top_k_num, N // 2))
1635 - intermediate_cache3 = _resize_cache(workspace13,
1636 - (num_tokens, top_k_num, K))
1639
1640 sorted_token_ids, expert_ids, num_tokens_post_padded = (
1641 moe_align_block_size(topk_ids, config['BLOCK_SIZE_M'],
@@ -1671,7 +1673,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1673
1674 invoke_fused_moe_kernel(qintermediate_cache2,
1675 w2,
1674 - intermediate_cache3,
1676 + output,
1677 a2q_scale,
1678 w2_scale,
1679 w2_zp,
@@ -1690,8 +1692,6 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1692 per_channel_quant=self.per_channel_quant,
1693 block_shape=self.block_shape)
1694
1693 - return intermediate_cache3
1694 -
1695
1696 def modular_triton_fused_moe(
1697 use_fp8_w8a8: bool,
vllm/model_executor/layers/fused_moe/modular_kernel.py
+143 -85
@@ -1,10 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from abc import ABC, abstractmethod
4 +from math import prod
5 from typing import Optional
6
7 import torch
8
9 +import vllm.envs as envs
10 +from vllm.model_executor.layers.fused_moe.utils import _resize_cache
11 +from vllm.utils import cdiv
12 +
13 #
14 # This file defines a set of base classes used to make MoE kernels more modular.
15 # The goal is to be able to utilize different communication mechanisms with
@@ -115,9 +120,9 @@ class FusedMoEPrepareAndFinalize(ABC):
120 - quantized + dispatched a.
121 - quantized + dispatched a1_scales.
122 - Optional tensor as big as number of local experts that contains the
118 - number of tokens assigned to each local expert.
123 + number of tokens assigned to each local expert.
124 - Optional dispatched expert topk IDs
120 - - Optional dispatched expert topk weight
125 + - Optional dispatched expert topk weight
126 """
127 raise NotImplementedError
128
@@ -159,7 +164,7 @@ class FusedMoEPrepareAndFinalize(ABC):
164 Some PrepareFinalize All2All implementations are batched. Meaning,
165 they can processes only as set of tokens at a time. This
166 function returns the batch size i.e the maximum number of tokens
162 - the implementation can process at a time.
167 + the implementation can process at a time.
168 Return None if there are no such restrictions.
169 """
170 raise NotImplementedError
@@ -171,6 +176,15 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
176 above.
177 """
178
179 + # TODO (bnell): make this return a CHUNK_SIZE or None instead?
180 + @abstractmethod
181 + def supports_chunking(self) -> bool:
182 + """
183 + A flag indicating whether or not this class supports activation
184 + chunking.
185 + """
186 + raise NotImplementedError
187 +
188 @abstractmethod
189 def workspace_shapes(
190 self,
@@ -181,19 +195,22 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
195 K: int,
196 topk: int,
197 num_experts: int,
184 - ) -> tuple[int, int, torch.dtype]:
198 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
199 """
186 - Compute the number of elements for the temporary outputs of the two
187 - gemms and activation in the fused expert function. Since the
188 - gemms are independent, the workspace for the first gemm can be shared
189 - with the workspace for the last gemm.
200 + Compute the shapes for the temporary and final outputs of the two gemms
201 + and activation in the fused expert function. Since the gemms are
202 + independent, the workspace for the first gemm can be shared with the
203 + workspace for the last gemm.
204
205 Returns a tuple of:
192 - - Number of workspace13 elements: must be large enough to hold the
206 + - workspace13 shape tuple: must be large enough to hold the
207 result of either expert gemm.
194 - - Number of workspace2 elements: must be large enough to hold the
208 + - workspace2 shape tuple: must be large enough to hold the
209 result of the activation function.
210 + - output shape tuple: must be exact size of the final gemm output.
211 - Workspace type: The dtype to use for the workspace tensors.
212 + - Note: in order for activation chunking to work, the first dimension
213 + of each tuple must be the number of tokens.
214 """
215 raise NotImplementedError
216
@@ -210,6 +227,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
227 @abstractmethod
228 def apply(
229 self,
230 + output: torch.Tensor,
231 hidden_states: torch.Tensor,
232 w1: torch.Tensor,
233 w2: torch.Tensor,
@@ -226,12 +244,13 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
244 workspace13: torch.Tensor,
245 workspace2: torch.Tensor,
246 expert_num_tokens: Optional[torch.Tensor],
229 - ) -> torch.Tensor:
247 + ):
248 """
249 This function computes the intermediate result of a Mixture of Experts
250 (MoE) layer using two sets of weights, w1 and w2.
251
252 Parameters:
253 + - output: (torch.Tensor): The unweighted, unreduced output tensor.
254 - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
255 layer.
256 - w1 (torch.Tensor): The first set of expert weights.
@@ -259,13 +278,20 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
278 function.
279 - expert_num_tokens: An optional tensor containing the number of tokens
280 assigned to each expert when using batched experts format input.
262 -
263 - Returns:
264 - - torch.Tensor: The unweighted, unreduced output tensor
281 """
282 raise NotImplementedError
283
284
285 +def _chunk_scales(scales: Optional[torch.Tensor], start: int,
286 + end: int) -> Optional[torch.Tensor]:
287 + if scales is not None:
288 + if scales.numel() == 1:
289 + return scales
290 + else:
291 + return scales[start:end]
292 + return None
293 +
294 +
295 class FusedMoEModularKernel(torch.nn.Module):
296 """
297 This class combines a FusedMoEPrepareAndFinalize instance and
@@ -288,61 +314,6 @@ class FusedMoEModularKernel(torch.nn.Module):
314 self.prepare_finalize = prepare_finalize
315 self.fused_experts = fused_experts
316
291 - def _do_fused_experts(
292 - self,
293 - a1: torch.Tensor, # input to forward fn
294 - a1q: torch.Tensor, # output of prepare fn
295 - w1: torch.Tensor,
296 - w2: torch.Tensor,
297 - topk_ids: torch.Tensor,
298 - expert_num_tokens: torch.Tensor,
299 - activation: str,
300 - global_num_experts: int,
301 - expert_map: Optional[torch.Tensor],
302 - w1_scale: Optional[torch.Tensor],
303 - w2_scale: Optional[torch.Tensor],
304 - w1_zp: Optional[torch.Tensor],
305 - w2_zp: Optional[torch.Tensor],
306 - a1q_scale: Optional[torch.Tensor],
307 - a2_scale: Optional[torch.Tensor]) -> torch.Tensor:
308 -
309 - _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
310 -
311 - # Use a1 here to decipher the correct workspace datatype
312 - workspace13_shape, workspace2_shape, workspace_dtype = (
313 - self.fused_experts.workspace_shapes(a1, a1q, M, N, K, top_k,
314 - global_num_experts))
315 -
316 - # We can reuse the memory between cache1 and cache3 because by the time
317 - # we need cache3, we're done with cache1
318 - workspace13 = torch.zeros(workspace13_shape,
319 - device=a1.device,
320 - dtype=workspace_dtype)
321 - workspace2 = torch.zeros(workspace2_shape,
322 - device=a1.device,
323 - dtype=workspace_dtype)
324 -
325 - fused_out = self.fused_experts.apply(
326 - a1q,
327 - w1,
328 - w2,
329 - topk_ids,
330 - activation=activation,
331 - global_num_experts=global_num_experts,
332 - expert_map=expert_map,
333 - w1_scale=w1_scale,
334 - w2_scale=w2_scale,
335 - w1_zp=w1_zp,
336 - w2_zp=w2_zp,
337 - a1q_scale=a1q_scale,
338 - a2_scale=a2_scale,
339 - workspace13=workspace13,
340 - workspace2=workspace2,
341 - expert_num_tokens=expert_num_tokens,
342 - )
343 -
344 - return fused_out
345 -
317 def forward(
318 self,
319 hidden_states: torch.Tensor,
@@ -408,12 +379,14 @@ class FusedMoEModularKernel(torch.nn.Module):
379 _expert_topk_weights) = self.prepare_finalize.prepare(
380 a1, a1_scale, a2_scale, topk_weights, topk_ids,
381 global_num_experts, expert_map, apply_router_weight_on_input)
382 +
383 # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
384 topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
385 topk_weights = (topk_weights if _expert_topk_weights is None else
386 _expert_topk_weights)
387
388 fused_out = None
389 +
390 if a1q.numel() == 0:
391 # This happens when none of the tokens from the all2all reach this
392 # EP rank. Also, note that this is only relevant for CUDAGraph
@@ -423,22 +396,107 @@ class FusedMoEModularKernel(torch.nn.Module):
396 # and can never run into the tensor.numel() == 0 case.
397 fused_out = torch.empty_like(a1q).to(dtype=a1.dtype)
398 else:
426 - fused_out = self._do_fused_experts(
427 - a1=a1,
428 - a1q=a1q,
429 - w1=w1,
430 - w2=w2,
431 - topk_ids=topk_ids,
432 - expert_num_tokens=expert_num_tokens,
433 - activation=activation,
434 - global_num_experts=global_num_experts,
435 - expert_map=expert_map,
436 - w1_scale=w1_scale,
437 - w2_scale=w2_scale,
438 - w1_zp=w1_zp,
439 - w2_zp=w2_zp,
440 - a1q_scale=a1q_scale,
441 - a2_scale=a2_scale)
399 + _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
400 +
401 + if self.fused_experts.supports_chunking():
402 + CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
403 + num_chunks = cdiv(M, CHUNK_SIZE)
404 + else:
405 + CHUNK_SIZE = M
406 + num_chunks = 1
407 +
408 + if num_chunks == 1:
409 + (workspace13_shape, workspace2_shape, fused_out_shape,
410 + workspace_dtype) = self.fused_experts.workspace_shapes(
411 + a1, a1q, M, N, K, top_k, global_num_experts)
412 + else:
413 + # Use the full M to get the final output shape.
414 + _, _, fused_out_shape, _ = (
415 + self.fused_experts.workspace_shapes(
416 + a1, a1q, M, N, K, top_k, global_num_experts))
417 + # Use the CHUNK_SIZE to get the workspace shapes.
418 + workspace13_shape, workspace2_shape, _, workspace_dtype = (
419 + self.fused_experts.workspace_shapes(
420 + a1, a1q, CHUNK_SIZE, N, K, top_k, global_num_experts))
421 +
422 + # We can reuse the memory between cache1 and cache3 because by the
423 + # time we need cache3, we're done with cache1.
424 + workspace13 = torch.zeros(prod(workspace13_shape),
425 + device=a1.device,
426 + dtype=workspace_dtype)
427 + workspace2 = torch.zeros(prod(workspace2_shape),
428 + device=a1.device,
429 + dtype=workspace_dtype)
430 +
431 + if num_chunks == 1:
432 + fused_out = _resize_cache(workspace13, fused_out_shape)
433 +
434 + self.fused_experts.apply(
435 + fused_out,
436 + a1q,
437 + w1,
438 + w2,
439 + topk_ids,
440 + activation=activation,
441 + global_num_experts=global_num_experts,
442 + expert_map=expert_map,
443 + w1_scale=w1_scale,
444 + w2_scale=w2_scale,
445 + w1_zp=w1_zp,
446 + w2_zp=w2_zp,
447 + a1q_scale=a1q_scale,
448 + a2_scale=a2_scale,
449 + workspace13=workspace13,
450 + workspace2=workspace2,
451 + expert_num_tokens=expert_num_tokens,
452 + )
453 + else:
454 + # The leading output dimension may not be equal to M, so
455 + # we compute output indices separately.
456 + M_out = fused_out_shape[0]
457 + assert M_out >= M
458 + factor = M_out // M
459 + assert factor > 0
460 + OUT_CHUNK_SIZE = CHUNK_SIZE * factor
461 +
462 + fused_out = torch.empty(fused_out_shape,
463 + device=a1q.device,
464 + dtype=workspace_dtype)
465 +
466 + assert cdiv(M_out, OUT_CHUNK_SIZE) == num_chunks, (
467 + f"{cdiv(M_out, OUT_CHUNK_SIZE)} == {num_chunks}")
468 +
469 + for chunk in range(num_chunks):
470 + begin_chunk_idx = chunk * CHUNK_SIZE
471 + end_chunk_idx = min((chunk + 1) * CHUNK_SIZE, M)
472 + begin_out_idx = chunk * OUT_CHUNK_SIZE
473 + end_out_idx = min((chunk + 1) * OUT_CHUNK_SIZE, M_out)
474 + curr_a1q = a1q[begin_chunk_idx:end_chunk_idx]
475 + curr_a1q_scale = _chunk_scales(a1q_scale, begin_chunk_idx,
476 + end_chunk_idx)
477 + curr_a2_scale = _chunk_scales(a2_scale, begin_chunk_idx,
478 + end_chunk_idx)
479 + curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
480 +
481 + self.fused_experts.apply(
482 + fused_out[begin_out_idx:end_out_idx],
483 + curr_a1q,
484 + w1,
485 + w2,
486 + curr_topk_ids,
487 + activation=activation,
488 + global_num_experts=global_num_experts,
489 + expert_map=expert_map,
490 + w1_scale=w1_scale,
491 + w2_scale=w2_scale,
492 + w1_zp=w1_zp,
493 + w2_zp=w2_zp,
494 + a1q_scale=curr_a1q_scale,
495 + a2_scale=curr_a2_scale,
496 + workspace13=workspace13,
497 + workspace2=workspace2,
498 + expert_num_tokens=expert_num_tokens,
499 + )
500
501 self.prepare_finalize.finalize(output, fused_out, topk_weights,
502 topk_ids, apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+35 -42
@@ -34,6 +34,12 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
34 self.deep_gemm_expert = DeepGemmExperts(
35 ) if self.allow_deep_gemm else None
36
37 + def supports_chunking(self) -> bool:
38 + dge = self.deep_gemm_expert
39 + te = self.triton_expert
40 + return ((dge is None or dge.supports_chunking())
41 + and (te is None or te.supports_chunking()))
42 +
43 def workspace_shapes(
44 self,
45 a: torch.Tensor,
@@ -43,7 +49,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
49 K: int,
50 topk: int,
51 num_experts: int,
46 - ) -> tuple[int, int, torch.dtype]:
52 + ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
53 # Note: the deep gemm workspaces are strictly larger than the triton
54 # workspaces so we can be pessimistic here and allocate for DeepGemm
55 # even if we fall back to triton later, e.g. if expert maps are set.
@@ -57,6 +63,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
63
64 def apply(
65 self,
66 + output: torch.Tensor,
67 hidden_states: torch.Tensor,
68 w1: torch.Tensor,
69 w2: torch.Tensor,
@@ -73,45 +80,31 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
80 workspace13: torch.Tensor,
81 workspace2: torch.Tensor,
82 expert_num_tokens: Optional[torch.Tensor],
76 - ) -> torch.Tensor:
83 + ):
84 N = w1.size(1)
78 - if (self.allow_deep_gemm and self.use_fp8_w8a8 and N > 512
79 - and _valid_deep_gemm(hidden_states, w1, w2)):
80 - assert self.deep_gemm_expert is not None
81 - return self.deep_gemm_expert.apply(
82 - hidden_states,
83 - w1,
84 - w2,
85 - topk_ids,
86 - activation,
87 - global_num_experts,
88 - expert_map,
89 - w1_scale,
90 - w2_scale,
91 - w1_zp,
92 - w2_zp,
93 - a1q_scale,
94 - a2_scale,
95 - workspace13,
96 - workspace2,
97 - expert_num_tokens,
98 - )
99 - else:
100 - return self.triton_expert.apply(
101 - hidden_states,
102 - w1,
103 - w2,
104 - topk_ids,
105 - activation,
106 - global_num_experts,
107 - expert_map,
108 - w1_scale,
109 - w2_scale,
110 - w1_zp,
111 - w2_zp,
112 - a1q_scale,
113 - a2_scale,
114 - workspace13,
115 - workspace2,
116 - expert_num_tokens,
117 - )
85 +
86 + use_deep_gemm = (self.allow_deep_gemm and self.use_fp8_w8a8 and N > 512
87 + and _valid_deep_gemm(hidden_states, w1, w2))
88 +
89 + experts = self.deep_gemm_expert if use_deep_gemm else self.triton_expert
90 + assert experts is not None
91 +
92 + experts.apply(
93 + output,
94 + hidden_states,
95 + w1,
96 + w2,
97 + topk_ids,
98 + activation,
99 + global_num_experts,
100 + expert_map,
101 + w1_scale,
102 + w2_scale,
103 + w1_zp,
104 + w2_zp,
105 + a1q_scale,
106 + a2_scale,
107 + workspace13,
108 + workspace2,
109 + expert_num_tokens,
110 + )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+5 -2
@@ -562,9 +562,12 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
562 (moe.num_experts + prepare_finalize.world_size - 1) //
563 prepare_finalize.world_size)
564 experts = CutlassExpertsFp8(
565 - max_experts_per_worker, moe.in_dtype,
565 + max_experts_per_worker,
566 + moe.in_dtype,
567 self.input_quant.strategy == QuantizationStrategy.TOKEN,
567 - self.weight_quant.strategy == QuantizationStrategy.CHANNEL)
568 + self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
569 + use_batched_format=True,
570 + )
571
572 if has_pplx and isinstance(
573 prepare_finalize,