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

[Kernel] Integrate batched/masked deepgemm kernel (#19111)

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

Varun Sundar Rabindranath committed Jun 4, 2025 at 17:59 UTC c3fd4d669a4446578ea83bf7c330112339d2d8e1
6 files changed +472 -51
tests/kernels/moe/deepep_utils.py
+4 -1
@@ -162,12 +162,14 @@ def make_deepep_ll_a2a(pg: ProcessGroup,
162 low_latency_mode=True,
163 num_qps_per_rank=deepep_ll_args.num_experts //
164 pgi.world_size)
165 +
166 return DeepEPLLPrepareAndFinalize(
167 buffer=buffer,
168 world_size=pgi.world_size,
169 dp_size=dp_size,
170 max_tokens_per_rank=deepep_ll_args.max_tokens_per_rank,
171 quant_dtype=q_dtype,
172 + block_shape=block_shape,
173 use_fp8_dispatch=deepep_ll_args.use_fp8_dispatch,
174 )
175
@@ -185,4 +187,5 @@ def make_deepep_a2a(pg: ProcessGroup,
187 block_shape)
188
189 assert deepep_ll_args is not None
188 - return make_deepep_ll_a2a(pg, pgi, dp_size, deepep_ll_args, q_dtype)
190 + return make_deepep_ll_a2a(pg, pgi, dp_size, deepep_ll_args, q_dtype,
191 + block_shape)
tests/kernels/moe/test_deepep_deepgemm_moe.py
+166 -24
@@ -1,6 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 """
3 Test DeepEP + DeepGEMM integration
4 +DeepGEMM are gemm kernels specialized for the
5 +fp8 block-quantized case.
6 """
7
8 import dataclasses
@@ -33,10 +35,14 @@ except ImportError:
35 if has_deep_ep:
36 from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
37 DeepEPHTPrepareAndFinalize)
38 + from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
39 + DeepEPLLPrepareAndFinalize)
40
37 - from .deepep_utils import DeepEPHTArgs, make_deepep_a2a
41 + from .deepep_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a
42
43 if has_deep_gemm:
44 + from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
45 + BatchedDeepGemmExperts)
46 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
47 DeepGemmExperts)
48
@@ -53,6 +59,13 @@ requires_deep_gemm = pytest.mark.skipif(
59 P = ParamSpec("P")
60
61
62 +def next_power_of_2(x):
63 + import math
64 + if x == 0:
65 + return 1
66 + return 2**math.ceil(math.log2(x))
67 +
68 +
69 def per_block_cast_to_fp8(
70 x: torch.Tensor,
71 block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
@@ -126,6 +139,9 @@ class TestConfig:
139 n: int
140 num_experts: int
141 block_size: list[int]
142 + # configs for testing low-latency kernels
143 + low_latency: bool
144 + use_fp8_dispatch: Optional[bool] = False
145
146
147 @dataclasses.dataclass
@@ -170,9 +186,43 @@ class TestTensors:
186 config=config)
187
188
173 -def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
174 - num_local_experts: int, q_dtype: Optional[torch.dtype],
175 - block_shape: list[int]) -> FusedMoEModularKernel:
189 +def make_ll_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
190 + max_tokens_per_rank: int, dp_size: int,
191 + hidden_size: int, q_dtype: Optional[torch.dtype],
192 + test_config: TestConfig) -> FusedMoEModularKernel:
193 +
194 + assert test_config.low_latency
195 + assert test_config.use_fp8_dispatch is not None
196 +
197 + a2a: DeepEPLLPrepareAndFinalize = make_deepep_a2a(
198 + pg=pg,
199 + pgi=pgi,
200 + dp_size=dp_size,
201 + deepep_ht_args=None,
202 + deepep_ll_args=DeepEPLLArgs(
203 + max_tokens_per_rank=max_tokens_per_rank,
204 + hidden_size=hidden_size,
205 + num_experts=test_config.num_experts,
206 + use_fp8_dispatch=test_config.use_fp8_dispatch),
207 + q_dtype=q_dtype,
208 + block_shape=test_config.block_size)
209 +
210 + fused_experts = BatchedDeepGemmExperts(max_num_tokens=max_tokens_per_rank,
211 + world_size=pgi.world_size,
212 + dp_size=dp_size,
213 + block_shape=test_config.block_size)
214 + mk = FusedMoEModularKernel(prepare_finalize=a2a,
215 + fused_experts=fused_experts)
216 + return mk
217 +
218 +
219 +def make_ht_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
220 + dp_size: int, num_local_experts: int,
221 + q_dtype: Optional[torch.dtype],
222 + test_config: TestConfig) -> FusedMoEModularKernel:
223 +
224 + assert not test_config.low_latency
225 + assert test_config.use_fp8_dispatch is None
226
227 a2a: DeepEPHTPrepareAndFinalize = make_deepep_a2a(
228 pg=pg,
@@ -181,7 +231,7 @@ def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
231 deepep_ht_args=DeepEPHTArgs(num_local_experts=num_local_experts),
232 deepep_ll_args=None,
233 q_dtype=q_dtype,
184 - block_shape=block_shape)
234 + block_shape=test_config.block_size)
235
236 fused_experts = DeepGemmExperts()
237 mk = FusedMoEModularKernel(prepare_finalize=a2a,
@@ -189,12 +239,42 @@ def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
239 return mk
240
241
192 -def deep_ep_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
193 - test_tensors: TestTensors, w1: torch.Tensor,
194 - w2: torch.Tensor, w1_scale: Optional[torch.Tensor],
195 - w2_scale: Optional[torch.Tensor],
196 - num_experts: int) -> torch.Tensor:
242 +def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
243 + num_local_experts: int,
244 + test_tensors: TestTensors) -> FusedMoEModularKernel:
245 +
246 + q_dtype = torch.float8_e4m3fn
247 + test_config = test_tensors.config
248 +
249 + mk: FusedMoEModularKernel
250 + # Make modular kernel
251 + if test_config.low_latency:
252 + max_tokens_per_rank = max(
253 + 64, next_power_of_2(test_tensors.rank_tokens.size(0)))
254 + hidden_size = test_tensors.rank_tokens.size(-1)
255 +
256 + mk = make_ll_modular_kernel(pg=pg,
257 + pgi=pgi,
258 + max_tokens_per_rank=max_tokens_per_rank,
259 + dp_size=dp_size,
260 + hidden_size=hidden_size,
261 + q_dtype=q_dtype,
262 + test_config=test_config)
263 + else:
264 + mk = make_ht_modular_kernel(pg, pgi, dp_size, num_local_experts,
265 + q_dtype, test_config)
266 +
267 + return mk
268 +
269 +
270 +def deepep_deepgemm_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
271 + dp_size: int, test_tensors: TestTensors,
272 + w1: torch.Tensor, w2: torch.Tensor,
273 + w1_scale: Optional[torch.Tensor],
274 + w2_scale: Optional[torch.Tensor]) -> torch.Tensor:
275
276 + test_config = test_tensors.config
277 + num_experts = test_config.num_experts
278 num_local_experts = w1.size(0)
279
280 def build_expert_map():
@@ -208,14 +288,17 @@ def deep_ep_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo, dp_size: int,
288 return expert_map.to(device=torch.cuda.current_device(),
289 dtype=torch.int32)
290
211 - q_dtype = torch.float8_e4m3fn
212 -
291 # Make modular kernel
292 mk: FusedMoEModularKernel = make_modular_kernel(
215 - pg, pgi, dp_size, num_local_experts, q_dtype,
216 - test_tensors.config.block_size)
293 + pg=pg,
294 + pgi=pgi,
295 + dp_size=dp_size,
296 + num_local_experts=num_local_experts,
297 + test_tensors=test_tensors)
298
218 - a1_scale = test_tensors.rank_token_scales
299 + # Low-Latency kernels can't dispatch scales.
300 + a1_scale = (None
301 + if test_config.low_latency else test_tensors.rank_token_scales)
302
303 out = mk.forward(hidden_states=test_tensors.rank_tokens,
304 w1=w1,
@@ -258,7 +341,7 @@ def triton_impl(a: torch.Tensor, topk_ids: torch.Tensor,
341 allow_deep_gemm=False)
342
343
261 -def _deep_ep_moe(
344 +def _test_deepep_deepgemm_moe(
345 pgi: ProcessGroupInfo,
346 dp_size: int,
347 config: TestConfig,
@@ -302,7 +385,7 @@ def _deep_ep_moe(
385 w1_scale_ep = w1_scale[e_start:e_end]
386 w2_scale_ep = w2_scale[e_start:e_end]
387
305 - deepep_moe = deep_ep_moe_impl(
388 + deepep_moe = deepep_deepgemm_moe_impl(
389 pg,
390 pgi,
391 dp_size,
@@ -311,7 +394,6 @@ def _deep_ep_moe(
394 w2_ep,
395 w1_scale_ep,
396 w2_scale_ep,
314 - config.num_experts,
397 )
398
399 torch.testing.assert_close(
@@ -335,15 +417,21 @@ MNKs = [
417 (222, 1024, 2048),
418 ]
419
420 +TOPKS = [2, 6]
421 +NUM_EXPERTS = [32]
422 +
423
424 @pytest.mark.parametrize("mnk", MNKs)
340 -@pytest.mark.parametrize("num_experts", [32])
341 -@pytest.mark.parametrize("topk", [2, 6])
425 +@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
426 +@pytest.mark.parametrize("topk", TOPKS)
427 @pytest.mark.parametrize("world_dp_size", [(2, 1)])
428 @requires_deep_ep
429 @requires_deep_gemm
345 -def test_deep_ep_moe(mnk: tuple[int, int, int], num_experts: int, topk: int,
346 - world_dp_size: tuple[int, int]):
430 +def test_ht_deepep_deepgemm_moe(mnk: tuple[int, int, int], num_experts: int,
431 + topk: int, world_dp_size: tuple[int, int]):
432 + """
433 + Tests for High-Throughput DeepEP + DeepGemm integration.
434 + """
435
436 m, n, k = mnk
437 current_platform.seed_everything(7)
@@ -354,6 +442,58 @@ def test_deep_ep_moe(mnk: tuple[int, int, int], num_experts: int, topk: int,
442 block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
443 block_size = [block_m, block_m]
444
445 + world_size, dp_size = world_dp_size
446 + config = TestConfig(topk=topk,
447 + m=m,
448 + k=k,
449 + n=n,
450 + num_experts=num_experts,
451 + block_size=block_size,
452 + low_latency=False,
453 + use_fp8_dispatch=None)
454 +
455 + w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights(
456 + num_experts, n, k, block_size)
457 +
458 + parallel_launch(world_size, _test_deepep_deepgemm_moe, dp_size, config, w1,
459 + w2, w1_scale, w2_scale)
460 +
461 +
462 +MNKs = [
463 + (1, 128, 2560),
464 + (2, 128, 2560),
465 + (3, 1024, 2560),
466 + (32, 128, 2560),
467 + (45, 512, 2560),
468 + (64, 1024, 2560),
469 + (222, 1024, 2560),
470 +]
471 +# Fix tests for USE_FP8_DISPATCH=True
472 +USE_FP8_DISPATCH = [False]
473 +
474 +
475 +@pytest.mark.parametrize("mnk", MNKs)
476 +@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
477 +@pytest.mark.parametrize("topk", TOPKS)
478 +@pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH)
479 +@pytest.mark.parametrize("block_size", [[128, 128]])
480 +@pytest.mark.parametrize("world_dp_size", [(2, 1)])
481 +@requires_deep_ep
482 +@requires_deep_gemm
483 +def test_ll_deepep_deepgemm_moe(mnk: tuple[int, int,
484 + int], num_experts: int, topk: int,
485 + use_fp8_dispatch: bool, block_size: list[int],
486 + world_dp_size: tuple[int, int]):
487 + """
488 + Tests for Low-Latency DeepEP + DeepGemm integration.
489 + """
490 +
491 + m, n, k = mnk
492 + current_platform.seed_everything(7)
493 +
494 + if topk > num_experts:
495 + pytest.skip(f"Skipping test: topk={topk} > E={num_experts}")
496 +
497 world_size, dp_size = world_dp_size
498 config = TestConfig(
499 topk=topk,
@@ -362,10 +502,12 @@ def test_deep_ep_moe(mnk: tuple[int, int, int], num_experts: int, topk: int,
502 n=n,
503 num_experts=num_experts,
504 block_size=block_size,
505 + low_latency=True,
506 + use_fp8_dispatch=use_fp8_dispatch,
507 )
508
509 w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights(
510 num_experts, n, k, block_size)
511
370 - parallel_launch(world_size, _deep_ep_moe, dp_size, config, w1, w2,
371 - w1_scale, w2_scale)
512 + parallel_launch(world_size, _test_deepep_deepgemm_moe, dp_size, config, w1,
513 + w2, w1_scale, w2_scale)
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py new
+124
@@ -0,0 +1,124 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import importlib.util
3 +from typing import Optional
4 +
5 +import torch
6 +
7 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm.logger import init_logger
9 +from vllm.model_executor.layers.fused_moe.utils import (
10 + _resize_cache, per_token_group_quant_fp8)
11 +
12 +logger = init_logger(__name__)
13 +
14 +has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
15 +
16 +
17 +class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
18 +
19 + # The Deep Gemm kernels only support block size of 128
20 + DEEPGEMM_BLOCK_SHAPE = 128
21 +
22 + def __init__(self, max_num_tokens: int, world_size: int, dp_size: int,
23 + block_shape: list[int]):
24 + """
25 + max_num_tokens: Maximum number of tokens from a DP Rank
26 + world_size: Number of EP ranks
27 + dp_size: Number of data-parallel ranks
28 + block_shape: Block quantization block shape
29 + """
30 + super().__init__()
31 + self.max_num_tokens = max_num_tokens
32 + self.world_size = world_size
33 + self.dp_size = dp_size
34 + self.block_shape = block_shape
35 +
36 + assert (len(self.block_shape) == 2 and all(
37 + [v == self.DEEPGEMM_BLOCK_SHAPE for v in self.block_shape]))
38 +
39 + def workspace_shapes(
40 + self,
41 + a: torch.Tensor,
42 + M: int,
43 + N: int,
44 + K: int,
45 + topk: int,
46 + num_experts: int,
47 + ) -> tuple[int, int, torch.dtype]:
48 + assert a.dim() == 2
49 + num_dp = self.world_size // self.dp_size
50 + max_num_tokens = a.size(
51 + 0) if self.max_num_tokens is None else self.max_num_tokens
52 + workspace13 = num_experts * max_num_tokens * num_dp * max(K, N)
53 + workspace2 = num_experts * max_num_tokens * num_dp * (N // 2)
54 + return (workspace13, workspace2, a.dtype)
55 +
56 + def apply(
57 + self,
58 + hidden_states: torch.Tensor,
59 + w1: torch.Tensor,
60 + w2: torch.Tensor,
61 + topk_ids: torch.Tensor,
62 + activation: str,
63 + global_num_experts: int,
64 + expert_map: Optional[torch.Tensor],
65 + w1_scale: Optional[torch.Tensor],
66 + w2_scale: Optional[torch.Tensor],
67 + w1_zp: Optional[torch.Tensor],
68 + w2_zp: Optional[torch.Tensor],
69 + a1q_scale: Optional[torch.Tensor],
70 + a2_scale: Optional[torch.Tensor],
71 + workspace13: torch.Tensor,
72 + workspace2: torch.Tensor,
73 + expert_num_tokens: Optional[torch.Tensor],
74 + ) -> torch.Tensor:
75 + import deep_gemm as dg
76 + assert hidden_states.ndim == 3
77 +
78 + a1q = hidden_states
79 + _, N, K = w1.size()
80 +
81 + if global_num_experts == -1:
82 + global_num_experts = w1.size(0)
83 +
84 + assert w2.size(1) == K
85 +
86 + E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
87 + hidden_states, w1, w2, topk_ids)
88 +
89 + workspace1 = _resize_cache(workspace13, (E, max_num_tokens, N))
90 + workspace2 = _resize_cache(workspace2, (E, max_num_tokens, N // 2))
91 + workspace3 = _resize_cache(workspace13, (E, max_num_tokens, K))
92 +
93 + # (from deepgemm docs) : A value hint (which is a value on CPU)
94 + # for the M expectation of each batch, correctly setting this value
95 + # may lead to better performance.
96 + expected_m = max_num_tokens
97 +
98 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a1q, a1q_scale),
99 + (w1, w1_scale),
100 + out=workspace1,
101 + masked_m=expert_num_tokens,
102 + expected_m=expected_m)
103 +
104 + # TODO (varun) [Optimization]: Use a batched version of activation.
105 + # Similarly for the quant below.
106 + self.activation(activation, workspace2, workspace1.view(-1, N))
107 +
108 + w2_hidden_size = workspace2.size(-1)
109 + workspace2 = workspace2.view(-1, w2_hidden_size)
110 +
111 + a2q_scale: Optional[torch.Tensor] = None
112 + a2q, a2q_scale = per_token_group_quant_fp8(workspace2,
113 + self.block_shape[1],
114 + column_major_scales=False)
115 + a2q = a2q.view(E, max_num_tokens, -1)
116 + a2q_scale = a2q_scale.view(E, max_num_tokens, -1)
117 +
118 + dg.m_grouped_gemm_fp8_fp8_bf16_nt_masked((a2q, a2q_scale),
119 + (w2, w2_scale),
120 + out=workspace3,
121 + masked_m=expert_num_tokens,
122 + expected_m=expected_m)
123 +
124 + return workspace3
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py new
+116
@@ -0,0 +1,116 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
3 +
4 +import torch
5 +
6 +import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 +from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8 + BatchedDeepGemmExperts)
9 +from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
10 + BatchedTritonExperts)
11 +
12 +
13 +class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
14 +
15 + def __init__(self,
16 + max_num_tokens: int,
17 + world_size: int,
18 + dp_size: int,
19 + use_fp8_w8a8: bool = False,
20 + use_int8_w8a8: bool = False,
21 + use_int8_w8a16: bool = False,
22 + use_int4_w4a16: bool = False,
23 + per_channel_quant: bool = False,
24 + block_shape: Optional[list[int]] = None,
25 + allow_deep_gemm: bool = False):
26 + super().__init__()
27 + assert not use_int8_w8a8, "NYI"
28 + assert not use_int8_w8a16, "NYI"
29 + assert not use_int4_w4a16, "NYI"
30 +
31 + self.max_num_tokens = max_num_tokens
32 + self.world_size = world_size
33 + self.dp_size = dp_size
34 + self.use_fp8_w8a8 = use_fp8_w8a8
35 + self.use_int8_w8a8 = use_int8_w8a8
36 + self.use_int8_w8a16 = use_int8_w8a16
37 + self.use_int4_w4a16 = use_int4_w4a16
38 + self.per_channel_quant = per_channel_quant
39 + self.block_shape = block_shape
40 + self.allow_deep_gemm = allow_deep_gemm
41 +
42 + # BatchedTritonKernel doesn't support block quantization
43 + # at the moment.
44 + self.batched_triton_experts = BatchedTritonExperts(
45 + max_num_tokens=self.max_num_tokens,
46 + use_fp8_w8a8=self.use_fp8_w8a8,
47 + use_int8_w8a8=self.use_int8_w8a8,
48 + use_int8_w8a16=self.use_int8_w8a16,
49 + use_int4_w4a16=self.use_int4_w4a16,
50 + per_channel_quant=self.per_channel_quant,
51 + block_shape=self.block_shape,
52 + world_size=self.world_size,
53 + dp_size=self.dp_size) if self.block_shape is None else None
54 +
55 + is_fp8_128_block_quantized = (self.use_fp8_w8a8
56 + and self.block_shape is not None
57 + and len(self.block_shape) == 2 and all(
58 + [b == 128
59 + for b in self.block_shape]))
60 + self.batched_deep_gemm_experts = BatchedDeepGemmExperts(
61 + max_num_tokens=self.max_num_tokens,
62 + world_size=self.world_size,
63 + dp_size=self.dp_size,
64 + block_shape=self.block_shape, # type: ignore[arg-type]
65 + ) if (self.allow_deep_gemm and is_fp8_128_block_quantized) else None
66 +
67 + def workspace_shapes(
68 + self,
69 + a: torch.Tensor,
70 + M: int,
71 + N: int,
72 + K: int,
73 + topk: int,
74 + num_experts: int,
75 + ) -> tuple[int, int, torch.dtype]:
76 + # Note: the deep gemm workspaces are strictly larger than the triton
77 + # workspaces so we can be pessimistic here and allocate for DeepGemm
78 + # even if we fall back to triton later, e.g. if expert maps are set.
79 + if self.allow_deep_gemm and self.batched_deep_gemm_experts is not None:
80 + return self.batched_deep_gemm_experts.workspace_shapes(
81 + a, M, N, K, topk, num_experts)
82 + else:
83 + assert self.batched_triton_experts is not None
84 + return self.batched_triton_experts.workspace_shapes(
85 + a, M, N, K, topk, num_experts)
86 +
87 + def apply(
88 + self,
89 + hidden_states: torch.Tensor,
90 + w1: torch.Tensor,
91 + w2: torch.Tensor,
92 + topk_ids: torch.Tensor,
93 + activation: str,
94 + global_num_experts: int,
95 + expert_map: Optional[torch.Tensor],
96 + w1_scale: Optional[torch.Tensor],
97 + w2_scale: Optional[torch.Tensor],
98 + w1_zp: Optional[torch.Tensor],
99 + w2_zp: Optional[torch.Tensor],
100 + a1q_scale: Optional[torch.Tensor],
101 + a2_scale: Optional[torch.Tensor],
102 + workspace13: torch.Tensor,
103 + workspace2: torch.Tensor,
104 + expert_num_tokens: Optional[torch.Tensor],
105 + ) -> torch.Tensor:
106 + use_batched_deep_gemm_experts = (self.allow_deep_gemm
107 + and self.batched_deep_gemm_experts
108 + is not None)
109 + experts = (self.batched_deep_gemm_experts
110 + if use_batched_deep_gemm_experts else
111 + self.batched_triton_experts)
112 + assert experts is not None
113 + return experts.apply(hidden_states, w1, w2, topk_ids, activation,
114 + global_num_experts, expert_map, w1_scale,
115 + w2_scale, w1_zp, w2_zp, a1q_scale, a2_scale,
116 + workspace13, workspace2, expert_num_tokens)
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+55 -21
@@ -1,5 +1,5 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -from typing import Optional
2 +from typing import Optional, Union
3
4 import deep_ep
5 import torch
@@ -65,6 +65,54 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
65 def topk_indices_dtype(self) -> Optional[torch.dtype]:
66 return torch.int64
67
68 + def _do_quant(
69 + self, x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
70 + a1_scale: Optional[torch.Tensor], a2_scale: Optional[torch.Tensor],
71 + a1_dtype: torch.dtype
72 + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
73 +
74 + block_k = self.block_shape[1] if self.block_shape is not None else None
75 + if self.use_fp8_dispatch:
76 + if block_k == DEEPEP_QUANT_BLOCK_SIZE:
77 + # DeepEP kernels did the quantization for us.
78 + x, x_scales = x
79 + return x, x_scales
80 +
81 + # Dequant to get back the tokens in the datatype we dispatched in.
82 + x_fp8, x_scales = x
83 + x = dequant_fp8(x_fp8, x_scales).to(dtype=a1_dtype)
84 +
85 + assert isinstance(x, torch.Tensor)
86 +
87 + # Check if there is a block_shape / or if we can infer the quantization
88 + # schemes from the scales.
89 + per_token_quant = None
90 + if all([v is None for v in [self.block_shape, a1_scale, a2_scale]
91 + ]) and self.quant_dtype is not None:
92 + # Quantization required despite none of the inputs suggesting
93 + # quantization. Fallback to per_token_dynamic quant.
94 + per_token_quant = True
95 + else:
96 + per_token_quant = ((self.block_shape is not None) or
97 + (a1_scale is not None and a1_scale.numel() != 1)
98 + or (a2_scale is not None
99 + and a2_scale.numel() != 1))
100 +
101 + num_experts, max_tokens, hidden_dim = x.size()
102 +
103 + # TODO (varun): Optimization - Use a batched version of quant
104 + x = x.view((-1, hidden_dim))
105 + x, x_scales = moe_kernel_quantize_input(x, a1_scale, self.quant_dtype,
106 + per_token_quant,
107 + self.block_shape)
108 + x = x.view((num_experts, -1, hidden_dim))
109 +
110 + if per_token_quant:
111 + assert x_scales is not None
112 + x_scales = x_scales.view(num_experts, max_tokens, -1)
113 +
114 + return x, x_scales
115 +
116 def prepare(
117 self,
118 a1: torch.Tensor,
@@ -87,11 +135,11 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
135 assert hidden_size % 128 == 0, \
136 "DeepEP kernels quantize the inputs in blocks of shape 128"
137
90 - # Quantize
91 - per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
138 + has_per_token_scales = a1_scale.numel(
139 + ) != 1 if a1_scale is not None else (
140 a2_scale.numel() != 1 if a2_scale is not None else False)
93 - assert not per_act_token, (
94 - "low_latency kernels don't support per-act-token quant")
141 + assert not has_per_token_scales, (
142 + "low_latency kernels doesn't support dispatching per-token scales")
143
144 if apply_router_weight_on_input:
145 topk = rank_topk_ids.size(1)
@@ -110,22 +158,8 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
158 async_finish=False,
159 return_recv_hook=False)
160
113 - if self.use_fp8_dispatch:
114 - # TODO (varun) : In the case of dynamic quantization, we could
115 - # probably skip the quant below and use the results directly.
116 - # Although note that the deepep quant is per token 128 elements.
117 - expert_x_fp8, expert_x_scales = expert_x
118 - expert_x = dequant_fp8(expert_x_fp8,
119 - expert_x_scales).to(dtype=a1.dtype)
120 -
121 - num_experts = expert_x.size(0)
122 - hidden_dim = expert_x.size(-1)
123 -
124 - expert_x = expert_x.view((-1, expert_x.size(-1)))
125 - expert_x, expert_x_scale = moe_kernel_quantize_input(
126 - expert_x, a1_scale, self.quant_dtype, per_act_token,
127 - self.block_shape)
128 - expert_x = expert_x.view((num_experts, -1, hidden_dim))
161 + expert_x, expert_x_scale = self._do_quant(expert_x, a1_scale, a2_scale,
162 + a1.dtype)
163
164 return (expert_x, expert_x_scale, expert_num_tokens, None, None)
165
vllm/model_executor/layers/quantization/fp8.py
+7 -5
@@ -771,21 +771,21 @@ class Fp8MoEMethod(FusedMoEMethodBase):
771
772 def select_gemm_impl(self, prepare_finalize):
773
774 - from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
775 - BatchedTritonExperts)
774 + from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
775 + BatchedTritonOrDeepGemmExperts)
776 from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
777 TritonOrDeepGemmExperts)
778
779 assert not self.use_marlin and not self.rocm_aiter_moe_enabled, (
780 "Marlin and ROCm AITER are not supported with all2all yet.")
781
782 - experts: Optional[Union[BatchedTritonExperts,
782 + experts: Optional[Union[BatchedTritonOrDeepGemmExperts,
783 TritonOrDeepGemmExperts]] = None
784 max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
785 use_batched_experts = max_num_tokens_per_rank is not None
786
787 if use_batched_experts:
788 - experts = BatchedTritonExperts(
788 + experts = BatchedTritonOrDeepGemmExperts(
789 max_num_tokens=max_num_tokens_per_rank,
790 world_size=prepare_finalize.world_size,
791 dp_size=prepare_finalize.dp_size,
@@ -793,7 +793,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
793 use_int8_w8a8=False,
794 use_int8_w8a16=False,
795 use_int4_w4a16=False,
796 - block_shape=None,
796 + per_channel_quant=False,
797 + block_shape=self.quant_config.weight_block_size,
798 + allow_deep_gemm=self.allow_deep_gemm,
799 )
800 else:
801 experts = TritonOrDeepGemmExperts(