@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 78fe77534b
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, 03:30 UTC

[Kernel] Enable fp8 support for pplx and BatchedTritonExperts. (#18864)

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

bnellnm committed Jul 3, 2025 at 17:55 UTC 78fe77534bd87f5356881f12d2a4ecbd8db29c41
25 files changed +1275 -661
tests/kernels/moe/parallel_utils.py
+3 -7
@@ -137,8 +137,7 @@ def make_deepep_ht_a2a(pg: ProcessGroup,
137 low_latency_mode=low_latency_mode,
138 num_qps_per_rank=num_qps_per_rank)
139 return DeepEPHTPrepareAndFinalize(buffer=buffer,
140 - world_size=pgi.world_size,
141 - rank=pgi.rank,
140 + num_dispatchers=pgi.world_size,
141 dp_size=dp_size,
142 rank_expert_offset=pgi.rank *
143 ht_args.num_local_experts)
@@ -146,7 +145,6 @@ def make_deepep_ht_a2a(pg: ProcessGroup,
145
146 def make_deepep_ll_a2a(pg: ProcessGroup,
147 pgi: ProcessGroupInfo,
149 - dp_size: int,
148 deepep_ll_args: DeepEPLLArgs,
149 q_dtype: Optional[torch.dtype] = None,
150 block_shape: Optional[list[int]] = None):
@@ -166,8 +164,7 @@ def make_deepep_ll_a2a(pg: ProcessGroup,
164
165 return DeepEPLLPrepareAndFinalize(
166 buffer=buffer,
169 - world_size=pgi.world_size,
170 - dp_size=dp_size,
167 + num_dispatchers=pgi.world_size,
168 max_tokens_per_rank=deepep_ll_args.max_tokens_per_rank,
169 use_fp8_dispatch=deepep_ll_args.use_fp8_dispatch,
170 )
@@ -186,5 +183,4 @@ def make_deepep_a2a(pg: ProcessGroup,
183 block_shape)
184
185 assert deepep_ll_args is not None
189 - return make_deepep_ll_a2a(pg, pgi, dp_size, deepep_ll_args, q_dtype,
190 - block_shape)
186 + return make_deepep_ll_a2a(pg, pgi, deepep_ll_args, q_dtype, block_shape)
tests/kernels/moe/test_batched_moe.py
+54 -30
@@ -10,7 +10,7 @@ import triton.language as tl
10
11 from tests.kernels.moe.utils import (batched_moe,
12 make_quantized_test_activations,
13 - make_test_weights, triton_moe)
13 + make_test_weights, naive_batched_moe)
14 from tests.kernels.quant_utils import native_batched_masked_quant_matmul
15 from tests.kernels.utils import torch_experts
16 from vllm.config import VllmConfig, set_current_vllm_config
@@ -33,12 +33,10 @@ MNK_FACTORS = [
33 (45, 512, 512),
34 (45, 1024, 128),
35 (45, 1024, 2048),
36 - (64, 128, 128),
36 (64, 512, 512),
37 (64, 1024, 2048),
38 (222, 128, 128),
39 (222, 128, 2048),
41 - (222, 512, 512),
40 (222, 1024, 128),
41 (222, 1024, 2048),
42 ]
@@ -95,11 +93,12 @@ class BatchedMMTensors:
93 @pytest.mark.parametrize("max_tokens_per_expert",
94 [32, 64, 128, 192, 224, 256, 512])
95 @pytest.mark.parametrize("K", [128, 256, 1024])
98 -@pytest.mark.parametrize("N", [128, 256, 512, 1024])
99 -@pytest.mark.parametrize("dtype",
100 - [torch.float32, torch.float16, torch.bfloat16])
101 -@pytest.mark.parametrize("block_shape", [None])
102 -@pytest.mark.parametrize("per_act_token_quant", [False])
96 +@pytest.mark.parametrize("N", [128, 256, 1024])
97 +@pytest.mark.parametrize(
98 + "dtype",
99 + [torch.float8_e4m3fn, torch.float32, torch.float16, torch.bfloat16])
100 +@pytest.mark.parametrize("block_shape", [None, [128, 128]])
101 +@pytest.mark.parametrize("per_act_token_quant", [False, True])
102 def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
103 N: int, dtype: torch.dtype,
104 block_shape: Optional[list[int]],
@@ -134,7 +133,8 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
133 in_dtype=act_dtype,
134 quant_dtype=quant_dtype,
135 block_shape=block_shape,
137 - per_act_token_quant=per_act_token_quant)
136 + per_act_token_quant=per_act_token_quant,
137 + )
138
139 B, B_q, B_scale, _, _, _ = make_test_weights(
140 num_experts,
@@ -143,6 +143,7 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
143 in_dtype=act_dtype,
144 quant_dtype=quant_dtype,
145 block_shape=block_shape,
146 + per_act_token_quant=per_act_token_quant,
147 )
148
149 out_shape = (num_experts, max_tokens_per_expert, N)
@@ -177,6 +178,7 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
178 "BLOCK_SIZE_N": 16,
179 "BLOCK_SIZE_K": 16 if dtype.itemsize > 1 else 32
180 },
181 + per_act_token_quant=per_act_token_quant,
182 block_shape=block_shape,
183 )
184
@@ -185,15 +187,13 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
187 B,
188 ref_output,
189 num_expert_tokens,
188 - None,
189 - None,
190 - None,
190 )
191
192 q_ref_output = native_batched_masked_quant_matmul(A_q, B_q, q_ref_output,
193 num_expert_tokens,
194 A_scale, B_scale,
196 - block_shape)
195 + block_shape,
196 + per_act_token_quant)
197
198 rtol, atol = {
199 torch.float16: (6e-2, 6e-2),
@@ -201,16 +201,17 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
201 torch.float32: (1e-2, 1e-2),
202 }[test_output.dtype]
203
204 - torch.testing.assert_close(ref_output, test_output, atol=atol, rtol=rtol)
204 + torch.testing.assert_close(ref_output, q_ref_output, atol=atol, rtol=rtol)
205 torch.testing.assert_close(test_output, q_ref_output, atol=atol, rtol=rtol)
206
207
208 @pytest.mark.parametrize(("m", "n", "k"), MNK_FACTORS)
209 @pytest.mark.parametrize("e", NUM_EXPERTS)
210 @pytest.mark.parametrize("topk", TOP_KS)
211 -@pytest.mark.parametrize("dtype", [torch.bfloat16])
212 -@pytest.mark.parametrize("per_act_token_quant", [False])
213 -@pytest.mark.parametrize("block_shape", [None])
211 +@pytest.mark.parametrize("dtype", [torch.float8_e4m3fn, torch.bfloat16])
212 +@pytest.mark.parametrize("per_act_token_quant", [False, True])
213 +@pytest.mark.parametrize("block_shape", [None, [128, 128]])
214 +@pytest.mark.parametrize("input_scales", [False])
215 def test_fused_moe_batched_experts(
216 m: int,
217 n: int,
@@ -220,15 +221,19 @@ def test_fused_moe_batched_experts(
221 dtype: torch.dtype,
222 per_act_token_quant: bool,
223 block_shape: Optional[list[int]],
224 + input_scales: bool,
225 ):
226 current_platform.seed_everything(7)
227
228 use_fp8_w8a8 = dtype == torch.float8_e4m3fn
229
230 + if topk > e:
231 + pytest.skip("topk > e")
232 +
233 if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
234 pytest.skip("Skip quantization test for non-quantized type")
235
231 - if per_act_token_quant and block_shape is not None or topk > e:
236 + if per_act_token_quant and block_shape is not None:
237 pytest.skip("Skip illegal quantization test.")
238
239 a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
@@ -241,16 +246,27 @@ def test_fused_moe_batched_experts(
246 act_dtype = dtype
247 quant_dtype = None
248
244 - _, w1, w1_s, _, w2, w2_s = make_test_weights(e,
245 - n,
246 - k,
247 - block_shape=block_shape,
248 - in_dtype=act_dtype,
249 - quant_dtype=quant_dtype)
249 + w1_16, w1, w1_s, w2_16, w2, w2_s = make_test_weights(
250 + e,
251 + n,
252 + k,
253 + block_shape=block_shape,
254 + in_dtype=act_dtype,
255 + quant_dtype=quant_dtype,
256 + per_act_token_quant=per_act_token_quant,
257 + )
258 +
259 + if input_scales and quant_dtype is not None:
260 + a1_scale = torch.tensor(1, device="cuda", dtype=torch.float32)
261 + a2_scale = torch.tensor(1, device="cuda", dtype=torch.float32)
262 + else:
263 + a1_scale = None
264 + a2_scale = None
265
266 with set_current_vllm_config(vllm_config):
267 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
253 - batched_output = batched_moe(
268 +
269 + baseline_output = torch_experts(
270 a,
271 w1,
272 w2,
@@ -258,11 +274,14 @@ def test_fused_moe_batched_experts(
274 topk_ids,
275 w1_scale=w1_s,
276 w2_scale=w2_s,
277 + a1_scale=a1_scale,
278 + a2_scale=a2_scale,
279 quant_dtype=quant_dtype,
280 per_act_token_quant=per_act_token_quant,
281 block_shape=block_shape,
282 )
265 - baseline_output = torch_experts(
283 +
284 + batched_output = naive_batched_moe(
285 a,
286 w1,
287 w2,
@@ -270,11 +289,14 @@ def test_fused_moe_batched_experts(
289 topk_ids,
290 w1_scale=w1_s,
291 w2_scale=w2_s,
292 + a1_scale=a1_scale,
293 + a2_scale=a2_scale,
294 quant_dtype=quant_dtype,
295 per_act_token_quant=per_act_token_quant,
275 - block_shape=block_shape)
296 + block_shape=block_shape,
297 + )
298
277 - triton_output = triton_moe(
299 + triton_output = batched_moe(
300 a,
301 w1,
302 w2,
@@ -282,14 +304,16 @@ def test_fused_moe_batched_experts(
304 topk_ids,
305 w1_scale=w1_s,
306 w2_scale=w2_s,
307 + a1_scale=a1_scale,
308 + a2_scale=a2_scale,
309 quant_dtype=quant_dtype,
310 per_act_token_quant=per_act_token_quant,
311 block_shape=block_shape,
312 )
313
290 - torch.testing.assert_close(triton_output,
314 + torch.testing.assert_close(batched_output,
315 baseline_output,
292 - atol=2e-2,
316 + atol=3e-2,
317 rtol=2e-2)
318
319 torch.testing.assert_close(triton_output,
tests/kernels/moe/test_deepep_deepgemm_moe.py
+1 -2
@@ -148,8 +148,7 @@ def make_ll_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
148
149 fused_experts = BatchedDeepGemmExperts(
150 max_num_tokens=max_tokens_per_rank,
151 - world_size=pgi.world_size,
152 - dp_size=dp_size,
151 + num_dispatchers=pgi.world_size // dp_size,
152 block_shape=test_config.block_size,
153 per_act_token_quant=test_config.per_act_token_quant)
154 mk = FusedMoEModularKernel(prepare_finalize=a2a,
tests/kernels/moe/test_deepep_moe.py
+3 -2
@@ -154,12 +154,13 @@ def make_modular_kernel(
154 deepep_ht_args = ht_args,
155 deepep_ll_args = ll_args)
156
157 + num_dispatchers = pgi.world_size // dp_size
158 +
159 if low_latency_mode:
160 assert not per_act_token_quant, "not supported in ll mode"
161 fused_experts = BatchedTritonExperts(
162 max_num_tokens=MAX_TOKENS_PER_RANK,
161 - world_size=pgi.world_size,
162 - dp_size=dp_size,
163 + num_dispatchers=num_dispatchers,
164 use_fp8_w8a8=is_quantized,
165 use_int8_w8a8=False,
166 use_int8_w8a16=False,
tests/kernels/moe/test_pplx_cutlass_moe.py
+44 -35
@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14 from vllm.model_executor.layers.fused_moe.modular_kernel import (
15 FusedMoEModularKernel)
16 from vllm.platforms import current_platform
17 +from vllm.utils import cdiv
18
19 from .parallel_utils import ProcessGroupInfo, parallel_launch
20
@@ -112,18 +113,21 @@ def pplx_cutlass_moe(
113 w2_scale = w2_scale.to(device)
114 a1_scale = a1_scale.to(device)
115
116 + assert num_experts % world_size == 0
117 + num_local_experts = cdiv(num_experts, world_size)
118 + num_dispatchers = pgi.world_size // dp_size
119 +
120 prepare_finalize = PplxPrepareAndFinalize(
121 ata,
117 - max_num_tokens,
118 - pgi.world_size,
119 - rank,
120 - dp_size,
121 - )
122 + max_num_tokens=max_num_tokens,
123 + num_local_experts=num_local_experts,
124 + num_dispatchers=num_dispatchers)
125
123 - experts = CutlassExpertsFp8((num_experts + world_size - 1) // world_size,
126 + experts = CutlassExpertsFp8(num_local_experts,
127 out_dtype,
128 per_act_token,
129 per_out_ch,
130 + num_dispatchers=num_dispatchers,
131 use_batched_format=True)
132
133 fused_cutlass_experts = FusedMoEModularKernel(
@@ -181,35 +185,40 @@ def _pplx_moe(
185 per_out_ch: bool,
186 use_internode: bool,
187 ):
184 - if use_internode:
185 - uid = nvshmem_get_unique_id(
186 - ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
187 - torch.distributed.broadcast(uid, src=0)
188 - nvshmem_init(uid, pgi.rank, pgi.world_size)
189 - else:
190 - group_ranks = list(range(pgi.world_size))
191 - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
192 - group_name = cpu_group.group_name
193 -
194 - with set_current_vllm_config(vllm_config):
195 - torch_output = torch_experts(a_full, w1_full, w2_full, topk_weights,
196 - topk_ids)
197 - pplx_output = pplx_cutlass_moe(pgi, dp_size, a, w1, w2, w1_scale,
198 - w2_scale, topk_weights, topk_ids,
199 - a1_scale, out_dtype, per_act_token,
200 - per_out_ch, group_name)
201 -
202 - torch_output = chunk_by_rank(torch_output, pgi.rank,
203 - pgi.world_size).to(pplx_output.device)
204 -
205 - # Uncomment if more debugging is needed
206 - # print("PPLX OUT:", pplx_output)
207 - # print("TORCH OUT:", torch_output)
208 -
209 - torch.testing.assert_close(pplx_output, torch_output, atol=0.05, rtol=0)
210 -
211 - if use_internode:
212 - nvshmem_finalize()
188 + try:
189 + if use_internode:
190 + uid = nvshmem_get_unique_id(
191 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
192 + torch.distributed.broadcast(uid, src=0)
193 + nvshmem_init(uid, pgi.rank, pgi.world_size)
194 + else:
195 + group_ranks = list(range(pgi.world_size))
196 + cpu_group = torch.distributed.new_group(group_ranks,
197 + backend="gloo")
198 + group_name = cpu_group.group_name
199 +
200 + with set_current_vllm_config(vllm_config):
201 + torch_output = torch_experts(a_full, w1_full, w2_full,
202 + topk_weights, topk_ids)
203 + pplx_output = pplx_cutlass_moe(pgi, dp_size, a, w1, w2, w1_scale,
204 + w2_scale, topk_weights, topk_ids,
205 + a1_scale, out_dtype, per_act_token,
206 + per_out_ch, group_name)
207 +
208 + torch_output = chunk_by_rank(torch_output, pgi.rank,
209 + pgi.world_size).to(pplx_output.device)
210 +
211 + # Uncomment if more debugging is needed
212 + # print("PPLX OUT:", pplx_output)
213 + # print("TORCH OUT:", torch_output)
214 +
215 + torch.testing.assert_close(pplx_output,
216 + torch_output,
217 + atol=0.05,
218 + rtol=0)
219 + finally:
220 + if use_internode:
221 + nvshmem_finalize()
222
223
224 @pytest.mark.parametrize("m", [2, 224])
tests/kernels/moe/test_pplx_moe.py
+488 -244
@@ -4,7 +4,10 @@
4
5 Run `pytest tests/kernels/test_pplx_moe.py`.
6 """
7 -from typing import Optional
7 +import itertools
8 +import textwrap
9 +import traceback
10 +from typing import Callable, Optional
11
12 import pytest
13 import torch
@@ -19,12 +22,13 @@ except ImportError:
22 has_pplx = False
23
24 from tests.kernels.moe.utils import make_test_weights, naive_batched_moe
25 +from tests.kernels.quant_utils import dequant
26 from tests.kernels.utils import torch_experts
27 from vllm.config import VllmConfig, set_current_vllm_config
28 from vllm.model_executor.layers.fused_moe import fused_topk, override_config
29 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
30 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
27 - BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
31 + BatchedTritonExperts)
32 from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config
33 from vllm.model_executor.layers.fused_moe.modular_kernel import (
34 FusedMoEModularKernel)
@@ -38,22 +42,22 @@ requires_pplx = pytest.mark.skipif(
42 reason="Requires PPLX kernels",
43 )
44
41 -PPLX_PREPARE_COMBOS = [(4, 128, 128), (32, 1024, 512), (64, 1024, 512),
42 - (222, 2048, 1024)]
43 -
44 -PPLX_MOE_COMBOS = [
45 - (1, 128, 128),
45 +PPLX_COMBOS = [
46 + # TODO: figure out why this fails, seems to be test problem
47 + #(1, 128, 128),
48 (2, 128, 512),
49 (3, 1024, 2048),
48 - (32, 128, 1024),
50 + (4, 128, 128),
51 + (32, 1024, 512),
52 (45, 512, 2048),
50 - (64, 1024, 1024),
51 - (222, 1024, 2048),
53 + (64, 1024, 512),
54 + (222, 2048, 1024),
55 + (256, 1408, 2048),
56 ]
57
58 NUM_EXPERTS = [8, 64]
55 -EP_SIZE = [1, 4]
59 TOP_KS = [1, 2, 6]
60 +DTYPES = [torch.float8_e4m3fn, torch.bfloat16]
61
62 vllm_config = VllmConfig()
63 vllm_config.scheduler_config.max_num_seqs = 128
@@ -169,9 +173,11 @@ def test_fused_moe_batched_experts(
173
174 with set_current_vllm_config(vllm_config):
175 topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
172 - baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids)
176 + baseline_output = torch_experts(a, w1, w2, topk_weight,
177 + topk_ids) # only for baseline
178 torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids)
174 - batched_output = naive_batched_moe(a, w1, w2, topk_weight, topk_ids)
179 + batched_output = naive_batched_moe(
180 + a, w1, w2, topk_weight, topk_ids) # pick torch_experts or this
181
182 torch.testing.assert_close(baseline_output,
183 torch_output,
@@ -183,6 +189,63 @@ def test_fused_moe_batched_experts(
189 rtol=0)
190
191
192 +def create_pplx_prepare_finalize(
193 + num_tokens: int,
194 + hidden_dim: int,
195 + topk: int,
196 + num_experts: int,
197 + rank: int,
198 + dp_size: int,
199 + world_size: int,
200 + in_dtype: torch.dtype,
201 + quant_dtype: Optional[torch.dtype],
202 + block_shape: Optional[list[int]],
203 + per_act_token_quant: bool,
204 + group_name: Optional[str],
205 +):
206 + from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
207 + PplxPrepareAndFinalize, pplx_hidden_dim_scale_bytes)
208 +
209 + max_num_tokens = max(rank_chunk(num_tokens, 0, world_size), 1)
210 + num_local_experts = rank_chunk(num_experts, 0, world_size)
211 +
212 + hidden_dim_bytes, scale_bytes = pplx_hidden_dim_scale_bytes(
213 + max_num_tokens,
214 + hidden_dim,
215 + in_dtype,
216 + quant_dtype,
217 + per_act_token_quant=per_act_token_quant,
218 + block_shape=block_shape,
219 + )
220 +
221 + args = dict(
222 + max_num_tokens=max_num_tokens,
223 + num_experts=num_experts,
224 + experts_per_token=topk,
225 + rank=rank,
226 + world_size=world_size,
227 + dp_size=dp_size,
228 + hidden_dim=hidden_dim,
229 + hidden_dim_bytes=hidden_dim_bytes,
230 + hidden_dim_scale_bytes=scale_bytes,
231 + )
232 +
233 + if group_name is None:
234 + ata = AllToAll.internode(**args)
235 + else:
236 + args["group_name"] = group_name
237 + ata = AllToAll.intranode(**args)
238 +
239 + prepare_finalize = PplxPrepareAndFinalize(
240 + ata,
241 + max_num_tokens=max_num_tokens,
242 + num_local_experts=num_local_experts,
243 + num_dispatchers=world_size // dp_size,
244 + )
245 +
246 + return prepare_finalize, ata
247 +
248 +
249 def rank_chunk(num: int, r: int, w: int) -> int:
250 rem = num % w
251 return (num // w) + (1 if r < rem else 0)
@@ -193,6 +256,35 @@ def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor:
256 return t[(r * chunk):(r + 1) * chunk]
257
258
259 +def maybe_chunk_by_rank(t: Optional[torch.Tensor], r: int,
260 + w: int) -> Optional[torch.Tensor]:
261 + if t is not None:
262 + return chunk_by_rank(t, r, w)
263 + else:
264 + return t
265 +
266 +
267 +def chunk_scales_by_rank(t: Optional[torch.Tensor], r: int,
268 + w: int) -> Optional[torch.Tensor]:
269 + if t is not None and t.numel() > 1:
270 + chunk = rank_chunk(t.shape[0], r, w)
271 + return t[(r * chunk):(r + 1) * chunk]
272 + else:
273 + return t
274 +
275 +
276 +def chunk_scales(t: Optional[torch.Tensor], start: int,
277 + end: int) -> Optional[torch.Tensor]:
278 + if t is not None and t.numel() > 1:
279 + return t[start:end]
280 + else:
281 + return t
282 +
283 +
284 +def dummy_work(a: torch.Tensor) -> torch.Tensor:
285 + return a * 1.1
286 +
287 +
288 def pplx_prepare_finalize(
289 pgi: ProcessGroupInfo,
290 dp_size: int,
@@ -200,11 +292,11 @@ def pplx_prepare_finalize(
292 topk_weight: torch.Tensor,
293 topk_ids: torch.Tensor,
294 num_experts: int,
295 + quant_dtype: Optional[torch.dtype],
296 + block_shape: Optional[list[int]],
297 + per_act_token_quant: bool,
298 group_name: Optional[str],
299 ) -> torch.Tensor:
205 - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
206 - PplxPrepareAndFinalize)
207 -
300 assert torch.cuda.current_device() == pgi.local_rank
301
302 topk = topk_ids.shape[1]
@@ -212,60 +304,66 @@ def pplx_prepare_finalize(
304 device = pgi.device
305 rank = pgi.rank
306 world_size = pgi.world_size
215 - max_num_tokens = rank_chunk(num_tokens, 0, world_size)
216 -
217 - args = dict(
218 - max_num_tokens=max_num_tokens,
219 - num_experts=num_experts,
220 - experts_per_token=topk,
221 - rank=rank,
222 - world_size=world_size,
223 - dp_size=dp_size,
224 - hidden_dim=hidden_dim,
225 - hidden_dim_bytes=hidden_dim * a.dtype.itemsize,
226 - hidden_dim_scale_bytes=0,
227 - )
228 -
229 - if group_name is None:
230 - ata = AllToAll.internode(**args)
231 - else:
232 - args["group_name"] = group_name
233 - ata = AllToAll.intranode(**args)
307
308 topk_ids = topk_ids.to(dtype=torch.uint32)
309
237 - prepare_finalize = PplxPrepareAndFinalize(
238 - ata,
239 - max_num_tokens,
240 - world_size,
310 + prepare_finalize, ata = create_pplx_prepare_finalize(
311 + num_tokens,
312 + hidden_dim,
313 + topk,
314 + num_experts,
315 rank,
316 dp_size,
317 + world_size,
318 + a.dtype,
319 + quant_dtype,
320 + block_shape,
321 + per_act_token_quant,
322 + group_name,
323 )
324
325 + assert a.shape[0] == topk_ids.shape[0]
326 +
327 a_chunk = chunk_by_rank(a, rank, world_size).to(device)
328 chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
329 chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
330
331 + assert a_chunk.shape[0] == chunk_topk_ids.shape[0]
332 +
333 + out = torch.full(
334 + a_chunk.shape,
335 + torch.nan,
336 + dtype=a.dtype,
337 + device=device,
338 + )
339 +
340 + if (quant_dtype is not None and not per_act_token_quant
341 + and block_shape is None):
342 + a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
343 + a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
344 + else:
345 + a1_scale = None
346 + a2_scale = None
347 +
348 b_a, b_a_scale, expert_num_tokens, _, _ = prepare_finalize.prepare(
349 a_chunk,
251 - None,
252 - None,
350 + a1_scale,
351 + a2_scale,
352 chunk_topk_weight,
353 chunk_topk_ids,
354 num_experts,
355 None,
356 False,
258 - FusedMoEQuantConfig(),
357 + FusedMoEQuantConfig(
358 + quant_dtype,
359 + per_act_token_quant,
360 + False,
361 + block_shape,
362 + ),
363 )
364
261 - b_a = b_a * 1.5
262 -
263 - out = torch.full(
264 - (max_num_tokens, hidden_dim),
265 - torch.nan,
266 - dtype=a.dtype,
267 - device=device,
268 - )
365 + b_a = dummy_work(
366 + dequant(b_a, b_a_scale, block_shape, per_act_token_quant, a.dtype))
367
368 prepare_finalize.finalize(
369 out,
@@ -291,70 +389,96 @@ def _pplx_prepare_finalize(
389 score: torch.Tensor,
390 topk: torch.Tensor,
391 num_experts: int,
392 + quant_dtype: Optional[torch.dtype],
393 + block_shape: Optional[list[int]],
394 + per_act_token_quant: bool,
395 use_internode: bool,
396 ):
296 - if use_internode:
297 - uid = nvshmem_get_unique_id(
298 - ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
299 - torch.distributed.broadcast(uid, src=0)
300 - nvshmem_init(uid, pgi.rank, pgi.world_size)
301 - group_name = None
302 - else:
303 - group_ranks = list(range(pgi.world_size))
304 - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
305 - group_name = cpu_group.group_name
306 -
307 - device = pgi.device
397 + try:
398 + if use_internode:
399 + uid = nvshmem_get_unique_id(
400 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
401 + torch.distributed.broadcast(uid, src=0)
402 + nvshmem_init(uid, pgi.rank, pgi.world_size)
403 + group_name = None
404 + else:
405 + group_ranks = list(range(pgi.world_size))
406 + cpu_group = torch.distributed.new_group(group_ranks,
407 + backend="gloo")
408 + group_name = cpu_group.group_name
409
309 - topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
310 - k = a.shape[1]
311 -
312 - a_rep = torch.repeat_interleave(a, topk, dim=0).to(device)
410 + topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
411 + m, k = a.shape
412
314 - torch_output = (a_rep.view(-1, topk, k) * 1.5 *
315 - topk_weight.view(-1, topk, 1).to(device)).sum(dim=1).to(
316 - a.dtype)
413 + a_rep = torch.repeat_interleave(dummy_work(a), topk, dim=0)
414
318 - pplx_output = pplx_prepare_finalize(pgi, dp_size, a, topk_weight, topk_ids,
319 - num_experts, group_name)
415 + torch_output = (a_rep.view(m, topk, k) *
416 + topk_weight.view(m, topk, 1).to(a_rep.dtype)).sum(
417 + dim=1)
418
321 - torch_output = chunk_by_rank(torch_output, pgi.rank,
322 - pgi.world_size).to(pplx_output.device)
419 + pplx_output = pplx_prepare_finalize(pgi, dp_size, a, topk_weight,
420 + topk_ids, num_experts, quant_dtype,
421 + block_shape, per_act_token_quant,
422 + group_name)
423
324 - torch.testing.assert_close(pplx_output, torch_output, atol=2e-2, rtol=0)
424 + torch_output = chunk_by_rank(torch_output, pgi.rank,
425 + pgi.world_size).to(pgi.device)
426
326 - if use_internode:
327 - nvshmem_finalize()
427 + torch.testing.assert_close(pplx_output,
428 + torch_output,
429 + atol=3e-2,
430 + rtol=3e-2)
431 + finally:
432 + if use_internode:
433 + nvshmem_finalize()
434
435
330 -# TODO (bnell): this test point does not work for odd M due to how the test is
331 -# written, not due to limitations of the pplx kernels. The pplx_moe
332 -# test below is able to deal with odd M.
333 -# TODO (bnell) add fp8 tests
334 -@pytest.mark.parametrize("mnk", PPLX_PREPARE_COMBOS)
436 +@pytest.mark.parametrize("mnk", PPLX_COMBOS)
437 @pytest.mark.parametrize("e", NUM_EXPERTS)
438 @pytest.mark.parametrize("topk", TOP_KS)
337 -@pytest.mark.parametrize("dtype", [torch.bfloat16])
439 +@pytest.mark.parametrize("dtype", DTYPES)
440 @pytest.mark.parametrize("world_dp_size", [[2, 1]])
441 +@pytest.mark.parametrize("per_act_token_quant", [False, True])
442 +@pytest.mark.parametrize("block_shape", [None, [128, 128]])
443 @pytest.mark.parametrize("use_internode", [False])
444 +@pytest.mark.optional
445 @requires_pplx
341 -def test_pplx_prepare_finalize(
446 +def test_pplx_prepare_finalize_slow(
447 mnk: tuple[int, int, int],
448 e: int,
449 topk: int,
450 dtype: torch.dtype,
451 world_dp_size: tuple[int, int],
452 + per_act_token_quant: bool,
453 + block_shape: Optional[list[int]],
454 use_internode: bool,
455 ):
456 + if dtype == torch.float8_e4m3fn:
457 + use_fp8_w8a8 = True
458 + act_dtype = torch.bfloat16
459 + quant_dtype = dtype
460 + else:
461 + use_fp8_w8a8 = False
462 + act_dtype = dtype
463 + quant_dtype = None
464 +
465 + if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
466 + pytest.skip("Skip quantization test for non-quantized type")
467 +
468 + if per_act_token_quant and block_shape is not None:
469 + pytest.skip("Skip illegal quantization combination")
470 +
471 current_platform.seed_everything(7)
472 m, n, k = mnk
473 world_size, dp_size = world_dp_size
474 device = "cuda"
353 - a = torch.randn((m, k), device=device, dtype=dtype) / 10
354 - score = torch.randn((m, e), device=device, dtype=dtype)
475 +
476 + a = torch.randn((m, k), device=device, dtype=act_dtype) / 10
477 + score = torch.randn((m, e), device=device, dtype=act_dtype)
478
479 parallel_launch(world_size, _pplx_prepare_finalize, dp_size, a, score,
357 - topk, e, use_internode)
480 + topk, e, quant_dtype, block_shape, per_act_token_quant,
481 + use_internode)
482
483
484 def pplx_moe(
@@ -369,84 +493,62 @@ def pplx_moe(
493 topk_ids: torch.Tensor,
494 w1_scale: Optional[torch.Tensor] = None,
495 w2_scale: Optional[torch.Tensor] = None,
372 - qtype: Optional[torch.dtype] = None,
496 + a1_scale: Optional[torch.Tensor] = None,
497 + a2_scale: Optional[torch.Tensor] = None,
498 + quant_dtype: Optional[torch.dtype] = None,
499 per_act_token_quant=False,
500 block_shape: Optional[list[int]] = None,
501 use_compile: bool = False,
502 use_cudagraphs: bool = True,
503 ) -> torch.Tensor:
378 - from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
379 - PplxPrepareAndFinalize, pplx_hidden_dim_scale_bytes)
504
381 - device = torch.device("cuda", rank)
382 - hidden_dim = a.shape[1]
505 + num_tokens, hidden_dim = a.shape
506 num_experts = w1.shape[0]
507 topk = topk_ids.shape[1]
385 - max_num_tokens = round_up(rank_chunk(a.shape[0], 0, world_size), 64)
508 + max_num_tokens = round_up(rank_chunk(a.shape[0], 0, world_size), 16)
509
387 - hidden_dim_bytes, scale_bytes = pplx_hidden_dim_scale_bytes(
388 - max_num_tokens,
510 + prepare_finalize, ata = create_pplx_prepare_finalize(
511 + num_tokens,
512 hidden_dim,
513 + topk,
514 + num_experts,
515 + rank,
516 + dp_size,
517 + world_size,
518 a.dtype,
391 - qtype,
392 - per_act_token_quant=per_act_token_quant,
393 - block_shape=block_shape,
519 + quant_dtype,
520 + block_shape,
521 + per_act_token_quant,
522 + group_name,
523 )
524
396 - args = dict(
397 - max_num_tokens=max_num_tokens,
398 - num_experts=num_experts,
399 - experts_per_token=topk,
400 - rank=rank,
401 - world_size=world_size,
402 - dp_size=dp_size,
403 - hidden_dim=hidden_dim,
404 - hidden_dim_bytes=hidden_dim_bytes,
405 - hidden_dim_scale_bytes=scale_bytes,
406 - )
407 -
408 - if group_name is None:
409 - ata = AllToAll.internode(**args)
410 - else:
411 - args["group_name"] = group_name
412 - ata = AllToAll.intranode(**args)
413 -
525 topk_ids = topk_ids.to(dtype=torch.uint32)
526
416 - prepare_finalize = PplxPrepareAndFinalize(
417 - ata,
418 - max_num_tokens,
419 - world_size,
420 - rank,
421 - dp_size,
527 + experts = BatchedTritonExperts(
528 + max_num_tokens=max_num_tokens,
529 + num_dispatchers=prepare_finalize.num_dispatchers(),
530 + use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
531 + block_shape=block_shape,
532 + per_act_token_quant=per_act_token_quant,
533 )
534
424 - experts = BatchedTritonExperts(max_num_tokens=max_num_tokens,
425 - world_size=world_size,
426 - dp_size=dp_size,
427 - use_fp8_w8a8=qtype == torch.float8_e4m3fn,
428 - block_shape=block_shape)
429 -
535 fused_experts = FusedMoEModularKernel(
536 prepare_finalize,
537 experts,
538 )
539
540 # Note: workers with the same dp_rank must use the exact same inputs.
436 - a_chunk = chunk_by_rank(a, rank, world_size).to(device)
437 - chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
438 - chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
541 + a_chunk = chunk_by_rank(a, rank, world_size)
542 + chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size)
543 + chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size)
544
545 # Chunking weights like this only works for batched format
441 - w1_chunk = chunk_by_rank(w1, rank, world_size).to(device)
442 - w2_chunk = chunk_by_rank(w2, rank, world_size).to(device)
443 -
444 - if w1_scale is not None:
445 - w1_scale_chunk = chunk_by_rank(w1_scale, rank, world_size).to(device)
446 - w2_scale_chunk = chunk_by_rank(w2_scale, rank, world_size).to(device)
447 - else:
448 - w1_scale_chunk = None
449 - w2_scale_chunk = None
546 + w1_chunk = chunk_by_rank(w1, rank, world_size)
547 + w2_chunk = chunk_by_rank(w2, rank, world_size)
548 + w1_scale_chunk = maybe_chunk_by_rank(w1_scale, rank, world_size)
549 + w2_scale_chunk = maybe_chunk_by_rank(w2_scale, rank, world_size)
550 + a1_scale_chunk = chunk_scales_by_rank(a1_scale, rank, world_size)
551 + a2_scale_chunk = chunk_scales_by_rank(a2_scale, rank, world_size)
552
553 # Note: for now use_compile will error out if the problem size is
554 # large enough to trigger chunking. I'm leaving the flag and
@@ -468,6 +570,8 @@ def pplx_moe(
570 chunk_topk_ids,
571 w1_scale=w1_scale_chunk,
572 w2_scale=w2_scale_chunk,
573 + a1_scale=a1_scale_chunk,
574 + a2_scale=a2_scale_chunk,
575 global_num_experts=num_experts)
576
577 if use_cudagraphs:
@@ -482,6 +586,8 @@ def pplx_moe(
586 chunk_topk_ids,
587 w1_scale=w1_scale_chunk,
588 w2_scale=w2_scale_chunk,
589 + a1_scale=a1_scale_chunk,
590 + a2_scale=a2_scale_chunk,
591 global_num_experts=num_experts)
592
593 torch.cuda.synchronize()
@@ -494,48 +600,6 @@ def pplx_moe(
600 return out
601
602
497 -def _batched_moe(pgi, dp_size, a, w1, w2, topk_weight, topk_ids):
498 - assert torch.cuda.current_device() == pgi.local_rank
499 -
500 - num_experts = w1.shape[0]
501 - device = pgi.device
502 - rank = pgi.rank
503 - world_size = pgi.world_size
504 - max_num_tokens = rank_chunk(a.shape[0], 0, world_size)
505 -
506 - prepare_finalize = BatchedPrepareAndFinalize(
507 - max_num_tokens=max_num_tokens,
508 - world_size=world_size,
509 - dp_size=dp_size,
510 - rank=rank,
511 - )
512 -
513 - experts = NaiveBatchedExperts(max_num_tokens=a.shape[0],
514 - world_size=1,
515 - dp_size=1)
516 -
517 - fused_experts = FusedMoEModularKernel(
518 - prepare_finalize,
519 - experts,
520 - )
521 -
522 - # Note: workers with the same dp_rank must use the exact same inputs.
523 - a_chunk = chunk_by_rank(a, rank, world_size).to(device)
524 - chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
525 - chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
526 -
527 - out = fused_experts(
528 - a_chunk,
529 - # Chunking weights like this only works for batched format
530 - chunk_by_rank(w1, rank, world_size).to(device),
531 - chunk_by_rank(w2, rank, world_size).to(device),
532 - chunk_topk_weight,
533 - chunk_topk_ids,
534 - global_num_experts=num_experts)
535 -
536 - return out
537 -
538 -
603 def _pplx_moe(
604 pgi: ProcessGroupInfo,
605 dp_size: int,
@@ -544,75 +608,130 @@ def _pplx_moe(
608 w2: torch.Tensor,
609 score: torch.Tensor,
610 topk: int,
611 + num_experts: int,
612 w1_s: Optional[torch.Tensor] = None,
613 w2_s: Optional[torch.Tensor] = None,
549 - qtype: Optional[torch.dtype] = None,
614 + quant_dtype: Optional[torch.dtype] = None,
615 per_act_token_quant: bool = False,
616 block_shape: Optional[list[int]] = None,
617 use_internode: bool = False,
618 ):
554 - if use_internode:
555 - uid = nvshmem_get_unique_id(
556 - ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
557 - torch.distributed.broadcast(uid, src=0)
558 - nvshmem_init(uid, pgi.rank, pgi.world_size)
559 - group_name = None
560 - else:
561 - group_ranks = list(range(pgi.world_size))
562 - cpu_group = torch.distributed.new_group(group_ranks, backend="gloo")
563 - group_name = cpu_group.group_name
564 -
565 - m, k = a.shape
566 - e, _, n = w2.shape
567 -
568 - moe_config = get_default_config(m, e, n, k, topk, a.dtype, False)
569 -
570 - device = torch.device("cuda", pgi.rank)
571 - a = a.to(device)
572 - w1 = w1.to(device)
573 - w2 = w2.to(device)
574 - w1_s = w1_s.to(device) if w1_s is not None else None
575 - w2_s = w2_s.to(device) if w2_s is not None else None
576 -
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_experts(a,
580 - w1,
581 - w2,
582 - topk_weight,
583 - topk_ids,
584 - w1_scale=w1_s,
585 - w2_scale=w2_s,
586 - quant_dtype=qtype,
587 - per_act_token_quant=per_act_token_quant,
588 - block_shape=block_shape)
589 - pplx_output = pplx_moe(group_name, pgi.rank, pgi.world_size, dp_size,
590 - a, w1, w2, topk_weight, topk_ids, w1_s, w2_s,
591 - qtype, per_act_token_quant, block_shape)
592 - # TODO (bnell): fix + re-enable
593 - #batched_output = _batched_moe(pgi, dp_size, a, w1, w2, topk_weight,
594 - # topk_ids)
595 -
596 - torch_output = chunk_by_rank(torch_output, pgi.rank,
597 - pgi.world_size).to(pplx_output.device)
598 -
599 - torch.testing.assert_close(pplx_output, torch_output, atol=2e-2, rtol=0)
600 - #torch.testing.assert_close(batched_output, torch_output, atol=2e-2, rtol=0)
601 -
602 - if use_internode:
603 - nvshmem_finalize()
604 -
605 -
606 -@pytest.mark.parametrize("mnk", PPLX_MOE_COMBOS)
619 + try:
620 + if use_internode:
621 + uid = nvshmem_get_unique_id(
622 + ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
623 + torch.distributed.broadcast(uid, src=0)
624 + nvshmem_init(uid, pgi.rank, pgi.world_size)
625 + group_name = None
626 + else:
627 + group_ranks = list(range(pgi.world_size))
628 + cpu_group = torch.distributed.new_group(group_ranks,
629 + backend="gloo")
630 + group_name = cpu_group.group_name
631 +
632 + m, k = a.shape
633 + e, _, n = w2.shape
634 +
635 + moe_config = get_default_config(m, e, n, k, topk, a.dtype, False)
636 +
637 + device = torch.device("cuda", pgi.rank)
638 + rank = pgi.rank
639 + world_size = pgi.world_size
640 +
641 + a = a.to(device)
642 + w1 = w1.to(device)
643 + w2 = w2.to(device)
644 + w1_s = w1_s.to(device) if w1_s is not None else None
645 + w2_s = w2_s.to(device) if w2_s is not None else None
646 +
647 + if (quant_dtype is not None and not per_act_token_quant
648 + and block_shape is None):
649 + a1_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
650 + a2_scale = torch.tensor(1.0, device="cuda", dtype=torch.float32)
651 + else:
652 + a1_scale = None
653 + a2_scale = None
654 +
655 + with set_current_vllm_config(vllm_config), override_config(moe_config):
656 + topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
657 +
658 + torch_output = torch_experts(
659 + a,
660 + w1,
661 + w2,
662 + topk_weight,
663 + topk_ids,
664 + w1_scale=w1_s,
665 + w2_scale=w2_s,
666 + a1_scale=a1_scale,
667 + a2_scale=a2_scale,
668 + quant_dtype=quant_dtype,
669 + per_act_token_quant=per_act_token_quant,
670 + block_shape=block_shape,
671 + )
672 +
673 + batched_output = naive_batched_moe(
674 + a,
675 + w1,
676 + w2,
677 + topk_weight,
678 + topk_ids,
679 + w1_scale=w1_s,
680 + w2_scale=w2_s,
681 + a1_scale=a1_scale,
682 + a2_scale=a2_scale,
683 + quant_dtype=quant_dtype,
684 + per_act_token_quant=per_act_token_quant,
685 + block_shape=block_shape,
686 + )
687 +
688 + pplx_output = pplx_moe(
689 + group_name,
690 + rank,
691 + world_size,
692 + dp_size,
693 + a,
694 + w1,
695 + w2,
696 + topk_weight,
697 + topk_ids,
698 + w1_scale=w1_s,
699 + w2_scale=w2_s,
700 + a1_scale=a1_scale,
701 + a2_scale=a2_scale,
702 + quant_dtype=quant_dtype,
703 + per_act_token_quant=per_act_token_quant,
704 + block_shape=block_shape,
705 + )
706 +
707 + chunked_batch_output = chunk_by_rank(
708 + batched_output, pgi.rank, pgi.world_size).to(pplx_output.device)
709 +
710 + torch.testing.assert_close(batched_output,
711 + torch_output,
712 + atol=3e-2,
713 + rtol=3e-2)
714 +
715 + torch.testing.assert_close(pplx_output,
716 + chunked_batch_output,
717 + atol=3e-2,
718 + rtol=3e-2)
719 + finally:
720 + if use_internode:
721 + nvshmem_finalize()
722 +
723 +
724 +@pytest.mark.parametrize("mnk", PPLX_COMBOS)
725 @pytest.mark.parametrize("e", NUM_EXPERTS)
726 @pytest.mark.parametrize("topk", TOP_KS)
609 -@pytest.mark.parametrize("dtype", [torch.bfloat16])
727 +@pytest.mark.parametrize("dtype", DTYPES)
728 @pytest.mark.parametrize("world_dp_size", [[2, 1]])
729 @pytest.mark.parametrize("per_act_token_quant", [False, True])
730 @pytest.mark.parametrize("block_shape", [None, [128, 128]])
731 @pytest.mark.parametrize("use_internode", [False])
732 +@pytest.mark.optional
733 @requires_pplx
615 -def test_pplx_moe(
734 +def test_pplx_moe_slow(
735 mnk: tuple[int, int, int],
736 e: int,
737 topk: int,
@@ -633,18 +752,143 @@ def test_pplx_moe(
752 use_fp8_w8a8 = False
753 quant_dtype = None
754
636 - if not use_fp8_w8a8 and per_act_token_quant and block_shape is not None:
755 + if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
756 pytest.skip("Skip quantization test for non-quantized type")
757
758 + if per_act_token_quant and block_shape is not None:
759 + pytest.skip("Skip illegal quantization combination")
760 +
761 a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
762 score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
763
642 - _, w1, w1_s, _, w2, w2_s = make_test_weights(e,
643 - n,
644 - k,
645 - quant_dtype=quant_dtype,
646 - block_shape=block_shape)
764 + _, w1, w1_s, _, w2, w2_s = make_test_weights(
765 + e,
766 + n,
767 + k,
768 + quant_dtype=quant_dtype,
769 + block_shape=block_shape,
770 + per_act_token_quant=per_act_token_quant,
771 + )
772
648 - parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk,
773 + parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk, e,
774 w1_s, w2_s, quant_dtype, per_act_token_quant, block_shape,
775 use_internode)
776 +
777 +
778 +def _pplx_test_loop(pgi: ProcessGroupInfo, dp_size: int, use_internode: bool,
779 + make_weights: bool, test_fn: Callable):
780 +
781 + def format_result(msg, ex=None):
782 + if ex is not None:
783 + x = str(ex)
784 + newx = x.strip(" \n\t")[:16]
785 + if len(newx) < len(x):
786 + newx = newx + " ..."
787 +
788 + prefix = "E\t"
789 + print(f"{textwrap.indent(traceback.format_exc(), prefix)}")
790 + print(f"FAILED {msg} - {newx}\n")
791 + else:
792 + print(f"PASSED {msg}")
793 +
794 + current_platform.seed_everything(7)
795 + combos = itertools.product(PPLX_COMBOS, NUM_EXPERTS, TOP_KS, DTYPES,
796 + [False, True], [None, [128, 128]])
797 + exceptions = []
798 + count = 0
799 + for mnk, e, topk, dtype, per_act_token_quant, block_shape in combos:
800 + count = count + 1
801 + m, n, k = mnk
802 +
803 + if dtype == torch.float8_e4m3fn:
804 + use_fp8_w8a8 = True
805 + quant_dtype = dtype
806 + else:
807 + use_fp8_w8a8 = False
808 + quant_dtype = None
809 +
810 + test_desc = (f"test_pplx_moe[mnk={mnk}, e={e}, topk={topk}, "
811 + f"dtype={dtype}, per_act_token={per_act_token_quant}, "
812 + f"block_shape={block_shape}")
813 +
814 + if not use_fp8_w8a8 and (per_act_token_quant
815 + or block_shape is not None):
816 + print(
817 + f"{test_desc} - Skip quantization test for non-quantized type."
818 + )
819 + continue
820 +
821 + if per_act_token_quant and block_shape is not None:
822 + print(f"{test_desc} - Skip illegal quantization combination.")
823 + continue
824 +
825 + a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
826 + score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
827 +
828 + args = dict()
829 + if make_weights:
830 + _, w1, w1_s, _, w2, w2_s = make_test_weights(
831 + e,
832 + n,
833 + k,
834 + quant_dtype=quant_dtype,
835 + block_shape=block_shape,
836 + per_act_token_quant=per_act_token_quant,
837 + )
838 + args["w1"] = w1
839 + args["w2"] = w2
840 + args["w1_s"] = w1_s
841 + args["w2_s"] = w2_s
842 +
843 + try:
844 + test_fn(
845 + pgi=pgi,
846 + dp_size=dp_size,
847 + a=a,
848 + score=score,
849 + topk=topk,
850 + num_experts=e,
851 + quant_dtype=quant_dtype,
852 + per_act_token_quant=per_act_token_quant,
853 + block_shape=block_shape,
854 + use_internode=use_internode,
855 + **args,
856 + )
857 + format_result(test_desc)
858 + except Exception as ex:
859 + format_result(test_desc, ex)
860 + exceptions.append(ex)
861 +
862 + if len(exceptions) > 0:
863 + raise RuntimeError(
864 + f"{len(exceptions)} of {count} tests failed in child process, "
865 + f"rank={pgi.rank}.")
866 + else:
867 + print(f"{count} of {count} tests passed in child process, "
868 + f"rank={pgi.rank}.")
869 +
870 +
871 +@pytest.mark.parametrize("world_dp_size", [[2, 1]])
872 +@pytest.mark.parametrize("use_internode", [False])
873 +@requires_pplx
874 +def test_pplx_prepare_finalize(
875 + world_dp_size: tuple[int, int],
876 + use_internode: bool,
877 +):
878 + current_platform.seed_everything(7)
879 + world_size, dp_size = world_dp_size
880 + parallel_launch(world_size * dp_size, _pplx_test_loop, dp_size,
881 + use_internode, False, _pplx_prepare_finalize)
882 +
883 +
884 +@pytest.mark.parametrize("world_dp_size", [[2, 1]])
885 +@pytest.mark.parametrize("use_internode", [False])
886 +@requires_pplx
887 +def test_pplx_moe(
888 + world_dp_size: tuple[int, int],
889 + use_internode: bool,
890 +):
891 + current_platform.seed_everything(7)
892 + world_size, dp_size = world_dp_size
893 + parallel_launch(world_size, _pplx_test_loop, dp_size, use_internode, True,
894 + _pplx_moe)
tests/kernels/moe/utils.py
+6 -8
@@ -63,13 +63,12 @@ def batched_moe(
63
64 fused_experts = FusedMoEModularKernel(
65 BatchedPrepareAndFinalize(max_num_tokens,
66 - world_size=1,
67 - dp_size=1,
66 + num_dispatchers=1,
67 + num_local_experts=w1.shape[0],
68 rank=0),
69 BatchedTritonExperts(
70 max_num_tokens=max_num_tokens,
71 - world_size=1,
72 - dp_size=1,
71 + num_dispatchers=1,
72 use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
73 per_act_token_quant=per_act_token_quant,
74 block_shape=block_shape,
@@ -105,13 +104,12 @@ def naive_batched_moe(
104
105 fused_experts = FusedMoEModularKernel(
106 BatchedPrepareAndFinalize(max_num_tokens,
108 - world_size=1,
109 - dp_size=1,
107 + num_dispatchers=1,
108 + num_local_experts=w1.shape[0],
109 rank=0),
110 NaiveBatchedExperts(
111 max_num_tokens=max_num_tokens,
113 - dp_size=1,
114 - world_size=1,
112 + num_dispatchers=1,
113 use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
114 per_act_token_quant=per_act_token_quant,
115 block_shape=block_shape,
tests/kernels/quant_utils.py
+18
@@ -277,6 +277,24 @@ def dequant(
277 return t.to(out_dtype)
278
279
280 +def batched_dequant(
281 + t: torch.Tensor,
282 + scale: Optional[torch.Tensor],
283 + block_shape: Optional[list[int]],
284 + per_act_token_quant: bool,
285 + out_dtype: Optional[torch.dtype] = torch.float32,
286 +) -> torch.Tensor:
287 + if scale is not None:
288 + assert t.shape[0] == scale.shape[0]
289 + out = torch.empty_like(t, dtype=out_dtype)
290 + for e in range(t.shape[0]):
291 + out[e] = dequant(t[e], scale[e], block_shape, per_act_token_quant,
292 + out_dtype)
293 + return out
294 +
295 + return t.to(out_dtype)
296 +
297 +
298 def native_batched_masked_quant_matmul(
299 A: torch.Tensor,
300 B: torch.Tensor,
tests/kernels/utils.py
+6 -4
@@ -1094,6 +1094,8 @@ def torch_experts(
1094 if expert_map is not None:
1095 topk_ids = expert_map[topk_ids]
1096
1097 + f32 = torch.float32
1098 +
1099 for i in range(num_experts):
1100 mask = topk_ids == i
1101 if mask.sum():
@@ -1109,7 +1111,8 @@ def torch_experts(
1111 out.dtype)
1112 tmp2 = SiluAndMul()(tmp1)
1113 tmp2, b_scale = moe_kernel_quantize_input(
1112 - tmp2, None, quant_dtype, per_act_token_quant, block_shape)
1114 + tmp2, a2_scale, quant_dtype, per_act_token_quant,
1115 + block_shape)
1116
1117 out[mask] = native_w8a8_block_matmul(tmp2, w2[i], b_scale,
1118 w2_scale[i], block_shape,
@@ -1117,7 +1120,6 @@ def torch_experts(
1120 else:
1121 assert (a_scale is not None and w1_scale is not None
1122 and w2_scale is not None)
1120 - f32 = torch.float32
1123 scales = a_scale if a_scale.numel() == 1 else a_scale[mask]
1124 tmp1 = a[mask].to(f32) * scales
1125 w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1)
@@ -1126,8 +1128,8 @@ def torch_experts(
1128 w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1)
1129 out[mask] = (tmp2 @ w2_dq).to(out.dtype)
1130
1129 - return (out.view(M, -1, w2.shape[1]) *
1130 - topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
1131 + return (out.view(M, -1, w2.shape[1]).to(f32) *
1132 + topk_weight.view(M, -1, 1)).sum(dim=1).to(out.dtype)
1133
1134
1135 def torch_moe(a: torch.Tensor,
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+6 -8
@@ -184,15 +184,14 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
184
185 def __init__(self,
186 max_num_tokens: int,
187 - world_size: int,
188 - dp_size: int,
187 + num_dispatchers: int,
188 block_shape: list[int],
189 per_act_token_quant=False):
190 """
191 max_num_tokens: Maximum number of tokens from a DP Rank
193 - world_size: Number of EP ranks
194 - dp_size: Number of data-parallel ranks
195 - block_shape: Block quantization block shape
192 + num_dispatchers: The number of DP dispatchers.
193 + block_shape: Block quantization block shape.
194 + per_act_token_quant: Per activation token quantization flag.
195 """
196 super().__init__(
197 FusedMoEQuantConfig(
@@ -202,8 +201,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
201 ))
202 assert self.block_shape == self.DEEPGEMM_BLOCK_SHAPE
203 self.max_num_tokens = max_num_tokens
205 - self.world_size = world_size
206 - self.dp_size = dp_size
204 + self.num_dispatchers = num_dispatchers
205
206 @property
207 def activation_formats(
@@ -233,7 +231,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
231 # FIXME (varun): We should be able to dispatch only from the leader
232 # DP ranks in the case of TP > 1. At the moment, all the Ranks
233 # end up sending their tokens. This needs to be fixed.
236 - num_dispatchers = self.world_size
234 + num_dispatchers = self.num_dispatchers
235 num_experts = local_num_experts
236 max_num_tokens = a.size(
237 0) if self.max_num_tokens is None else self.max_num_tokens
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+11 -23
@@ -15,8 +15,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
15
16 def __init__(self,
17 max_num_tokens: int,
18 - world_size: int,
19 - dp_size: int,
18 + num_dispatchers: int,
19 use_fp8_w8a8: bool = False,
20 use_int8_w8a8: bool = False,
21 use_int8_w8a16: bool = False,
@@ -37,35 +36,28 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
36 block_shape=block_shape,
37 per_act_token_quant=per_act_token_quant,
38 ))
40 - self.max_num_tokens = max_num_tokens
41 - self.world_size = world_size
42 - self.dp_size = dp_size
39 self.allow_deep_gemm = allow_deep_gemm
40
45 - # BatchedTritonKernel doesn't support block quantization
46 - # at the moment.
41 self.batched_triton_experts = BatchedTritonExperts(
48 - max_num_tokens=self.max_num_tokens,
49 - world_size=self.world_size,
50 - dp_size=self.dp_size,
42 + max_num_tokens=max_num_tokens,
43 + num_dispatchers=num_dispatchers,
44 use_fp8_w8a8=use_fp8_w8a8,
45 use_int8_w8a8=use_int8_w8a8,
46 use_int8_w8a16=use_int8_w8a16,
47 use_int4_w4a16=use_int4_w4a16,
48 per_act_token_quant=self.per_act_token_quant,
49 block_shape=self.block_shape,
57 - ) if self.block_shape is None else None
50 + )
51
59 - is_fp8_128_block_quantized = (
60 - use_fp8_w8a8 and self.block_shape
61 - == BatchedDeepGemmExperts.DEEPGEMM_BLOCK_SHAPE)
52 + self.allow_deep_gemm = (allow_deep_gemm and use_fp8_w8a8
53 + and self.block_shape
54 + == BatchedDeepGemmExperts.DEEPGEMM_BLOCK_SHAPE)
55
56 self.batched_deep_gemm_experts = BatchedDeepGemmExperts(
64 - max_num_tokens=self.max_num_tokens,
65 - world_size=self.world_size,
66 - dp_size=self.dp_size,
57 + max_num_tokens=max_num_tokens,
58 + num_dispatchers=num_dispatchers,
59 block_shape=self.block_shape, # type: ignore[arg-type]
68 - ) if (self.allow_deep_gemm and is_fp8_128_block_quantized) else None
60 + ) if self.allow_deep_gemm else None
61
62 assert (self.batched_deep_gemm_experts is not None
63 or self.batched_triton_experts is not None)
@@ -138,12 +130,8 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
130 workspace2: torch.Tensor,
131 expert_num_tokens: Optional[torch.Tensor],
132 ):
141 - use_batched_deep_gemm_experts = (self.allow_deep_gemm
142 - and self.batched_deep_gemm_experts
143 - is not None)
133 experts = (self.batched_deep_gemm_experts
145 - if use_batched_deep_gemm_experts else
146 - self.batched_triton_experts)
134 + if self.allow_deep_gemm else self.batched_triton_experts)
135 assert experts is not None
136 experts.apply(output, hidden_states, w1, w2, topk_ids, activation,
137 global_num_experts, expert_map, w1_scale, w2_scale,
vllm/model_executor/layers/fused_moe/config.py
+55 -9
@@ -14,6 +14,7 @@ from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
14 from vllm.logger import init_logger
15 from vllm.model_executor.layers.quantization.base_config import (
16 QuantizationConfig)
17 +from vllm.utils import cdiv
18
19 logger = init_logger(__name__)
20
@@ -68,6 +69,57 @@ class FusedMoEQuantConfig:
69 # TODO: add col major flag?
70 # add detailed quant info for input, intermediates, weights, etc?
71
72 + def __post_init__(self):
73 + assert (not self.per_act_token_quant
74 + or self.block_shape is None), "illegal quantization"
75 +
76 + @property
77 + def is_quantized(self) -> bool:
78 + return self.quant_dtype is not None
79 +
80 + @property
81 + def is_per_act_token(self) -> bool:
82 + return self.per_act_token_quant
83 +
84 + @property
85 + def is_block_quantized(self) -> bool:
86 + return self.block_shape is not None
87 +
88 + @property
89 + def is_per_tensor(self) -> bool:
90 + return not self.per_act_token_quant and self.block_shape is None
91 +
92 + def scale_shape(
93 + self,
94 + max_tokens: int,
95 + hidden_dim: int,
96 + ) -> Optional[tuple[int, int]]:
97 + if self.is_quantized:
98 + if self.is_block_quantized:
99 + assert self.block_shape is not None
100 + _, block_k = self.block_shape
101 + k_tiles = cdiv(hidden_dim, block_k)
102 + return (max_tokens, k_tiles)
103 + elif self.is_per_act_token:
104 + return (max_tokens, 1)
105 + else:
106 + return (1, 1)
107 + else:
108 + return None
109 +
110 + def batched_scale_shape(
111 + self,
112 + num_experts: int,
113 + max_tokens: int,
114 + hidden_dim: int,
115 + ) -> Optional[tuple[int, int, int]]:
116 + if self.is_quantized:
117 + scale_shape = self.scale_shape(max_tokens, hidden_dim)
118 + assert scale_shape is not None
119 + return (num_experts, *scale_shape)
120 + else:
121 + return None
122 +
123 @staticmethod
124 def make(
125 use_fp8_w8a8: bool = False,
@@ -109,7 +161,6 @@ class FusedMoEParallelConfig:
161 tp_rank: int
162 dp_rank: int
163 ep_rank: int
112 - world_size: int
164
165 use_ep: bool # whether to use EP or not
166
@@ -133,7 +184,7 @@ class FusedMoEParallelConfig:
184 and envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency")
185
186 @staticmethod
136 - def make(tp_size_: int, dp_size_: int, world_size_: int,
187 + def make(tp_size_: int, dp_size_: int,
188 vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
189 """
190 Determine MoE parallel configuration. Based on the input tp_size_,
@@ -144,7 +195,6 @@ class FusedMoEParallelConfig:
195 tp_size_ (int): tp_size passed into the FusedMoE constructor.
196 dp_size_ (int): dp_size passed into the FusedMoE constructor.
197 ep_size_ (int): ep_size passed into the FusedMoE constructor.
147 - world_size_ (int): the world size of the current All2All manager.
198 vllm_parallel_config (ParallelConfig): vllm's parallel config
199 object.
200
@@ -223,7 +273,6 @@ class FusedMoEParallelConfig:
273 dp_rank=dp_rank,
274 ep_size=1,
275 ep_rank=0,
226 - world_size=world_size_,
276 use_ep=False)
277 # DP + EP / TP + EP / DP + TP + EP
278 assert use_ep
@@ -237,7 +286,6 @@ class FusedMoEParallelConfig:
286 dp_rank=dp_rank,
287 ep_size=ep_size,
288 ep_rank=ep_rank,
240 - world_size=world_size_,
289 use_ep=True)
290
291
@@ -263,6 +311,8 @@ class FusedMoEConfig:
311 logger.debug("Using FusedMoEConfig::max_num_tokens=%d",
312 self.max_num_tokens)
313
314 + assert self.max_num_tokens > 0
315 +
316 @property
317 def quant_dtype(self) -> Optional[torch.dtype]:
318 if self.quant_config is not None:
@@ -303,10 +353,6 @@ class FusedMoEConfig:
353 def ep_size(self):
354 return self.moe_parallel_config.ep_size
355
306 - @property
307 - def world_size(self):
308 - return self.moe_parallel_config.world_size
309 -
356 @property
357 def tp_rank(self):
358 return self.moe_parallel_config.tp_rank
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+12 -6
@@ -41,10 +41,7 @@ def run_cutlass_moe_fp8(
41 assert w2_scale is not None
42 assert w1.dtype == torch.float8_e4m3fn
43 assert w2.dtype == torch.float8_e4m3fn
44 - if expert_num_tokens is None:
45 - assert a1q.size(1) == w1.size(2), "Hidden size mismatch w1"
46 - else:
47 - assert a1q.size(2) == w1.size(2), "Hidden size mismatch w1"
44 + assert a1q.size(-1) == w1.size(2), "Hidden size mismatch w1"
45 assert w1.size(1) == w2.size(2) * 2, "Hidden size mismatch w2"
46 assert w1_scale.dim() == 1 or w1_scale.size(
47 1) == 1 or w1_scale.shape[1] == w1.size(1), "W1 scale shape mismatch"
@@ -178,6 +175,8 @@ def run_cutlass_moe_fp8(
175 c2 = _resize_cache(workspace2, (M * topk, N))
176 c3 = _resize_cache(workspace13, (M * topk, K))
177
178 + c1.fill_(0)
179 +
180 ops.cutlass_moe_mm(c1, a1q, w1, a1q_scale, w1_scale, expert_offsets,
181 problem_sizes1, ab_strides1, ab_strides1, c_strides1,
182 per_act_token, per_out_ch)
@@ -213,6 +212,7 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
212 per_act_token_quant: bool,
213 per_out_ch_quant: bool,
214 block_shape: Optional[list[int]] = None,
215 + num_dispatchers: Optional[int] = None,
216 use_batched_format: bool = False,
217 ):
218 super().__init__(
@@ -223,7 +223,9 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
223 block_shape=block_shape,
224 ))
225 assert max_experts_per_worker > 0
226 + assert not use_batched_format or num_dispatchers is not None
227 self.max_experts_per_worker = max_experts_per_worker
228 + self.num_dispatchers = num_dispatchers
229 self.out_dtype = out_dtype
230 self.use_batched_format = use_batched_format
231
@@ -260,8 +262,12 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
262 output: tuple[int, ...] = ()
263 if self.use_batched_format:
264 padded_M = aq.size(1)
263 - workspace1 = (self.max_experts_per_worker, padded_M, max(N, K))
264 - workspace2 = (self.max_experts_per_worker, padded_M, (N // 2))
265 + num_dp = self.num_dispatchers
266 + assert num_dp is not None
267 + workspace1 = (self.max_experts_per_worker, padded_M * num_dp,
268 + max(N, K))
269 + workspace2 = (self.max_experts_per_worker, padded_M * num_dp,
270 + (N // 2))
271 output = (self.max_experts_per_worker, padded_M, K)
272 else:
273 workspace1 = (M * topk, max(2 * N, K))
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+6 -17
@@ -16,12 +16,11 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
16 Prepare/Finalize using DeepEP High-Throughput kernels.
17 """
18
19 - def __init__(self, buffer: deep_ep.Buffer, world_size: int, rank: int,
19 + def __init__(self, buffer: deep_ep.Buffer, num_dispatchers: int,
20 dp_size: int, rank_expert_offset: int):
21 super().__init__()
22 self.buffer = buffer
23 - self.world_size = world_size
24 - self.rank = rank
23 + self.num_dispatchers_ = num_dispatchers
24 self.dp_size = dp_size
25 self.rank_expert_offset = rank_expert_offset
26 # The dispatch function returns a handle that the combine function
@@ -32,6 +31,9 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
31 # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164
32 self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]
33
34 + def num_dispatchers(self) -> int:
35 + return self.num_dispatchers_
36 +
37 @property
38 def activation_format(self) -> mk.FusedMoEActivationFormat:
39 return mk.FusedMoEActivationFormat.Standard
@@ -136,20 +138,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
138 "apply_router_weight_on_input is only implemented for topk=1")
139 a1 = a1 * topk_weights.to(a1.dtype)
140
139 - # Check if there is a block_shape / or if we can infer the quantization
140 - # schemes from the scales.
141 - per_token_quant = None
142 - if all([
143 - x is None
144 - for x in [quant_config.block_shape, a1_scale, a2_scale]
145 - ]) and quant_config.quant_dtype is not None:
146 - # Quantization required despite none of the inputs suggesting
147 - # quantization. Fallback to per_token_dynamic quant.
148 - per_token_quant = True
149 - else:
150 - per_token_quant = False
151 -
152 - if per_token_quant:
141 + if quant_config.per_act_token_quant:
142 a1q, a1q_scale = moe_kernel_quantize_input(
143 a1,
144 a1_scale,
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+7 -8
@@ -7,7 +7,7 @@ import torch
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9 from vllm.model_executor.layers.fused_moe.utils import (
10 - maybe_fix_scales, moe_kernel_quantize_input)
10 + moe_kernel_quantize_input, normalize_batched_scales_shape)
11
12 # DeepEP kernels quantize dispatch inputs in 128 element chunks.
13 DEEPEP_QUANT_BLOCK_SIZE = 128
@@ -42,20 +42,21 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
42 def __init__(self,
43 buffer: deep_ep.Buffer,
44 max_tokens_per_rank: int,
45 - world_size: int,
46 - dp_size: int,
45 + num_dispatchers: int,
46 use_fp8_dispatch: bool = False):
47 super().__init__()
48
49 self.buffer = buffer
50 self.max_tokens_per_rank = max_tokens_per_rank
52 - self.world_size = world_size
53 - self.dp_size = dp_size
51 self.use_fp8_dispatch = use_fp8_dispatch
52 # The dispatch function returns a handle that the combine function
53 # requires. We store the handle here so it is available to the
54 # combine function.
55 self.handle = None
56 + self.num_dispatchers_ = num_dispatchers
57 +
58 + def num_dispatchers(self) -> int:
59 + return self.num_dispatchers_
60
61 @property
62 def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -91,8 +92,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
92
93 assert isinstance(x, torch.Tensor)
94
94 - assert not per_act_token_quant
95 -
95 num_experts, max_tokens, hidden_dim = x.size()
96
97 # TODO (varun): Optimization - Use a batched version of quant
@@ -104,7 +103,7 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
103
104 if quant_dtype is not None:
105 assert x_scales is not None
107 - x_scales = maybe_fix_scales(x_scales, num_experts)
106 + x_scales = normalize_batched_scales_shape(x_scales, num_experts)
107
108 return x, x_scales
109
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+391 -192
@@ -12,42 +12,49 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12 from vllm.model_executor.layers.fused_moe.fused_moe import (
13 get_config_dtype_str, try_get_optimal_moe_config)
14 from vllm.model_executor.layers.fused_moe.utils import (
15 - _resize_cache, moe_kernel_quantize_input)
15 + _resize_cache, moe_kernel_quantize_input, normalize_batched_scales_shape,
16 + normalize_scales_shape)
17 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 + group_broadcast)
19
20
21 @triton.jit
22 def moe_mmk(
20 - a_ptrs,
21 - b_ptrs,
22 - K,
23 - expert_id,
24 - a_scale_ptr,
25 - b_scale_ptr,
26 - # The stride variables represent how much to increase the ptr by when
27 - # moving by 1 element in a particular dimension. E.g. `stride_am` is
28 - # how much to increase `a_ptr` by to get the element one row down
29 - # (A has M rows).
30 - stride_ak,
31 - stride_bk,
32 - stride_asm,
33 - stride_ask,
34 - stride_bse,
35 - stride_bsk,
36 - stride_bsn,
37 - # Offsets and masks
38 - offs_m,
39 - offs_n,
40 - mask_m,
41 - # Block size for block-wise quantization
42 - group_n: tl.constexpr,
43 - group_k: tl.constexpr,
44 - # Meta-parameters
45 - BLOCK_M: tl.constexpr,
46 - BLOCK_N: tl.constexpr,
47 - BLOCK_K: tl.constexpr,
48 - compute_type: tl.constexpr,
49 - use_w8a8: tl.constexpr,
50 - use_w8a16: tl.constexpr):
23 + a_ptrs,
24 + b_ptrs,
25 + K,
26 + expert_id,
27 + a_scale_ptr,
28 + b_scale_ptr,
29 + # The stride variables represent how much to increase the ptr by when
30 + # moving by 1 element in a particular dimension. E.g. `stride_am` is
31 + # how much to increase `a_ptr` by to get the element one row down
32 + # (A has M rows).
33 + stride_ak: tl.int64,
34 + stride_bk: tl.int64,
35 + stride_ase: tl.int64,
36 + stride_asm: tl.int64,
37 + stride_ask: tl.int64,
38 + stride_bse: tl.int64,
39 + stride_bsk: tl.int64,
40 + stride_bsn: tl.int64,
41 + # Offsets and masks
42 + offs_m,
43 + offs_n,
44 + offs_bn,
45 + mask_m,
46 + # Block size for block-wise quantization
47 + group_n: tl.constexpr,
48 + group_k: tl.constexpr,
49 + # Meta-parameters
50 + BLOCK_M: tl.constexpr,
51 + BLOCK_N: tl.constexpr,
52 + BLOCK_K: tl.constexpr,
53 + compute_type: tl.constexpr,
54 + use_w8a8: tl.constexpr,
55 + use_w8a16: tl.constexpr,
56 + per_act_token_quant: tl.constexpr,
57 +):
58
59 offs_k = tl.arange(0, BLOCK_K)
60
@@ -60,13 +67,22 @@ def moe_mmk(
67 # block-wise
68 if group_k > 0 and group_n > 0:
69 a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
63 - offs_bsn = offs_n // group_n
64 - b_scale_ptrs = (b_scale_ptr + expert_id * stride_bse +
65 - offs_bsn * stride_bsn)
70 + offs_bsn = offs_bn // group_n
71 + b_scale_ptrs = b_scale_ptr + offs_bsn * stride_bsn
72 +
73 + # per act token
74 + elif per_act_token_quant:
75 + # Load per-token scale for activations
76 + a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
77 + a_scale = tl.load(a_scale_ptrs, mask=mask_m, other=0.0)[:, None]
78 +
79 + b_scale_ptrs = b_scale_ptr + offs_bn[None, :] * stride_bsn
80 + b_scale = tl.load(b_scale_ptrs)
81 +
82 # tensor-wise
83 else:
84 a_scale = tl.load(a_scale_ptr)
69 - b_scale = tl.load(b_scale_ptr + expert_id)
85 + b_scale = tl.load(b_scale_ptr)
86
87 # -----------------------------------------------------------
88 # Iterate to compute a block of the C matrix.
@@ -96,13 +112,11 @@ def moe_mmk(
112 accumulator += tl.dot(a, b) * a_scale[:,
113 None] * b_scale[None, :]
114 else:
99 - if use_w8a8:
100 - # acc used to enable fp8_fast_accum
101 - accumulator = tl.dot(a, b, acc=accumulator)
102 - else:
103 - accumulator += tl.dot(a, b)
115 + # acc used to enable fp8_fast_accum
116 + accumulator = tl.dot(a, b, acc=accumulator)
117 else:
118 accumulator += tl.dot(a, b)
119 +
120 # Advance the ptrs to the next K block.
121 a_ptrs += BLOCK_K * stride_ak
122 b_ptrs += BLOCK_K * stride_bk
@@ -122,47 +136,53 @@ def moe_mmk(
136
137 @triton.jit
138 def expert_triton_kernel(
125 - a_ptr, #[max_tokens, K]
126 - b_ptr, #[K, N]
127 - c_ptr, #[max_tokens, N]
128 - expert_id,
129 - compute_type: tl.constexpr,
130 - # Dimensions
131 - M,
132 - N,
133 - K,
134 - # Quantization data
135 - a_scale_ptr,
136 - b_scale_ptr,
137 - b_zp_ptr,
138 - # strides
139 - stride_am,
140 - stride_ak,
141 - stride_bk,
142 - stride_bn,
143 - stride_cm,
144 - stride_cn,
145 - stride_asm,
146 - stride_ask,
147 - stride_bse,
148 - stride_bsk,
149 - stride_bsn,
150 - # Blockwise quantization data
151 - group_n,
152 - group_k,
153 - # Quantization schemes
154 - use_fp8_w8a8: tl.constexpr,
155 - use_int8_w8a16: tl.constexpr,
156 - # Kernel config
157 - BLOCK_M: tl.constexpr,
158 - BLOCK_N: tl.constexpr,
159 - BLOCK_K: tl.constexpr):
139 + a_ptr, #[max_tokens, K]
140 + b_ptr, #[K, N]
141 + c_ptr, #[max_tokens, N]
142 + expert_id,
143 + compute_type: tl.constexpr,
144 + # Dimensions
145 + M,
146 + N,
147 + K,
148 + # Quantization data
149 + a_scale_ptr,
150 + b_scale_ptr,
151 + b_zp_ptr,
152 + # strides
153 + stride_am: tl.int64,
154 + stride_ak: tl.int64,
155 + stride_bk: tl.int64,
156 + stride_bn: tl.int64,
157 + stride_cm: tl.int64,
158 + stride_cn: tl.int64,
159 + stride_ase: tl.int64,
160 + stride_asm: tl.int64,
161 + stride_ask: tl.int64,
162 + stride_bse: tl.int64,
163 + stride_bsk: tl.int64,
164 + stride_bsn: tl.int64,
165 + # offsets
166 + offs_bn,
167 + # Blockwise quantization data
168 + group_n,
169 + group_k,
170 + # Quantization schemes
171 + use_fp8_w8a8: tl.constexpr,
172 + use_int8_w8a16: tl.constexpr,
173 + per_act_token_quant: tl.constexpr,
174 + # Kernel config
175 + BLOCK_M: tl.constexpr,
176 + BLOCK_N: tl.constexpr,
177 + BLOCK_K: tl.constexpr,
178 +):
179
180 offs_m = tl.arange(0, BLOCK_M)
181 offs_n = tl.arange(0, BLOCK_N) % N
182 offs_k = tl.arange(0, BLOCK_K)
183 mask_m = offs_m < M
184
185 + # Make grids of a + b pointers
186 a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
187 b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
188
@@ -179,6 +199,7 @@ def expert_triton_kernel(
199 # (A has M rows).
200 stride_ak,
201 stride_bk,
202 + stride_ase,
203 stride_asm,
204 stride_ask,
205 stride_bse,
@@ -187,6 +208,7 @@ def expert_triton_kernel(
208 # Offsets and masks
209 offs_m,
210 offs_n,
211 + offs_bn,
212 mask_m,
213 # Block size for block-wise quantization
214 group_n,
@@ -197,7 +219,8 @@ def expert_triton_kernel(
219 BLOCK_K,
220 compute_type,
221 use_fp8_w8a8,
200 - use_int8_w8a16)
222 + use_int8_w8a16,
223 + per_act_token_quant)
224
225 # store in C
226 offs_cn = tl.arange(0, BLOCK_N)
@@ -208,53 +231,57 @@ def expert_triton_kernel(
231
232 @triton.jit
233 def batched_triton_kernel(
211 - a_ptr, # [E, max_num_tokens, K]
212 - b_ptr, # [E, K, N]
213 - c_ptr, # [E, max_num_tokens, N]
214 - expert_num_tokens, # [E]
215 - compute_type: tl.constexpr,
216 - # Dimensions
217 - max_num_tokens,
218 - K,
219 - N,
220 - # Quantization data
221 - a_scale_ptr,
222 - b_scale_ptr,
223 - b_zp_ptr,
224 - # The stride variables represent how much to increase the ptr by when
225 - # moving by 1 element in a particular dimension. E.g. `stride_am` is
226 - # how much to increase `a_ptr` by to get the element one row down
227 - # (A has M rows).
228 - stride_ae,
229 - stride_am,
230 - stride_ak,
231 - stride_be,
232 - stride_bk,
233 - stride_bn,
234 - stride_ce,
235 - stride_cm,
236 - stride_cn,
237 - stride_asm,
238 - stride_ask,
239 - stride_bse,
240 - stride_bsk,
241 - stride_bsn,
242 - # Blockwise quantization data
243 - group_n: tl.constexpr,
244 - group_k: tl.constexpr,
245 - # Quantization schemes
246 - use_fp8_w8a8: tl.constexpr,
247 - use_int8_w8a16: tl.constexpr,
248 - # Kernel config
249 - BLOCK_M: tl.constexpr,
250 - BLOCK_N: tl.constexpr,
251 - BLOCK_K: tl.constexpr):
234 + a_ptr, # [E, max_num_tokens, K]
235 + b_ptr, # [E, K, N]
236 + c_ptr, # [E, max_num_tokens, N]
237 + expert_num_tokens, # [E]
238 + compute_type: tl.constexpr,
239 + # Dimensions
240 + max_num_tokens,
241 + K,
242 + N,
243 + # Quantization data
244 + a_scale_ptr,
245 + b_scale_ptr,
246 + b_zp_ptr,
247 + # The stride variables represent how much to increase the ptr by when
248 + # moving by 1 element in a particular dimension. E.g. `stride_am` is
249 + # how much to increase `a_ptr` by to get the element one row down
250 + # (A has M rows).
251 + stride_ae: tl.int64,
252 + stride_am: tl.int64,
253 + stride_ak: tl.int64,
254 + stride_be: tl.int64,
255 + stride_bk: tl.int64,
256 + stride_bn: tl.int64,
257 + stride_ce: tl.int64,
258 + stride_cm: tl.int64,
259 + stride_cn: tl.int64,
260 + stride_ase: tl.int64,
261 + stride_asm: tl.int64,
262 + stride_ask: tl.int64,
263 + stride_bse: tl.int64,
264 + stride_bsk: tl.int64,
265 + stride_bsn: tl.int64,
266 + # Blockwise quantization data
267 + group_n: tl.constexpr,
268 + group_k: tl.constexpr,
269 + # Quantization schemes
270 + use_fp8_w8a8: tl.constexpr,
271 + use_int8_w8a16: tl.constexpr,
272 + per_act_token_quant: tl.constexpr,
273 + # Kernel config
274 + BLOCK_M: tl.constexpr,
275 + BLOCK_N: tl.constexpr,
276 + BLOCK_K: tl.constexpr,
277 +):
278 expert_id = tl.program_id(axis=0)
279 e_num_tokens = tl.load(expert_num_tokens + expert_id)
280 if e_num_tokens == 0:
281 # Early exit
282 return
283
284 + # axis 1 is M_blocks * N_blocks
285 pid_mn = tl.program_id(axis=1)
286 #num_pid_m = tl.cdiv(max_num_tokens, BLOCK_M)
287 num_pid_n = tl.cdiv(N, BLOCK_N)
@@ -275,6 +302,16 @@ def batched_triton_kernel(
302 c_ptr = (c_ptr + expert_id * stride_ce + cta_m_start * stride_cm +
303 cta_n_start * stride_cn)
304
305 + offs_bn = (pid_n * BLOCK_N + tl.arange(0, BLOCK_N).to(tl.int64)) % N
306 +
307 + if use_fp8_w8a8:
308 + a_scale_ptr = a_scale_ptr + expert_id * stride_ase
309 + b_scale_ptr = b_scale_ptr + expert_id * stride_bse
310 +
311 + # block-wise
312 + if group_k > 0 and group_n > 0 or per_act_token_quant:
313 + a_scale_ptr = a_scale_ptr + cta_m_start * stride_asm
314 +
315 expert_triton_kernel(
316 a_ptr,
317 b_ptr,
@@ -294,17 +331,21 @@ def batched_triton_kernel(
331 stride_bn,
332 stride_cm,
333 stride_cn,
334 + stride_ase,
335 stride_asm,
336 stride_ask,
337 stride_bse,
338 stride_bsk,
339 stride_bsn,
340 + # offsets
341 + offs_bn,
342 # Blockwise quantization data
343 group_n,
344 group_k,
345 # Quantization schemes
346 use_fp8_w8a8,
347 use_int8_w8a16,
348 + per_act_token_quant,
349 # Kernel config
350 BLOCK_M,
351 BLOCK_N,
@@ -326,6 +367,7 @@ def invoke_moe_batched_triton_kernel(
367 use_int8_w8a16: bool,
368 use_int4_w4a16: bool,
369 config: dict[str, int],
370 + per_act_token_quant: bool,
371 block_shape: Optional[list[int]] = None):
372
373 assert not use_int4_w4a16
@@ -340,6 +382,42 @@ def invoke_moe_batched_triton_kernel(
382 grid = (expert_num_tokens.size(0), triton.cdiv(max_num_tokens, BLOCK_M) *
383 triton.cdiv(B.size(1), BLOCK_N))
384
385 + A_scale = normalize_batched_scales_shape(A_scale,
386 + expert_num_tokens.shape[0])
387 +
388 + if B_scale is not None and B_scale.ndim == 1:
389 + assert B_scale.numel() == expert_num_tokens.shape[0]
390 + B_scale = B_scale.view(-1, 1, 1)
391 +
392 + assert A_scale is None or A_scale.ndim == 3, (
393 + f"{0 if A_scale is None else A_scale.shape}")
394 + assert B_scale is None or B_scale.ndim == 1 or B_scale.ndim == 3, (
395 + f"{0 if B_scale is None else B_scale.shape}")
396 +
397 + if B_scale is not None:
398 + if B_scale.ndim == 1:
399 + stride_bse = 1
400 + stride_bsk = 0
401 + stride_bsn = 0
402 + else:
403 + stride_bse = B_scale.stride(0)
404 + stride_bsk = B_scale.stride(2)
405 + stride_bsn = B_scale.stride(1)
406 +
407 + else:
408 + stride_bse = 0
409 + stride_bsk = 0
410 + stride_bsn = 0
411 +
412 + if A_scale is not None:
413 + stride_ase = A_scale.stride(0)
414 + stride_asm = A_scale.stride(1)
415 + stride_ask = A_scale.stride(2)
416 + else:
417 + stride_ase = 0
418 + stride_asm = 0
419 + stride_ask = 0
420 +
421 batched_triton_kernel[grid](
422 A,
423 B,
@@ -364,17 +442,19 @@ def invoke_moe_batched_triton_kernel(
442 C.stride(0),
443 C.stride(1),
444 C.stride(2),
367 - A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0,
368 - A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0,
369 - B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0,
370 - B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0,
371 - B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0,
445 + stride_ase,
446 + stride_asm,
447 + stride_ask,
448 + stride_bse,
449 + stride_bsk,
450 + stride_bsn,
451 # Blockwise quantization data
452 0 if block_shape is None else block_shape[0],
453 0 if block_shape is None else block_shape[1],
454 # Quantization schemes
455 use_fp8_w8a8,
456 use_int8_w8a16,
457 + per_act_token_quant,
458 # Kernel config
459 BLOCK_M=BLOCK_M,
460 BLOCK_N=BLOCK_N,
@@ -391,15 +471,15 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
471 def __init__(
472 self,
473 max_num_tokens: int,
394 - world_size: int,
395 - dp_size: int,
474 + num_local_experts: int,
475 + num_dispatchers: int,
476 rank: int,
477 ):
478 super().__init__()
399 - self.world_size = world_size
400 - self.dp_size = dp_size
401 - self.rank = rank
479 self.max_num_tokens = max_num_tokens
480 + self.num_local_experts = num_local_experts
481 + self.rank = rank
482 + self.num_dispatchers_ = num_dispatchers
483
484 @property
485 def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -411,6 +491,9 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
491 def topk_indices_dtype(self) -> Optional[torch.dtype]:
492 return None
493
494 + def num_dispatchers(self) -> int:
495 + return self.num_dispatchers_
496 +
497 def prepare(
498 self,
499 a1: torch.Tensor,
@@ -442,9 +525,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
525 dtype=torch.int,
526 device=a1.device)
527
445 - assert num_experts % self.world_size == 0
446 -
447 - num_local_experts = num_experts // self.world_size
528 + num_local_experts = self.num_local_experts
529
530 if quant_config.quant_dtype is None:
531 b_type = a1.dtype
@@ -456,21 +537,53 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
537 dtype=b_type,
538 device=a1.device)
539
459 - b_a1_scale = None
540 + if quant_config.is_quantized:
541 + scale_shape = quant_config.batched_scale_shape(
542 + num_local_experts, self.max_num_tokens, hidden_dim)
543
461 - assert quant_config.quant_dtype is None, "quantization NYI"
544 + b_a1_scale = torch.empty(scale_shape,
545 + dtype=torch.float32,
546 + device=a1.device)
547 + else:
548 + assert a1_scale is None
549 + b_a1_scale = None
550
551 first_expert = num_local_experts * self.rank
552 last_expert = first_expert + num_local_experts
553
554 + a1_scale = normalize_scales_shape(a1_scale)
555 + a2_scale = normalize_scales_shape(a2_scale)
556 +
557 for expert_id in range(first_expert, last_expert):
558 topks = torch.any(topk_ids == expert_id, dim=1).flatten()
559 rows = torch.count_nonzero(topks.flatten())
560 if rows == 0:
561 continue
562 idx = expert_id - first_expert
472 - b_a1[idx, :rows, :] = a1[:topks.numel()][topks]
563 tokens_per_expert[idx] = rows
564 + rhs = a1[:topks.numel()][topks]
565 + if quant_config.quant_dtype is not None:
566 + if a1_scale is not None:
567 + if quant_config.is_per_act_token:
568 + rhs_a1_scale = a1_scale[:topks.numel()][topks]
569 + else:
570 + rhs_a1_scale = a1_scale
571 + else:
572 + rhs_a1_scale = None
573 + b_a1[idx, :rows, :], b_s = moe_kernel_quantize_input(
574 + rhs,
575 + rhs_a1_scale,
576 + quant_config.quant_dtype,
577 + quant_config.per_act_token_quant,
578 + quant_config.block_shape,
579 + )
580 + assert b_s is not None
581 + if quant_config.is_per_act_token:
582 + b_a1_scale[idx, :rows] = b_s[:rows]
583 + else:
584 + b_a1_scale[idx, :b_s.shape[0]] = b_s
585 + else:
586 + b_a1[idx, :rows, :] = rhs
587
588 assert b_a1_scale is None or b_a1_scale.ndim == 3
589
@@ -514,8 +627,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
627 def __init__(
628 self,
629 max_num_tokens: int,
517 - world_size: int,
518 - dp_size: int,
630 + num_dispatchers: int,
631 use_fp8_w8a8: bool = False,
632 use_int8_w8a8: bool = False,
633 use_int8_w8a16: bool = False,
@@ -532,13 +644,11 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
644 per_act_token_quant=per_act_token_quant,
645 block_shape=block_shape,
646 ))
535 - assert not use_fp8_w8a8, "NYI"
647 assert not use_int8_w8a8, "NYI"
648 assert not use_int8_w8a16, "NYI"
649 assert not use_int4_w4a16, "NYI"
650 self.max_num_tokens = max_num_tokens
540 - self.world_size = world_size
541 - self.dp_size = dp_size
651 + self.num_dispatchers = num_dispatchers
652
653 @property
654 def activation_formats(
@@ -565,11 +675,21 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
675 local_num_experts: int,
676 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
677 assert a.dim() == 2
568 - num_dp = self.dp_size
678 + num_dp = self.num_dispatchers
679 num_experts = local_num_experts
680 workspace13 = (num_experts, self.max_num_tokens * num_dp, K)
681 workspace2 = (self.max_num_tokens * num_dp, N)
572 - return (workspace13, workspace2, workspace13, a.dtype)
682 + output = workspace13
683 + return (workspace13, workspace2, output, a.dtype)
684 +
685 + def dequant(self, t: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
686 + assert self.quant_config.is_quantized
687 + f32 = torch.float32
688 + if (self.quant_config.is_per_act_token
689 + or self.quant_config.is_per_tensor):
690 + return t.to(f32) * scale
691 + else:
692 + return t.to(f32) * group_broadcast(scale, t.shape)
693
694 def apply(
695 self,
@@ -612,9 +732,95 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
732 continue
733
734 tmp = _resize_cache(workspace2, (num, N))
615 - input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
616 - self.activation(activation, tmp, input)
617 - output[expert, :num, :] = tmp @ w2[expert].transpose(0, 1)
735 +
736 + if self.quant_config.is_quantized:
737 + assert a1q_scale is not None and w1_scale is not None
738 + input = self.dequant(hidden_states[expert, :, :],
739 + a1q_scale[expert])
740 + w1_dq = self.dequant(w1[expert], w1_scale[expert])
741 + input = input[:num] @ w1_dq.transpose(0, 1)
742 + else:
743 + input = hidden_states[expert, :num, :] @ w1[expert].transpose(
744 + 0, 1)
745 +
746 + self.activation(activation, tmp, input.to(tmp.dtype))
747 +
748 + if self.quant_config.is_quantized:
749 + assert w2_scale is not None
750 + w2_dq = self.dequant(w2[expert], w2_scale[expert])
751 + else:
752 + w2_dq = w2[expert]
753 +
754 + output[expert, :num, :] = tmp @ w2_dq.transpose(0, 1).to(tmp.dtype)
755 +
756 +
757 +def batched_moe_kernel_quantize_input(
758 + A: torch.Tensor,
759 + A_scale: Optional[torch.Tensor],
760 + num_tokens: int,
761 + E: int,
762 + N: int,
763 + expert_num_tokens: torch.Tensor,
764 + qtype: Optional[torch.dtype],
765 + per_act_token_quant: bool,
766 + block_shape: Optional[list[int]] = None,
767 +) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
768 + if (torch.compiler.is_compiling()
769 + or torch.cuda.is_current_stream_capturing()):
770 + # Note: this does a bunch of extra work because expert_num_tokens is
771 + # ignored but it does support torch.compile + cudagraphs.
772 + hidden_dim = A.size(-1)
773 + assert A_scale is None or A_scale.ndim <= 2, (
774 + f"{A_scale.shape if A_scale is not None else None}")
775 + A_q, A_q_scale = moe_kernel_quantize_input(A.view(-1,
776 + hidden_dim), A_scale,
777 + qtype, per_act_token_quant,
778 + block_shape)
779 + A_q = A_q.view(E, -1, hidden_dim)
780 + A_q_scale = normalize_batched_scales_shape(A_q_scale, E)
781 +
782 + return A_q, A_q_scale
783 + elif qtype is None:
784 + return A, normalize_batched_scales_shape(A_scale, E)
785 + else:
786 + A_q = torch.empty_like(A, dtype=qtype)
787 +
788 + if per_act_token_quant:
789 + assert block_shape is None
790 + scale_shape = (E, num_tokens, 1)
791 + elif block_shape is not None:
792 + _, block_k = block_shape
793 + k_tiles = (A.shape[-1] + block_k - 1) // block_k
794 + scale_shape = (E, num_tokens, k_tiles)
795 + else:
796 + scale_shape = (E, 1, 1)
797 +
798 + A_q_scale = torch.zeros(scale_shape,
799 + dtype=torch.float32,
800 + device=A.device)
801 +
802 + num_experts = expert_num_tokens.numel()
803 +
804 + A_scale = normalize_batched_scales_shape(A_scale, num_experts)
805 +
806 + for e in range(E):
807 + num_tokens = int(expert_num_tokens[e].item())
808 + if num_tokens > 0:
809 + if A_scale is not None:
810 + scales = A_scale[e, :min(num_tokens, A_scale.shape[1])]
811 + else:
812 + scales = None
813 + A_q[e, :num_tokens], tmp_scale = moe_kernel_quantize_input(
814 + A[e, :num_tokens],
815 + scales,
816 + qtype,
817 + per_act_token_quant,
818 + block_shape,
819 + )
820 + assert tmp_scale is not None
821 + A_q_scale[e, :tmp_scale.shape[0]] = tmp_scale
822 +
823 + return A_q, A_q_scale
824
825
826 class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
@@ -627,8 +833,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
833 def __init__(
834 self,
835 max_num_tokens: int,
630 - world_size: int,
631 - dp_size: int,
836 + num_dispatchers: int,
837 use_fp8_w8a8: bool = False,
838 use_int8_w8a8: bool = False,
839 use_int8_w8a16: bool = False,
@@ -648,17 +853,14 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
853 assert not use_int8_w8a8, "NYI"
854 assert not use_int8_w8a16, "NYI"
855 assert not use_int4_w4a16, "NYI"
856 + assert max_num_tokens > 0
857 + assert num_dispatchers > 0
858 self.use_fp8_w8a8 = use_fp8_w8a8
859 self.use_int8_w8a8 = use_int8_w8a8
860 self.use_int4_w4a16 = use_int4_w4a16
861 self.use_int8_w8a16 = use_int8_w8a16
862 self.max_num_tokens = max_num_tokens
656 - self.world_size = world_size
657 - self.dp_size = dp_size
658 - assert world_size > 0
659 - assert dp_size > 0
660 - assert dp_size <= world_size
661 - assert max_num_tokens > 0
863 + self.num_dispatchers = num_dispatchers
864
865 @property
866 def activation_formats(
@@ -685,7 +887,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
887 local_num_experts: int,
888 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
889 assert a.dim() == 2
688 - num_dp = self.world_size
890 + num_dp = self.num_dispatchers
891 num_experts = local_num_experts
892 max_num_tokens = self.max_num_tokens
893 workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
@@ -772,51 +974,48 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
974 if self.use_fp8_w8a8:
975 intermediate_cache1.fill_(0)
976
977 + a1q_scale = normalize_batched_scales_shape(a1q_scale, E)
978 +
979 # MM1
776 - invoke_moe_batched_triton_kernel(A=hidden_states,
777 - B=w1,
778 - C=intermediate_cache1,
779 - expert_num_tokens=expert_num_tokens,
780 - compute_type=compute_type,
781 - A_scale=a1q_scale,
782 - B_scale=w1_scale,
783 - B_zp=w1_zp,
784 - use_fp8_w8a8=self.use_fp8_w8a8,
785 - use_int8_w8a16=self.use_int8_w8a16,
786 - use_int4_w4a16=self.use_int4_w4a16,
787 - config=config,
788 - block_shape=self.block_shape)
980 + invoke_moe_batched_triton_kernel(
981 + A=hidden_states,
982 + B=w1,
983 + C=intermediate_cache1,
984 + expert_num_tokens=expert_num_tokens,
985 + compute_type=compute_type,
986 + A_scale=a1q_scale,
987 + B_scale=w1_scale,
988 + B_zp=w1_zp,
989 + use_fp8_w8a8=self.use_fp8_w8a8,
990 + use_int8_w8a16=self.use_int8_w8a16,
991 + use_int4_w4a16=self.use_int4_w4a16,
992 + config=config,
993 + per_act_token_quant=self.per_act_token_quant,
994 + block_shape=self.block_shape)
995
996 intermediate_cache2.fill_(0)
997
792 - # TODO: would be nice to use expert_num_tokens here to reduce
793 - # garbage compute
998 + # TODO (bnell): use triton utility from batched deep gemm.
999 self.activation(activation, intermediate_cache2.view(-1, N // 2),
1000 intermediate_cache1.view(-1, N))
1001
797 - ic2_hidden_size = intermediate_cache2.size(-1)
798 - intermediate_cache2 = intermediate_cache2.view(-1, ic2_hidden_size)
799 -
800 - qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
801 - A=intermediate_cache2,
802 - A_scale=a2_scale,
803 - quant_dtype=self.quant_dtype,
1002 + qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
1003 + intermediate_cache2, a2_scale, max_num_tokens, E, N,
1004 + expert_num_tokens, self.quant_dtype, self.per_act_token_quant,
1005 + self.block_shape)
1006 +
1007 + invoke_moe_batched_triton_kernel(
1008 + A=qintermediate_cache2,
1009 + B=w2,
1010 + C=output,
1011 + expert_num_tokens=expert_num_tokens,
1012 + compute_type=compute_type,
1013 + A_scale=a2q_scale,
1014 + B_scale=w2_scale,
1015 + B_zp=w2_zp,
1016 + use_fp8_w8a8=self.use_fp8_w8a8,
1017 + use_int8_w8a16=self.use_int8_w8a16,
1018 + use_int4_w4a16=self.use_int4_w4a16,
1019 + config=config,
1020 per_act_token_quant=self.per_act_token_quant,
1021 block_shape=self.block_shape)
806 -
807 - qintermediate_cache2 = qintermediate_cache2.view(
808 - (E, -1, ic2_hidden_size))
809 -
810 - invoke_moe_batched_triton_kernel(A=qintermediate_cache2,
811 - B=w2,
812 - C=output,
813 - expert_num_tokens=expert_num_tokens,
814 - compute_type=compute_type,
815 - A_scale=a2q_scale,
816 - B_scale=w2_scale,
817 - B_zp=w2_zp,
818 - use_fp8_w8a8=self.use_fp8_w8a8,
819 - use_int8_w8a16=self.use_int8_w8a16,
820 - use_int4_w4a16=self.use_int4_w4a16,
821 - config=config,
822 - block_shape=self.block_shape)
vllm/model_executor/layers/fused_moe/fused_moe.py
+2
@@ -1127,6 +1127,8 @@ def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:
1127 return torch_vllm_outplace_fused_experts
1128
1129
1130 +# TODO (bnell): replace this with modular op. Can get rid of inplace/outplace
1131 +# torch ops.
1132 def fused_experts(hidden_states: torch.Tensor,
1133 w1: torch.Tensor,
1134 w2: torch.Tensor,
vllm/model_executor/layers/fused_moe/layer.py
+12 -22
@@ -14,7 +14,6 @@ import vllm.envs as envs
14 from vllm.config import get_current_vllm_config
15 from vllm.distributed import (get_dp_group, get_ep_group,
16 get_tensor_model_parallel_world_size,
17 - get_world_group,
17 tensor_model_parallel_all_reduce)
18 from vllm.distributed.eplb.eplb_state import EplbState
19 from vllm.forward_context import ForwardContext, get_forward_context
@@ -114,6 +113,9 @@ class FusedMoEMethodBase(QuantizeMethodBase):
113 hidden_dim_scale_bytes=hidden_scale_bytes,
114 )
115
116 + num_dispatchers = (all2all_manager.world_size //
117 + all2all_manager.tp_group.world_size)
118 +
119 # Intranode pplx a2a takes a group name while internode does not.
120 if not all2all_manager.internode:
121 all_to_all_args[
@@ -124,10 +126,8 @@ class FusedMoEMethodBase(QuantizeMethodBase):
126 prepare_finalize = PplxPrepareAndFinalize(
127 handle,
128 max_num_tokens=moe.max_num_tokens,
127 - world_size=all2all_manager.world_size,
128 - rank=all2all_manager.rank,
129 - # dp_size actually means tp_size, bug in pplx kernels
130 - dp_size=all2all_manager.tp_group.world_size,
129 + num_local_experts=moe.num_local_experts,
130 + num_dispatchers=num_dispatchers,
131 )
132 elif moe.use_deepep_ht_kernels:
133 assert moe.dp_size == all2all_manager.dp_world_size
@@ -136,16 +136,13 @@ class FusedMoEMethodBase(QuantizeMethodBase):
136 handle = all2all_manager.get_handle(all_to_all_args)
137 prepare_finalize = DeepEPHTPrepareAndFinalize(
138 handle,
139 - world_size=all2all_manager.world_size,
140 - rank=all2all_manager.rank,
139 + num_dispatchers=all2all_manager.world_size,
140 dp_size=all2all_manager.dp_world_size,
141 rank_expert_offset=all2all_manager.rank *
142 moe.num_local_experts,
143 )
144
145 elif moe.use_deepep_ll_kernels:
147 - assert moe.dp_size == all2all_manager.dp_world_size
148 -
146 all_to_all_args = dict(
147 max_num_tokens_per_dp_rank=moe.max_num_tokens,
148 token_hidden_size=moe.hidden_dim,
@@ -168,8 +165,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
165 prepare_finalize = DeepEPLLPrepareAndFinalize(
166 handle,
167 max_tokens_per_rank=moe.max_num_tokens,
171 - world_size=all2all_manager.world_size,
172 - dp_size=all2all_manager.dp_world_size,
168 + num_dispatchers=all2all_manager.world_size,
169 use_fp8_dispatch=use_fp8_dispatch,
170 )
171
@@ -245,18 +241,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
241
242 assert self.fused_experts == fused_experts
243
248 - all2all_manager = get_ep_group().device_communicator.all2all_manager
249 - assert all2all_manager is not None
250 -
244 if (prepare_finalize.activation_format ==
245 FusedMoEActivationFormat.BatchedExperts):
246 logger.debug("BatchedTritonExperts %s", self.moe)
254 - assert self.moe.dp_size == all2all_manager.dp_world_size
247 return BatchedTritonExperts(
248 max_num_tokens=self.moe.max_num_tokens,
257 - world_size=all2all_manager.world_size,
258 - # dp_size actually means tp_size, bug in pplx kernels
259 - dp_size=all2all_manager.tp_group.world_size,
249 + num_dispatchers=prepare_finalize.num_dispatchers(),
250 )
251 else:
252 logger.debug("TritonExperts %s", self.moe)
@@ -652,14 +642,12 @@ class FusedMoE(torch.nn.Module):
642 get_tensor_model_parallel_world_size())
643 dp_size_ = (dp_size
644 if dp_size is not None else get_dp_group().world_size)
655 - world_size_ = get_world_group().world_size
645
646 vllm_config = get_current_vllm_config()
647 self.moe_parallel_config: FusedMoEParallelConfig = (
648 FusedMoEParallelConfig.make(
649 tp_size_=tp_size_,
650 dp_size_=dp_size_,
662 - world_size_=world_size_,
651 vllm_parallel_config=vllm_config.parallel_config))
652
653 self.global_num_experts = num_experts + num_redundant_experts
@@ -1186,9 +1174,9 @@ class FusedMoE(torch.nn.Module):
1174 logical_replica_count: Optional[torch.Tensor] = None,
1175 ) -> tuple[torch.Tensor, torch.Tensor]:
1176 """
1189 - Route the input hidden states to the top-k experts based on the
1177 + Route the input hidden states to the top-k experts based on the
1178 router logits.
1191 -
1179 +
1180 Returns:
1181 (topk_weights, topk_ids) (tuple[torch.Tensor, torch.Tensor]):
1182 The weights and *global physical* expert ids of the top-k experts.
@@ -1299,6 +1287,8 @@ class FusedMoE(torch.nn.Module):
1287
1288 topk_ids = topk_ids.to(dtype=indices_type)
1289
1290 + assert topk_ids.dtype == indices_type or indices_type is None
1291 +
1292 return topk_weights, topk_ids
1293
1294 def must_reduce_shared_expert_outputs(self) -> bool:
vllm/model_executor/layers/fused_moe/modular_kernel.py
+4
@@ -193,6 +193,10 @@ class FusedMoEPrepareAndFinalize(ABC):
193 """
194 raise NotImplementedError
195
196 + @abstractmethod
197 + def num_dispatchers(self) -> int:
198 + raise NotImplementedError
199 +
200
201 class FusedMoEPermuteExpertsUnpermute(ABC):
202 """
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+61 -34
@@ -8,7 +8,7 @@ import torch
8 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10 from vllm.model_executor.layers.fused_moe.utils import (
11 - moe_kernel_quantize_input)
11 + _validate_scale_shape, moe_kernel_quantize_input)
12 from vllm.utils import cdiv, round_up
13
14
@@ -32,16 +32,16 @@ def pplx_hidden_dim_scale_bytes(
32 elem_size = torch.float32.itemsize
33
34 if per_act_token_quant:
35 - # per-token
35 + # per-token (M x 1)
36 assert block_shape is None
37 hidden_scale_bytes = elem_size
38 elif block_shape is not None:
39 - # per-group
39 + # per-group (M x K_tiles)
40 block_size = block_shape[1]
41 num_blocks = cdiv(hidden_dim, block_size)
42 hidden_scale_bytes = num_blocks * elem_size
43 else:
44 - # per-tensor
44 + # per-tensor (1 x 1)
45 hidden_scale_bytes = elem_size
46 else:
47 hidden_dim_bytes = hidden_dim * in_dtype.itemsize
@@ -53,25 +53,22 @@ def pplx_hidden_dim_scale_bytes(
53 )
54
55
56 -# The max_num_tokens, world_size and dp_size must be the same
57 -# as the ones used to create the AllToAll.
56 class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
57
58 def __init__(
59 self,
60 a2a: pplx.AllToAll,
61 max_num_tokens: int,
64 - world_size: int,
65 - rank: int,
66 - dp_size: int,
62 + num_local_experts: int,
63 + num_dispatchers: int,
64 ):
65 super().__init__()
66 assert max_num_tokens > 0
67 + assert num_local_experts > 0
68 self.a2a = a2a
69 self.max_num_tokens = max_num_tokens
72 - self.world_size = world_size
73 - self.rank = rank
74 - self.dp_size = dp_size
70 + self.num_local_experts = num_local_experts
71 + self.num_dispatchers_ = num_dispatchers
72
73 @property
74 def activation_format(self) -> mk.FusedMoEActivationFormat:
@@ -83,6 +80,9 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
80 def topk_indices_dtype(self) -> Optional[torch.dtype]:
81 return torch.uint32
82
83 + def num_dispatchers(self) -> int:
84 + return self.num_dispatchers_
85 +
86 def prepare(
87 self,
88 a1: torch.Tensor,
@@ -120,42 +120,64 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
120 per_act_token_quant=quant_config.per_act_token_quant,
121 block_shape=quant_config.block_shape)
122
123 + _validate_scale_shape(a1q, a1q_scale, quant_config.per_act_token_quant,
124 + quant_config.block_shape)
125 +
126 if a1q_scale is not None:
124 - if a1q_scale.numel() == 1:
125 - orig_a_scale_block_shape = 1
126 - else:
127 - orig_a_scale_block_shape = a1q_scale.shape[-1]
128 - a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols)
127 + scalar_scales = a1q_scale.numel() == 1
128 +
129 + # pplx requires 2-d scales even for scalar scales
130 + if a1q_scale.dim() <= 1:
131 + assert scalar_scales
132 + a1q_scale = a1q_scale.view(1, 1)
133 +
134 + orig_a_scale_block_shape = a1q_scale.shape[-1]
135
130 - # rem_experts need to be 0 for pplx to work properly.
131 - rem_experts = num_experts % self.world_size
132 - assert rem_experts == 0
133 - num_local_experts = ((num_experts // self.world_size) +
134 - (1 if self.rank < rem_experts else 0))
136 + if not quant_config.is_block_quantized:
137 + # TODO (bnell): use group_broadcast instead?
138 + a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols)
139 +
140 + assert a1q_scale is None or a1q_scale.ndim == 2, \
141 + f"{0 if a1q_scale is None else (a1q_scale.ndim, a1q_scale.shape)}"
142
143 expert_num_tokens = torch.empty(
137 - num_local_experts,
144 + self.num_local_experts,
145 dtype=torch.int32,
146 device=device,
147 )
148
142 - num_dp = self.world_size // self.dp_size
149 expert_x = torch.empty(
144 - (num_local_experts, self.max_num_tokens * num_dp, hidden_dim),
150 + (self.num_local_experts,
151 + self.max_num_tokens * self.num_dispatchers(), hidden_dim),
152 dtype=a1q.dtype,
153 device=device,
154 )
155
156 expert_x_scale: Optional[torch.Tensor] = None
157 if a1q.dtype.itemsize == 1:
151 - block_size = (quant_config.block_shape[1]
152 - if quant_config.block_shape is not None else 1)
158 + if quant_config.is_per_act_token:
159 + # (M x 1) -> (E x M x K)
160 + final_dim = expert_x.size(2)
161 + elif quant_config.is_per_tensor:
162 + # (1 x 1) -> (E x 1 x 1)
163 + final_dim = 1
164 + else:
165 + # (M x K_tiles) -> (E x M x K_tiles)
166 + assert quant_config.block_shape is not None
167 + num_blocks = cdiv(expert_x.size(2),
168 + quant_config.block_shape[1])
169 + final_dim = num_blocks
170 +
171 + expert_x_scale_shape = (
172 + self.num_local_experts,
173 + expert_x.size(1),
174 + round_up(final_dim, 4) # round up for alignment
175 + )
176 +
177 expert_x_scale = torch.empty(
154 - (num_local_experts, expert_x.size(1),
155 - round_up(
156 - (expert_x.size(2) + block_size - 1) // block_size, 4)),
178 + expert_x_scale_shape,
179 dtype=torch.float32,
158 - device=device,
180 + device=expert_x.device,
181 )
182
183 # This argument is optional, defaults to indices.size(0)
@@ -171,8 +193,10 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
193 indices=topk_ids,
194 bound_m=bound_m,
195 )
196 +
197 if expert_x_scale is not None:
198 expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape]
199 + assert expert_x_scale.ndim == 3
200
201 return expert_x, expert_x_scale, expert_num_tokens, None, None
202
@@ -184,13 +208,16 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
208 topk_ids: torch.Tensor,
209 apply_router_weight_on_input: bool,
210 ) -> None:
187 - num_tokens = output.size(0) # M
211 # This argument is optional
212 # There's not much point setting this unless it is != topk_ids.size(0)
213 bound_m: Optional[torch.Tensor] = None
214
192 - assert topk_ids.size(0) == num_tokens, (
193 - f"{topk_ids.size(0)} == {num_tokens}")
215 + # TODO (bnell): fails in test_pplx_moe.py, figure out what's going on
216 + #num_tokens = output.size(0) # M
217 + #assert topk_ids.size(0) == num_tokens, (
218 + # f"{topk_ids.size(0)} == {num_tokens}")
219 + assert topk_ids.size() == topk_weights.size(), (
220 + f"{topk_ids.size()} == {topk_weights.size()}")
221 assert output.size(0) <= self.max_num_tokens, (
222 f"{output.size(0)} <= {self.max_num_tokens}")
223 assert output.size(1) == fused_expert_output.size(-1)
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+3
@@ -24,6 +24,9 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
24 def topk_indices_dtype(self) -> Optional[torch.dtype]:
25 return None
26
27 + def num_dispatchers(self) -> int:
28 + return 1
29 +
30 def prepare(
31 self,
32 a1: torch.Tensor,
vllm/model_executor/layers/fused_moe/utils.py
+34 -3
@@ -99,9 +99,20 @@ def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
99 return m[idx, ...]
100
101
102 -# TODO(bnell): better name
103 -def maybe_fix_scales(scales: Optional[torch.Tensor],
104 - num_experts: int) -> Optional[torch.Tensor]:
102 +def normalize_scales_shape(
103 + scales: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
104 + if scales is not None:
105 + if scales.numel() == 1:
106 + scales = scales.view(1, 1)
107 + else:
108 + scales = scales.view(-1, scales.size(-1))
109 + return scales
110 +
111 +
112 +def normalize_batched_scales_shape(
113 + scales: Optional[torch.Tensor],
114 + num_experts: int,
115 +) -> Optional[torch.Tensor]:
116 if scales is not None and scales.ndim < 3:
117 if scales.numel() == 1:
118 scales = scales.view(1)
@@ -111,3 +122,23 @@ def maybe_fix_scales(scales: Optional[torch.Tensor],
122 scales = scales.view(num_experts, -1, scales.size(-1))
123
124 return scales
125 +
126 +
127 +def _validate_scale_shape(
128 + a: torch.Tensor,
129 + a_scale: Optional[torch.Tensor],
130 + per_act_token_quant: bool,
131 + block_shape: Optional[list[int]],
132 +) -> None:
133 + if a_scale is None:
134 + return
135 +
136 + if not per_act_token_quant and block_shape is None:
137 + assert a_scale.numel() == 1, f"{a_scale.shape}"
138 + elif per_act_token_quant:
139 + assert a_scale.shape[0] == a.shape[0] and a_scale.shape[1] == 1, (
140 + f"{a_scale.shape[0]} == {a.shape[0]} and {a_scale.shape[1]} == 1")
141 + else:
142 + assert block_shape is not None
143 + expected = (a.shape[0], cdiv(a.shape[1], block_shape[1]))
144 + assert a_scale.shape == expected, f"{a_scale.shape} == {expected}"
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+46 -2
@@ -573,6 +573,41 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
573 from vllm.model_executor.layers.fused_moe import fused_experts
574 self.fused_experts_func = fused_experts
575
576 + def select_gemm_impl(
577 + self,
578 + prepare_finalize: FusedMoEPrepareAndFinalize,
579 + moe: FusedMoEConfig,
580 + ) -> FusedMoEPermuteExpertsUnpermute:
581 + from vllm.model_executor.layers.fused_moe import TritonExperts
582 + from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
583 + BatchedTritonExperts)
584 +
585 + assert not self.rocm_aiter_moe_enabled and not self.use_marlin
586 +
587 + logger.debug("BatchedTritonExperts(%s)", self.__class__.__name__)
588 +
589 + if (prepare_finalize.activation_format ==
590 + FusedMoEActivationFormat.BatchedExperts):
591 + max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank(
592 + )
593 + assert max_num_tokens_per_rank is not None
594 +
595 + return BatchedTritonExperts(
596 + max_num_tokens=max_num_tokens_per_rank,
597 + num_dispatchers=prepare_finalize.num_dispatchers(),
598 + use_fp8_w8a8=True,
599 + block_shape=self.quant_config.weight_block_size,
600 + per_act_token_quant=(
601 + self.input_quant.strategy == QuantizationStrategy.TOKEN),
602 + )
603 + else:
604 + return TritonExperts(
605 + use_fp8_w8a8=True,
606 + block_shape=self.quant_config.weight_block_size,
607 + per_act_token_quant=(
608 + self.input_quant.strategy == QuantizationStrategy.TOKEN),
609 + )
610 +
611 def apply(
612 self,
613 layer: torch.nn.Module,
@@ -610,7 +645,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
645 num_expert_group=num_expert_group,
646 custom_routing_function=custom_routing_function,
647 scoring_func=scoring_func,
613 - e_score_correction_bias=e_score_correction_bias)
648 + e_score_correction_bias=e_score_correction_bias,
649 + indices_type=self.topk_indices_dtype,
650 + )
651
652 if self.rocm_aiter_moe_enabled:
653 return self.rocm_aiter_fused_experts_func(
@@ -832,18 +869,25 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
869 use_batched_format = (prepare_finalize.activation_format ==
870 FusedMoEActivationFormat.BatchedExperts)
871
872 + num_dispatchers = prepare_finalize.num_dispatchers()
873 +
874 num_experts = (moe.num_local_experts
875 if use_batched_format else moe.num_experts)
876
877 + logger.debug("CutlassExpertsFp8(%s)", self.__class__.__name__)
878 +
879 experts = CutlassExpertsFp8(
880 num_experts,
881 moe.in_dtype,
882 self.input_quant.strategy == QuantizationStrategy.TOKEN,
883 self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
884 + num_dispatchers=num_dispatchers,
885 use_batched_format=use_batched_format,
886 )
887
846 - self.disable_expert_map = not experts.supports_expert_map()
888 + self.disable_expert_map = (num_dispatchers > 1
889 + or not experts.supports_expert_map())
890 +
891 return experts
892
893 def apply(
vllm/model_executor/layers/quantization/fp8.py
+1 -4
@@ -802,10 +802,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
802 self.quant_config.weight_block_size, False)
803 return BatchedTritonOrDeepGemmExperts(
804 max_num_tokens=max_num_tokens_per_rank,
805 - world_size=prepare_finalize.
806 - world_size, # type: ignore [attr-defined]
807 - dp_size=prepare_finalize.
808 - dp_size, # type: ignore [attr-defined]
805 + num_dispatchers=prepare_finalize.num_dispatchers(),
806 use_fp8_w8a8=True,
807 block_shape=self.quant_config.weight_block_size,
808 per_act_token_quant=False,
vllm/model_executor/models/qwen3_moe.py
+1 -1
@@ -135,7 +135,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
135 router_logits, _ = self.gate(hidden_states)
136 final_hidden_states = self.experts(hidden_states=hidden_states,
137 router_logits=router_logits)
138 - final_hidden_states = final_hidden_states
138 +
139 if self.tp_size > 1:
140 final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
141 final_hidden_states)