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] DeepEP dispatch-combine kernel integration (#18434)
Signed-off-by: Varun <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Jun 3, 2025 at 15:30 UTC
fa98d77773c649de05a4bda9847682c80287aa36
23 files changed
+1952
-124
csrc/moe/topk_softmax_kernels.cu
+14
-2
@@ -516,9 +516,8 @@ void topk_softmax(
516
topk,
517
stream);
518
}
519
- else
519
+ else if (topk_indices.scalar_type() == at::ScalarType::UInt32)
520
{
521
- assert(topk_indices.scalar_type() == at::ScalarType::UInt32);
521
vllm::moe::topkGatingSoftmaxKernelLauncher(
522
gating_output.data_ptr<float>(),
523
topk_weights.data_ptr<float>(),
@@ -530,4 +529,17 @@ void topk_softmax(
529
topk,
530
stream);
531
}
532
+ else {
533
+ assert(topk_indices.scalar_type() == at::ScalarType::Int64);
534
+ vllm::moe::topkGatingSoftmaxKernelLauncher(
535
+ gating_output.data_ptr<float>(),
536
+ topk_weights.data_ptr<float>(),
537
+ topk_indices.data_ptr<int64_t>(),
538
+ token_expert_indices.data_ptr<int>(),
539
+ softmax_workspace.data_ptr<float>(),
540
+ num_tokens,
541
+ num_experts,
542
+ topk,
543
+ stream);
544
+ }
545
}
tests/kernels/moe/__init__.py
tests/kernels/moe/deepep_utils.py
new
+188
@@ -0,0 +1,188 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+DeepEP test utilities
4
+"""
5
+import dataclasses
6
+import importlib
7
+import traceback
8
+from typing import Callable, Optional
9
+
10
+import torch
11
+from torch.distributed import ProcessGroup
12
+from torch.multiprocessing import (
13
+ spawn) # pyright: ignore[reportPrivateImportUsage]
14
+from typing_extensions import Concatenate, ParamSpec
15
+
16
+has_deep_ep = importlib.util.find_spec("deep_ep") is not None
17
+if has_deep_ep:
18
+ from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
19
+ DeepEPHTPrepareAndFinalize)
20
+ from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
21
+ DeepEPLLPrepareAndFinalize)
22
+
23
+## Parallel Processes Utils
24
+
25
+P = ParamSpec("P")
26
+
27
+
28
+@dataclasses.dataclass
29
+class ProcessGroupInfo:
30
+ world_size: int
31
+ world_local_size: int
32
+ rank: int
33
+ node_rank: int
34
+ local_rank: int
35
+ device: torch.device
36
+
37
+
38
+def _worker_parallel_launch(
39
+ local_rank: int,
40
+ world_size: int,
41
+ world_local_size: int,
42
+ node_rank: int,
43
+ init_method: str,
44
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
45
+ *args: P.args,
46
+ **kwargs: P.kwargs,
47
+) -> None:
48
+ rank = node_rank * world_local_size + local_rank
49
+ torch.cuda.set_device(local_rank)
50
+ device = torch.device("cuda", local_rank)
51
+ torch.distributed.init_process_group(
52
+ backend="cpu:gloo,cuda:nccl",
53
+ init_method=init_method,
54
+ rank=rank,
55
+ world_size=world_size,
56
+ device_id=device,
57
+ )
58
+ barrier = torch.tensor([rank], device=device)
59
+ torch.distributed.all_reduce(barrier)
60
+
61
+ try:
62
+ worker(
63
+ ProcessGroupInfo(
64
+ world_size=world_size,
65
+ world_local_size=world_local_size,
66
+ rank=rank,
67
+ node_rank=node_rank,
68
+ local_rank=local_rank,
69
+ device=device,
70
+ ),
71
+ *args,
72
+ **kwargs,
73
+ )
74
+ except Exception as ex:
75
+ print(ex)
76
+ traceback.print_exc()
77
+ raise
78
+ finally:
79
+ torch.distributed.destroy_process_group()
80
+
81
+
82
+def parallel_launch(
83
+ world_size: int,
84
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
85
+ *args: P.args,
86
+ **kwargs: P.kwargs,
87
+) -> None:
88
+ assert not kwargs
89
+ spawn(
90
+ _worker_parallel_launch,
91
+ args=(
92
+ world_size,
93
+ world_size,
94
+ 0,
95
+ "tcp://localhost:29500",
96
+ worker,
97
+ ) + args,
98
+ nprocs=world_size,
99
+ join=True,
100
+ )
101
+
102
+
103
+## DeepEP specific utils
104
+
105
+
106
+@dataclasses.dataclass
107
+class DeepEPHTArgs:
108
+ num_local_experts: int
109
+
110
+
111
+@dataclasses.dataclass
112
+class DeepEPLLArgs:
113
+ max_tokens_per_rank: int
114
+ hidden_size: int
115
+ num_experts: int
116
+ use_fp8_dispatch: bool
117
+
118
+
119
+def make_deepep_ht_a2a(pg: ProcessGroup,
120
+ pgi: ProcessGroupInfo,
121
+ dp_size: int,
122
+ ht_args: DeepEPHTArgs,
123
+ q_dtype: Optional[torch.dtype] = None,
124
+ block_shape: Optional[list[int]] = None):
125
+
126
+ import deep_ep
127
+
128
+ # high throughput a2a
129
+ num_nvl_bytes = 1024 * 1024 * 1024 # 1GB
130
+ num_rdma_bytes, low_latency_mode, num_qps_per_rank = 0, False, 1
131
+ buffer = deep_ep.Buffer(group=pg,
132
+ num_nvl_bytes=num_nvl_bytes,
133
+ num_rdma_bytes=num_rdma_bytes,
134
+ low_latency_mode=low_latency_mode,
135
+ num_qps_per_rank=num_qps_per_rank)
136
+ return DeepEPHTPrepareAndFinalize(buffer=buffer,
137
+ world_size=pgi.world_size,
138
+ rank=pgi.rank,
139
+ dp_size=dp_size,
140
+ rank_expert_offset=pgi.rank *
141
+ ht_args.num_local_experts,
142
+ quant_dtype=q_dtype,
143
+ block_shape=block_shape)
144
+
145
+
146
+def make_deepep_ll_a2a(pg: ProcessGroup,
147
+ pgi: ProcessGroupInfo,
148
+ dp_size: int,
149
+ deepep_ll_args: DeepEPLLArgs,
150
+ q_dtype: Optional[torch.dtype] = None,
151
+ block_shape: Optional[list[int]] = None):
152
+
153
+ import deep_ep
154
+
155
+ # low-latency a2a
156
+ num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
157
+ deepep_ll_args.max_tokens_per_rank, deepep_ll_args.hidden_size,
158
+ pgi.world_size, deepep_ll_args.num_experts)
159
+
160
+ buffer = deep_ep.Buffer(group=pg,
161
+ num_rdma_bytes=num_rdma_bytes,
162
+ low_latency_mode=True,
163
+ num_qps_per_rank=deepep_ll_args.num_experts //
164
+ pgi.world_size)
165
+ return DeepEPLLPrepareAndFinalize(
166
+ buffer=buffer,
167
+ world_size=pgi.world_size,
168
+ dp_size=dp_size,
169
+ max_tokens_per_rank=deepep_ll_args.max_tokens_per_rank,
170
+ quant_dtype=q_dtype,
171
+ use_fp8_dispatch=deepep_ll_args.use_fp8_dispatch,
172
+ )
173
+
174
+
175
+def make_deepep_a2a(pg: ProcessGroup,
176
+ pgi: ProcessGroupInfo,
177
+ dp_size: int,
178
+ deepep_ht_args: Optional[DeepEPHTArgs],
179
+ deepep_ll_args: Optional[DeepEPLLArgs],
180
+ q_dtype: Optional[torch.dtype] = None,
181
+ block_shape: Optional[list[int]] = None):
182
+ if deepep_ht_args is not None:
183
+ assert deepep_ll_args is None
184
+ return make_deepep_ht_a2a(pg, pgi, dp_size, deepep_ht_args, q_dtype,
185
+ block_shape)
186
+
187
+ assert deepep_ll_args is not None
188
+ return make_deepep_ll_a2a(pg, pgi, dp_size, deepep_ll_args, q_dtype)
tests/kernels/moe/test_deepep_deepgemm_moe.py
new
+371
@@ -0,0 +1,371 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Test DeepEP + DeepGEMM integration
4
+"""
5
+
6
+import dataclasses
7
+import importlib
8
+from typing import Optional
9
+
10
+import pytest
11
+import torch.distributed
12
+from torch.distributed import ProcessGroup
13
+from typing_extensions import ParamSpec
14
+
15
+from vllm.config import VllmConfig, set_current_vllm_config
16
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
17
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
18
+ FusedMoEModularKernel)
19
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
20
+ per_token_group_quant_fp8)
21
+from vllm.platforms import current_platform
22
+
23
+from .deepep_utils import ProcessGroupInfo, parallel_launch
24
+
25
+has_deep_ep = importlib.util.find_spec("deep_ep") is not None
26
+
27
+try:
28
+ import deep_gemm
29
+ has_deep_gemm = True
30
+except ImportError:
31
+ has_deep_gemm = False
32
+
33
+if has_deep_ep:
34
+ from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
35
+ DeepEPHTPrepareAndFinalize)
36
+
37
+ from .deepep_utils import DeepEPHTArgs, make_deepep_a2a
38
+
39
+if has_deep_gemm:
40
+ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
41
+ DeepGemmExperts)
42
+
43
+requires_deep_ep = pytest.mark.skipif(
44
+ not has_deep_ep,
45
+ reason="Requires deep_ep kernels",
46
+)
47
+
48
+requires_deep_gemm = pytest.mark.skipif(
49
+ not has_deep_gemm,
50
+ reason="Requires deep_gemm kernels",
51
+)
52
+
53
+P = ParamSpec("P")
54
+
55
+
56
+def per_block_cast_to_fp8(
57
+ x: torch.Tensor,
58
+ block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
59
+ assert x.dim() == 2
60
+ m, n = x.shape
61
+ x_padded = torch.zeros(
62
+ (deep_gemm.ceil_div(m, 128) * 128,
63
+ deep_gemm.ceil_div(n, block_size_n) * block_size_n),
64
+ dtype=x.dtype,
65
+ device=x.device)
66
+ x_padded[:m, :n] = x
67
+ x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
68
+ x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
69
+ x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
70
+ x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
71
+ scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
72
+ return x_scaled_sub, scales
73
+
74
+
75
+def make_block_quant_fp8_weights(
76
+ e: int,
77
+ n: int,
78
+ k: int,
79
+ block_size: list[int],
80
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
81
+ """
82
+ Return weights w1, w2, w1q, w2q, w1_scale, w2_scale
83
+ """
84
+ dtype = torch.bfloat16
85
+
86
+ fp8_info = torch.finfo(torch.float8_e4m3fn)
87
+ fp8_max, fp8_min = fp8_info.max, fp8_info.min
88
+
89
+ w1_bf16 = torch.randn((e, 2 * n, k), dtype=dtype) / 10
90
+ w1_bf16 = w1_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
91
+
92
+ w2_bf16 = torch.randn((e, k, n), dtype=dtype) / 10
93
+ w2_bf16 = w2_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
94
+
95
+ block_n, block_k = block_size[0], block_size[1]
96
+ n_tiles_w1 = ((2 * n) + block_n - 1) // block_n
97
+ k_tiles_w1 = (k + block_k - 1) // block_k
98
+ n_tiles_w2 = (k + block_n - 1) // block_n
99
+ k_tiles_w2 = (n + block_k - 1) // block_k
100
+
101
+ w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn)
102
+ w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn)
103
+
104
+ w1_s = torch.empty((e, n_tiles_w1, k_tiles_w1),
105
+ device="cuda",
106
+ dtype=torch.float32)
107
+ w2_s = torch.empty((e, n_tiles_w2, k_tiles_w2),
108
+ device="cuda",
109
+ dtype=torch.float32)
110
+
111
+ assert w1_s.shape == (e, (2 * n + 127) // 128, (k + 127) // 128)
112
+ assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
113
+
114
+ for i in range(e):
115
+ w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
116
+ w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
117
+
118
+ return w1, w2, w1_s, w2_s
119
+
120
+
121
+@dataclasses.dataclass
122
+class TestConfig:
123
+ topk: int
124
+ m: int
125
+ k: int
126
+ n: int
127
+ num_experts: int
128
+ block_size: list[int]
129
+
130
+
131
+@dataclasses.dataclass
132
+class TestTensors:
133
+ rank_tokens: torch.Tensor # all ranks make this many tokens
134
+ rank_token_scales: Optional[torch.Tensor]
135
+ topk: torch.Tensor
136
+ topk_weights: torch.Tensor
137
+ config: TestConfig
138
+
139
+ @staticmethod
140
+ def make(config: TestConfig, rank) -> "TestTensors":
141
+
142
+ dtype = torch.bfloat16
143
+ topk, m, k, block_size = (config.topk, config.m, config.k,
144
+ config.block_size)
145
+
146
+ fp8_info = torch.finfo(torch.float8_e4m3fn)
147
+ fp8_max, fp8_min = fp8_info.max, fp8_info.min
148
+
149
+ rank_tokens = torch.randn(
150
+ (m, k), device=torch.cuda.current_device(), dtype=dtype) / 10.0
151
+ rank_tokens = rank_tokens.clamp(min=fp8_min, max=fp8_max)
152
+
153
+ block_k = block_size[1]
154
+ _, rank_token_scales = per_token_group_quant_fp8(rank_tokens, block_k)
155
+
156
+ topk_ids = torch.randint(
157
+ low=0,
158
+ high=config.num_experts,
159
+ size=(m, topk),
160
+ device=torch.cuda.current_device()).to(dtype=torch.int64)
161
+
162
+ topk_weights = torch.randn(topk_ids.shape,
163
+ dtype=torch.float32,
164
+ device=torch.cuda.current_device())
165
+
166
+ return TestTensors(rank_tokens=rank_tokens,
167
+ rank_token_scales=rank_token_scales,
168
+ topk=topk_ids,
169
+ topk_weights=topk_weights,
170
+ config=config)
171
+
172
+
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:
176
+
177
+ a2a: DeepEPHTPrepareAndFinalize = make_deepep_a2a(
178
+ pg=pg,
179
+ pgi=pgi,
180
+ dp_size=dp_size,
181
+ deepep_ht_args=DeepEPHTArgs(num_local_experts=num_local_experts),
182
+ deepep_ll_args=None,
183
+ q_dtype=q_dtype,
184
+ block_shape=block_shape)
185
+
186
+ fused_experts = DeepGemmExperts()
187
+ mk = FusedMoEModularKernel(prepare_finalize=a2a,
188
+ fused_experts=fused_experts)
189
+ return mk
190
+
191
+
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:
197
+
198
+ num_local_experts = w1.size(0)
199
+
200
+ def build_expert_map():
201
+ num_local_experts = w1.size(0)
202
+ expert_map = torch.full((num_experts, ),
203
+ fill_value=-1,
204
+ dtype=torch.int32)
205
+ s = pgi.rank * num_local_experts
206
+ e = s + num_local_experts
207
+ expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
208
+ return expert_map.to(device=torch.cuda.current_device(),
209
+ dtype=torch.int32)
210
+
211
+ q_dtype = torch.float8_e4m3fn
212
+
213
+ # Make modular kernel
214
+ mk: FusedMoEModularKernel = make_modular_kernel(
215
+ pg, pgi, dp_size, num_local_experts, q_dtype,
216
+ test_tensors.config.block_size)
217
+
218
+ a1_scale = test_tensors.rank_token_scales
219
+
220
+ out = mk.forward(hidden_states=test_tensors.rank_tokens,
221
+ w1=w1,
222
+ w2=w2,
223
+ topk_weights=test_tensors.topk_weights,
224
+ topk_ids=test_tensors.topk,
225
+ inplace=False,
226
+ activation="silu",
227
+ global_num_experts=num_experts,
228
+ expert_map=build_expert_map(),
229
+ w1_scale=w1_scale,
230
+ w2_scale=w2_scale,
231
+ w1_zp=None,
232
+ w2_zp=None,
233
+ a1_scale=a1_scale,
234
+ a2_scale=None,
235
+ apply_router_weight_on_input=False)
236
+ return out
237
+
238
+
239
+def triton_impl(a: torch.Tensor, topk_ids: torch.Tensor,
240
+ topk_weights: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
241
+ w1_scale: torch.Tensor, w2_scale: torch.Tensor,
242
+ a1_scale: torch.Tensor, block_shape: list[int]):
243
+
244
+ return fused_experts(
245
+ hidden_states=a,
246
+ w1=w1,
247
+ w2=w2,
248
+ topk_weights=topk_weights,
249
+ topk_ids=topk_ids,
250
+ inplace=False,
251
+ use_fp8_w8a8=True,
252
+ w1_scale=w1_scale,
253
+ w2_scale=w2_scale,
254
+ a1_scale=a1_scale,
255
+ block_shape=block_shape,
256
+ # Make sure this is set to False so we
257
+ # dont end up comparing the same implementation.
258
+ allow_deep_gemm=False)
259
+
260
+
261
+def _deep_ep_moe(
262
+ pgi: ProcessGroupInfo,
263
+ dp_size: int,
264
+ config: TestConfig,
265
+ w1: torch.Tensor,
266
+ w2: torch.Tensor,
267
+ w1_scale: torch.Tensor,
268
+ w2_scale: torch.Tensor,
269
+):
270
+ current_platform.seed_everything(pgi.rank)
271
+
272
+ w1 = w1.to(device=torch.cuda.current_device())
273
+ w2 = w2.to(device=torch.cuda.current_device())
274
+ w1_scale = w1_scale.to(device=torch.cuda.current_device())
275
+ w2_scale = w2_scale.to(device=torch.cuda.current_device())
276
+
277
+ pg = torch.distributed.new_group(list(range(pgi.world_size)))
278
+ test_tensors = TestTensors.make(config, pgi.rank)
279
+ block_shape = [
280
+ w1.size(1) // w1_scale.size(1),
281
+ w1.size(2) // w1_scale.size(2)
282
+ ]
283
+
284
+ with set_current_vllm_config(VllmConfig()):
285
+ # Reference
286
+ triton_moe = triton_impl(a=test_tensors.rank_tokens,
287
+ topk_ids=test_tensors.topk,
288
+ topk_weights=test_tensors.topk_weights,
289
+ w1=w1,
290
+ w2=w2,
291
+ w1_scale=w1_scale,
292
+ w2_scale=w2_scale,
293
+ a1_scale=test_tensors.rank_token_scales,
294
+ block_shape=block_shape)
295
+
296
+ # Slice experts for this rank.
297
+ num_local_experts = config.num_experts // pgi.world_size
298
+ e_start = num_local_experts * pgi.rank
299
+ e_end = e_start + num_local_experts
300
+ w1_ep = w1[e_start:e_end]
301
+ w2_ep = w2[e_start:e_end]
302
+ w1_scale_ep = w1_scale[e_start:e_end]
303
+ w2_scale_ep = w2_scale[e_start:e_end]
304
+
305
+ deepep_moe = deep_ep_moe_impl(
306
+ pg,
307
+ pgi,
308
+ dp_size,
309
+ test_tensors,
310
+ w1_ep,
311
+ w2_ep,
312
+ w1_scale_ep,
313
+ w2_scale_ep,
314
+ config.num_experts,
315
+ )
316
+
317
+ torch.testing.assert_close(
318
+ triton_moe,
319
+ deepep_moe,
320
+ atol=6e-2,
321
+ rtol=6e-2,
322
+ )
323
+
324
+
325
+MNKs = [
326
+ (8, 128, 128),
327
+ (8, 128, 512),
328
+ (8, 512, 512),
329
+ (3, 1024, 2048),
330
+ (32, 128, 1024),
331
+ (45, 512, 2048),
332
+ (64, 1024, 1024),
333
+ (129, 128, 256),
334
+ (129, 1024, 2048),
335
+ (222, 1024, 2048),
336
+]
337
+
338
+
339
+@pytest.mark.parametrize("mnk", MNKs)
340
+@pytest.mark.parametrize("num_experts", [32])
341
+@pytest.mark.parametrize("topk", [2, 6])
342
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
343
+@requires_deep_ep
344
+@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]):
347
+
348
+ m, n, k = mnk
349
+ current_platform.seed_everything(7)
350
+
351
+ if topk > num_experts:
352
+ pytest.skip(f"Skipping test: topk={topk} > E={num_experts}")
353
+
354
+ block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
355
+ block_size = [block_m, block_m]
356
+
357
+ world_size, dp_size = world_dp_size
358
+ config = TestConfig(
359
+ topk=topk,
360
+ m=m,
361
+ k=k,
362
+ n=n,
363
+ num_experts=num_experts,
364
+ block_size=block_size,
365
+ )
366
+
367
+ w1, w2, w1_scale, w2_scale = make_block_quant_fp8_weights(
368
+ num_experts, n, k, block_size)
369
+
370
+ parallel_launch(world_size, _deep_ep_moe, dp_size, config, w1, w2,
371
+ w1_scale, w2_scale)
tests/kernels/moe/test_deepep_moe.py
new
+459
@@ -0,0 +1,459 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+Test deepep dispatch-combine logic
4
+"""
5
+
6
+import dataclasses
7
+import importlib
8
+from typing import Optional, Union
9
+
10
+import pytest
11
+import torch.distributed
12
+from torch.distributed import ProcessGroup
13
+
14
+from vllm import _custom_ops as ops
15
+from vllm.config import VllmConfig, set_current_vllm_config
16
+from vllm.model_executor.layers.activation import SiluAndMul
17
+from vllm.model_executor.layers.fused_moe import TritonExperts
18
+from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
19
+ BatchedTritonExperts)
20
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
21
+ FusedMoEModularKernel)
22
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
23
+ per_token_group_quant_fp8)
24
+from vllm.platforms import current_platform
25
+
26
+from .deepep_utils import ProcessGroupInfo, parallel_launch
27
+
28
+has_deep_ep = importlib.util.find_spec("deep_ep") is not None
29
+
30
+if has_deep_ep:
31
+ from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
32
+ DeepEPHTPrepareAndFinalize)
33
+ from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
34
+ DeepEPLLPrepareAndFinalize)
35
+
36
+ from .deepep_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a
37
+
38
+requires_deep_ep = pytest.mark.skipif(
39
+ not has_deep_ep,
40
+ reason="Requires deep_ep kernels",
41
+)
42
+
43
+MAX_TOKENS_PER_RANK = 64
44
+
45
+
46
+def make_weights(
47
+ e, n, k, dtype
48
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
49
+ """
50
+ Return weights w1, w2, w1_scale, w2_scale
51
+ """
52
+ if dtype in [torch.float16, torch.bfloat16]:
53
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
54
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
55
+ return w1, w2, None, None
56
+
57
+ # per-out-channel weight quantization
58
+ assert dtype == torch.float8_e4m3fn
59
+ w1 = torch.empty((e, 2 * n, k), device="cuda", dtype=torch.float16)
60
+ w2 = torch.empty((e, k, n), device="cuda", dtype=torch.float16)
61
+
62
+ n_b_scales = 2 * n
63
+ k_b_scales = k
64
+ w1_q = torch.empty_like(w1, dtype=dtype)
65
+ w2_q = torch.empty_like(w2, dtype=dtype)
66
+ w1_scale = torch.empty((e, n_b_scales, 1),
67
+ device="cuda",
68
+ dtype=torch.float32)
69
+ w2_scale = torch.empty((e, k_b_scales, 1),
70
+ device="cuda",
71
+ dtype=torch.float32)
72
+ for expert in range(e):
73
+ w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
74
+ w1[expert], use_per_token_if_dynamic=True)
75
+ w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
76
+ w2[expert], use_per_token_if_dynamic=True)
77
+ return w1_q, w2_q, w1_scale, w2_scale
78
+
79
+
80
+@dataclasses.dataclass
81
+class TestConfig:
82
+ dtype: torch.dtype
83
+ topk: int
84
+ m: int
85
+ k: int
86
+ n: int
87
+ num_experts: int
88
+
89
+
90
+@dataclasses.dataclass
91
+class TestTensors:
92
+ rank_tokens: torch.Tensor # all ranks make this many tokens
93
+ rank_token_scales: Optional[torch.Tensor]
94
+ topk: torch.Tensor
95
+ topk_weights: torch.Tensor
96
+ config: TestConfig
97
+
98
+ @staticmethod
99
+ def make(config: TestConfig, low_latency_mode: bool) -> "TestTensors":
100
+ # TODO (varun) - check that float16 works ?
101
+ assert config.dtype in [torch.bfloat16, torch.float8_e4m3fn]
102
+ token_dtype = (torch.bfloat16 if config.dtype == torch.float8_e4m3fn
103
+ else config.dtype)
104
+ rank_tokens = torch.randn(
105
+ (config.m, config.k), device="cuda", dtype=token_dtype) / 10
106
+ rank_token_scales = None
107
+ if config.dtype == torch.float8_e4m3fn:
108
+ # low_latency_mode kernels dont support per-token quant.
109
+ _, rank_token_scales = ops.scaled_fp8_quant(
110
+ rank_tokens, use_per_token_if_dynamic=not low_latency_mode)
111
+
112
+ topk = torch.randint(low=0,
113
+ high=config.num_experts,
114
+ size=(config.m, config.topk),
115
+ device="cuda").to(dtype=torch.int64)
116
+ topk_weights = torch.randn(topk.shape,
117
+ dtype=torch.float32,
118
+ device="cuda")
119
+ return TestTensors(rank_tokens=rank_tokens,
120
+ rank_token_scales=rank_token_scales,
121
+ topk=topk,
122
+ topk_weights=topk_weights,
123
+ config=config)
124
+
125
+
126
+def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
127
+ low_latency_mode: bool, hidden_size: int, dp_size: int,
128
+ num_experts: int, num_local_experts: int,
129
+ q_dtype: Optional[torch.dtype],
130
+ use_fp8_dispatch: bool) -> FusedMoEModularKernel:
131
+
132
+ is_quantized = q_dtype is not None
133
+
134
+ ht_args: Optional[DeepEPHTArgs] = None
135
+ ll_args: Optional[DeepEPLLArgs] = None
136
+
137
+ if low_latency_mode:
138
+ ll_args = DeepEPLLArgs(max_tokens_per_rank=MAX_TOKENS_PER_RANK,
139
+ hidden_size=hidden_size,
140
+ num_experts=num_experts,
141
+ use_fp8_dispatch=use_fp8_dispatch)
142
+ else:
143
+ assert not use_fp8_dispatch, (
144
+ "FP8 Dispatch is valid only for low-latency kernels")
145
+ ht_args = DeepEPHTArgs(num_local_experts=num_local_experts)
146
+
147
+ a2a : Union[DeepEPHTPrepareAndFinalize, DeepEPLLPrepareAndFinalize] = \
148
+ make_deepep_a2a(pg = pg,
149
+ pgi = pgi,
150
+ dp_size = dp_size,
151
+ q_dtype = q_dtype,
152
+ block_shape = None,
153
+ deepep_ht_args = ht_args,
154
+ deepep_ll_args = ll_args)
155
+
156
+ if low_latency_mode:
157
+ fused_experts = BatchedTritonExperts(
158
+ max_num_tokens=MAX_TOKENS_PER_RANK,
159
+ world_size=pgi.world_size,
160
+ dp_size=dp_size,
161
+ use_fp8_w8a8=is_quantized,
162
+ use_int8_w8a8=False,
163
+ use_int8_w8a16=False,
164
+ use_int4_w4a16=False)
165
+ else:
166
+ fused_experts = TritonExperts(use_fp8_w8a8=is_quantized,
167
+ use_int8_w8a8=False,
168
+ use_int8_w8a16=False,
169
+ use_int4_w4a16=False,
170
+ per_channel_quant=False)
171
+
172
+ mk = FusedMoEModularKernel(prepare_finalize=a2a,
173
+ fused_experts=fused_experts)
174
+ return mk
175
+
176
+
177
+def deep_ep_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
178
+ low_latency_mode: bool, dp_size: int,
179
+ test_tensors: TestTensors, w1: torch.Tensor,
180
+ w2: torch.Tensor, w1_scale: Optional[torch.Tensor],
181
+ w2_scale: Optional[torch.Tensor], num_experts: int,
182
+ use_fp8_dispatch: bool) -> torch.Tensor:
183
+
184
+ num_local_experts = w1.size(0)
185
+
186
+ def build_expert_map():
187
+ num_local_experts = w1.size(0)
188
+ expert_map = torch.full((num_experts, ),
189
+ fill_value=-1,
190
+ dtype=torch.int32)
191
+ s = pgi.rank * num_local_experts
192
+ e = s + num_local_experts
193
+ expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
194
+ return expert_map.to(device=torch.cuda.current_device(),
195
+ dtype=torch.int32)
196
+
197
+ hidden_size = test_tensors.rank_tokens.size(1)
198
+ is_quantized = w1.dtype == torch.float8_e4m3fn
199
+ q_dtype = None
200
+ if is_quantized:
201
+ q_dtype = torch.float8_e4m3fn
202
+
203
+ # Make modular kernel
204
+ mk: FusedMoEModularKernel = make_modular_kernel(pg, pgi, low_latency_mode,
205
+ hidden_size, dp_size,
206
+ num_experts,
207
+ num_local_experts, q_dtype,
208
+ use_fp8_dispatch)
209
+
210
+ out_hidden_states = torch.empty_like(test_tensors.rank_tokens)
211
+ total_num_tokens = test_tensors.rank_tokens.size(0)
212
+
213
+ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
214
+ rank_tokens_chunk = test_tensors.rank_tokens[chunk_start:chunk_end]
215
+ topk_weights_chunk = test_tensors.topk_weights[chunk_start:chunk_end]
216
+ topk_chunk = test_tensors.topk[chunk_start:chunk_end]
217
+ rank_token_scales_chunk = test_tensors.rank_token_scales
218
+ if rank_token_scales_chunk is not None and rank_token_scales_chunk.size(
219
+ 0) == total_num_tokens:
220
+ # per act token
221
+ rank_token_scales_chunk = rank_token_scales_chunk[
222
+ chunk_start:chunk_end]
223
+
224
+ out = mk.forward(hidden_states=rank_tokens_chunk,
225
+ w1=w1,
226
+ w2=w2,
227
+ topk_weights=topk_weights_chunk,
228
+ topk_ids=topk_chunk,
229
+ inplace=False,
230
+ activation="silu",
231
+ global_num_experts=num_experts,
232
+ expert_map=build_expert_map(),
233
+ w1_scale=w1_scale,
234
+ w2_scale=w2_scale,
235
+ w1_zp=None,
236
+ w2_zp=None,
237
+ a1_scale=rank_token_scales_chunk,
238
+ a2_scale=None,
239
+ apply_router_weight_on_input=False)
240
+
241
+ if not skip_result_store:
242
+ out_hidden_states[chunk_start:chunk_end, :].copy_(
243
+ out, non_blocking=True)
244
+
245
+ max_num_tokens_per_dp = (MAX_TOKENS_PER_RANK
246
+ if low_latency_mode else total_num_tokens)
247
+
248
+ for chunk_start_ in range(0, total_num_tokens, max_num_tokens_per_dp):
249
+ chunk_start = chunk_start_
250
+ chunk_end = min(chunk_start + max_num_tokens_per_dp, total_num_tokens)
251
+ # clamp start and end
252
+ chunk_start = min(chunk_start, total_num_tokens - 1)
253
+ chunk_end = min(chunk_end, total_num_tokens)
254
+
255
+ process_chunk(chunk_start,
256
+ chunk_end,
257
+ skip_result_store=chunk_start_ >= total_num_tokens)
258
+
259
+ return out_hidden_states
260
+
261
+
262
+def torch_moe_impl(test_tensors: TestTensors, w1: torch.Tensor,
263
+ w2: torch.Tensor, w1_scale: Optional[torch.Tensor],
264
+ w2_scale: Optional[torch.Tensor], using_fp8_dispatch: bool):
265
+
266
+ a, topk_ids, topk_weights = (test_tensors.rank_tokens, test_tensors.topk,
267
+ test_tensors.topk_weights)
268
+ if using_fp8_dispatch:
269
+ # The DeepEP implementation is requested to dispatch using FP8.
270
+ # For numerical stability for testing, emulate the fp8 dispatch by
271
+ # blockwise quant and de-quant.
272
+ a = test_tensors.rank_tokens
273
+ aq, aq_scale = per_token_group_quant_fp8(a, 128)
274
+ a = (aq.view(-1, 128).to(torch.float32) * aq_scale.view(-1, 1)).view(
275
+ a.shape).to(a.dtype)
276
+
277
+ is_quantized = w1.dtype == torch.float8_e4m3fn
278
+ a_dtype = a.dtype
279
+ if is_quantized:
280
+ w1 = w1.to(dtype=torch.float32) * w1_scale
281
+ w2 = w2.to(dtype=torch.float32) * w2_scale
282
+ a = a.to(dtype=torch.float32)
283
+
284
+ m, _ = a.shape
285
+ topk = topk_ids.size(1)
286
+ out = torch.zeros_like(a)
287
+
288
+ for i in range(m):
289
+ a_i = a[i]
290
+ o_i = out[i]
291
+ for j in range(topk):
292
+ e = topk_ids[i][j]
293
+ e_w = topk_weights[i][j]
294
+ w1_e = w1[e]
295
+ w2_e = w2[e]
296
+ o_i += (SiluAndMul()
297
+ (a_i @ w1_e.transpose(0, 1)) @ w2_e.transpose(0, 1)) * e_w
298
+
299
+ if is_quantized:
300
+ out = out.to(dtype=a_dtype)
301
+
302
+ return out
303
+
304
+
305
+def _deep_ep_moe(
306
+ pgi: ProcessGroupInfo,
307
+ low_latency_mode: bool,
308
+ dp_size: int,
309
+ config: TestConfig,
310
+ w1: torch.Tensor,
311
+ w2: torch.Tensor,
312
+ w1_scale: Optional[torch.Tensor],
313
+ w2_scale: Optional[torch.Tensor],
314
+ use_fp8_dispatch: bool,
315
+):
316
+
317
+ if not low_latency_mode:
318
+ assert not use_fp8_dispatch, (
319
+ "FP8 dispatch interface is available only in low-latency mode")
320
+
321
+ is_quantized = w1.dtype == torch.float8_e4m3fn
322
+ w1 = w1.to(device=torch.cuda.current_device())
323
+ w2 = w2.to(device=torch.cuda.current_device())
324
+ if is_quantized:
325
+ w1_scale = w1_scale.to( # type: ignore
326
+ device=torch.cuda.current_device())
327
+ w2_scale = w2_scale.to( # type: ignore
328
+ device=torch.cuda.current_device())
329
+
330
+ pg = torch.distributed.new_group(list(range(pgi.world_size)))
331
+ test_tensors = TestTensors.make(config, low_latency_mode)
332
+
333
+ with set_current_vllm_config(VllmConfig()):
334
+ # Reference
335
+ torch_combined = torch_moe_impl(test_tensors, w1, w2, w1_scale,
336
+ w2_scale, use_fp8_dispatch)
337
+
338
+ # Splice experts for this rank.
339
+ num_local_experts = config.num_experts // pgi.world_size
340
+ e_start = num_local_experts * pgi.rank
341
+ e_end = e_start + num_local_experts
342
+ w1_ep = w1[e_start:e_end]
343
+ w2_ep = w2[e_start:e_end]
344
+
345
+ w1_scale_ep, w2_scale_ep = None, None
346
+ if is_quantized:
347
+ w1_scale_ep = w1_scale[e_start:e_end] # type: ignore
348
+ w2_scale_ep = w2_scale[e_start:e_end] # type: ignore
349
+ deepep_combined = deep_ep_moe_impl(
350
+ pg,
351
+ pgi,
352
+ low_latency_mode,
353
+ dp_size,
354
+ test_tensors,
355
+ w1_ep,
356
+ w2_ep,
357
+ w1_scale_ep,
358
+ w2_scale_ep,
359
+ config.num_experts,
360
+ use_fp8_dispatch,
361
+ )
362
+
363
+ torch.testing.assert_close(
364
+ torch_combined,
365
+ deepep_combined,
366
+ atol=6e-2,
367
+ rtol=6e-2,
368
+ )
369
+
370
+
371
+MNKs = [
372
+ (1, 128, 128),
373
+ (2, 128, 512),
374
+ (3, 1024, 2048),
375
+ (32, 128, 1024),
376
+ (45, 512, 2048),
377
+ (64, 1024, 1024),
378
+ (222, 1024, 2048),
379
+]
380
+
381
+DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
382
+
383
+
384
+@pytest.mark.parametrize("dtype", DTYPES)
385
+@pytest.mark.parametrize("mnk", MNKs)
386
+@pytest.mark.parametrize("num_experts", [32])
387
+@pytest.mark.parametrize("topk", [6])
388
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
389
+@requires_deep_ep
390
+def test_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],
391
+ num_experts: int, topk: int, world_dp_size: tuple[int,
392
+ int]):
393
+ low_latency_mode = False
394
+ use_fp8_dispatch = False
395
+ m, n, k = mnk
396
+
397
+ current_platform.seed_everything(7)
398
+ world_size, dp_size = world_dp_size
399
+ config = TestConfig(dtype=dtype,
400
+ topk=topk,
401
+ m=m,
402
+ k=k,
403
+ n=n,
404
+ num_experts=num_experts)
405
+
406
+ w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
407
+
408
+ parallel_launch(world_size, _deep_ep_moe, low_latency_mode, dp_size,
409
+ config, w1, w2, w1_scale, w2_scale, use_fp8_dispatch)
410
+
411
+
412
+MNKs = [
413
+ (1, 128, 2560),
414
+ (2, 128, 2560),
415
+ (3, 1024, 2560),
416
+ (32, 128, 2560),
417
+ (45, 512, 2560),
418
+ (64, 1024, 2560),
419
+ (222, 1024, 2560),
420
+]
421
+DTYPES = [torch.float8_e4m3fn, torch.bfloat16]
422
+USE_FP8_DISPATCH = [True, False]
423
+
424
+
425
+@pytest.mark.parametrize("dtype", DTYPES)
426
+@pytest.mark.parametrize("mnk", MNKs)
427
+@pytest.mark.parametrize("num_experts", [32])
428
+@pytest.mark.parametrize("topk", [6])
429
+@pytest.mark.parametrize("world_dp_size", [(2, 1)])
430
+@pytest.mark.parametrize("use_fp8_dispatch", USE_FP8_DISPATCH)
431
+@requires_deep_ep
432
+def test_low_latency_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],
433
+ num_experts: int, topk: int,
434
+ world_dp_size: tuple[int, int],
435
+ use_fp8_dispatch: bool):
436
+
437
+ low_latency_mode = True
438
+ m, n, k = mnk
439
+
440
+ if (low_latency_mode
441
+ and k not in DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES):
442
+ pytest.skip(
443
+ f"Skipping test as hidden size {k} is not in list of supported "
444
+ f"hidden sizes {DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES}"
445
+ )
446
+
447
+ current_platform.seed_everything(7)
448
+ world_size, dp_size = world_dp_size
449
+ config = TestConfig(dtype=dtype,
450
+ topk=topk,
451
+ m=m,
452
+ k=k,
453
+ n=n,
454
+ num_experts=num_experts)
455
+
456
+ w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
457
+
458
+ parallel_launch(world_size, _deep_ep_moe, low_latency_mode, dp_size,
459
+ config, w1, w2, w1_scale, w2_scale, use_fp8_dispatch)
vllm/config.py
+2
@@ -1856,6 +1856,8 @@ class ParallelConfig:
1856
factors.append(self.pipeline_parallel_size)
1857
factors.append(self.tensor_parallel_size)
1858
factors.append(self.enable_expert_parallel)
1859
+ factors.append(self.data_parallel_size)
1860
+ factors.append(envs.VLLM_ALL2ALL_BACKEND)
1861
return hashlib.sha256(str(factors).encode()).hexdigest()
1862
1863
def __post_init__(self) -> None:
vllm/distributed/device_communicators/all2all.py
+145
-1
@@ -1,7 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import importlib.util
4
-from typing import TYPE_CHECKING
4
+from typing import TYPE_CHECKING, Any
5
6
import torch
7
import torch.distributed as dist
@@ -129,3 +129,147 @@ class PPLXAll2AllManager(All2AllManagerBase):
129
from pplx_kernels.nvshmem import nvshmem_finalize
130
logger.debug("PPLX NVSHMEM finalize")
131
nvshmem_finalize()
132
+
133
+
134
+class DeepEPAll2AllManagerBase(All2AllManagerBase):
135
+ """
136
+ All2All communication based on DeepEP High-Throughput kernels.
137
+ """
138
+
139
+ def __init__(self, cpu_group):
140
+ has_deepep = importlib.util.find_spec("deep_ep") is not None
141
+ assert has_deepep, "DeepEP kernels not found. Please follow https://github.com/vllm-project/vllm/blob/main/tools/ep_kernels/README.md to install DeepEP kernels." # noqa
142
+ super().__init__(cpu_group)
143
+ self.handle_cache = Cache()
144
+
145
+ # This is the DeepEP default. Stick to it till we can establish
146
+ # reasonable defaults based on profiling.
147
+ self.num_sms = 20
148
+
149
+ def get_handle(self, kwargs):
150
+ raise NotImplementedError
151
+
152
+ def dispatch(self, hidden_states: torch.Tensor,
153
+ router_logits: torch.Tensor):
154
+ raise NotImplementedError
155
+
156
+ def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
157
+ raise NotImplementedError
158
+
159
+ def destroy(self):
160
+ pass
161
+
162
+
163
+class DeepEPHTAll2AllManager(DeepEPAll2AllManagerBase):
164
+ """
165
+ All2All communication based on DeepEP High-Throughput kernels.
166
+ """
167
+
168
+ def __init__(self, cpu_group):
169
+ super().__init__(cpu_group)
170
+
171
+ def _make_all2all_kwargs(self) -> dict[Any, Any]:
172
+ # Defaults for internode and intranode are taken from DeepEP tests.
173
+ num_nvl_bytes = 1024 * 1024 * 1024
174
+ num_rdma_bytes = None
175
+ num_qps_per_rank = None
176
+
177
+ if self.internode:
178
+ num_rdma_bytes = 1024 * 1024 * 1024
179
+ num_qps_per_rank = self.num_sms // 2
180
+ else:
181
+ assert self.intranode
182
+ num_rdma_bytes = 0
183
+ num_qps_per_rank = 1
184
+
185
+ assert num_rdma_bytes is not None
186
+ assert num_qps_per_rank is not None
187
+ return dict(group=self.cpu_group,
188
+ num_nvl_bytes=num_nvl_bytes,
189
+ num_rdma_bytes=num_rdma_bytes,
190
+ low_latency_mode=False,
191
+ num_qps_per_rank=num_qps_per_rank)
192
+
193
+ def get_handle(self, kwargs):
194
+
195
+ assert len(kwargs) == 0, (
196
+ "DeepEPHTAll2AllManager expects no arguments. All the required "
197
+ "args are computed in the Manager itself.")
198
+
199
+ import deep_ep
200
+ buffer_kwargs = self._make_all2all_kwargs()
201
+ logger.debug("DeepEP all2all args %s", buffer_kwargs)
202
+ handle: deep_ep.Buffer = self.handle_cache.get_or_create(
203
+ buffer_kwargs, deep_ep.Buffer)
204
+ # It is dangerous to set num sms outside this function. num_sms is not
205
+ # a part of the hash-key that identifies this object. If we are in a
206
+ # situation where we make objects with different num_sms, the hash key
207
+ # in get_or_create must be updated.
208
+ handle.set_num_sms(self.num_sms)
209
+ return handle
210
+
211
+
212
+class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
213
+ """
214
+ All2All communication based on DeepEP Low-Latency kernels.
215
+ """
216
+
217
+ def __init__(self, cpu_group):
218
+ super().__init__(cpu_group)
219
+
220
+ def _make_all2all_kwargs(
221
+ self,
222
+ max_num_tokens_per_dp_rank: int,
223
+ token_hidden_size: int,
224
+ num_ep_ranks: int,
225
+ num_global_experts: int,
226
+ num_local_experts: int,
227
+ ) -> dict[Any, Any]:
228
+ """
229
+ max_num_tokens_per_dp_rank : the maximum number of tokens a DP rank
230
+ can dispatch all the ranks must hold the same value.
231
+ token_hidden_size: the hidden dimension of each token.
232
+ num_ep_ranks: the number of EP group ranks.
233
+ num_global_experts: Number of experts in the model.
234
+ num_local_experts: Number of experts in an EP rank.
235
+ """
236
+ import deep_ep
237
+
238
+ # Defaults for internode and intranode are taken from DeepEP tests.
239
+ num_nvl_bytes = 1024 * 1024 * 1024
240
+ num_qps_per_rank = num_local_experts
241
+ num_rdma_bytes = None
242
+
243
+ if self.internode:
244
+ num_rdma_bytes = 1024 * 1024 * 1024
245
+ else:
246
+ assert self.intranode
247
+ num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
248
+ num_max_dispatch_tokens_per_rank=max_num_tokens_per_dp_rank,
249
+ hidden=token_hidden_size,
250
+ num_ranks=num_ep_ranks,
251
+ num_experts=num_global_experts)
252
+
253
+ assert num_rdma_bytes is not None
254
+ return dict(group=self.cpu_group,
255
+ num_nvl_bytes=num_nvl_bytes,
256
+ num_rdma_bytes=num_rdma_bytes,
257
+ low_latency_mode=True,
258
+ num_qps_per_rank=num_qps_per_rank)
259
+
260
+ def get_handle(self, kwargs):
261
+ """
262
+ The kwargs for DeepEPLLAll2AllManager is dictated by
263
+ _make_all2all_kwargs.
264
+ """
265
+ import deep_ep
266
+ buffer_kwargs = self._make_all2all_kwargs(**kwargs)
267
+ logger.debug("DeepEP all2all args %s", buffer_kwargs)
268
+ handle: deep_ep.Buffer = self.handle_cache.get_or_create(
269
+ buffer_kwargs, deep_ep.Buffer)
270
+ # It is dangerous to set num sms outside this function. num_sms is not
271
+ # a part of the hash-key that identifies this object. If we are in a
272
+ # situation where we make objects with different num_sms, the hash key
273
+ # in get_or_create must be updated.
274
+ handle.set_num_sms(self.num_sms)
275
+ return handle
vllm/distributed/device_communicators/cuda_communicator.py
+8
@@ -67,6 +67,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
67
from .all2all import PPLXAll2AllManager
68
self.all2all_manager = PPLXAll2AllManager(self.cpu_group)
69
logger.info("Using PPLX all2all manager.")
70
+ elif all2all_backend == "deepep_high_throughput":
71
+ from .all2all import DeepEPHTAll2AllManager
72
+ self.all2all_manager = DeepEPHTAll2AllManager(self.cpu_group)
73
+ logger.info("Using DeepEP High-Throughput all2all manager.")
74
+ elif all2all_backend == "deepep_low_latency":
75
+ from .all2all import DeepEPLLAll2AllManager
76
+ self.all2all_manager = DeepEPLLAll2AllManager(self.cpu_group)
77
+ logger.info("Using DeepEP Low-Latency all2all manager.")
78
else:
79
raise ValueError(f"Unknown all2all backend: {all2all_backend}")
80
vllm/envs.py
+2
@@ -826,6 +826,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
826
# Available options:
827
# - "naive": naive all2all implementation using all-reduce
828
# - "pplx": use pplx kernels
829
+ # - "deepep_high_throughput", use deepep high-throughput kernels
830
+ # - "deepep_low_latency", use deepep low-latency kernels
831
"VLLM_ALL2ALL_BACKEND":
832
lambda: os.getenv("VLLM_ALL2ALL_BACKEND", "naive"),
833
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+18
-14
@@ -12,8 +12,8 @@ from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
12
_moe_permute)
13
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
14
MoEPrepareAndFinalizeNoEP)
15
-from vllm.model_executor.layers.fused_moe.utils import (_fp8_quantize,
16
- _resize_cache)
15
+from vllm.model_executor.layers.fused_moe.utils import (
16
+ _resize_cache, per_token_group_quant_fp8)
17
from vllm.utils import round_up
18
19
logger = init_logger(__name__)
@@ -34,10 +34,8 @@ def _valid_deep_gemm_shape(M: int, N: int, K: int):
34
return align <= M and N % align == 0 and K % align == 0
35
36
37
-def _valid_deep_gemm(hidden_states: torch.Tensor,
38
- w1: torch.Tensor,
39
- w2: torch.Tensor,
40
- expert_map: Optional[torch.Tensor] = None) -> bool:
37
+def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
38
+ w2: torch.Tensor) -> bool:
39
"""
40
Check if the given problem size is supported by the DeepGemm grouped
41
gemm kernel. All of M, N, K and the quantization block_shape must be
@@ -47,10 +45,6 @@ def _valid_deep_gemm(hidden_states: torch.Tensor,
45
logger.debug("DeepGemm disabled: deep_gemm not available.")
46
return False
47
50
- if expert_map is not None:
51
- logger.debug("DeepGemm disabled: expert map NYI.")
52
- return False
53
-
48
M = hidden_states.size(0)
49
_, K, N = w2.size()
50
if not _valid_deep_gemm_shape(M, N, K):
@@ -116,7 +110,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
110
a1q = hidden_states
111
_, N, K = w1.size()
112
119
- assert global_num_experts != -1
113
+ if global_num_experts == -1:
114
+ global_num_experts = w1.size(0)
115
+
116
assert w2.size(1) == K
117
118
a1q, a1q_scale, _, expert_ids, inv_perm = _moe_permute(
@@ -128,6 +124,14 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
124
self.block_shape[0],
125
)
126
127
+ if expert_map is not None:
128
+ # DeepGemm (Grouped Contiguous) kernel needs a valid B index
129
+ # for all rows of A. To that effect, simply compute with
130
+ # the 0th weight matrix.
131
+ # Note that this relies on the fact that corresponding topk
132
+ # weights would be 0 during weight multiplication.
133
+ expert_ids = torch.where(expert_ids == -1, 0, expert_ids)
134
+
135
# Note: M_sum is different than the pre-permuted shape of a1q.
136
M_sum = a1q.size(0)
137
workspace1 = _resize_cache(workspace13, (M_sum, N))
@@ -140,9 +144,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
144
self.activation(activation, workspace2, workspace1.view(-1, N))
145
146
a2q_scale: Optional[torch.Tensor] = None
143
-
144
- a2q, a2q_scale = _fp8_quantize(workspace2, a2_scale, False,
145
- self.block_shape)
147
+ a2q, a2q_scale = per_token_group_quant_fp8(workspace2,
148
+ self.block_shape[1],
149
+ column_major_scales=True)
150
151
dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
152
(a2q, a2q_scale), (w2, w2_scale), workspace3, expert_ids)
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
new
+236
@@ -0,0 +1,236 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import deep_ep
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.utils import (
9
+ moe_kernel_quantize_input)
10
+
11
+
12
+class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
13
+ """
14
+ Prepare/Finalize using DeepEP High-Throughput kernels.
15
+ """
16
+
17
+ def __init__(self,
18
+ buffer: deep_ep.Buffer,
19
+ world_size: int,
20
+ rank: int,
21
+ dp_size: int,
22
+ rank_expert_offset: int,
23
+ quant_dtype: Optional[torch.dtype] = None,
24
+ block_shape: Optional[list[int]] = None):
25
+ super().__init__()
26
+ self.buffer = buffer
27
+ self.world_size = world_size
28
+ self.rank = rank
29
+ self.dp_size = dp_size
30
+ self.rank_expert_offset = rank_expert_offset
31
+ self.quant_dtype = quant_dtype
32
+ self.block_shape = block_shape
33
+ # The dispatch function returns a handle that the combine function
34
+ # requires. We store the handle here so it is available to the
35
+ # combine function.
36
+ self.handle = None
37
+
38
+ # From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164
39
+ self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]
40
+
41
+ def max_num_tokens_per_rank(self) -> Optional[int]:
42
+ return None
43
+
44
+ def topk_indices_dtype(self) -> Optional[torch.dtype]:
45
+ return torch.int64
46
+
47
+ def _get_dispatch_config(self) -> Optional[deep_ep.Config]:
48
+ if self.dp_size not in self.available_rank_configs:
49
+ return None
50
+ return deep_ep.Buffer.get_dispatch_config(self.dp_size)
51
+
52
+ def _get_combine_config(self) -> Optional[deep_ep.Config]:
53
+ if self.dp_size not in self.available_rank_configs:
54
+ return None
55
+ return deep_ep.Buffer.get_combine_config(self.dp_size)
56
+
57
+ def _do_quant(self, tokens: torch.Tensor,
58
+ token_scales: Optional[torch.Tensor], per_act_token: bool):
59
+ tokens, token_scales = moe_kernel_quantize_input(
60
+ tokens, token_scales, self.quant_dtype, per_act_token,
61
+ self.block_shape)
62
+ return tokens, token_scales
63
+
64
+ def _do_dispatch(self, tokens: torch.Tensor,
65
+ token_scales: Optional[torch.Tensor],
66
+ rank_topk_ids: torch.Tensor,
67
+ rank_topk_weights: torch.Tensor, num_experts: int):
68
+
69
+ has_scales = token_scales is not None
70
+
71
+ (num_tokens_per_rank, num_tokens_per_rdma_rank, expert_num_tokens,
72
+ is_token_in_rank, event) = self.buffer.get_dispatch_layout(
73
+ topk_idx=rank_topk_ids,
74
+ num_experts=num_experts,
75
+ previous_event=None,
76
+ async_finish=False,
77
+ allocate_on_comm_stream=False)
78
+
79
+ token_data = tokens
80
+ if has_scales:
81
+ token_data = (tokens, token_scales)
82
+
83
+ (
84
+ token_data, expert_topk_ids, expert_topk_weights,
85
+ expert_num_tokens_per_expert_list, self.handle, event
86
+ ) = self.buffer.dispatch(
87
+ x=token_data,
88
+ handle=None,
89
+ num_tokens_per_rank=num_tokens_per_rank,
90
+ num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
91
+ is_token_in_rank=is_token_in_rank,
92
+ num_tokens_per_expert=expert_num_tokens,
93
+ topk_idx=rank_topk_ids,
94
+ topk_weights=rank_topk_weights,
95
+ # expert_alignment rounds the number of tokens per expert
96
+ # to this value.
97
+ expert_alignment=1,
98
+ config=self._get_dispatch_config(),
99
+ previous_event=None,
100
+ async_finish=False,
101
+ allocate_on_comm_stream=False)
102
+
103
+ if has_scales:
104
+ expert_x, expert_x_scale = token_data
105
+ else:
106
+ expert_x, expert_x_scale = token_data, None
107
+
108
+ # The existing MOE kernels assume that all entries of topk_ids are
109
+ # valid. To that effect, set the -1s in expert_topk_ids to some expert
110
+ # outside this rank so the expert_map can remap it to -1 when safe.
111
+ # With Expert Parallel, the experts are divided amongst the rank
112
+ # sequentially. For rank 0, set it to num_experts - 1 and for all other
113
+ # ranks set it to 0 as we know that expert_map will have a -1 in those
114
+ # regions for those ranks.
115
+ #
116
+ # DeepEP's topk_ids output refers to the local experts directly. Offset
117
+ # the topk_ids to move it back to the global experts space so it aligns
118
+ # with existing vLLM interfaces.
119
+ expert_topk_ids = torch.where(
120
+ expert_topk_ids == -1,
121
+ num_experts - 1 if self.rank_expert_offset == 0 else 0,
122
+ expert_topk_ids + self.rank_expert_offset)
123
+
124
+ return (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,
125
+ expert_topk_weights)
126
+
127
+ def prepare(
128
+ self,
129
+ a1: torch.Tensor,
130
+ a1_scale: Optional[torch.Tensor],
131
+ a2_scale: Optional[torch.Tensor],
132
+ rank_topk_weights: torch.Tensor,
133
+ rank_topk_ids: torch.Tensor,
134
+ num_experts: int,
135
+ expert_map: Optional[torch.Tensor],
136
+ apply_router_weight_on_input: bool,
137
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
138
+ Optional[torch.Tensor], Optional[torch.Tensor]]:
139
+
140
+ if apply_router_weight_on_input:
141
+ topk = rank_topk_ids.size(1)
142
+ # TODO: this only works for topK=1, will need to update for topK>1
143
+ assert topk == 1, (
144
+ "apply_router_weight_on_input is only implemented for topk=1")
145
+ a1 = a1 * rank_topk_weights.to(a1.dtype)
146
+
147
+ # Check if there is a block_shape / or if we can infer the quantization
148
+ # schemes from the scales.
149
+ per_token_quant = None
150
+ if all([x is None for x in [self.block_shape, a1_scale, a2_scale]
151
+ ]) and self.quant_dtype is not None:
152
+ # Quantization required despite none of the inputs suggesting
153
+ # quantization. Fallback to per_token_dynamic quant.
154
+ per_token_quant = True
155
+ else:
156
+ per_token_quant = ((self.block_shape is not None) or
157
+ (a1_scale is not None and a1_scale.numel() != 1)
158
+ or (a2_scale is not None
159
+ and a2_scale.numel() != 1))
160
+
161
+ if per_token_quant:
162
+ a1q, a1q_scale = self._do_quant(a1, a1_scale, per_act_token=True)
163
+ (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,
164
+ expert_topk_weights) = self._do_dispatch(
165
+ tokens=a1q,
166
+ token_scales=a1q_scale,
167
+ rank_topk_ids=rank_topk_ids,
168
+ rank_topk_weights=rank_topk_weights,
169
+ num_experts=num_experts)
170
+ else:
171
+ # DeepEP kernels only support dispatching per-token-quant
172
+ # quantization. dispatch in bfloat16.
173
+ (expert_x, _, expert_num_tokens, expert_topk_ids,
174
+ expert_topk_weights) = self._do_dispatch(
175
+ tokens=a1,
176
+ token_scales=None,
177
+ rank_topk_ids=rank_topk_ids,
178
+ rank_topk_weights=rank_topk_weights,
179
+ num_experts=num_experts)
180
+ # quantize now
181
+ expert_x_scale = None
182
+ if expert_x.numel() != 0:
183
+ expert_x, expert_x_scale = self._do_quant(expert_x,
184
+ a1_scale,
185
+ per_act_token=False)
186
+
187
+ return (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,
188
+ expert_topk_weights)
189
+
190
+ def _apply_weights_and_reduce(self, num_tokens: int,
191
+ fused_expert_output: torch.Tensor,
192
+ topk_weights: torch.Tensor,
193
+ apply_router_weight_on_input: bool,
194
+ output_dtype: torch.dtype):
195
+
196
+ if fused_expert_output.ndim == 2:
197
+ hidden_dim = fused_expert_output.size(-1)
198
+ fused_expert_output = fused_expert_output.view(
199
+ num_tokens, -1, hidden_dim)
200
+
201
+ if not apply_router_weight_on_input:
202
+ # The DeepEP combine kernels don't do the topk weight
203
+ # multiplication. We multiply the weights locally.
204
+ fused_expert_output = fused_expert_output.to(torch.float32)
205
+ fused_expert_output = fused_expert_output * topk_weights.view(
206
+ fused_expert_output.size(0), -1, 1)
207
+ fused_expert_output = fused_expert_output.to(output_dtype)
208
+
209
+ return fused_expert_output.sum(dim=1).to(output_dtype)
210
+
211
+ def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
212
+ topk_weights: torch.Tensor, topk_ids: torch.Tensor,
213
+ apply_router_weight_on_input: bool) -> None:
214
+
215
+ assert self.handle is not None
216
+
217
+ # fused_expert_output can have 0 tokens - This happens when none of the
218
+ # tokens from the all2all reach this EP rank.
219
+ if fused_expert_output.numel() != 0:
220
+ fused_expert_output = self._apply_weights_and_reduce(
221
+ num_tokens=topk_ids.size(0),
222
+ fused_expert_output=fused_expert_output,
223
+ topk_weights=topk_weights,
224
+ apply_router_weight_on_input=apply_router_weight_on_input,
225
+ output_dtype=output.dtype)
226
+
227
+ combined_x, _, event = self.buffer.combine(
228
+ x=fused_expert_output,
229
+ handle=self.handle,
230
+ topk_weights=None,
231
+ config=self._get_combine_config(),
232
+ previous_event=None,
233
+ async_finish=False,
234
+ allocate_on_comm_stream=False)
235
+ # Respect inplace outputs.
236
+ output.copy_(combined_x, non_blocking=True)
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
new
+152
@@ -0,0 +1,152 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import deep_ep
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.utils import (
9
+ moe_kernel_quantize_input)
10
+
11
+# DeepEP kernels quantize dispatch inputs in 128 element chunks.
12
+DEEPEP_QUANT_BLOCK_SIZE = 128
13
+
14
+
15
+def dequant_fp8(expert_x_fp8: torch.Tensor,
16
+ expert_x_scales: torch.Tensor) -> torch.Tensor:
17
+ """
18
+ Return dequantized tensor in fp32
19
+ """
20
+ # TODO (varun) : Optimize leverage num_tokens_per_expert counts
21
+ assert expert_x_fp8.is_contiguous()
22
+ expert_x_scales = expert_x_scales.contiguous()
23
+ num_experts = expert_x_fp8.size(0)
24
+
25
+ expert_x_fp32 = expert_x_fp8.to(torch.float32).view(
26
+ num_experts, -1, DEEPEP_QUANT_BLOCK_SIZE)
27
+ expert_x_scales = expert_x_scales.view(num_experts, -1, 1)
28
+ return (expert_x_fp32 * expert_x_scales).view(expert_x_fp8.shape)
29
+
30
+
31
+class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
32
+ """
33
+ Prepare/Finalize using DeepEP low-latency kernels.
34
+ """
35
+
36
+ # DeepEP low-latency kernels are compiled only for certain
37
+ # specific hidden sizes.
38
+ SUPPORTED_HIDDEN_SIZES = [2560, 4096, 5120, 7168]
39
+
40
+ def __init__(self,
41
+ buffer: deep_ep.Buffer,
42
+ world_size: int,
43
+ dp_size: int,
44
+ max_tokens_per_rank: int,
45
+ quant_dtype: Optional[torch.dtype] = None,
46
+ block_shape: Optional[list[int]] = None,
47
+ use_fp8_dispatch: bool = False):
48
+ super().__init__()
49
+
50
+ self.buffer = buffer
51
+ self.world_size = world_size
52
+ self.dp_size = dp_size
53
+ self.quant_dtype = quant_dtype
54
+ self.block_shape = block_shape
55
+ self.max_tokens_per_rank = max_tokens_per_rank
56
+ self.use_fp8_dispatch = use_fp8_dispatch
57
+ # The dispatch function returns a handle that the combine function
58
+ # requires. We store the handle here so it is available to the
59
+ # combine function.
60
+ self.handle = None
61
+
62
+ def max_num_tokens_per_rank(self) -> Optional[int]:
63
+ return self.max_tokens_per_rank
64
+
65
+ def topk_indices_dtype(self) -> Optional[torch.dtype]:
66
+ return torch.int64
67
+
68
+ def prepare(
69
+ self,
70
+ a1: torch.Tensor,
71
+ a1_scale: Optional[torch.Tensor],
72
+ a2_scale: Optional[torch.Tensor],
73
+ rank_topk_weights: torch.Tensor,
74
+ rank_topk_ids: torch.Tensor,
75
+ num_experts: int,
76
+ expert_map: Optional[torch.Tensor],
77
+ apply_router_weight_on_input: bool,
78
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
79
+ Optional[torch.Tensor], Optional[torch.Tensor]]:
80
+
81
+ hidden_size = a1.size(1)
82
+ assert hidden_size in self.SUPPORTED_HIDDEN_SIZES, \
83
+ (f"Hidden Size {hidden_size} not in supported list of hidden sizes"
84
+ f"{self.SUPPORTED_HIDDEN_SIZES}")
85
+
86
+ if self.use_fp8_dispatch:
87
+ assert hidden_size % 128 == 0, \
88
+ "DeepEP kernels quantize the inputs in blocks of shape 128"
89
+
90
+ # Quantize
91
+ per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
92
+ 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")
95
+
96
+ if apply_router_weight_on_input:
97
+ topk = rank_topk_ids.size(1)
98
+ # TODO: this only works for topK=1, will need to update for topK>1
99
+ assert topk == 1, (
100
+ "apply_router_weight_on_input is only implemented for topk=1")
101
+ a1 = a1 * rank_topk_weights.to(a1.dtype)
102
+
103
+ # Dispatch
104
+ expert_x, expert_num_tokens, self.handle, event, hook = \
105
+ self.buffer.low_latency_dispatch(a1,
106
+ rank_topk_ids,
107
+ self.max_tokens_per_rank,
108
+ num_experts,
109
+ use_fp8=self.use_fp8_dispatch,
110
+ async_finish=False,
111
+ return_recv_hook=False)
112
+
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))
129
+
130
+ return (expert_x, expert_x_scale, expert_num_tokens, None, None)
131
+
132
+ def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
133
+ topk_weights: torch.Tensor, topk_ids: torch.Tensor,
134
+ apply_router_weight_on_input: bool) -> None:
135
+
136
+ assert self.handle is not None
137
+
138
+ combine_topk_weights = topk_weights
139
+ if apply_router_weight_on_input:
140
+ # weights have already been applied.
141
+ combine_topk_weights = torch.ones_like(topk_weights)
142
+
143
+ # TODO (varun) : Enable zero copy mode
144
+ _, event, hook = self.buffer.low_latency_combine(
145
+ fused_expert_output,
146
+ topk_ids,
147
+ combine_topk_weights,
148
+ self.handle,
149
+ async_finish=False,
150
+ zero_copy=False,
151
+ return_recv_hook=False,
152
+ out=output)
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+37
-20
@@ -10,7 +10,8 @@ import triton.language as tl
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm.model_executor.layers.fused_moe.fused_moe import (
12
get_config_dtype_str, try_get_optimal_moe_config)
13
-from vllm.model_executor.layers.fused_moe.utils import _resize_cache
13
+from vllm.model_executor.layers.fused_moe.utils import (
14
+ _resize_cache, moe_kernel_quantize_input)
15
16
17
@triton.jit
@@ -397,6 +398,12 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
398
self.rank = rank
399
self.max_num_tokens = max_num_tokens
400
401
+ def max_num_tokens_per_rank(self) -> Optional[int]:
402
+ return self.max_num_tokens
403
+
404
+ def topk_indices_dtype(self) -> Optional[torch.dtype]:
405
+ return None
406
+
407
def prepare(
408
self,
409
a1: torch.Tensor,
@@ -407,7 +414,8 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
414
num_experts: int,
415
expert_map: Optional[torch.Tensor],
416
apply_router_weight_on_input: bool,
410
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
417
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
418
+ Optional[torch.Tensor], Optional[torch.Tensor]]:
419
assert a1.dim() == 2
420
assert topk_ids.dim() == 2
421
assert topk_ids.size(0) == a1.size(0)
@@ -450,7 +458,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
458
first_expert, :rows, :] = a1[:topks.numel()][topks]
459
tokens_per_expert[expert_id - first_expert] = rows
460
453
- return b_a1, a1_scale, tokens_per_expert
461
+ return b_a1, a1_scale, tokens_per_expert, None, None
462
463
def finalize(
464
self,
@@ -601,6 +609,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
609
use_int8_w8a8: bool = False,
610
use_int8_w8a16: bool = False,
611
use_int4_w4a16: bool = False,
612
+ per_channel_quant: bool = False,
613
block_shape: Optional[list[int]] = None,
614
world_size: int = 1,
615
dp_size: int = 1,
@@ -611,12 +620,15 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
620
self.use_int4_w4a16 = use_int4_w4a16
621
self.use_int8_w8a16 = use_int8_w8a16
622
self.block_shape = block_shape
623
+ self.per_channel_quant = per_channel_quant
624
self.max_num_tokens = max_num_tokens
615
- assert not use_int8_w8a8, "NYI"
616
- assert not use_int4_w4a16, "NYI"
625
self.world_size = world_size
626
self.dp_size = dp_size
627
628
+ assert not use_int8_w8a8, "NYI"
629
+ assert not use_int4_w4a16, "NYI"
630
+ assert self.block_shape is None, "NYI"
631
+
632
def workspace_shapes(
633
self,
634
a: torch.Tensor,
@@ -670,8 +682,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
682
torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
683
]
684
673
- # TODO: num_tokens -> max_num_tokens?
674
- E, num_tokens, N, K, top_k_num = mk._moe_problem_size(
685
+ E, max_num_tokens, N, K, top_k_num = mk._moe_problem_size(
686
hidden_states, w1, w2, topk_ids)
687
688
assert w1.size(0) == E
@@ -687,7 +698,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
698
w2.size(),
699
top_k_num,
700
config_dtype,
690
- num_tokens,
701
+ max_num_tokens,
702
block_shape=self.block_shape,
703
)
704
@@ -706,10 +717,12 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
717
#print(f"shape: E={E}, M={num_tokens}, N={N}, K={K}, top_k={top_k_num}")
718
# We can reuse the memory between these because by the time we need
719
# cache3, we're done with cache1
709
- intermediate_cache1 = _resize_cache(workspace13, (E, num_tokens, N))
720
+ intermediate_cache1 = _resize_cache(workspace13,
721
+ (E, max_num_tokens, N))
722
intermediate_cache2 = _resize_cache(workspace2,
711
- (E, num_tokens, N // 2))
712
- intermediate_cache3 = _resize_cache(workspace13, (E, num_tokens, K))
723
+ (E, max_num_tokens, N // 2))
724
+ intermediate_cache3 = _resize_cache(workspace13,
725
+ (E, max_num_tokens, K))
726
727
# MM1
728
invoke_moe_batched_triton_kernel(A=hidden_states,
@@ -731,15 +744,20 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
744
self.activation(activation, intermediate_cache2.view(-1, N // 2),
745
intermediate_cache1.view(-1, N))
746
734
- #qintermediate_cache2 = intermediate_cache2
735
- a2q_scale = a2_scale
736
- # TODO (varun) : support w8a8
737
- assert not self.use_fp8_w8a8
738
- #if self.use_fp8_w8a8:
739
- # qintermediate_cache2, a2q_scale = _fp8_quantize(
740
- # intermediate_cache2, a2_scale, self.block_shape)
747
+ ic2_hidden_size = intermediate_cache2.size(-1)
748
+ intermediate_cache2 = intermediate_cache2.view(-1, ic2_hidden_size)
749
+
750
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
751
+ A=intermediate_cache2,
752
+ A_scale=a2_scale,
753
+ qtype=torch.float8_e4m3fn if self.use_fp8_w8a8 else None,
754
+ per_channel_quant=self.per_channel_quant,
755
+ block_shape=self.block_shape)
756
742
- invoke_moe_batched_triton_kernel(A=intermediate_cache2,
757
+ qintermediate_cache2 = qintermediate_cache2.view(
758
+ (E, -1, ic2_hidden_size))
759
+
760
+ invoke_moe_batched_triton_kernel(A=qintermediate_cache2,
761
B=w2,
762
C=intermediate_cache3,
763
expert_num_tokens=expert_num_tokens,
@@ -752,5 +770,4 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
770
use_int4_w4a16=self.use_int4_w4a16,
771
config=config,
772
block_shape=self.block_shape)
755
-
773
return intermediate_cache3
vllm/model_executor/layers/fused_moe/fused_moe.py
+1
-1
@@ -1164,7 +1164,7 @@ def fused_experts(hidden_states: torch.Tensor,
1164
# permute/unpermute ops are available.
1165
N = w1.shape[1]
1166
if (allow_deep_gemm and use_fp8_w8a8 and N > 512
1167
- and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):
1167
+ and _valid_deep_gemm(hidden_states, w1, w2)):
1168
assert apply_router_weight_on_input is False
1169
return deep_gemm_moe_fp8(
1170
hidden_states=hidden_states,
vllm/model_executor/layers/fused_moe/layer.py
+120
-28
@@ -5,7 +5,7 @@ import importlib
5
from abc import abstractmethod
6
from dataclasses import dataclass
7
from enum import Enum
8
-from typing import Callable, Optional
8
+from typing import Callable, Optional, Union
9
10
import torch
11
import torch.nn.functional as F
@@ -30,16 +30,19 @@ from vllm.platforms.interface import CpuArchEnum
30
from vllm.utils import direct_register_custom_op
31
32
has_pplx = importlib.util.find_spec("pplx_kernels") is not None
33
+has_deepep = importlib.util.find_spec("deep_ep") is not None
34
35
if current_platform.is_cuda_alike():
35
- from .fused_batched_moe import (BatchedPrepareAndFinalize,
36
- BatchedTritonExperts)
36
+ from .fused_batched_moe import BatchedTritonExperts
37
from .fused_moe import TritonExperts, fused_experts
38
from .modular_kernel import (FusedMoEModularKernel,
39
FusedMoEPermuteExpertsUnpermute,
40
FusedMoEPrepareAndFinalize)
41
if has_pplx:
42
from .pplx_prepare_finalize import PplxPrepareAndFinalize
43
+ if has_deepep:
44
+ from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
45
+ from .deepep_ll_prepare_finalize import DeepEPLLPrepareAndFinalize
46
else:
47
fused_experts = None # type: ignore
48
FusedMoEPermuteExpertsUnpermute = None # type: ignore
@@ -71,10 +74,24 @@ class FusedMoEParallelConfig:
74
75
use_ep: bool # whether to use EP or not
76
77
+ @property
78
+ def use_all2all_kernels(self):
79
+ return self.dp_size > 1 and self.use_ep
80
+
81
@property
82
def use_pplx_kernels(self):
76
- return self.dp_size > 1 and self.use_ep and \
77
- envs.VLLM_ALL2ALL_BACKEND == "pplx"
83
+ return (self.use_all2all_kernels
84
+ and envs.VLLM_ALL2ALL_BACKEND == "pplx")
85
+
86
+ @property
87
+ def use_deepep_ht_kernels(self):
88
+ return (self.use_all2all_kernels
89
+ and envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput")
90
+
91
+ @property
92
+ def use_deepep_ll_kernels(self):
93
+ return (self.use_all2all_kernels
94
+ and envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency")
95
96
@staticmethod
97
def make(tp_size_: int, dp_size_: int,
@@ -231,6 +248,14 @@ class MoEConfig:
248
def use_pplx_kernels(self):
249
return self.moe_parallel_config.use_pplx_kernels
250
251
+ @property
252
+ def use_deepep_ht_kernels(self):
253
+ return self.moe_parallel_config.use_deepep_ht_kernels
254
+
255
+ @property
256
+ def use_deepep_ll_kernels(self):
257
+ return self.moe_parallel_config.use_deepep_ll_kernels
258
+
259
260
class FusedMoeWeightScaleSupported(Enum):
261
TENSOR = "tensor"
@@ -252,7 +277,16 @@ class FusedMoEMethodBase(QuantizeMethodBase):
277
all2all_manager = get_ep_group().device_communicator.all2all_manager
278
assert all2all_manager is not None
279
255
- prepare_finalize = None
280
+ quant_dtype = None
281
+ act_quant_block_size = None
282
+ from vllm.model_executor.layers.quantization.fp8 import Fp8Config
283
+ if isinstance(quant_config, Fp8Config):
284
+ act_quant_block_size = quant_config.weight_block_size
285
+ quant_dtype = torch.float8_e4m3fn
286
+
287
+ prepare_finalize: Optional[Union[PplxPrepareAndFinalize,
288
+ DeepEPHTPrepareAndFinalize,
289
+ DeepEPLLPrepareAndFinalize]] = None
290
if moe.use_pplx_kernels:
291
all_to_all_args = dict(
292
max_num_tokens=moe.max_num_tokens,
@@ -288,8 +322,49 @@ class FusedMoEMethodBase(QuantizeMethodBase):
322
dp_size=all2all_manager.tp_group.world_size,
323
quant_dtype=moe.in_dtype,
324
)
325
+ elif moe.use_deepep_ht_kernels:
326
+ assert moe.dp_size == all2all_manager.dp_world_size
327
+
328
+ all_to_all_args = dict()
329
+ handle = all2all_manager.get_handle(all_to_all_args)
330
+ prepare_finalize = DeepEPHTPrepareAndFinalize(
331
+ handle,
332
+ world_size=all2all_manager.world_size,
333
+ rank=all2all_manager.rank,
334
+ dp_size=all2all_manager.dp_world_size,
335
+ rank_expert_offset=all2all_manager.rank *
336
+ moe.num_local_experts,
337
+ quant_dtype=quant_dtype,
338
+ block_shape=act_quant_block_size,
339
+ )
340
+
341
+ elif moe.use_deepep_ll_kernels:
342
+ assert moe.dp_size == all2all_manager.dp_world_size
343
344
+ all_to_all_args = dict(
345
+ max_num_tokens_per_dp_rank=moe.max_num_tokens,
346
+ token_hidden_size=moe.hidden_dim,
347
+ num_ep_ranks=all2all_manager.world_size,
348
+ num_global_experts=moe.num_experts,
349
+ num_local_experts=moe.num_experts //
350
+ all2all_manager.world_size)
351
+ handle = all2all_manager.get_handle(all_to_all_args)
352
+
353
+ # Note (varun): Whether to use FP8 dispatch or not needs some
354
+ # profiling. Turning it off for now.
355
+ prepare_finalize = DeepEPLLPrepareAndFinalize(
356
+ handle,
357
+ world_size=all2all_manager.world_size,
358
+ dp_size=all2all_manager.dp_world_size,
359
+ max_tokens_per_rank=moe.max_num_tokens,
360
+ quant_dtype=quant_dtype,
361
+ block_shape=act_quant_block_size,
362
+ use_fp8_dispatch=False,
363
+ )
364
+
365
+ self.topk_indices_dtype = None
366
if prepare_finalize is not None:
367
+ self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
368
experts = self.select_gemm_impl(prepare_finalize)
369
self.fused_experts = FusedMoEModularKernel(
370
prepare_finalize,
@@ -297,7 +372,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
372
)
373
374
def select_gemm_impl(
300
- self, prepare_finalize: Optional[FusedMoEPrepareAndFinalize]
375
+ self, prepare_finalize: FusedMoEPrepareAndFinalize
376
) -> FusedMoEPermuteExpertsUnpermute:
377
# based on the all2all implementation, select the appropriate
378
# gemm implementation
@@ -334,6 +409,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
409
def __init__(self, moe: MoEConfig):
410
super().__init__()
411
self.fused_experts = fused_experts # type: ignore
412
+ self.topk_indices_dtype = None
413
self.moe = moe
414
415
self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
@@ -343,8 +419,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
419
else:
420
self.rocm_aiter_fused_experts = None # type: ignore
421
346
- def select_gemm_impl(
347
- self, prepare_finalize: Optional[FusedMoEPrepareAndFinalize]):
422
+ def select_gemm_impl(self, prepare_finalize: FusedMoEPrepareAndFinalize):
423
424
assert self.fused_experts == fused_experts
425
@@ -353,11 +428,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
428
429
experts: Optional[FusedMoEPermuteExpertsUnpermute] = None
430
356
- if isinstance(prepare_finalize,
357
- (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):
431
+ use_batched_experts = prepare_finalize.max_num_tokens_per_rank(
432
+ ) is not None
433
+ if use_batched_experts:
434
logger.debug("BatchedTritonExperts %s", self.moe)
435
+ assert self.moe.dp_size == all2all_manager.dp_world_size
436
experts = BatchedTritonExperts(
360
- max_num_tokens=MOE_DP_CHUNK_SIZE,
437
+ max_num_tokens=self.moe.max_num_tokens,
438
world_size=all2all_manager.world_size,
439
# dp_size actually means tp_size, bug in pplx kernels
440
dp_size=all2all_manager.tp_group.world_size,
@@ -366,6 +443,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
443
use_int8_w8a16=False,
444
use_int4_w4a16=False,
445
block_shape=None,
446
+ per_channel_quant=False,
447
)
448
else:
449
logger.debug("TritonExperts %s", self.moe)
@@ -494,6 +572,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
572
apply_router_weight_on_input: bool = False,
573
activation: str = "silu",
574
) -> torch.Tensor:
575
+
576
topk_weights, topk_ids = FusedMoE.select_experts(
577
hidden_states=x,
578
router_logits=router_logits,
@@ -505,7 +584,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
584
custom_routing_function=custom_routing_function,
585
scoring_func=scoring_func,
586
e_score_correction_bias=e_score_correction_bias,
508
- indices_type=torch.uint32 if self.moe.use_pplx_kernels else None)
587
+ indices_type=self.topk_indices_dtype)
588
589
if self.rocm_aiter_moe_enabled:
590
assert expert_map is None
@@ -806,11 +885,8 @@ class FusedMoE(torch.nn.Module):
885
# Note: get_quant_method will look at the layer's local_num_experts
886
# for heuristic purposes, so it must be initialized first.
887
quant_method: Optional[QuantizeMethodBase] = None
809
-
810
- if quant_config is None:
811
- quant_method = UnquantizedFusedMoEMethod(moe)
812
- else:
813
- quant_method = quant_config.get_quant_method(self, prefix)
888
+ quant_method = (UnquantizedFusedMoEMethod(moe) if quant_config is None
889
+ else quant_config.get_quant_method(self, prefix))
890
891
assert quant_method is not None
892
assert isinstance(quant_method, FusedMoEMethodBase)
@@ -836,7 +912,8 @@ class FusedMoE(torch.nn.Module):
912
# Chunked all2all staging tensor
913
self.batched_hidden_states: Optional[torch.Tensor] = None
914
self.batched_router_logits: Optional[torch.Tensor] = None
839
- if self.moe_parallel_config.use_pplx_kernels:
915
+ if (self.moe_parallel_config.use_pplx_kernels
916
+ or self.moe_parallel_config.use_deepep_ll_kernels):
917
act_dtype = vllm_config.model_config.dtype
918
self.batched_hidden_states = torch.zeros(
919
(MOE_DP_CHUNK_SIZE, self.hidden_size),
@@ -880,6 +957,14 @@ class FusedMoE(torch.nn.Module):
957
def use_pplx_kernels(self):
958
return self.moe_parallel_config.use_pplx_kernels
959
960
+ @property
961
+ def use_deepep_ht_kernels(self):
962
+ return self.moe_parallel_config.use_deepep_ht_kernels
963
+
964
+ @property
965
+ def use_deepep_ll_kernels(self):
966
+ return self.moe_parallel_config.use_deepep_ll_kernels
967
+
968
def _load_per_tensor_weight_scale(self, shard_id: str,
969
param: torch.nn.Parameter,
970
loaded_weight: torch.Tensor,
@@ -1210,19 +1295,21 @@ class FusedMoE(torch.nn.Module):
1295
When just tensor-parallel is used, it is not required to reduce
1296
the shared_experts results immediately. Instead we reduce at the
1297
once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
1213
- With EP and the pplx kernels - this is no longer viable as all
1298
+ With EP and all2all kernels - this is no longer viable as all
1299
GPU ranks in DP, produce the complete set of hidden_states.
1300
Therefore it is required that we reduce the shared_experts output
1301
early.
1302
"""
1218
- return self.use_pplx_kernels
1303
+ return (self.use_pplx_kernels or self.use_deepep_ht_kernels
1304
+ or self.use_deepep_ll_kernels)
1305
1306
def maybe_all_reduce_tensor_model_parallel(
1307
self, final_hidden_states: torch.Tensor):
1308
"""
1309
The pplx combine kernel reduces across GPU ranks by default.
1310
"""
1225
- if self.use_pplx_kernels:
1311
+ if (self.use_pplx_kernels or self.use_deepep_ht_kernels
1312
+ or self.use_deepep_ll_kernels):
1313
return final_hidden_states
1314
else:
1315
return tensor_model_parallel_all_reduce(final_hidden_states)
@@ -1289,7 +1376,7 @@ class FusedMoE(torch.nn.Module):
1376
1377
ctx = get_forward_context()
1378
max_tokens_across_dp = ctx.dp_metadata.max_tokens_across_dp_cpu
1292
- moe_dp_chunk_size_per_rank = MOE_DP_CHUNK_SIZE
1379
+ moe_dp_chunk_size_per_rank = self.moe_config.max_num_tokens
1380
1381
num_tokens = full_hidden_states.size(0)
1382
for chunk_start_ in range(0, max_tokens_across_dp,
@@ -1310,12 +1397,17 @@ class FusedMoE(torch.nn.Module):
1397
def forward_impl(self, hidden_states: torch.Tensor,
1398
router_logits: torch.Tensor):
1399
assert self.quant_method is not None
1313
- if self.moe_parallel_config.use_pplx_kernels:
1400
+ if (self.moe_parallel_config.use_pplx_kernels
1401
+ or self.moe_parallel_config.use_deepep_ll_kernels):
1402
return self.forward_impl_chunked(hidden_states, router_logits)
1403
1316
- if self.dp_size > 1:
1404
+ do_naive_dispatch_combine: bool = (
1405
+ self.dp_size > 1
1406
+ and not self.moe_parallel_config.use_deepep_ht_kernels)
1407
+ if do_naive_dispatch_combine:
1408
hidden_states, router_logits = get_ep_group().dispatch(
1409
hidden_states, router_logits)
1410
+
1411
# Matrix multiply.
1412
final_hidden_states = self.quant_method.apply(
1413
layer=self,
@@ -1335,12 +1427,12 @@ class FusedMoE(torch.nn.Module):
1427
apply_router_weight_on_input=self.apply_router_weight_on_input,
1428
)
1429
1338
- if self.dp_size > 1:
1430
+ if do_naive_dispatch_combine:
1431
final_hidden_states = get_ep_group().combine(final_hidden_states)
1432
1433
if self.reduce_results and (self.tp_size > 1 or self.ep_size > 1):
1342
- # Default set to False. (May have to add shared expert outputs.)
1343
- final_hidden_states = tensor_model_parallel_all_reduce(
1434
+ # Default set to False. (May have to add shared expert outputs.
1435
+ final_hidden_states = self.maybe_all_reduce_tensor_model_parallel(
1436
final_hidden_states)
1437
1438
return final_hidden_states
vllm/model_executor/layers/fused_moe/modular_kernel.py
+121
-41
@@ -94,7 +94,8 @@ class FusedMoEPrepareAndFinalize(ABC):
94
num_experts: int,
95
expert_map: Optional[torch.Tensor],
96
apply_router_weight_on_input: bool,
97
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
97
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
98
+ Optional[torch.Tensor], Optional[torch.Tensor]]:
99
"""
100
Perform any quantization (and/or) dispatching needed
101
for this kernel.
@@ -113,6 +114,10 @@ class FusedMoEPrepareAndFinalize(ABC):
114
Returns a tuple of:
115
- quantized + dispatched a.
116
- quantized + dispatched a1_scales.
117
+ - Optional tensor as big as number of local experts that contains the
118
+ number of tokens assigned to each local expert.
119
+ - Optional dispatched expert topk IDs
120
+ - Optional dispatched expert topk weight
121
"""
122
raise NotImplementedError
123
@@ -138,6 +143,27 @@ class FusedMoEPrepareAndFinalize(ABC):
143
"""
144
raise NotImplementedError
145
146
+ @abstractmethod
147
+ def topk_indices_dtype(self) -> Optional[torch.dtype]:
148
+ """
149
+ The PrepareFinalize All2All implementations generally constrain the
150
+ dtype of the topk_ids they support. This function returns the
151
+ required topk indices dtype so it can be respected.
152
+ Return None if there are no such restrictions.
153
+ """
154
+ raise NotImplementedError
155
+
156
+ @abstractmethod
157
+ def max_num_tokens_per_rank(self) -> Optional[int]:
158
+ """
159
+ Some PrepareFinalize All2All implementations are batched. Meaning,
160
+ they can processes only as set of tokens at a time. This
161
+ function returns the batch size i.e the maximum number of tokens
162
+ the implementation can process at a time.
163
+ Return None if there are no such restrictions.
164
+ """
165
+ raise NotImplementedError
166
+
167
168
class FusedMoEPermuteExpertsUnpermute(ABC):
169
"""
@@ -261,6 +287,61 @@ class FusedMoEModularKernel(torch.nn.Module):
287
self.prepare_finalize = prepare_finalize
288
self.fused_experts = fused_experts
289
290
+ def _do_fused_experts(
291
+ self,
292
+ a1: torch.Tensor, # input to forward fn
293
+ a1q: torch.Tensor, # output of prepare fn
294
+ w1: torch.Tensor,
295
+ w2: torch.Tensor,
296
+ topk_ids: torch.Tensor,
297
+ expert_num_tokens: torch.Tensor,
298
+ activation: str,
299
+ global_num_experts: int,
300
+ expert_map: Optional[torch.Tensor],
301
+ w1_scale: Optional[torch.Tensor],
302
+ w2_scale: Optional[torch.Tensor],
303
+ w1_zp: Optional[torch.Tensor],
304
+ w2_zp: Optional[torch.Tensor],
305
+ a1q_scale: Optional[torch.Tensor],
306
+ a2_scale: Optional[torch.Tensor]) -> torch.Tensor:
307
+
308
+ _, M, N, K, top_k = _moe_problem_size(a1q, w1, w2, topk_ids)
309
+
310
+ # Use a1 here to decipher the correct workspace datatype
311
+ workspace13_shape, workspace2_shape, workspace_dtype = (
312
+ self.fused_experts.workspace_shapes(a1, M, N, K, top_k,
313
+ global_num_experts))
314
+
315
+ # We can reuse the memory between cache1 and cache3 because by the time
316
+ # we need cache3, we're done with cache1
317
+ workspace13 = torch.zeros(workspace13_shape,
318
+ device=a1.device,
319
+ dtype=workspace_dtype)
320
+ workspace2 = torch.zeros(workspace2_shape,
321
+ device=a1.device,
322
+ dtype=workspace_dtype)
323
+
324
+ fused_out = self.fused_experts.apply(
325
+ a1q,
326
+ w1,
327
+ w2,
328
+ topk_ids,
329
+ activation=activation,
330
+ global_num_experts=global_num_experts,
331
+ expert_map=expert_map,
332
+ w1_scale=w1_scale,
333
+ w2_scale=w2_scale,
334
+ w1_zp=w1_zp,
335
+ w2_zp=w2_zp,
336
+ a1q_scale=a1q_scale,
337
+ a2_scale=a2_scale,
338
+ workspace13=workspace13,
339
+ workspace2=workspace2,
340
+ expert_num_tokens=expert_num_tokens,
341
+ )
342
+
343
+ return fused_out
344
+
345
def forward(
346
self,
347
hidden_states: torch.Tensor,
@@ -315,49 +396,48 @@ class FusedMoEModularKernel(torch.nn.Module):
396
Returns:
397
- torch.Tensor: The output tensor after applying the MoE layer.
398
"""
318
- a1 = hidden_states
319
- E, M, N, K, top_k = _moe_problem_size(a1, w1, w2, topk_ids)
320
-
321
- if global_num_experts == -1:
322
- global_num_experts = E
399
400
+ a1 = hidden_states
401
output = a1 if inplace else torch.zeros_like(a1)
402
326
- workspace13_shape, workspace2_shape, workspace_dtype = (
327
- self.fused_experts.workspace_shapes(a1, M, N, K, top_k,
328
- global_num_experts))
329
-
330
- # We can reuse the memory between cache1 and cache3 because by the time
331
- # we need cache3, we're done with cache1
332
- workspace13 = torch.zeros(workspace13_shape,
333
- device=a1.device,
334
- dtype=workspace_dtype)
335
- workspace2 = torch.zeros(workspace2_shape,
336
- device=a1.device,
337
- dtype=workspace_dtype)
338
-
339
- a1q, a1q_scale, expert_num_tokens = self.prepare_finalize.prepare(
340
- a1, a1_scale, a2_scale, topk_weights, topk_ids, global_num_experts,
341
- expert_map, apply_router_weight_on_input)
342
-
343
- fused_out = self.fused_experts.apply(
344
- a1q,
345
- w1,
346
- w2,
347
- topk_ids,
348
- activation=activation,
349
- global_num_experts=global_num_experts,
350
- expert_map=expert_map,
351
- w1_scale=w1_scale,
352
- w2_scale=w2_scale,
353
- w1_zp=w1_zp,
354
- w2_zp=w2_zp,
355
- a1q_scale=a1q_scale,
356
- a2_scale=a2_scale,
357
- workspace13=workspace13,
358
- workspace2=workspace2,
359
- expert_num_tokens=expert_num_tokens,
360
- )
403
+ if global_num_experts == -1:
404
+ global_num_experts = w1.size(0)
405
+
406
+ (a1q, a1q_scale, expert_num_tokens, _expert_topk_ids,
407
+ _expert_topk_weights) = self.prepare_finalize.prepare(
408
+ a1, a1_scale, a2_scale, topk_weights, topk_ids,
409
+ global_num_experts, expert_map, apply_router_weight_on_input)
410
+ # Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
411
+ topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
412
+ topk_weights = (topk_weights if _expert_topk_weights is None else
413
+ _expert_topk_weights)
414
+
415
+ fused_out = None
416
+ if a1q.numel() == 0:
417
+ # This happens when none of the tokens from the all2all reach this
418
+ # EP rank. Also, note that this is only relevant for CUDAGraph
419
+ # incompatible all2all kernels like the DeepEP high-throughput
420
+ # kernels. CUDAGraph compatible all2all kernels like the pplx
421
+ # kernels and the DeepEP low-latency kernels are always batched
422
+ # and can never run into the tensor.numel() == 0 case.
423
+ fused_out = torch.empty_like(a1q).to(dtype=a1.dtype)
424
+ else:
425
+ fused_out = self._do_fused_experts(
426
+ a1=a1,
427
+ a1q=a1q,
428
+ w1=w1,
429
+ w2=w2,
430
+ topk_ids=topk_ids,
431
+ expert_num_tokens=expert_num_tokens,
432
+ activation=activation,
433
+ global_num_experts=global_num_experts,
434
+ expert_map=expert_map,
435
+ w1_scale=w1_scale,
436
+ w2_scale=w2_scale,
437
+ w1_zp=w1_zp,
438
+ w2_zp=w2_zp,
439
+ a1q_scale=a1q_scale,
440
+ a2_scale=a2_scale)
441
442
self.prepare_finalize.finalize(output, fused_out, topk_weights,
443
topk_ids, apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+3
-2
@@ -25,7 +25,7 @@ def _moe_permute(
25
"""
26
top_k_num = curr_topk_ids.size(1)
27
28
- tokens_in_chunk = curr_hidden_states.sizze(0)
28
+ tokens_in_chunk = curr_hidden_states.size(0)
29
30
sorted_token_ids, expert_ids, num_tokens_post_padded = (
31
moe_align_block_size(curr_topk_ids,
@@ -37,11 +37,12 @@ def _moe_permute(
37
inv_perm: Optional[torch.Tensor] = None
38
39
num_tokens = top_k_num * tokens_in_chunk
40
- sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
40
expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
41
inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
42
43
# Permute according to sorted token ids.
44
+ sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
45
+
46
curr_hidden_states = _fp8_perm(curr_hidden_states,
47
sorted_token_ids // top_k_num)
48
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+9
-2
@@ -32,6 +32,12 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
32
self.dp_size = dp_size
33
self.quant_dtype = quant_dtype
34
35
+ def max_num_tokens_per_rank(self) -> Optional[int]:
36
+ return self.max_num_tokens
37
+
38
+ def topk_indices_dtype(self) -> Optional[torch.dtype]:
39
+ return torch.uint32
40
+
41
def prepare(
42
self,
43
a1: torch.Tensor,
@@ -42,7 +48,8 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
48
num_experts: int,
49
expert_map: Optional[torch.Tensor],
50
apply_router_weight_on_input: bool,
45
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
51
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
52
+ Optional[torch.Tensor], Optional[torch.Tensor]]:
53
num_tokens = a1.size(0) # M
54
hidden_dim = a1.size(-1) # K
55
@@ -115,7 +122,7 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
122
bound_m=bound_m,
123
)
124
118
- return expert_x, expert_x_scale, expert_num_tokens
125
+ return expert_x, expert_x_scale, expert_num_tokens, None, None
126
127
def finalize(
128
self,
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+10
-2
@@ -24,6 +24,12 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
24
self.block_shape = block_shape
25
self.quant_dtype = quant_dtype
26
27
+ def max_num_tokens_per_rank(self) -> Optional[int]:
28
+ return None
29
+
30
+ def topk_indices_dtype(self) -> Optional[torch.dtype]:
31
+ return None
32
+
33
def prepare(
34
self,
35
a1: torch.Tensor,
@@ -34,7 +40,9 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
40
num_experts: int,
41
expert_map: Optional[torch.Tensor],
42
apply_router_weight_on_input: bool = False,
37
- ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
43
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
44
+ Optional[torch.Tensor], Optional[torch.Tensor]]:
45
+
46
if apply_router_weight_on_input:
47
topk = topk_ids.size(1)
48
# TODO: this only works for topK=1, will need to update for topK>1
@@ -47,7 +55,7 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
55
self.per_channel_quant,
56
self.block_shape)
57
50
- return a1q, a1q_scale, None
58
+ return a1q, a1q_scale, None, None, None
59
60
def finalize(
61
self,
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+5
-2
@@ -29,9 +29,10 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
29
per_channel_quant=per_channel_quant,
30
block_shape=block_shape,
31
block_m=block_m)
32
- self.deep_gemm_expert = DeepGemmExperts()
32
self.allow_deep_gemm = allow_deep_gemm
33
self.use_fp8_w8a8 = use_fp8_w8a8
34
+ self.deep_gemm_expert = DeepGemmExperts(
35
+ ) if self.allow_deep_gemm else None
36
37
def workspace_shapes(
38
self,
@@ -46,6 +47,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
47
# workspaces so we can be pessimistic here and allocate for DeepGemm
48
# even if we fall back to triton later, e.g. if expert maps are set.
49
if self.allow_deep_gemm and _valid_deep_gemm_shape(M, N, K):
50
+ assert self.deep_gemm_expert is not None
51
return self.deep_gemm_expert.workspace_shapes(
52
a, M, N, K, topk, num_experts)
53
else:
@@ -73,7 +75,8 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
75
) -> torch.Tensor:
76
N = w1.size(1)
77
if (self.allow_deep_gemm and self.use_fp8_w8a8 and N > 512
76
- and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):
78
+ and _valid_deep_gemm(hidden_states, w1, w2)):
79
+ assert self.deep_gemm_expert is not None
80
return self.deep_gemm_expert.apply(
81
hidden_states,
82
w1,
vllm/model_executor/layers/fused_moe/utils.py
+2
-2
@@ -18,8 +18,8 @@ def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
18
Shrink the given tensor and apply the given view to it. This is
19
used to resize the intermediate fused_moe caches.
20
"""
21
- assert prod(
22
- v) <= x.numel(), f"{prod(v)} <= {x.numel()}" # CUDAGRAPH unfriendly?
21
+ assert prod(v) <= x.numel(
22
+ ), f"{v} ({prod(v)}) <= {x.shape} ({x.numel()})" # CUDAGRAPH unfriendly?
23
return x.flatten()[:prod(v)].view(*v)
24
25
vllm/model_executor/layers/quantization/fp8.py
+34
-7
@@ -3,7 +3,7 @@
3
4
import functools
5
import importlib.util
6
-from typing import Any, Callable, Optional
6
+from typing import Any, Callable, Optional, Union
7
8
import torch
9
import torch.nn.functional as F
@@ -452,6 +452,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
452
if envs.VLLM_USE_DEEP_GEMM:
453
if not has_deep_gemm:
454
logger.warning_once("Failed to import DeepGemm kernels.")
455
+ elif not self.block_quant:
456
+ logger.warning_once("Model is not block quantized. Not using "
457
+ " DeepGemm kernels")
458
elif (current_platform.is_cuda()
459
and current_platform.has_device_capability(90)):
460
logger.info_once("Using DeepGemm kernels for Fp8MoEMethod.")
@@ -460,8 +463,10 @@ class Fp8MoEMethod(FusedMoEMethodBase):
463
logger.warning_once(
464
"DeepGemm not supported on the current platform.")
465
466
+ self.topk_indices_dtype = None
467
self.fused_experts = functools.partial( # type: ignore
468
fused_experts,
469
+ use_fp8_w8a8=True,
470
block_shape=self.quant_config.weight_block_size,
471
allow_deep_gemm=self.allow_deep_gemm)
472
@@ -765,18 +770,39 @@ class Fp8MoEMethod(FusedMoEMethodBase):
770
del layer.w2_input_scale
771
772
def select_gemm_impl(self, prepare_finalize):
773
+
774
+ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
775
+ BatchedTritonExperts)
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
774
- experts = TritonOrDeepGemmExperts(
775
- use_fp8_w8a8=True,
776
- block_shape=self.quant_config.weight_block_size,
777
- allow_deep_gemm=self.allow_deep_gemm,
778
- )
782
+ experts: Optional[Union[BatchedTritonExperts,
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(
789
+ max_num_tokens=max_num_tokens_per_rank,
790
+ world_size=prepare_finalize.world_size,
791
+ dp_size=prepare_finalize.dp_size,
792
+ use_fp8_w8a8=True,
793
+ use_int8_w8a8=False,
794
+ use_int8_w8a16=False,
795
+ use_int4_w4a16=False,
796
+ block_shape=None,
797
+ )
798
+ else:
799
+ experts = TritonOrDeepGemmExperts(
800
+ use_fp8_w8a8=True,
801
+ block_shape=self.quant_config.weight_block_size,
802
+ allow_deep_gemm=self.allow_deep_gemm,
803
+ )
804
+
805
+ assert experts is not None
806
return experts
807
808
def apply(
@@ -797,6 +823,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
823
apply_router_weight_on_input: bool = False,
824
activation: str = "silu",
825
) -> torch.Tensor:
826
+
827
topk_weights, topk_ids = FusedMoE.select_experts(
828
hidden_states=x,
829
router_logits=router_logits,
@@ -808,6 +835,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
835
custom_routing_function=custom_routing_function,
836
scoring_func=scoring_func,
837
e_score_correction_bias=e_score_correction_bias,
838
+ indices_type=self.topk_indices_dtype,
839
)
840
841
if self.rocm_aiter_moe_enabled:
@@ -855,7 +883,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
883
topk_ids=topk_ids,
884
inplace=True,
885
activation=activation,
858
- use_fp8_w8a8=True,
886
global_num_experts=global_num_experts,
887
apply_router_weight_on_input=apply_router_weight_on_input,
888
expert_map=expert_map,
vllm/platforms/cuda.py
+15
@@ -154,6 +154,21 @@ class CudaPlatformBase(Platform):
154
logger.info(
155
"Forcing kv cache block size to 64 for FlashMLA backend.")
156
157
+ if (envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput"
158
+ and parallel_config.data_parallel_size > 1
159
+ and vllm_config.compilation_config.use_cudagraph):
160
+ logger.info(
161
+ "Data Parallel: Forcing enforce eager to be True since DP "
162
+ "with DeepEP high-throughput kernels are not CUDA Graph "
163
+ "compatible. The DeepEP low-latency kernels are CUDA Graph "
164
+ "compatible. Set the all_to_all backend to deepep_low_latency "
165
+ "to use those kernels instead.")
166
+ vllm_config.compilation_config.use_cudagraph = False
167
+ vllm_config.model_config.enforce_eager = True
168
+ # TODO (varun): Turning this ON gives incorrect results for the
169
+ # Deepseek-V2-lite model.
170
+ vllm_config.compilation_config.use_inductor = False
171
+
172
@classmethod
173
def get_current_memory_usage(cls,
174
device: Optional[torch.types.Device] = None