A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 03:30 UTC
[Kernels] MoE refactor (#19636)
Signed-off-by: Bill Nell <bnell@redhat.com> Signed-off-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: ElizaWszola <ewszola@redhat.com>
bnellnm committed
Jul 2, 2025 at 09:08 UTC
c1909e7e8ccd2037e76536a8e726120c85d3754e
36 files changed
+2713
-1599
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py
+8
-3
@@ -113,6 +113,7 @@ def bench_run(
113
w2_scale: torch.Tensor,
114
topk_weights: torch.Tensor,
115
topk_ids: torch.Tensor,
116
+ per_act_token: bool,
117
num_repeats: int,
118
):
119
for _ in range(num_repeats):
@@ -124,7 +125,8 @@ def bench_run(
125
topk_ids,
126
w1_scale,
127
w2_scale,
127
- a1_scale=a_scale,
128
+ per_act_token,
129
+ a1_scale=None,
130
)
131
132
def run_cutlass_from_graph(
@@ -148,7 +150,8 @@ def bench_run(
150
topk_ids,
151
w1_scale,
152
w2_scale,
151
- a1_scale=a_scale,
153
+ per_act_token,
154
+ a1_scale=None,
155
)
156
157
def run_triton_from_graph(
@@ -227,6 +230,7 @@ def bench_run(
230
"w2_q": w2_q,
231
"w1_scale": w1_scale,
232
"w2_scale": w2_scale,
233
+ "per_act_token": per_act_token,
234
# cuda graph params
235
"cutlass_graph": cutlass_graph,
236
"triton_graph": triton_graph,
@@ -287,12 +291,13 @@ def bench_run(
291
w2_scale,
292
topk_weights,
293
topk_ids,
294
+ per_act_token,
295
num_warmup,
296
)
297
298
results.append(
299
benchmark.Timer(
295
- stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, num_runs)", # noqa: E501
300
+ stmt="run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, per_act_token, num_runs)", # noqa: E501
301
globals=globals,
302
label=label,
303
sub_label=sub_label,
tests/kernels/moe/parallel_utils.py
new
+190
@@ -0,0 +1,190 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""
3
+DeepEP test utilities
4
+"""
5
+import dataclasses
6
+import importlib
7
+import os
8
+import traceback
9
+from typing import Callable, Optional
10
+
11
+import torch
12
+from torch.distributed import ProcessGroup
13
+from torch.multiprocessing import (
14
+ spawn) # pyright: ignore[reportPrivateImportUsage]
15
+from typing_extensions import Concatenate, ParamSpec
16
+
17
+from vllm.utils import get_open_port
18
+
19
+has_deep_ep = importlib.util.find_spec("deep_ep") is not None
20
+if has_deep_ep:
21
+ from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
22
+ DeepEPHTPrepareAndFinalize)
23
+ from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
24
+ DeepEPLLPrepareAndFinalize)
25
+
26
+## Parallel Processes Utils
27
+
28
+P = ParamSpec("P")
29
+
30
+
31
+@dataclasses.dataclass
32
+class ProcessGroupInfo:
33
+ world_size: int
34
+ world_local_size: int
35
+ rank: int
36
+ node_rank: int
37
+ local_rank: int
38
+ device: torch.device
39
+
40
+
41
+def _worker_parallel_launch(
42
+ local_rank: int,
43
+ world_size: int,
44
+ world_local_size: int,
45
+ node_rank: int,
46
+ init_method: str,
47
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
48
+ *args: P.args,
49
+ **kwargs: P.kwargs,
50
+) -> None:
51
+ rank = node_rank * world_local_size + local_rank
52
+ torch.cuda.set_device(local_rank)
53
+ device = torch.device("cuda", local_rank)
54
+ torch.distributed.init_process_group(
55
+ backend="cpu:gloo,cuda:nccl",
56
+ init_method=init_method,
57
+ rank=rank,
58
+ world_size=world_size,
59
+ device_id=device,
60
+ )
61
+ barrier = torch.tensor([rank], device=device)
62
+ torch.distributed.all_reduce(barrier)
63
+
64
+ try:
65
+ worker(
66
+ ProcessGroupInfo(
67
+ world_size=world_size,
68
+ world_local_size=world_local_size,
69
+ rank=rank,
70
+ node_rank=node_rank,
71
+ local_rank=local_rank,
72
+ device=device,
73
+ ),
74
+ *args,
75
+ **kwargs,
76
+ )
77
+ except Exception as ex:
78
+ print(ex)
79
+ traceback.print_exc()
80
+ raise
81
+ finally:
82
+ torch.distributed.destroy_process_group()
83
+
84
+
85
+def parallel_launch(
86
+ world_size: int,
87
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
88
+ *args: P.args,
89
+ **kwargs: P.kwargs,
90
+) -> None:
91
+ assert not kwargs
92
+ spawn(
93
+ _worker_parallel_launch,
94
+ args=(
95
+ world_size,
96
+ world_size,
97
+ 0,
98
+ f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
99
+ worker,
100
+ ) + args,
101
+ nprocs=world_size,
102
+ join=True,
103
+ )
104
+
105
+
106
+## DeepEP specific utils
107
+
108
+
109
+@dataclasses.dataclass
110
+class DeepEPHTArgs:
111
+ num_local_experts: int
112
+
113
+
114
+@dataclasses.dataclass
115
+class DeepEPLLArgs:
116
+ max_tokens_per_rank: int
117
+ hidden_size: int
118
+ num_experts: int
119
+ use_fp8_dispatch: bool
120
+
121
+
122
+def make_deepep_ht_a2a(pg: ProcessGroup,
123
+ pgi: ProcessGroupInfo,
124
+ dp_size: int,
125
+ ht_args: DeepEPHTArgs,
126
+ q_dtype: Optional[torch.dtype] = None,
127
+ block_shape: Optional[list[int]] = None):
128
+
129
+ import deep_ep
130
+
131
+ # high throughput a2a
132
+ num_nvl_bytes = 1024 * 1024 * 1024 # 1GB
133
+ num_rdma_bytes, low_latency_mode, num_qps_per_rank = 0, False, 1
134
+ buffer = deep_ep.Buffer(group=pg,
135
+ num_nvl_bytes=num_nvl_bytes,
136
+ num_rdma_bytes=num_rdma_bytes,
137
+ low_latency_mode=low_latency_mode,
138
+ num_qps_per_rank=num_qps_per_rank)
139
+ return DeepEPHTPrepareAndFinalize(buffer=buffer,
140
+ world_size=pgi.world_size,
141
+ rank=pgi.rank,
142
+ dp_size=dp_size,
143
+ rank_expert_offset=pgi.rank *
144
+ ht_args.num_local_experts)
145
+
146
+
147
+def make_deepep_ll_a2a(pg: ProcessGroup,
148
+ pgi: ProcessGroupInfo,
149
+ dp_size: int,
150
+ deepep_ll_args: DeepEPLLArgs,
151
+ q_dtype: Optional[torch.dtype] = None,
152
+ block_shape: Optional[list[int]] = None):
153
+
154
+ import deep_ep
155
+
156
+ # low-latency a2a
157
+ num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
158
+ deepep_ll_args.max_tokens_per_rank, deepep_ll_args.hidden_size,
159
+ pgi.world_size, deepep_ll_args.num_experts)
160
+
161
+ buffer = deep_ep.Buffer(group=pg,
162
+ num_rdma_bytes=num_rdma_bytes,
163
+ low_latency_mode=True,
164
+ num_qps_per_rank=deepep_ll_args.num_experts //
165
+ pgi.world_size)
166
+
167
+ return DeepEPLLPrepareAndFinalize(
168
+ buffer=buffer,
169
+ world_size=pgi.world_size,
170
+ dp_size=dp_size,
171
+ max_tokens_per_rank=deepep_ll_args.max_tokens_per_rank,
172
+ use_fp8_dispatch=deepep_ll_args.use_fp8_dispatch,
173
+ )
174
+
175
+
176
+def make_deepep_a2a(pg: ProcessGroup,
177
+ pgi: ProcessGroupInfo,
178
+ dp_size: int,
179
+ deepep_ht_args: Optional[DeepEPHTArgs],
180
+ deepep_ll_args: Optional[DeepEPLLArgs],
181
+ q_dtype: Optional[torch.dtype] = None,
182
+ block_shape: Optional[list[int]] = None):
183
+ if deepep_ht_args is not None:
184
+ assert deepep_ll_args is None
185
+ return make_deepep_ht_a2a(pg, pgi, dp_size, deepep_ht_args, q_dtype,
186
+ block_shape)
187
+
188
+ assert deepep_ll_args is not None
189
+ return make_deepep_ll_a2a(pg, pgi, dp_size, deepep_ll_args, q_dtype,
190
+ block_shape)
tests/kernels/moe/test_batched_moe.py
+218
-35
@@ -2,18 +2,59 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
+from typing import Optional
6
7
import pytest
8
import torch
9
import triton.language as tl
10
11
+from tests.kernels.moe.utils import (batched_moe,
12
+ make_quantized_test_activations,
13
+ make_test_weights, triton_moe)
14
+from tests.kernels.quant_utils import native_batched_masked_quant_matmul
15
+from tests.kernels.utils import torch_experts
16
+from vllm.config import VllmConfig, set_current_vllm_config
17
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
18
invoke_moe_batched_triton_kernel)
19
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
20
+from vllm.platforms import current_platform
21
+
22
+MNK_FACTORS = [
23
+ (1, 128, 128),
24
+ (1, 128, 2048),
25
+ (1, 512, 512),
26
+ (1, 1024, 128),
27
+ (1, 1024, 2048),
28
+ (32, 128, 128),
29
+ (32, 512, 512),
30
+ (32, 1024, 2048),
31
+ (45, 128, 128),
32
+ (45, 128, 2048),
33
+ (45, 512, 512),
34
+ (45, 1024, 128),
35
+ (45, 1024, 2048),
36
+ (64, 128, 128),
37
+ (64, 512, 512),
38
+ (64, 1024, 2048),
39
+ (222, 128, 128),
40
+ (222, 128, 2048),
41
+ (222, 512, 512),
42
+ (222, 1024, 128),
43
+ (222, 1024, 2048),
44
+]
45
+NUM_EXPERTS = [8, 64]
46
+TOP_KS = [1, 2, 6]
47
+
48
+vllm_config = VllmConfig()
49
+vllm_config.scheduler_config.max_num_seqs = 128
50
+vllm_config.scheduler_config.max_model_len = 8192
51
52
53
@dataclass
54
class BatchedMMConfig:
16
- dtype: torch.dtype
55
+ in_dtype: torch.dtype
56
+ quant_dtype: Optional[torch.dtype]
57
+ out_dtype: torch.dtype
58
num_experts: int
59
max_tokens_per_expert: int
60
K: int
@@ -32,79 +73,127 @@ class BatchedMMTensors:
73
A = torch.randn(
74
(config.num_experts, config.max_tokens_per_expert, config.K),
75
device="cuda",
35
- dtype=config.dtype) / 10
76
+ dtype=config.in_dtype) / 10
77
B = torch.randn((config.num_experts, config.N, config.K),
78
device="cuda",
38
- dtype=config.dtype)
79
+ dtype=config.in_dtype)
80
C = torch.zeros(
81
(config.num_experts, config.max_tokens_per_expert, config.N),
82
device="cuda",
42
- dtype=config.dtype)
83
+ dtype=config.out_dtype)
84
+
85
num_expert_tokens = torch.randint(low=0,
86
high=config.max_tokens_per_expert,
87
size=(config.num_experts, ),
88
device="cuda",
89
dtype=torch.int32)
48
- return BatchedMMTensors(A, B, C, num_expert_tokens)
49
-
90
51
-def ref_impl(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
52
- num_expert_tokens: torch.Tensor) -> torch.Tensor:
53
-
54
- num_expert_tokens_cpu = num_expert_tokens.clone()
55
- num_expert_tokens_cpu = num_expert_tokens_cpu.to(device="cpu")
56
- num_experts = num_expert_tokens.size(0)
57
-
58
- for e in range(num_experts):
59
- num_tokens = num_expert_tokens_cpu[e]
60
- C[e, :num_tokens, :] = A[e, :num_tokens, :] @ B[e].transpose(0, 1)
61
-
62
- return C
91
+ return BatchedMMTensors(A, B, C, num_expert_tokens)
92
93
65
-@pytest.mark.parametrize("num_experts", [16, 32])
94
+@pytest.mark.parametrize("num_experts", [8, 16, 32])
95
@pytest.mark.parametrize("max_tokens_per_expert",
96
[32, 64, 128, 192, 224, 256, 512])
97
@pytest.mark.parametrize("K", [128, 256, 1024])
98
@pytest.mark.parametrize("N", [128, 256, 512, 1024])
99
@pytest.mark.parametrize("dtype",
100
[torch.float32, torch.float16, torch.bfloat16])
101
+@pytest.mark.parametrize("block_shape", [None])
102
+@pytest.mark.parametrize("per_act_token_quant", [False])
103
def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
73
- N: int, dtype: torch.dtype):
104
+ N: int, dtype: torch.dtype,
105
+ block_shape: Optional[list[int]],
106
+ per_act_token_quant: bool):
107
+ current_platform.seed_everything(7)
108
75
- config = BatchedMMConfig(dtype, num_experts, max_tokens_per_expert, K, N)
76
- tensors = BatchedMMTensors.make_tensors(config)
109
+ use_fp8_w8a8 = dtype == torch.float8_e4m3fn
110
78
- test_output = tensors.C
79
- ref_output = test_output.clone()
111
+ if (per_act_token_quant or block_shape is not None) and not use_fp8_w8a8:
112
+ pytest.skip("Don't test blocking for non-quantized types.")
113
+
114
+ if per_act_token_quant and block_shape is not None:
115
+ pytest.skip("Skip illegal quantization test.")
116
+
117
+ if dtype.itemsize == 1:
118
+ act_dtype = torch.bfloat16
119
+ quant_dtype = dtype
120
+ else:
121
+ act_dtype = dtype
122
+ quant_dtype = None
123
+
124
+ num_expert_tokens = torch.randint(low=0,
125
+ high=max_tokens_per_expert,
126
+ size=(num_experts, ),
127
+ device="cuda",
128
+ dtype=torch.int32)
129
+
130
+ A, A_q, A_scale = make_quantized_test_activations(
131
+ num_experts,
132
+ max_tokens_per_expert,
133
+ K,
134
+ in_dtype=act_dtype,
135
+ quant_dtype=quant_dtype,
136
+ block_shape=block_shape,
137
+ per_act_token_quant=per_act_token_quant)
138
+
139
+ B, B_q, B_scale, _, _, _ = make_test_weights(
140
+ num_experts,
141
+ N // 2,
142
+ K,
143
+ in_dtype=act_dtype,
144
+ quant_dtype=quant_dtype,
145
+ block_shape=block_shape,
146
+ )
147
+
148
+ out_shape = (num_experts, max_tokens_per_expert, N)
149
+ test_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda")
150
+ ref_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda")
151
+ q_ref_output = torch.zeros(out_shape, dtype=act_dtype, device="cuda")
152
153
compute_tl_dtype = {
154
torch.float16: tl.float16,
155
torch.bfloat16: tl.bfloat16,
156
torch.float32: tl.float32
157
}[test_output.dtype]
158
+
159
+ assert A_q.dtype == B_q.dtype
160
+
161
invoke_moe_batched_triton_kernel(
87
- tensors.A,
88
- tensors.B,
162
+ A_q,
163
+ B_q,
164
test_output,
90
- tensors.num_expert_tokens,
165
+ num_expert_tokens,
166
compute_tl_dtype,
167
# Quantization data
93
- None,
94
- None,
168
+ A_scale,
169
+ B_scale,
170
None,
171
# Quantization schemes
97
- False,
172
+ use_fp8_w8a8,
173
False,
174
False,
175
config={
176
"BLOCK_SIZE_M": 16,
177
"BLOCK_SIZE_N": 16,
103
- "BLOCK_SIZE_K": 16
104
- })
178
+ "BLOCK_SIZE_K": 16 if dtype.itemsize > 1 else 32
179
+ },
180
+ block_shape=block_shape,
181
+ )
182
106
- ref_output = ref_impl(tensors.A, tensors.B, ref_output,
107
- tensors.num_expert_tokens)
183
+ ref_output = native_batched_masked_quant_matmul(
184
+ A,
185
+ B,
186
+ ref_output,
187
+ num_expert_tokens,
188
+ None,
189
+ None,
190
+ None,
191
+ )
192
+
193
+ q_ref_output = native_batched_masked_quant_matmul(A_q, B_q, q_ref_output,
194
+ num_expert_tokens,
195
+ A_scale, B_scale,
196
+ block_shape)
197
198
rtol, atol = {
199
torch.float16: (6e-2, 6e-2),
@@ -112,4 +201,98 @@ def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
201
torch.float32: (1e-2, 1e-2),
202
}[test_output.dtype]
203
115
- torch.testing.assert_close(test_output, ref_output, atol=atol, rtol=rtol)
204
+ torch.testing.assert_close(ref_output, test_output, atol=atol, rtol=rtol)
205
+ torch.testing.assert_close(test_output, q_ref_output, atol=atol, rtol=rtol)
206
+
207
+
208
+@pytest.mark.parametrize(("m", "n", "k"), MNK_FACTORS)
209
+@pytest.mark.parametrize("e", NUM_EXPERTS)
210
+@pytest.mark.parametrize("topk", TOP_KS)
211
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
212
+@pytest.mark.parametrize("per_act_token_quant", [False])
213
+@pytest.mark.parametrize("block_shape", [None])
214
+def test_fused_moe_batched_experts(
215
+ m: int,
216
+ n: int,
217
+ k: int,
218
+ e: int,
219
+ topk: int,
220
+ dtype: torch.dtype,
221
+ per_act_token_quant: bool,
222
+ block_shape: Optional[list[int]],
223
+):
224
+ current_platform.seed_everything(7)
225
+
226
+ use_fp8_w8a8 = dtype == torch.float8_e4m3fn
227
+
228
+ if not use_fp8_w8a8 and (per_act_token_quant or block_shape is not None):
229
+ pytest.skip("Skip quantization test for non-quantized type")
230
+
231
+ if per_act_token_quant and block_shape is not None or topk > e:
232
+ pytest.skip("Skip illegal quantization test.")
233
+
234
+ a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
235
+ score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
236
+
237
+ if dtype.itemsize == 1:
238
+ act_dtype = torch.bfloat16
239
+ quant_dtype = dtype
240
+ else:
241
+ act_dtype = dtype
242
+ quant_dtype = None
243
+
244
+ _, w1, w1_s, _, w2, w2_s = make_test_weights(e,
245
+ n,
246
+ k,
247
+ block_shape=block_shape,
248
+ in_dtype=act_dtype,
249
+ quant_dtype=quant_dtype)
250
+
251
+ with set_current_vllm_config(vllm_config):
252
+ topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
253
+ batched_output = batched_moe(
254
+ a,
255
+ w1,
256
+ w2,
257
+ topk_weight,
258
+ topk_ids,
259
+ w1_scale=w1_s,
260
+ w2_scale=w2_s,
261
+ quant_dtype=quant_dtype,
262
+ per_act_token_quant=per_act_token_quant,
263
+ block_shape=block_shape,
264
+ )
265
+ baseline_output = torch_experts(
266
+ a,
267
+ w1,
268
+ w2,
269
+ topk_weight,
270
+ topk_ids,
271
+ w1_scale=w1_s,
272
+ w2_scale=w2_s,
273
+ quant_dtype=quant_dtype,
274
+ per_act_token_quant=per_act_token_quant,
275
+ block_shape=block_shape)
276
+
277
+ triton_output = triton_moe(
278
+ a,
279
+ w1,
280
+ w2,
281
+ topk_weight,
282
+ topk_ids,
283
+ w1_scale=w1_s,
284
+ w2_scale=w2_s,
285
+ quant_dtype=quant_dtype,
286
+ per_act_token_quant=per_act_token_quant,
287
+ block_shape=block_shape,
288
+ )
289
+
290
+ torch.testing.assert_close(triton_output,
291
+ baseline_output,
292
+ atol=2e-2,
293
+ rtol=2e-2)
294
+
295
+ torch.testing.assert_close(triton_output,
296
+ batched_output,
297
+ atol=2e-2,
298
+ rtol=2e-2)
tests/kernels/moe/test_block_fp8.py
new
+296
@@ -0,0 +1,296 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.kernels.moe.utils import make_test_weights
8
+from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
9
+ native_w8a8_block_matmul)
10
+from vllm.config import VllmConfig, set_current_vllm_config
11
+from vllm.model_executor.layers.activation import SiluAndMul
12
+from vllm.model_executor.layers.fused_moe import fused_experts
13
+from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
14
+ _valid_deep_gemm_shape, deep_gemm_moe_fp8)
15
+from vllm.model_executor.layers.fused_moe.fused_moe import (
16
+ fused_topk, modular_triton_fused_moe)
17
+from vllm.platforms import current_platform
18
+
19
+dg_available = False
20
+try:
21
+ import deep_gemm
22
+ dg_available = True
23
+except ImportError:
24
+ pass
25
+
26
+if current_platform.get_device_capability() < (9, 0):
27
+ pytest.skip("FP8 Triton requires CUDA 9.0 or higher",
28
+ allow_module_level=True)
29
+
30
+vllm_config = VllmConfig()
31
+vllm_config.scheduler_config.max_num_seqs = 128
32
+vllm_config.scheduler_config.max_model_len = 8192
33
+
34
+# Test configurations
35
+DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
36
+# Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8
37
+# and its hidden size is 7168.
38
+MNK_FACTORS = [
39
+ (1, 128, 128),
40
+ (1, 512, 512),
41
+ (1, 128, 7168),
42
+ (1, 1024, 7168),
43
+ (1, 4608, 128),
44
+ (1, 4608, 512),
45
+ (1, 4608, 7168),
46
+ (83, 128, 128),
47
+ (83, 512, 512),
48
+ (83, 1024, 7168),
49
+ (83, 4608, 512),
50
+ (83, 4608, 7168),
51
+ (128, 128, 128),
52
+ (128, 512, 512),
53
+ (128, 1024, 7168),
54
+ (128, 4608, 512),
55
+ (128, 4608, 7168),
56
+ (2048, 128, 128),
57
+ (2048, 1024, 7168),
58
+ (2048, 4608, 512),
59
+ (2048, 4608, 7168),
60
+ (8192, 128, 128),
61
+ (8192, 512, 512),
62
+ (8192, 128, 7168),
63
+ (8192, 1024, 7168),
64
+ (8192, 4608, 512),
65
+ (8192, 4608, 7168),
66
+]
67
+
68
+MNK_FACTORS_DG = [
69
+ (128, 128, 128),
70
+ (128, 512, 512),
71
+ (128, 128, 7168),
72
+ (128, 1024, 7168),
73
+ (128, 4608, 128),
74
+ (128, 4608, 512),
75
+ (128, 4608, 7168),
76
+ (192, 128, 128),
77
+ (192, 512, 512),
78
+ (192, 1024, 7168),
79
+ (192, 4608, 512),
80
+ (192, 4608, 7168),
81
+ (1335, 128, 128),
82
+ (1335, 1024, 7168),
83
+ (1335, 4608, 512),
84
+ (1335, 4608, 7168),
85
+ (2048, 128, 128),
86
+ (2048, 512, 512),
87
+ (2048, 128, 7168),
88
+ (2048, 1024, 7168),
89
+ (2048, 4608, 128),
90
+ (2048, 4608, 512),
91
+ (2048, 4608, 7168),
92
+]
93
+
94
+BLOCK_SIZE = [[128, 128]]
95
+E = [2, 8, 16] # [128, 256]
96
+TOP_KS = [1, 2, 6]
97
+SEEDS = [0]
98
+
99
+
100
+def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids,
101
+ block_shape):
102
+ """Fused moe with block-wise quantization using native torch."""
103
+ B, D = a.shape
104
+ topk = topk_ids.size(1)
105
+ a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
106
+ out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
107
+
108
+ topk_weight = topk_weight.view(-1)
109
+ topk_ids = topk_ids.view(-1)
110
+
111
+ _, block_k = block_shape[0], block_shape[1]
112
+ a_q, a_s = native_per_token_group_quant_fp8(a, block_k)
113
+ a_q = a_q.to(torch.float32)
114
+ for i in range(w1.shape[0]):
115
+ mask = topk_ids == i
116
+ if mask.sum():
117
+ inter_out = native_w8a8_block_matmul(a_q[mask],
118
+ w1[i],
119
+ a_s[mask],
120
+ w1_s[i],
121
+ block_shape,
122
+ output_dtype=a.dtype)
123
+ act_out = SiluAndMul().forward_native(inter_out)
124
+ act_out_q, act_out_s = native_per_token_group_quant_fp8(
125
+ act_out, block_k)
126
+ out[mask] = native_w8a8_block_matmul(act_out_q,
127
+ w2[i],
128
+ act_out_s,
129
+ w2_s[i],
130
+ block_shape,
131
+ output_dtype=a.dtype)
132
+ return (out.view(B, -1, w2.shape[1]) *
133
+ topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
134
+
135
+
136
+# Skip all tests if CUDA is not available
137
+pytest.importorskip("torch.cuda")
138
+
139
+
140
+@pytest.fixture(autouse=True)
141
+def setup_cuda():
142
+ torch.set_default_device("cuda")
143
+
144
+
145
+@pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS)
146
+@pytest.mark.parametrize("E", E)
147
+@pytest.mark.parametrize("topk", TOP_KS)
148
+@pytest.mark.parametrize("block_size", BLOCK_SIZE)
149
+@pytest.mark.parametrize("dtype", DTYPES)
150
+@pytest.mark.parametrize("seed", SEEDS)
151
+@torch.inference_mode()
152
+def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed,
153
+ monkeypatch):
154
+ if topk > E:
155
+ pytest.skip(f"Skipping test; topk={topk} > E={E}")
156
+
157
+ torch.manual_seed(seed)
158
+
159
+ monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", "2048")
160
+
161
+ a = torch.randn((M, K), dtype=dtype) / 10
162
+ score = torch.randn((M, E), dtype=dtype)
163
+
164
+ _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
165
+ N,
166
+ K,
167
+ dtype,
168
+ torch.float8_e4m3fn,
169
+ per_act_token_quant=False,
170
+ block_shape=block_size)
171
+
172
+ m_fused_moe = modular_triton_fused_moe(use_fp8_w8a8=True,
173
+ use_int8_w8a8=False,
174
+ use_int8_w8a16=False,
175
+ use_int4_w4a16=False,
176
+ per_act_token_quant=False,
177
+ block_shape=block_size)
178
+
179
+ topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
180
+
181
+ # Set the context to avoid lots of warning spam.
182
+ with set_current_vllm_config(vllm_config):
183
+ ref_out = torch_w8a8_block_fp8_moe(
184
+ a,
185
+ w1,
186
+ w2,
187
+ w1_s,
188
+ w2_s,
189
+ topk_weights,
190
+ topk_ids,
191
+ block_size,
192
+ )
193
+
194
+ out = fused_experts(
195
+ a,
196
+ w1,
197
+ w2,
198
+ topk_weights,
199
+ topk_ids,
200
+ use_fp8_w8a8=True,
201
+ w1_scale=w1_s,
202
+ w2_scale=w2_s,
203
+ block_shape=block_size,
204
+ )
205
+
206
+ m_out = m_fused_moe(
207
+ a,
208
+ w1,
209
+ w2,
210
+ topk_weights,
211
+ topk_ids,
212
+ w1_scale=w1_s,
213
+ w2_scale=w2_s,
214
+ )
215
+
216
+ # 0.039 only needed for [40000-4608-7168-2-1-block_size852-dtype852-0]
217
+ tol = 0.035 if M < 40000 else 0.039
218
+ torch.testing.assert_close(out, ref_out, atol=tol, rtol=tol)
219
+ torch.testing.assert_close(m_out, ref_out, atol=tol, rtol=tol)
220
+
221
+
222
+@pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS_DG)
223
+@pytest.mark.parametrize("E", E)
224
+@pytest.mark.parametrize("topk", TOP_KS)
225
+@pytest.mark.parametrize("seed", SEEDS)
226
+@pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
227
+@torch.inference_mode()
228
+def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
229
+ monkeypatch):
230
+ if topk > E:
231
+ pytest.skip(f"Skipping test: topk={topk} > E={E}")
232
+
233
+ if not _valid_deep_gemm_shape(M, N, K):
234
+ pytest.skip(f"Skipping test: invalid size m={M}, n={N}, k={K}")
235
+
236
+ chunk_size = 1024
237
+
238
+ torch.manual_seed(seed)
239
+
240
+ monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
241
+
242
+ block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
243
+ block_size = [block_m, block_m]
244
+ dtype = torch.bfloat16
245
+
246
+ a = torch.randn((M, K), dtype=dtype) / 10
247
+ score = torch.randn((M, E), dtype=dtype)
248
+
249
+ _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
250
+ N,
251
+ K,
252
+ dtype,
253
+ torch.float8_e4m3fn,
254
+ per_act_token_quant=False,
255
+ block_shape=block_size)
256
+
257
+ # Note: for now use_compile will error out if the problem size is
258
+ # large enough to trigger chunking. I'm leaving the flag and
259
+ # setup code in case we are able to revisit this later.
260
+ use_compile = False
261
+
262
+ use_cudagraph = (chunk_size < M and N >= 1024 and K >= 1024
263
+ and current_platform.is_cuda_alike())
264
+
265
+ topk_weights, topk_ids, _ = fused_topk(a, score.float(), topk, False)
266
+
267
+ # Set the context to avoid lots of warning spam.
268
+ with set_current_vllm_config(vllm_config):
269
+ ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weights,
270
+ topk_ids, block_size)
271
+
272
+ if use_compile:
273
+ deep_gemm_moe_fp8_fn = torch.compile(deep_gemm_moe_fp8,
274
+ backend="inductor",
275
+ fullgraph=True)
276
+ torch._dynamo.mark_dynamic(a, 0)
277
+ torch._dynamo.mark_dynamic(topk_weights, 0)
278
+ torch._dynamo.mark_dynamic(topk_ids, 0)
279
+ else:
280
+ deep_gemm_moe_fp8_fn = deep_gemm_moe_fp8
281
+
282
+ out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights,
283
+ topk_ids)
284
+
285
+ if use_cudagraph:
286
+ out.fill_(0)
287
+ stream = torch.cuda.Stream()
288
+ graph = torch.cuda.CUDAGraph()
289
+ with torch.cuda.graph(graph, stream=stream):
290
+ out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights,
291
+ topk_ids)
292
+ torch.cuda.synchronize()
293
+ graph.replay()
294
+ torch.cuda.synchronize()
295
+
296
+ torch.testing.assert_close(out, ref_out, atol=0.035, rtol=0.035)
tests/kernels/moe/test_block_int8.py
new
+147
@@ -0,0 +1,147 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import pytest
5
+import torch
6
+
7
+from tests.kernels.moe.utils import make_test_weights
8
+from tests.kernels.quant_utils import (native_per_token_group_quant_int8,
9
+ native_w8a8_block_matmul)
10
+from vllm.config import VllmConfig, set_current_vllm_config
11
+from vllm.model_executor.layers.activation import SiluAndMul
12
+from vllm.model_executor.layers.fused_moe import fused_moe
13
+from vllm.platforms import current_platform
14
+
15
+if current_platform.get_device_capability() < (7, 0):
16
+ pytest.skip("INT8 Triton requires CUDA 7.0 or higher",
17
+ allow_module_level=True)
18
+
19
+vllm_config = VllmConfig()
20
+vllm_config.scheduler_config.max_num_seqs = 128
21
+vllm_config.scheduler_config.max_model_len = 8192
22
+
23
+DTYPES = [torch.half, torch.bfloat16]
24
+
25
+MNK_FACTORS = [
26
+ (1, 128, 128),
27
+ (1, 512, 512),
28
+ (1, 128, 7168),
29
+ (1, 1024, 7168),
30
+ (1, 4096, 128),
31
+ (1, 4096, 512),
32
+ (1, 4096, 7168),
33
+ (33, 128, 128),
34
+ (33, 512, 512),
35
+ (33, 128, 7168),
36
+ (33, 1024, 7168),
37
+ (33, 4096, 128),
38
+ (33, 4096, 512),
39
+ (33, 4096, 7168),
40
+ (128, 128, 128),
41
+ (128, 512, 512),
42
+ (128, 1024, 7168),
43
+ (128, 4096, 512),
44
+ (128, 4096, 7168),
45
+ (222, 128, 128),
46
+ (222, 512, 512),
47
+ (222, 1024, 7168),
48
+ (222, 4096, 512),
49
+ (222, 4096, 7168),
50
+ (2048, 128, 128),
51
+ (2048, 1024, 7168),
52
+ (2048, 4096, 512),
53
+ (2048, 4096, 7168),
54
+]
55
+
56
+E = [8, 24]
57
+TOP_KS = [2, 6]
58
+# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
59
+BLOCK_SIZE = [[128, 128]]
60
+SEEDS = [0]
61
+
62
+
63
+# For test
64
+def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
65
+ """This function performs fused moe with block-wise quantization using
66
+ native torch."""
67
+ B, D = a.shape
68
+ a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
69
+ out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
70
+ score = torch.softmax(score, dim=-1, dtype=torch.float32)
71
+ topk_weight, topk_ids = torch.topk(score, topk)
72
+ topk_weight = topk_weight.view(-1)
73
+ topk_ids = topk_ids.view(-1)
74
+
75
+ _, block_k = block_shape[0], block_shape[1]
76
+ a_q, a_s = native_per_token_group_quant_int8(a, block_k)
77
+ for i in range(w1.shape[0]):
78
+ mask = topk_ids == i
79
+ if mask.sum():
80
+ inter_out = native_w8a8_block_matmul(a_q[mask],
81
+ w1[i],
82
+ a_s[mask],
83
+ w1_s[i],
84
+ block_shape,
85
+ output_dtype=a.dtype)
86
+ act_out = SiluAndMul().forward_native(inter_out)
87
+ act_out_q, act_out_s = native_per_token_group_quant_int8(
88
+ act_out, block_k)
89
+ act_out = act_out.to(torch.float32)
90
+ out[mask] = native_w8a8_block_matmul(act_out_q,
91
+ w2[i],
92
+ act_out_s,
93
+ w2_s[i],
94
+ block_shape,
95
+ output_dtype=a.dtype)
96
+ return (out.view(B, -1, w2.shape[1]) *
97
+ topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
98
+
99
+
100
+@pytest.fixture(autouse=True, scope="module")
101
+def setup_cuda():
102
+ """Sets the default CUDA device for all tests in this module."""
103
+ torch.set_default_device("cuda")
104
+
105
+
106
+@pytest.mark.parametrize(("M", "N", "K"), MNK_FACTORS)
107
+@pytest.mark.parametrize("E", E)
108
+@pytest.mark.parametrize("topk", TOP_KS)
109
+@pytest.mark.parametrize("block_size", BLOCK_SIZE)
110
+@pytest.mark.parametrize("dtype", DTYPES)
111
+@pytest.mark.parametrize("seed", SEEDS)
112
+@torch.inference_mode()
113
+def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
114
+ """Tests the fused_moe kernel with W8A8 INT8 block quantization against a
115
+ native torch reference."""
116
+ torch.manual_seed(seed)
117
+
118
+ a = torch.randn((M, K), dtype=dtype) / 10
119
+ score = torch.randn((M, E), dtype=dtype)
120
+
121
+ _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
122
+ N,
123
+ K,
124
+ dtype,
125
+ torch.int8,
126
+ per_act_token_quant=False,
127
+ block_shape=block_size)
128
+
129
+ # Set the context to avoid lots of warning spam.
130
+ with set_current_vllm_config(vllm_config):
131
+ out = fused_moe(
132
+ a,
133
+ w1,
134
+ w2,
135
+ score,
136
+ topk,
137
+ renormalize=False,
138
+ use_int8_w8a8=True,
139
+ w1_scale=w1_s,
140
+ w2_scale=w2_s,
141
+ block_shape=block_size,
142
+ )
143
+ ref_out = torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk,
144
+ block_size)
145
+
146
+ # Check results
147
+ torch.testing.assert_close(out, ref_out, atol=0.065, rtol=0.065)
tests/kernels/moe/test_cutlass_moe.py
+13
-9
@@ -97,11 +97,9 @@ class MOETensors8Bit(MOETensors):
97
n_b_scales = 2 * n if per_out_channel else 1
98
k_b_scales = k if per_out_channel else 1
99
# Get the right scale for tests.
100
- _, a_scale = ops.scaled_fp8_quant(
101
- moe_tensors_fp16.a, use_per_token_if_dynamic=per_act_token)
102
- a_q, _ = ops.scaled_fp8_quant(moe_tensors_fp16.a,
103
- a_scale,
104
- use_per_token_if_dynamic=per_act_token)
100
+ a_q, a_scale = ops.scaled_fp8_quant(
101
+ moe_tensors_fp16.a, None, use_per_token_if_dynamic=per_act_token)
102
+
103
w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=q_dtype)
104
w2_q = torch.empty((e, k, n), device="cuda", dtype=q_dtype)
105
@@ -187,6 +185,7 @@ def run_with_expert_maps(num_experts: int, num_local_experts: int,
185
def run_8_bit(moe_tensors: MOETensors8Bit,
186
topk_weights: torch.Tensor,
187
topk_ids: torch.Tensor,
188
+ per_act_token: bool,
189
num_local_experts: Optional[int] = None) -> torch.Tensor:
190
assert not any([
191
t is None for t in [
@@ -203,7 +202,8 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
202
'topk_ids': topk_ids,
203
'w1_scale': moe_tensors.w1_scale,
204
'w2_scale': moe_tensors.w2_scale,
206
- 'a1_scale': moe_tensors.a_scale
205
+ 'per_act_token': per_act_token,
206
+ 'a1_scale': None #moe_tensors.a_scale
207
}
208
209
num_experts = moe_tensors.w1.size(0)
@@ -254,11 +254,13 @@ def test_cutlass_moe_8_bit_no_graph(
254
triton_output = fused_experts(mt.a_d, mt.w1_d, mt.w2_d, topk_weights,
255
topk_ids)
256
257
- cutlass_output = run_8_bit(mt, topk_weights, topk_ids)
257
+ cutlass_output = run_8_bit(mt, topk_weights, topk_ids, per_act_token)
258
259
+ # Note 5.5 only needed for larger problem sizes, 5 works ok for
260
+ # the rest.
261
torch.testing.assert_close(triton_output,
262
cutlass_output,
261
- atol=5e-2,
263
+ atol=5.5e-2,
264
rtol=1e-2)
265
266
@@ -303,7 +305,8 @@ def test_cutlass_moe_8_bit_cuda_graph(
305
stream = torch.cuda.Stream()
306
graph = torch.cuda.CUDAGraph()
307
with torch.cuda.graph(graph, stream=stream):
306
- cutlass_output = run_8_bit(mt, topk_weights, topk_ids)
308
+ cutlass_output = run_8_bit(mt, topk_weights, topk_ids,
309
+ per_act_token)
310
311
torch.cuda.synchronize()
312
graph.replay()
@@ -359,6 +362,7 @@ def test_cutlass_moe_8_bit_EP(
362
cutlass_output = run_8_bit(mt,
363
topk_weights,
364
topk_ids,
365
+ per_act_token,
366
num_local_experts=e // ep_size)
367
368
torch.testing.assert_close(triton_output,
tests/kernels/moe/test_deepep_deepgemm_moe.py
+28
-74
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
"""
3
-Test DeepEP + DeepGEMM integration
3
+Test DeepEP + DeepGEMM integration
4
DeepGEMM are gemm kernels specialized for the
5
fp8 block-quantized case.
6
"""
@@ -17,12 +17,11 @@ from vllm.config import VllmConfig, set_current_vllm_config
17
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
18
from vllm.model_executor.layers.fused_moe.modular_kernel import (
19
FusedMoEModularKernel)
20
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21
- per_token_group_quant_fp8)
20
from vllm.platforms import current_platform
21
from vllm.utils import has_deep_ep, has_deep_gemm
22
25
-from .utils import ProcessGroupInfo, parallel_launch
23
+from .parallel_utils import ProcessGroupInfo, parallel_launch
24
+from .utils import make_test_weights
25
26
if has_deep_ep():
27
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
@@ -30,10 +29,9 @@ if has_deep_ep():
29
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
30
DeepEPLLPrepareAndFinalize)
31
33
- from .utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a
32
+ from .parallel_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a
33
34
if has_deep_gemm():
36
- import deep_gemm
35
36
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
37
BatchedDeepGemmExperts)
@@ -60,25 +58,6 @@ def next_power_of_2(x):
58
return 2**math.ceil(math.log2(x))
59
60
63
-def per_block_cast_to_fp8(
64
- x: torch.Tensor,
65
- block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
66
- assert x.dim() == 2
67
- m, n = x.shape
68
- x_padded = torch.zeros(
69
- (deep_gemm.ceil_div(m, 128) * 128,
70
- deep_gemm.ceil_div(n, block_size_n) * block_size_n),
71
- dtype=x.dtype,
72
- device=x.device)
73
- x_padded[:m, :n] = x
74
- x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
75
- x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
76
- x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
77
- x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
78
- scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
79
- return x_scaled_sub, scales
80
-
81
-
61
def make_block_quant_fp8_weights(
62
e: int,
63
n: int,
@@ -86,43 +65,11 @@ def make_block_quant_fp8_weights(
65
block_size: list[int],
66
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
67
"""
89
- Return weights w1, w2, w1q, w2q, w1_scale, w2_scale
68
+ Return weights w1q, w2q, w1_scale, w2_scale
69
"""
91
- dtype = torch.bfloat16
92
-
93
- fp8_info = torch.finfo(torch.float8_e4m3fn)
94
- fp8_max, fp8_min = fp8_info.max, fp8_info.min
95
-
96
- w1_bf16 = torch.randn((e, 2 * n, k), dtype=dtype) / 10
97
- w1_bf16 = w1_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
98
-
99
- w2_bf16 = torch.randn((e, k, n), dtype=dtype) / 10
100
- w2_bf16 = w2_bf16.clamp(min=fp8_min, max=fp8_max).to(dtype=dtype)
101
-
102
- block_n, block_k = block_size[0], block_size[1]
103
- n_tiles_w1 = ((2 * n) + block_n - 1) // block_n
104
- k_tiles_w1 = (k + block_k - 1) // block_k
105
- n_tiles_w2 = (k + block_n - 1) // block_n
106
- k_tiles_w2 = (n + block_k - 1) // block_k
107
-
108
- w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn)
109
- w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn)
110
-
111
- w1_s = torch.empty((e, n_tiles_w1, k_tiles_w1),
112
- device="cuda",
113
- dtype=torch.float32)
114
- w2_s = torch.empty((e, n_tiles_w2, k_tiles_w2),
115
- device="cuda",
116
- dtype=torch.float32)
117
-
118
- assert w1_s.shape == (e, (2 * n + 127) // 128, (k + 127) // 128)
119
- assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
120
-
121
- for i in range(e):
122
- w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
123
- w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
124
-
125
- return w1, w2, w1_s, w2_s
70
+ w1, w1q, w1_scale, w2, w2q, w2_scale = make_test_weights(
71
+ e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_size)
72
+ return w1q, w2q, w1_scale, w2_scale
73
74
75
@dataclasses.dataclass
@@ -132,6 +79,7 @@ class TestConfig:
79
k: int
80
n: int
81
num_experts: int
82
+ per_act_token_quant: bool
83
block_size: list[int]
84
# configs for testing low-latency kernels
85
low_latency: bool
@@ -150,8 +98,7 @@ class TestTensors:
98
def make(config: TestConfig, rank) -> "TestTensors":
99
100
dtype = torch.bfloat16
153
- topk, m, k, block_size = (config.topk, config.m, config.k,
154
- config.block_size)
101
+ topk, m, k = (config.topk, config.m, config.k)
102
103
fp8_info = torch.finfo(torch.float8_e4m3fn)
104
fp8_max, fp8_min = fp8_info.max, fp8_info.min
@@ -159,9 +106,7 @@ class TestTensors:
106
rank_tokens = torch.randn(
107
(m, k), device=torch.cuda.current_device(), dtype=dtype) / 10.0
108
rank_tokens = rank_tokens.clamp(min=fp8_min, max=fp8_max)
162
-
163
- block_k = block_size[1]
164
- _, rank_token_scales = per_token_group_quant_fp8(rank_tokens, block_k)
109
+ rank_token_scales = None
110
111
topk_ids = torch.randint(
112
low=0,
@@ -201,10 +146,12 @@ def make_ll_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
146
q_dtype=q_dtype,
147
block_shape=test_config.block_size)
148
204
- fused_experts = BatchedDeepGemmExperts(max_num_tokens=max_tokens_per_rank,
205
- world_size=pgi.world_size,
206
- dp_size=dp_size,
207
- block_shape=test_config.block_size)
149
+ fused_experts = BatchedDeepGemmExperts(
150
+ max_num_tokens=max_tokens_per_rank,
151
+ world_size=pgi.world_size,
152
+ dp_size=dp_size,
153
+ block_shape=test_config.block_size,
154
+ per_act_token_quant=test_config.per_act_token_quant)
155
mk = FusedMoEModularKernel(prepare_finalize=a2a,
156
fused_experts=fused_experts)
157
return mk
@@ -426,6 +373,7 @@ def test_ht_deepep_deepgemm_moe(mnk: tuple[int, int, int], num_experts: int,
373
"""
374
Tests for High-Throughput DeepEP + DeepGemm integration.
375
"""
376
+ import deep_gemm
377
378
m, n, k = mnk
379
current_platform.seed_everything(7)
@@ -442,6 +390,7 @@ def test_ht_deepep_deepgemm_moe(mnk: tuple[int, int, int], num_experts: int,
390
k=k,
391
n=n,
392
num_experts=num_experts,
393
+ per_act_token_quant=False,
394
block_size=block_size,
395
low_latency=False,
396
use_fp8_dispatch=None)
@@ -474,10 +423,14 @@ USE_FP8_DISPATCH = [False]
423
@pytest.mark.parametrize("world_dp_size", [(2, 1)])
424
@requires_deep_ep
425
@requires_deep_gemm
477
-def test_ll_deepep_deepgemm_moe(mnk: tuple[int, int,
478
- int], num_experts: int, topk: int,
479
- use_fp8_dispatch: bool, block_size: list[int],
480
- world_dp_size: tuple[int, int]):
426
+def test_ll_deepep_deepgemm_moe(
427
+ mnk: tuple[int, int, int],
428
+ num_experts: int,
429
+ topk: int,
430
+ use_fp8_dispatch: bool,
431
+ block_size: list[int],
432
+ world_dp_size: tuple[int, int],
433
+):
434
"""
435
Tests for Low-Latency DeepEP + DeepGemm integration.
436
"""
@@ -495,6 +448,7 @@ def test_ll_deepep_deepgemm_moe(mnk: tuple[int, int,
448
k=k,
449
n=n,
450
num_experts=num_experts,
451
+ per_act_token_quant=False,
452
block_size=block_size,
453
low_latency=True,
454
use_fp8_dispatch=use_fp8_dispatch,
tests/kernels/moe/test_deepep_moe.py
+69
-37
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
23
from vllm.platforms import current_platform
24
from vllm.utils import has_deep_ep
25
26
-from .utils import ProcessGroupInfo, parallel_launch
26
+from .parallel_utils import ProcessGroupInfo, parallel_launch
27
28
if has_deep_ep():
29
from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
@@ -31,7 +31,7 @@ if has_deep_ep():
31
from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
32
DeepEPLLPrepareAndFinalize)
33
34
- from .utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a
34
+ from .parallel_utils import DeepEPHTArgs, DeepEPLLArgs, make_deepep_a2a
35
36
requires_deep_ep = pytest.mark.skipif(
37
not has_deep_ep(),
@@ -102,10 +102,6 @@ class TestTensors:
102
rank_tokens = torch.randn(
103
(config.m, config.k), device="cuda", dtype=token_dtype) / 10
104
rank_token_scales = None
105
- if config.dtype == torch.float8_e4m3fn:
106
- # low_latency_mode kernels dont support per-token quant.
107
- _, rank_token_scales = ops.scaled_fp8_quant(
108
- rank_tokens, use_per_token_if_dynamic=not low_latency_mode)
105
106
topk = torch.randint(low=0,
107
high=config.num_experts,
@@ -121,11 +117,18 @@ class TestTensors:
117
config=config)
118
119
124
-def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
125
- low_latency_mode: bool, hidden_size: int, dp_size: int,
126
- num_experts: int, num_local_experts: int,
127
- q_dtype: Optional[torch.dtype],
128
- use_fp8_dispatch: bool) -> FusedMoEModularKernel:
120
+def make_modular_kernel(
121
+ pg: ProcessGroup,
122
+ pgi: ProcessGroupInfo,
123
+ low_latency_mode: bool,
124
+ hidden_size: int,
125
+ dp_size: int,
126
+ num_experts: int,
127
+ num_local_experts: int,
128
+ q_dtype: Optional[torch.dtype],
129
+ use_fp8_dispatch: bool,
130
+ per_act_token_quant: bool,
131
+) -> FusedMoEModularKernel:
132
133
is_quantized = q_dtype is not None
134
@@ -152,6 +155,7 @@ def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
155
deepep_ll_args = ll_args)
156
157
if low_latency_mode:
158
+ assert not per_act_token_quant, "not supported in ll mode"
159
fused_experts = BatchedTritonExperts(
160
max_num_tokens=MAX_TOKENS_PER_RANK,
161
world_size=pgi.world_size,
@@ -159,25 +163,37 @@ def make_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
163
use_fp8_w8a8=is_quantized,
164
use_int8_w8a8=False,
165
use_int8_w8a16=False,
162
- use_int4_w4a16=False)
166
+ use_int4_w4a16=False,
167
+ per_act_token_quant=False,
168
+ )
169
else:
164
- fused_experts = TritonExperts(use_fp8_w8a8=is_quantized,
165
- use_int8_w8a8=False,
166
- use_int8_w8a16=False,
167
- use_int4_w4a16=False,
168
- per_channel_quant=False)
170
+ fused_experts = TritonExperts(
171
+ use_fp8_w8a8=is_quantized,
172
+ use_int8_w8a8=False,
173
+ use_int8_w8a16=False,
174
+ use_int4_w4a16=False,
175
+ per_act_token_quant=per_act_token_quant,
176
+ )
177
178
mk = FusedMoEModularKernel(prepare_finalize=a2a,
179
fused_experts=fused_experts)
180
return mk
181
182
175
-def deep_ep_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
176
- low_latency_mode: bool, dp_size: int,
177
- test_tensors: TestTensors, w1: torch.Tensor,
178
- w2: torch.Tensor, w1_scale: Optional[torch.Tensor],
179
- w2_scale: Optional[torch.Tensor], num_experts: int,
180
- use_fp8_dispatch: bool) -> torch.Tensor:
183
+def deep_ep_moe_impl(
184
+ pg: ProcessGroup,
185
+ pgi: ProcessGroupInfo,
186
+ low_latency_mode: bool,
187
+ dp_size: int,
188
+ test_tensors: TestTensors,
189
+ w1: torch.Tensor,
190
+ w2: torch.Tensor,
191
+ w1_scale: Optional[torch.Tensor],
192
+ w2_scale: Optional[torch.Tensor],
193
+ num_experts: int,
194
+ use_fp8_dispatch: bool,
195
+ per_act_token_quant: bool,
196
+) -> torch.Tensor:
197
198
num_local_experts = w1.size(0)
199
@@ -199,11 +215,9 @@ def deep_ep_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
215
q_dtype = torch.float8_e4m3fn
216
217
# Make modular kernel
202
- mk: FusedMoEModularKernel = make_modular_kernel(pg, pgi, low_latency_mode,
203
- hidden_size, dp_size,
204
- num_experts,
205
- num_local_experts, q_dtype,
206
- use_fp8_dispatch)
218
+ mk: FusedMoEModularKernel = make_modular_kernel(
219
+ pg, pgi, low_latency_mode, hidden_size, dp_size, num_experts,
220
+ num_local_experts, q_dtype, use_fp8_dispatch, per_act_token_quant)
221
222
out_hidden_states = torch.empty_like(test_tensors.rank_tokens)
223
total_num_tokens = test_tensors.rank_tokens.size(0)
@@ -257,9 +271,15 @@ def deep_ep_moe_impl(pg: ProcessGroup, pgi: ProcessGroupInfo,
271
return out_hidden_states
272
273
260
-def torch_moe_impl(test_tensors: TestTensors, w1: torch.Tensor,
261
- w2: torch.Tensor, w1_scale: Optional[torch.Tensor],
262
- w2_scale: Optional[torch.Tensor], using_fp8_dispatch: bool):
274
+def torch_moe_impl(
275
+ test_tensors: TestTensors,
276
+ w1: torch.Tensor,
277
+ w2: torch.Tensor,
278
+ w1_scale: Optional[torch.Tensor],
279
+ w2_scale: Optional[torch.Tensor],
280
+ using_fp8_dispatch: bool,
281
+ per_act_token_quant: bool,
282
+):
283
284
a, topk_ids, topk_weights = (test_tensors.rank_tokens, test_tensors.topk,
285
test_tensors.topk_weights)
@@ -267,6 +287,7 @@ def torch_moe_impl(test_tensors: TestTensors, w1: torch.Tensor,
287
# The DeepEP implementation is requested to dispatch using FP8.
288
# For numerical stability for testing, emulate the fp8 dispatch by
289
# blockwise quant and de-quant.
290
+ assert not per_act_token_quant
291
a = test_tensors.rank_tokens
292
aq, aq_scale = per_token_group_quant_fp8(a, 128)
293
a = (aq.view(-1, 128).to(torch.float32) * aq_scale.view(-1, 1)).view(
@@ -310,6 +331,7 @@ def _deep_ep_moe(
331
w1_scale: Optional[torch.Tensor],
332
w2_scale: Optional[torch.Tensor],
333
use_fp8_dispatch: bool,
334
+ per_act_token_quant: bool,
335
):
336
337
if not low_latency_mode:
@@ -331,7 +353,8 @@ def _deep_ep_moe(
353
with set_current_vllm_config(VllmConfig()):
354
# Reference
355
torch_combined = torch_moe_impl(test_tensors, w1, w2, w1_scale,
334
- w2_scale, use_fp8_dispatch)
356
+ w2_scale, use_fp8_dispatch,
357
+ per_act_token_quant)
358
359
# Splice experts for this rank.
360
num_local_experts = config.num_experts // pgi.world_size
@@ -356,6 +379,7 @@ def _deep_ep_moe(
379
w2_scale_ep,
380
config.num_experts,
381
use_fp8_dispatch,
382
+ per_act_token_quant,
383
)
384
385
torch.testing.assert_close(
@@ -384,10 +408,16 @@ DTYPES = [torch.bfloat16, torch.float8_e4m3fn]
408
@pytest.mark.parametrize("num_experts", [32])
409
@pytest.mark.parametrize("topk", [6])
410
@pytest.mark.parametrize("world_dp_size", [(2, 1)])
411
+@pytest.mark.parametrize("per_act_token_quant", [False, True])
412
@requires_deep_ep
388
-def test_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],
389
- num_experts: int, topk: int, world_dp_size: tuple[int,
390
- int]):
413
+def test_deep_ep_moe(
414
+ dtype: torch.dtype,
415
+ mnk: tuple[int, int, int],
416
+ num_experts: int,
417
+ topk: int,
418
+ world_dp_size: tuple[int, int],
419
+ per_act_token_quant: bool,
420
+):
421
low_latency_mode = False
422
use_fp8_dispatch = False
423
m, n, k = mnk
@@ -404,7 +434,8 @@ def test_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],
434
w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
435
436
parallel_launch(world_size, _deep_ep_moe, low_latency_mode, dp_size,
407
- config, w1, w2, w1_scale, w2_scale, use_fp8_dispatch)
437
+ config, w1, w2, w1_scale, w2_scale, use_fp8_dispatch,
438
+ per_act_token_quant)
439
440
441
MNKs = [
@@ -454,4 +485,5 @@ def test_low_latency_deep_ep_moe(dtype: torch.dtype, mnk: tuple[int, int, int],
485
w1, w2, w1_scale, w2_scale = make_weights(num_experts, n, k, dtype)
486
487
parallel_launch(world_size, _deep_ep_moe, low_latency_mode, dp_size,
457
- config, w1, w2, w1_scale, w2_scale, use_fp8_dispatch)
488
+ config, w1, w2, w1_scale, w2_scale, use_fp8_dispatch,
489
+ False)
tests/kernels/moe/test_moe.py
+13
-1
@@ -17,6 +17,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
17
import vllm.model_executor.layers.fused_moe # noqa
18
from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
19
from vllm.config import VllmConfig, set_current_vllm_config
20
+from vllm.distributed.parallel_state import init_distributed_environment
21
from vllm.forward_context import set_forward_context
22
from vllm.model_executor.layers.fused_moe import fused_moe
23
from vllm.model_executor.layers.fused_moe.fused_moe import (
@@ -142,6 +143,10 @@ def test_fused_moe(
143
# Setup test data
144
#
145
146
+ #
147
+ # Setup test data
148
+ #
149
+
150
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
151
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
152
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
@@ -169,7 +174,7 @@ def test_fused_moe(
174
use_int8_w8a8=False,
175
use_int8_w8a16=False,
176
use_int4_w4a16=False,
172
- per_channel_quant=False,
177
+ per_act_token_quant=False,
178
block_shape=None)
179
180
def m_fused_moe(
@@ -365,6 +370,13 @@ def test_mixtral_moe(dtype: torch.dtype, padding: bool, use_rocm_aiter: bool,
370
if dtype == torch.float32:
371
pytest.skip("AITER ROCm test skip for float32")
372
373
+ monkeypatch.setenv('RANK', "0")
374
+ monkeypatch.setenv('LOCAL_RANK', "0")
375
+ monkeypatch.setenv('WORLD_SIZE', "1")
376
+ monkeypatch.setenv('MASTER_ADDR', 'localhost')
377
+ monkeypatch.setenv('MASTER_PORT', '12345')
378
+ init_distributed_environment()
379
+
380
# Instantiate our and huggingface's MoE blocks
381
vllm_config.compilation_config.static_forward_context = dict()
382
with (set_current_vllm_config(vllm_config),
tests/kernels/moe/test_nvfp4_moe.py
+1
-1
@@ -14,7 +14,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
14
from vllm.platforms import current_platform
15
16
if not current_platform.has_device_capability(100):
17
- pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
17
+ pytest.skip("Nvfp4 Requires compute capability of 10 or above.",
18
allow_module_level=True)
19
20
MNK_FACTORS = [
tests/kernels/moe/test_pplx_cutlass_moe.py
+2
-4
@@ -15,7 +15,7 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
15
FusedMoEModularKernel)
16
from vllm.platforms import current_platform
17
18
-from .utils import ProcessGroupInfo, parallel_launch
18
+from .parallel_utils import ProcessGroupInfo, parallel_launch
19
20
try:
21
from pplx_kernels import AllToAll
@@ -93,7 +93,7 @@ def pplx_cutlass_moe(
93
num_experts=num_experts,
94
experts_per_token=topk,
95
rank=rank,
96
- world_size=pgi.world_size,
96
+ world_size=world_size,
97
dp_size=dp_size,
98
hidden_dim=hidden_dim,
99
hidden_dim_bytes=hidden_dim, # because a.dtype.itemsize == 1
@@ -118,8 +118,6 @@ def pplx_cutlass_moe(
118
pgi.world_size,
119
rank,
120
dp_size,
121
- quant_dtype=torch.float8_e4m3fn,
122
- per_act_token=per_act_token,
121
)
122
123
experts = CutlassExpertsFp8((num_experts + world_size - 1) // world_size,
tests/kernels/moe/test_pplx_moe.py
+96
-49
@@ -18,18 +18,20 @@ try:
18
except ImportError:
19
has_pplx = False
20
21
+from tests.kernels.moe.utils import make_test_weights, naive_batched_moe
22
from tests.kernels.utils import torch_experts
23
from vllm.config import VllmConfig, set_current_vllm_config
23
-from vllm.model_executor.layers.fused_moe import override_config
24
+from vllm.model_executor.layers.fused_moe import fused_topk, override_config
25
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
26
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
25
- BatchedExperts, BatchedPrepareAndFinalize, BatchedTritonExperts)
26
-from vllm.model_executor.layers.fused_moe.fused_moe import (fused_topk,
27
- get_default_config)
27
+ BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
28
+from vllm.model_executor.layers.fused_moe.fused_moe import get_default_config
29
from vllm.model_executor.layers.fused_moe.modular_kernel import (
30
FusedMoEModularKernel)
31
from vllm.platforms import current_platform
32
+from vllm.utils import round_up
33
32
-from .utils import ProcessGroupInfo, parallel_launch
34
+from .parallel_utils import ProcessGroupInfo, parallel_launch
35
36
requires_pplx = pytest.mark.skipif(
37
not has_pplx,
@@ -144,25 +146,6 @@ def torch_batched_moe(
146
return torch_finalize(out, topk_weight, topk_ids)
147
148
147
-def batched_moe(
148
- a: torch.Tensor,
149
- w1: torch.Tensor,
150
- w2: torch.Tensor,
151
- topk_weight: torch.Tensor,
152
- topk_ids: torch.Tensor,
153
-) -> torch.Tensor:
154
- num_experts = w1.shape[0]
155
-
156
- fused_experts = FusedMoEModularKernel(
157
- BatchedPrepareAndFinalize(max_num_tokens=a.shape[0],
158
- world_size=1,
159
- dp_size=1,
160
- rank=0),
161
- BatchedExperts(max_num_tokens=a.shape[0], dp_size=1, world_size=1))
162
-
163
- return fused_experts(a, w1, w2, topk_weight, topk_ids, num_experts)
164
-
165
-
149
@pytest.mark.parametrize("m", [1, 33, 64, 222])
150
@pytest.mark.parametrize("n", [128, 1024, 2048])
151
@pytest.mark.parametrize("k", [128, 512, 1024])
@@ -188,7 +171,7 @@ def test_fused_moe_batched_experts(
171
topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
172
baseline_output = torch_experts(a, w1, w2, topk_weight, topk_ids)
173
torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids)
191
- batched_output = batched_moe(a, w1, w2, topk_weight, topk_ids)
174
+ batched_output = naive_batched_moe(a, w1, w2, topk_weight, topk_ids)
175
176
torch.testing.assert_close(baseline_output,
177
torch_output,
@@ -226,7 +209,6 @@ def pplx_prepare_finalize(
209
210
topk = topk_ids.shape[1]
211
num_tokens, hidden_dim = a.shape
229
- block_size = 128
212
device = pgi.device
213
rank = pgi.rank
214
world_size = pgi.world_size
@@ -241,9 +223,7 @@ def pplx_prepare_finalize(
223
dp_size=dp_size,
224
hidden_dim=hidden_dim,
225
hidden_dim_bytes=hidden_dim * a.dtype.itemsize,
244
- hidden_dim_scale_bytes=(0 if a.dtype.itemsize != 1 else
245
- ((hidden_dim + block_size - 1) // block_size *
246
- torch.float32.itemsize)),
226
+ hidden_dim_scale_bytes=0,
227
)
228
229
if group_name is None:
@@ -260,7 +240,6 @@ def pplx_prepare_finalize(
240
world_size,
241
rank,
242
dp_size,
263
- a.dtype,
243
)
244
245
a_chunk = chunk_by_rank(a, rank, world_size).to(device)
@@ -276,6 +255,7 @@ def pplx_prepare_finalize(
255
num_experts,
256
None,
257
False,
258
+ FusedMoEQuantConfig(),
259
)
260
261
b_a = b_a * 1.5
@@ -350,6 +330,7 @@ def _pplx_prepare_finalize(
330
# TODO (bnell): this test point does not work for odd M due to how the test is
331
# written, not due to limitations of the pplx kernels. The pplx_moe
332
# test below is able to deal with odd M.
333
+# TODO (bnell) add fp8 tests
334
@pytest.mark.parametrize("mnk", PPLX_PREPARE_COMBOS)
335
@pytest.mark.parametrize("e", NUM_EXPERTS)
336
@pytest.mark.parametrize("topk", TOP_KS)
@@ -386,18 +367,31 @@ def pplx_moe(
367
w2: torch.Tensor,
368
topk_weight: torch.Tensor,
369
topk_ids: torch.Tensor,
370
+ w1_scale: Optional[torch.Tensor] = None,
371
+ w2_scale: Optional[torch.Tensor] = None,
372
+ qtype: Optional[torch.dtype] = None,
373
+ per_act_token_quant=False,
374
+ block_shape: Optional[list[int]] = None,
375
use_compile: bool = False,
376
use_cudagraphs: bool = True,
377
) -> torch.Tensor:
378
from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
393
- PplxPrepareAndFinalize)
379
+ PplxPrepareAndFinalize, pplx_hidden_dim_scale_bytes)
380
381
device = torch.device("cuda", rank)
382
hidden_dim = a.shape[1]
383
num_experts = w1.shape[0]
398
- block_size = 128
384
topk = topk_ids.shape[1]
400
- max_num_tokens = rank_chunk(a.shape[0], 0, world_size)
385
+ max_num_tokens = round_up(rank_chunk(a.shape[0], 0, world_size), 64)
386
+
387
+ hidden_dim_bytes, scale_bytes = pplx_hidden_dim_scale_bytes(
388
+ max_num_tokens,
389
+ hidden_dim,
390
+ a.dtype,
391
+ qtype,
392
+ per_act_token_quant=per_act_token_quant,
393
+ block_shape=block_shape,
394
+ )
395
396
args = dict(
397
max_num_tokens=max_num_tokens,
@@ -407,10 +401,8 @@ def pplx_moe(
401
world_size=world_size,
402
dp_size=dp_size,
403
hidden_dim=hidden_dim,
410
- hidden_dim_bytes=hidden_dim * a.dtype.itemsize,
411
- hidden_dim_scale_bytes=(0 if a.dtype.itemsize != 1 else
412
- ((hidden_dim + block_size - 1) // block_size *
413
- torch.float32.itemsize)),
404
+ hidden_dim_bytes=hidden_dim_bytes,
405
+ hidden_dim_scale_bytes=scale_bytes,
406
)
407
408
if group_name is None:
@@ -429,9 +421,11 @@ def pplx_moe(
421
dp_size,
422
)
423
432
- experts = BatchedTritonExperts(max_num_tokens=a.shape[0],
424
+ experts = BatchedTritonExperts(max_num_tokens=max_num_tokens,
425
world_size=world_size,
434
- dp_size=dp_size)
426
+ dp_size=dp_size,
427
+ use_fp8_w8a8=qtype == torch.float8_e4m3fn,
428
+ block_shape=block_shape)
429
430
fused_experts = FusedMoEModularKernel(
431
prepare_finalize,
@@ -447,6 +441,13 @@ def pplx_moe(
441
w1_chunk = chunk_by_rank(w1, rank, world_size).to(device)
442
w2_chunk = chunk_by_rank(w2, rank, world_size).to(device)
443
444
+ if w1_scale is not None:
445
+ w1_scale_chunk = chunk_by_rank(w1_scale, rank, world_size).to(device)
446
+ w2_scale_chunk = chunk_by_rank(w2_scale, rank, world_size).to(device)
447
+ else:
448
+ w1_scale_chunk = None
449
+ w2_scale_chunk = None
450
+
451
# Note: for now use_compile will error out if the problem size is
452
# large enough to trigger chunking. I'm leaving the flag and
453
# setup code in case we are able to revisit this later.
@@ -465,6 +466,8 @@ def pplx_moe(
466
w2_chunk,
467
chunk_topk_weight,
468
chunk_topk_ids,
469
+ w1_scale=w1_scale_chunk,
470
+ w2_scale=w2_scale_chunk,
471
global_num_experts=num_experts)
472
473
if use_cudagraphs:
@@ -477,6 +480,8 @@ def pplx_moe(
480
w2_chunk,
481
chunk_topk_weight,
482
chunk_topk_ids,
483
+ w1_scale=w1_scale_chunk,
484
+ w2_scale=w2_scale_chunk,
485
global_num_experts=num_experts)
486
487
torch.cuda.synchronize()
@@ -505,9 +510,9 @@ def _batched_moe(pgi, dp_size, a, w1, w2, topk_weight, topk_ids):
510
rank=rank,
511
)
512
508
- experts = BatchedExperts(max_num_tokens=a.shape[0],
509
- world_size=1,
510
- dp_size=1)
513
+ experts = NaiveBatchedExperts(max_num_tokens=a.shape[0],
514
+ world_size=1,
515
+ dp_size=1)
516
517
fused_experts = FusedMoEModularKernel(
518
prepare_finalize,
@@ -539,7 +544,12 @@ def _pplx_moe(
544
w2: torch.Tensor,
545
score: torch.Tensor,
546
topk: int,
542
- use_internode: bool,
547
+ w1_s: Optional[torch.Tensor] = None,
548
+ w2_s: Optional[torch.Tensor] = None,
549
+ qtype: Optional[torch.dtype] = None,
550
+ per_act_token_quant: bool = False,
551
+ block_shape: Optional[list[int]] = None,
552
+ use_internode: bool = False,
553
):
554
if use_internode:
555
uid = nvshmem_get_unique_id(
@@ -557,11 +567,28 @@ def _pplx_moe(
567
568
moe_config = get_default_config(m, e, n, k, topk, a.dtype, False)
569
570
+ device = torch.device("cuda", pgi.rank)
571
+ a = a.to(device)
572
+ w1 = w1.to(device)
573
+ w2 = w2.to(device)
574
+ w1_s = w1_s.to(device) if w1_s is not None else None
575
+ w2_s = w2_s.to(device) if w2_s is not None else None
576
+
577
with set_current_vllm_config(vllm_config), override_config(moe_config):
578
topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
562
- torch_output = torch_experts(a, w1, w2, topk_weight, topk_ids)
579
+ torch_output = torch_experts(a,
580
+ w1,
581
+ w2,
582
+ topk_weight,
583
+ topk_ids,
584
+ w1_scale=w1_s,
585
+ w2_scale=w2_s,
586
+ quant_dtype=qtype,
587
+ per_act_token_quant=per_act_token_quant,
588
+ block_shape=block_shape)
589
pplx_output = pplx_moe(group_name, pgi.rank, pgi.world_size, dp_size,
564
- a, w1, w2, topk_weight, topk_ids)
590
+ a, w1, w2, topk_weight, topk_ids, w1_s, w2_s,
591
+ qtype, per_act_token_quant, block_shape)
592
# TODO (bnell): fix + re-enable
593
#batched_output = _batched_moe(pgi, dp_size, a, w1, w2, topk_weight,
594
# topk_ids)
@@ -581,6 +608,8 @@ def _pplx_moe(
608
@pytest.mark.parametrize("topk", TOP_KS)
609
@pytest.mark.parametrize("dtype", [torch.bfloat16])
610
@pytest.mark.parametrize("world_dp_size", [[2, 1]])
611
+@pytest.mark.parametrize("per_act_token_quant", [False, True])
612
+@pytest.mark.parametrize("block_shape", [None, [128, 128]])
613
@pytest.mark.parametrize("use_internode", [False])
614
@requires_pplx
615
def test_pplx_moe(
@@ -589,15 +618,33 @@ def test_pplx_moe(
618
topk: int,
619
dtype: torch.dtype,
620
world_dp_size: tuple[int, int],
621
+ per_act_token_quant: bool,
622
+ block_shape: Optional[list[int]],
623
use_internode: bool,
624
):
625
current_platform.seed_everything(7)
626
m, n, k = mnk
627
world_size, dp_size = world_dp_size
597
- a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
598
- w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
599
- w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
600
- score = torch.randn((m, e), device="cuda", dtype=dtype)
628
+
629
+ if dtype == torch.float8_e4m3fn:
630
+ use_fp8_w8a8 = True
631
+ quant_dtype = dtype
632
+ else:
633
+ use_fp8_w8a8 = False
634
+ quant_dtype = None
635
+
636
+ if not use_fp8_w8a8 and per_act_token_quant and block_shape is not None:
637
+ pytest.skip("Skip quantization test for non-quantized type")
638
+
639
+ a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) / 10
640
+ score = torch.randn((m, e), device="cuda", dtype=torch.bfloat16)
641
+
642
+ _, w1, w1_s, _, w2, w2_s = make_test_weights(e,
643
+ n,
644
+ k,
645
+ quant_dtype=quant_dtype,
646
+ block_shape=block_shape)
647
648
parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk,
649
+ w1_s, w2_s, quant_dtype, per_act_token_quant, block_shape,
650
use_internode)
tests/kernels/moe/utils.py
+242
-187
@@ -1,194 +1,249 @@
1
# SPDX-License-Identifier: Apache-2.0
2
-"""
3
-DeepEP test utilities
4
-"""
5
-import dataclasses
6
-import importlib
7
-import os
8
-import traceback
9
-from typing import Callable, Optional
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
5
import torch
12
-from torch.distributed import ProcessGroup
13
-from torch.multiprocessing import (
14
- spawn) # pyright: ignore[reportPrivateImportUsage]
15
-from typing_extensions import Concatenate, ParamSpec
16
-
17
-from vllm.utils import get_open_port
18
-
19
-has_deep_ep = importlib.util.find_spec("deep_ep") is not None
20
-if has_deep_ep:
21
- from vllm.model_executor.layers.fused_moe.deepep_ht_prepare_finalize import ( # noqa: E501
22
- DeepEPHTPrepareAndFinalize)
23
- from vllm.model_executor.layers.fused_moe.deepep_ll_prepare_finalize import ( # noqa: E501
24
- DeepEPLLPrepareAndFinalize)
25
-
26
-## Parallel Processes Utils
27
-
28
-P = ParamSpec("P")
29
-
30
-
31
-@dataclasses.dataclass
32
-class ProcessGroupInfo:
33
- world_size: int
34
- world_local_size: int
35
- rank: int
36
- node_rank: int
37
- local_rank: int
38
- device: torch.device
39
-
40
-
41
-def _worker_parallel_launch(
42
- local_rank: int,
43
- world_size: int,
44
- world_local_size: int,
45
- node_rank: int,
46
- init_method: str,
47
- worker: Callable[Concatenate[ProcessGroupInfo, P], None],
48
- *args: P.args,
49
- **kwargs: P.kwargs,
50
-) -> None:
51
- rank = node_rank * world_local_size + local_rank
52
- torch.cuda.set_device(local_rank)
53
- device = torch.device("cuda", local_rank)
54
- torch.distributed.init_process_group(
55
- backend="cpu:gloo,cuda:nccl",
56
- init_method=init_method,
57
- rank=rank,
58
- world_size=world_size,
59
- device_id=device,
60
- )
61
- barrier = torch.tensor([rank], device=device)
62
- torch.distributed.all_reduce(barrier)
63
-
64
- try:
65
- worker(
66
- ProcessGroupInfo(
67
- world_size=world_size,
68
- world_local_size=world_local_size,
69
- rank=rank,
70
- node_rank=node_rank,
71
- local_rank=local_rank,
72
- device=device,
73
- ),
74
- *args,
75
- **kwargs,
76
- )
77
- except Exception as ex:
78
- print(ex)
79
- traceback.print_exc()
80
- raise
81
- finally:
82
- torch.distributed.destroy_process_group()
83
-
84
-
85
-def parallel_launch(
86
- world_size: int,
87
- worker: Callable[Concatenate[ProcessGroupInfo, P], None],
88
- *args: P.args,
89
- **kwargs: P.kwargs,
90
-) -> None:
91
- assert not kwargs
92
- spawn(
93
- _worker_parallel_launch,
94
- args=(
95
- world_size,
96
- world_size,
97
- 0,
98
- f"tcp://{os.getenv('LOCALHOST', 'localhost')}:{get_open_port()}",
99
- worker,
100
- ) + args,
101
- nprocs=world_size,
102
- join=True,
103
- )
104
-
6
106
-## DeepEP specific utils
107
-
108
-
109
-@dataclasses.dataclass
110
-class DeepEPHTArgs:
111
- num_local_experts: int
112
-
113
-
114
-@dataclasses.dataclass
115
-class DeepEPLLArgs:
116
- max_tokens_per_rank: int
117
- hidden_size: int
118
- num_experts: int
119
- use_fp8_dispatch: bool
120
-
121
-
122
-def make_deepep_ht_a2a(pg: ProcessGroup,
123
- pgi: ProcessGroupInfo,
124
- dp_size: int,
125
- ht_args: DeepEPHTArgs,
126
- q_dtype: Optional[torch.dtype] = None,
127
- block_shape: Optional[list[int]] = None):
128
-
129
- import deep_ep
130
-
131
- # high throughput a2a
132
- num_nvl_bytes = 1024 * 1024 * 1024 # 1GB
133
- num_rdma_bytes, low_latency_mode, num_qps_per_rank = 0, False, 1
134
- buffer = deep_ep.Buffer(group=pg,
135
- num_nvl_bytes=num_nvl_bytes,
136
- num_rdma_bytes=num_rdma_bytes,
137
- low_latency_mode=low_latency_mode,
138
- num_qps_per_rank=num_qps_per_rank)
139
- return DeepEPHTPrepareAndFinalize(buffer=buffer,
140
- world_size=pgi.world_size,
141
- rank=pgi.rank,
142
- dp_size=dp_size,
143
- rank_expert_offset=pgi.rank *
144
- ht_args.num_local_experts,
145
- quant_dtype=q_dtype,
146
- block_shape=block_shape)
147
-
148
-
149
-def make_deepep_ll_a2a(pg: ProcessGroup,
150
- pgi: ProcessGroupInfo,
151
- dp_size: int,
152
- deepep_ll_args: DeepEPLLArgs,
153
- q_dtype: Optional[torch.dtype] = None,
154
- block_shape: Optional[list[int]] = None):
155
-
156
- import deep_ep
157
-
158
- # low-latency a2a
159
- num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
160
- deepep_ll_args.max_tokens_per_rank, deepep_ll_args.hidden_size,
161
- pgi.world_size, deepep_ll_args.num_experts)
162
-
163
- buffer = deep_ep.Buffer(group=pg,
164
- num_rdma_bytes=num_rdma_bytes,
165
- low_latency_mode=True,
166
- num_qps_per_rank=deepep_ll_args.num_experts //
167
- pgi.world_size)
168
-
169
- return DeepEPLLPrepareAndFinalize(
170
- buffer=buffer,
171
- world_size=pgi.world_size,
172
- dp_size=dp_size,
173
- max_tokens_per_rank=deepep_ll_args.max_tokens_per_rank,
174
- quant_dtype=q_dtype,
175
- block_shape=block_shape,
176
- use_fp8_dispatch=deepep_ll_args.use_fp8_dispatch,
7
+import vllm._custom_ops as ops
8
+from tests.kernels.quant_utils import (per_block_cast_to_fp8,
9
+ per_block_cast_to_int8)
10
+from vllm.model_executor.layers.fused_moe import fused_experts
11
+from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
12
+ BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
13
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
14
+ FusedMoEModularKernel)
15
+from vllm.model_executor.layers.fused_moe.utils import (
16
+ moe_kernel_quantize_input)
17
+from vllm.utils import round_up
18
+
19
+
20
+def triton_moe(
21
+ a: torch.Tensor,
22
+ w1: torch.Tensor,
23
+ w2: torch.Tensor,
24
+ topk_weight: torch.Tensor,
25
+ topk_ids: torch.Tensor,
26
+ w1_scale: Optional[torch.Tensor] = None,
27
+ w2_scale: Optional[torch.Tensor] = None,
28
+ a1_scale: Optional[torch.Tensor] = None,
29
+ a2_scale: Optional[torch.Tensor] = None,
30
+ quant_dtype: Optional[torch.dtype] = None,
31
+ per_act_token_quant=False,
32
+ block_shape: Optional[list[int]] = None,
33
+) -> torch.Tensor:
34
+ return fused_experts(a,
35
+ w1,
36
+ w2,
37
+ topk_weight,
38
+ topk_ids,
39
+ w1_scale=w1_scale,
40
+ w2_scale=w2_scale,
41
+ a1_scale=a1_scale,
42
+ a2_scale=a2_scale,
43
+ per_channel_quant=per_act_token_quant,
44
+ use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
45
+ block_shape=block_shape)
46
+
47
+
48
+def batched_moe(
49
+ a: torch.Tensor,
50
+ w1: torch.Tensor,
51
+ w2: torch.Tensor,
52
+ topk_weight: torch.Tensor,
53
+ topk_ids: torch.Tensor,
54
+ w1_scale: Optional[torch.Tensor] = None,
55
+ w2_scale: Optional[torch.Tensor] = None,
56
+ a1_scale: Optional[torch.Tensor] = None,
57
+ a2_scale: Optional[torch.Tensor] = None,
58
+ quant_dtype: Optional[torch.dtype] = None,
59
+ per_act_token_quant: bool = False,
60
+ block_shape: Optional[list[int]] = None,
61
+) -> torch.Tensor:
62
+ max_num_tokens = round_up(a.shape[0], 64)
63
+
64
+ fused_experts = FusedMoEModularKernel(
65
+ BatchedPrepareAndFinalize(max_num_tokens,
66
+ world_size=1,
67
+ dp_size=1,
68
+ rank=0),
69
+ BatchedTritonExperts(
70
+ max_num_tokens=max_num_tokens,
71
+ world_size=1,
72
+ dp_size=1,
73
+ use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
74
+ per_act_token_quant=per_act_token_quant,
75
+ block_shape=block_shape,
76
+ ),
77
)
78
79
+ return fused_experts(a,
80
+ w1,
81
+ w2,
82
+ topk_weight,
83
+ topk_ids,
84
+ w1_scale=w1_scale,
85
+ w2_scale=w2_scale,
86
+ a1_scale=a1_scale,
87
+ a2_scale=a2_scale)
88
+
89
+
90
+def naive_batched_moe(
91
+ a: torch.Tensor,
92
+ w1: torch.Tensor,
93
+ w2: torch.Tensor,
94
+ topk_weight: torch.Tensor,
95
+ topk_ids: torch.Tensor,
96
+ w1_scale: Optional[torch.Tensor] = None,
97
+ w2_scale: Optional[torch.Tensor] = None,
98
+ a1_scale: Optional[torch.Tensor] = None,
99
+ a2_scale: Optional[torch.Tensor] = None,
100
+ quant_dtype: Optional[torch.dtype] = None,
101
+ per_act_token_quant: bool = False,
102
+ block_shape: Optional[list[int]] = None,
103
+) -> torch.Tensor:
104
+ max_num_tokens = round_up(a.shape[0], 64)
105
+
106
+ fused_experts = FusedMoEModularKernel(
107
+ BatchedPrepareAndFinalize(max_num_tokens,
108
+ world_size=1,
109
+ dp_size=1,
110
+ rank=0),
111
+ NaiveBatchedExperts(
112
+ max_num_tokens=max_num_tokens,
113
+ dp_size=1,
114
+ world_size=1,
115
+ use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
116
+ per_act_token_quant=per_act_token_quant,
117
+ block_shape=block_shape,
118
+ ),
119
+ )
120
180
-def make_deepep_a2a(pg: ProcessGroup,
181
- pgi: ProcessGroupInfo,
182
- dp_size: int,
183
- deepep_ht_args: Optional[DeepEPHTArgs],
184
- deepep_ll_args: Optional[DeepEPLLArgs],
185
- q_dtype: Optional[torch.dtype] = None,
186
- block_shape: Optional[list[int]] = None):
187
- if deepep_ht_args is not None:
188
- assert deepep_ll_args is None
189
- return make_deepep_ht_a2a(pg, pgi, dp_size, deepep_ht_args, q_dtype,
190
- block_shape)
191
-
192
- assert deepep_ll_args is not None
193
- return make_deepep_ll_a2a(pg, pgi, dp_size, deepep_ll_args, q_dtype,
194
- block_shape)
121
+ return fused_experts(a,
122
+ w1,
123
+ w2,
124
+ topk_weight,
125
+ topk_ids,
126
+ w1_scale=w1_scale,
127
+ w2_scale=w2_scale,
128
+ a1_scale=a1_scale,
129
+ a2_scale=a2_scale)
130
+
131
+
132
+def chunk_scales(scales: Optional[torch.Tensor], start: int,
133
+ end: int) -> Optional[torch.Tensor]:
134
+ if scales is not None:
135
+ if scales.numel() == 1:
136
+ return scales
137
+ else:
138
+ return scales[start:end]
139
+ return None
140
+
141
+
142
+def make_quantized_test_activations(
143
+ E: int,
144
+ m: int,
145
+ k: int,
146
+ in_dtype: torch.dtype,
147
+ quant_dtype: Optional[torch.dtype] = None,
148
+ block_shape: Optional[list[int]] = None,
149
+ per_act_token_quant: bool = False,
150
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
151
+ a = torch.randn((E, m, k), device="cuda", dtype=in_dtype) / 10
152
+ a_q = a
153
+ a_scale = None
154
+
155
+ if quant_dtype is not None:
156
+ assert (quant_dtype == torch.float8_e4m3fn
157
+ or quant_dtype == torch.int8), "only fp8/int8 supported"
158
+ a_q = torch.zeros_like(a, dtype=quant_dtype)
159
+ a_scale_l = [None] * E
160
+ for e in range(E):
161
+ a_q[e], a_scale_l[e] = moe_kernel_quantize_input(
162
+ a[e], None, quant_dtype, per_act_token_quant, block_shape)
163
+ a_scale = torch.stack(a_scale_l)
164
+
165
+ if not per_act_token_quant and block_shape is None:
166
+ a_scale = a_scale.view(E, 1, 1)
167
+
168
+ return a, a_q, a_scale
169
+
170
+
171
+def moe_quantize_weights(
172
+ w: torch.Tensor,
173
+ w_s: Optional[torch.Tensor],
174
+ quant_dtype: Optional[torch.dtype],
175
+ per_token_quant: bool,
176
+ block_shape: Optional[list[int]],
177
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
178
+ assert (quant_dtype == torch.float8_e4m3fn
179
+ or quant_dtype == torch.int8), "only fp8/int8 supported"
180
+
181
+ if block_shape is not None:
182
+ assert not per_token_quant
183
+ if quant_dtype == torch.int8:
184
+ w, w_s = per_block_cast_to_int8(w, block_shape)
185
+ else:
186
+ w, w_s = per_block_cast_to_fp8(w, block_shape)
187
+ else:
188
+ if quant_dtype == torch.int8:
189
+ w, w_s = ops.scaled_int8_quant(
190
+ w, w_s, use_per_token_if_dynamic=per_token_quant)
191
+ else:
192
+ w, w_s = ops.scaled_fp8_quant(
193
+ w, w_s, use_per_token_if_dynamic=per_token_quant)
194
+
195
+ return w, w_s
196
+
197
+
198
+def make_test_weight(
199
+ e: int,
200
+ rows: int,
201
+ cols: int,
202
+ in_dtype: torch.dtype = torch.bfloat16,
203
+ quant_dtype: Optional[torch.dtype] = None,
204
+ block_shape: Optional[list[int]] = None,
205
+ per_act_token_quant: bool = False,
206
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
207
+ w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15
208
+
209
+ if quant_dtype is not None:
210
+ w_l = [None] * e
211
+ w_s_l = [None] * e
212
+ for idx in range(e):
213
+ w_l[idx], w_s_l[idx] = moe_quantize_weights(
214
+ w_16[idx], None, quant_dtype, per_act_token_quant, block_shape)
215
+
216
+ w = torch.stack(w_l)
217
+ w_s = torch.stack(w_s_l)
218
+ if w_s.ndim == 2:
219
+ assert w_s.shape[-1] == 1
220
+ w_s = w_s.view(-1, 1, 1)
221
+
222
+ if block_shape is not None:
223
+ block_n, block_k = block_shape
224
+ n_tiles = (rows + block_n - 1) // block_n
225
+ k_tiles = (cols + block_k - 1) // block_k
226
+ assert w_s.shape == (e, n_tiles, k_tiles)
227
+ else:
228
+ w = w_16
229
+ w_s = None
230
+
231
+ return w_16, w, w_s
232
+
233
+
234
+def make_test_weights(
235
+ e: int,
236
+ n: int,
237
+ k: int,
238
+ in_dtype: torch.dtype = torch.bfloat16,
239
+ quant_dtype: Optional[torch.dtype] = None,
240
+ block_shape: Optional[list[int]] = None,
241
+ per_act_token_quant: bool = False,
242
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor,
243
+ torch.Tensor, Optional[torch.Tensor]]:
244
+ return (
245
+ *make_test_weight(e, 2 * n, k, in_dtype, quant_dtype, block_shape,
246
+ per_act_token_quant),
247
+ *make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape,
248
+ per_act_token_quant),
249
+ )
tests/kernels/quant_utils.py
+166
-8
@@ -5,7 +5,10 @@ from typing import Optional, Union
5
6
import torch
7
8
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
9
+ group_broadcast)
10
from vllm.platforms import current_platform
11
+from vllm.utils import round_up
12
13
# Using the default value (240.0) from pytorch will cause accuracy
14
# issue on dynamic quantization models. Here use 224.0 for rocm.
@@ -94,9 +97,15 @@ def ref_dynamic_per_tensor_fp8_quant(x: torch.tensor) \
97
return ref_out, ref_scale.view((1, ))
98
99
97
-def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
98
- As: torch.Tensor, Bs: torch.Tensor, block_size,
99
- output_dtype):
100
+def native_w8a8_block_matmul(
101
+ A: torch.Tensor,
102
+ B: torch.Tensor,
103
+ As: torch.Tensor,
104
+ Bs: torch.Tensor,
105
+ block_size: list[int],
106
+ output_dtype: torch.dtype,
107
+ compute_type: torch.dtype = torch.float32,
108
+) -> torch.Tensor:
109
"""This function performs matrix multiplication with block-wise
110
quantization using native torch.
111
It is agnostic to the input data type and can be used for both int8 and
@@ -106,8 +115,8 @@ def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
115
`Bs` (float32).
116
The output is returned in the specified `output_dtype`.
117
"""
109
- A = A.to(torch.float32)
110
- B = B.to(torch.float32)
118
+ A = A.to(compute_type)
119
+ B = B.to(compute_type)
120
assert A.shape[-1] == B.shape[-1]
121
assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2
122
assert len(block_size) == 2
@@ -122,11 +131,11 @@ def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
131
As = As.reshape(M, As.shape[-1])
132
n_tiles = (N + block_n - 1) // block_n
133
k_tiles = (K + block_k - 1) // block_k
125
- assert n_tiles == Bs.shape[0]
126
- assert k_tiles == Bs.shape[1]
134
+ assert n_tiles == Bs.shape[0], f"{n_tiles} == {Bs.shape[0]}"
135
+ assert k_tiles == Bs.shape[1], f"{k_tiles} == {Bs.shape[1]}"
136
137
C_shape = (M, N)
129
- C = torch.zeros(C_shape, dtype=torch.float32, device=A.device)
138
+ C = torch.zeros(C_shape, dtype=compute_type, device=A.device)
139
140
A_tiles = [
141
A[:, i * block_k:min((i + 1) * block_k, K)] for i in range(k_tiles)
@@ -152,3 +161,152 @@ def native_w8a8_block_matmul(A: torch.Tensor, B: torch.Tensor,
161
162
C = C.reshape(origin_C_shape).to(output_dtype)
163
return C
164
+
165
+
166
+def native_per_token_group_quant_fp8(x,
167
+ group_size,
168
+ eps=1e-10,
169
+ dtype=torch.float8_e4m3fn):
170
+ """Function to perform per-token-group quantization on an input tensor
171
+ `x` using native torch."""
172
+ assert x.shape[-1] % group_size == 0, ("the last dimension of `x` must "
173
+ "be divisible by `group_size`")
174
+ assert x.is_contiguous(), "`x` is not contiguous"
175
+
176
+ finfo = torch.finfo(dtype)
177
+ fp8_min = finfo.min
178
+ fp8_max = finfo.max
179
+
180
+ x_ = x.reshape(x.numel() // group_size, group_size)
181
+ amax = x_.abs().max(dim=-1,
182
+ keepdim=True)[0].clamp(min=eps).to(torch.float32)
183
+ x_s = amax / fp8_max
184
+ x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
185
+ x_q = x_q.reshape(x.shape)
186
+ x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))
187
+
188
+ return x_q, x_s
189
+
190
+
191
+def native_per_token_group_quant_int8(x,
192
+ group_size,
193
+ eps=1e-10,
194
+ dtype=torch.int8):
195
+ """Function to perform per-token-group quantization on an input tensor
196
+ `x` using native torch.
197
+
198
+ It converts the tensor values into int8 values and returns the
199
+ quantized tensor along with the scaling factor used for quantization.
200
+ """
201
+ assert (x.shape[-1] % group_size == 0
202
+ ), "the last dimension of `x` must be divisible by `group_size`"
203
+ assert x.is_contiguous(), "`x` is not contiguous"
204
+
205
+ iinfo = torch.iinfo(dtype)
206
+ int8_min = iinfo.min
207
+ int8_max = iinfo.max
208
+
209
+ x_ = x.reshape(x.numel() // group_size, group_size)
210
+ # Use float32 for scale calculation for stability
211
+ amax = x_.abs().max(dim=-1,
212
+ keepdim=True)[0].clamp(min=eps).to(torch.float32)
213
+ x_s = amax / int8_max
214
+ x_q = (x_.to(torch.float32) / x_s).round().clamp(
215
+ min=int8_min, max=int8_max).to(dtype) # Round before clamping
216
+ x_q = x_q.reshape(x.shape)
217
+ x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))
218
+
219
+ return x_q, x_s
220
+
221
+
222
+DEFAULT_BLOCK_SHAPE = [128, 128]
223
+
224
+
225
+def per_block_cast_to_fp8(
226
+ x: torch.Tensor,
227
+ block_shape: list[int] = DEFAULT_BLOCK_SHAPE,
228
+) -> tuple[torch.Tensor, torch.Tensor]:
229
+ block_m, block_n = block_shape
230
+ assert x.dim() == 2
231
+ m, n = x.shape
232
+ x_padded = torch.zeros((round_up(m, block_m), round_up(n, block_n)),
233
+ dtype=x.dtype,
234
+ device=x.device)
235
+ x_padded[:m, :n] = x
236
+ x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
237
+ x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
238
+ x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
239
+ x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
240
+ scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
241
+ return x_scaled_sub, scales
242
+
243
+
244
+def per_block_cast_to_int8(
245
+ x: torch.Tensor,
246
+ block_shape: list[int] = DEFAULT_BLOCK_SHAPE,
247
+) -> tuple[torch.Tensor, torch.Tensor]:
248
+ block_m, block_n = block_shape
249
+ assert x.dim() == 2
250
+ m, n = x.shape
251
+ x_padded = torch.zeros((round_up(m, block_m), round_up(n, block_n)),
252
+ dtype=x.dtype,
253
+ device=x.device)
254
+ x_padded[:m, :n] = x
255
+ x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
256
+ x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
257
+ x_scaled = (x_view * (256.0 / x_amax)).to(torch.int8)
258
+ x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
259
+ scales = (x_amax / 256.0).view(x_view.size(0), x_view.size(2))
260
+ return x_scaled_sub, scales
261
+
262
+
263
+def dequant(
264
+ t: torch.Tensor,
265
+ scale: Optional[torch.Tensor],
266
+ block_shape: Optional[list[int]],
267
+ per_act_token_quant: bool,
268
+ out_dtype: Optional[torch.dtype] = torch.float32,
269
+) -> torch.Tensor:
270
+ if scale is not None:
271
+ f32 = torch.float32
272
+ if per_act_token_quant or block_shape is None:
273
+ return (t.to(f32) * scale).to(out_dtype)
274
+ else:
275
+ return (t.to(f32) * group_broadcast(scale, t.shape)).to(out_dtype)
276
+ else:
277
+ return t.to(out_dtype)
278
+
279
+
280
+def native_batched_masked_quant_matmul(
281
+ A: torch.Tensor,
282
+ B: torch.Tensor,
283
+ C: torch.Tensor,
284
+ num_expert_tokens: torch.Tensor,
285
+ A_scale: Optional[torch.Tensor] = None,
286
+ B_scale: Optional[torch.Tensor] = None,
287
+ block_shape: Optional[list[int]] = None,
288
+ per_act_token_quant: bool = False,
289
+) -> torch.Tensor:
290
+ num_expert_tokens_cpu = num_expert_tokens.clone()
291
+ num_expert_tokens_cpu = num_expert_tokens_cpu.to(device="cpu")
292
+ num_experts = num_expert_tokens.size(0)
293
+
294
+ for e in range(num_experts):
295
+ num_tokens = num_expert_tokens_cpu[e]
296
+ if A.dtype.itemsize == 1 and block_shape is not None:
297
+ assert A_scale is not None and B_scale is not None
298
+ tmp = native_w8a8_block_matmul(A[e], B[e], A_scale[e], B_scale[e],
299
+ block_shape, C.dtype)
300
+ C[e, :num_tokens, :] = tmp[:num_tokens, :]
301
+ elif A.dtype.itemsize == 1 and block_shape is None:
302
+ assert A_scale is not None and B_scale is not None
303
+ A_dq = dequant(A[e], A_scale[e], block_shape, per_act_token_quant)
304
+ B_dq = dequant(B[e], B_scale[e], block_shape, per_act_token_quant)
305
+ C[e, :num_tokens, :] = (
306
+ A_dq[:num_tokens] @ B_dq.transpose(0, 1)).to(C.dtype)
307
+ else:
308
+ assert A_scale is None
309
+ assert B_scale is None
310
+ C[e, :num_tokens, :] = A[e, :num_tokens, :] @ B[e].transpose(0, 1)
311
+
312
+ return C
tests/kernels/quantization/test_block_fp8.py
+4
-367
@@ -7,16 +7,10 @@ import itertools
7
import pytest
8
import torch
9
10
-from tests.kernels.quant_utils import native_w8a8_block_matmul
11
-from vllm.config import VllmConfig, set_current_vllm_config
12
-from vllm.model_executor.layers.activation import SiluAndMul
13
-from vllm.model_executor.layers.fused_moe import fused_moe
14
-from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
15
- _valid_deep_gemm_shape, deep_gemm_moe_fp8)
16
-from vllm.model_executor.layers.fused_moe.fused_moe import (
17
- fused_topk, modular_triton_fused_moe)
18
-from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
19
- moe_align_block_size)
10
+from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
11
+ native_w8a8_block_matmul,
12
+ per_block_cast_to_fp8)
13
+from vllm.config import VllmConfig
14
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
15
per_token_group_quant_fp8, w8a8_block_fp8_matmul)
16
from vllm.platforms import current_platform
@@ -46,78 +40,10 @@ N = [128, 512, 7168, 7748, 13824]
40
K = [256, 3884, 4096, 13824, 16384]
41
# Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8
42
# and its hidden size is 7168.
49
-M_moe = [1, 2, 7, 83, 128, 2048, 1024 * 128]
50
-M_moe_dg = [128, 192, 1335, 2048]
51
-N_moe = [128, 256, 1024, 4608] # [13824]
52
-K_moe = [256, 512, 7168] # [13824]
43
BLOCK_SIZE = [[128, 128]]
54
-E = [2, 8, 16, 24] # [128, 256]
55
-TOP_KS = [1, 2, 6]
44
OUT_DTYPES = [torch.bfloat16] # [torch.float32, torch.half, torch.bfloat16]
45
SEEDS = [0]
46
59
-
60
-def native_per_token_group_quant_fp8(x,
61
- group_size,
62
- eps=1e-10,
63
- dtype=torch.float8_e4m3fn):
64
- """Function to perform per-token-group quantization on an input tensor
65
- `x` using native torch."""
66
- assert x.shape[-1] % group_size == 0, ("the last dimension of `x` cannot "
67
- "be divisible by `group_size`")
68
- assert x.is_contiguous(), "`x` is not contiguous"
69
-
70
- finfo = torch.finfo(dtype)
71
- fp8_min = finfo.min
72
- fp8_max = finfo.max
73
-
74
- x_ = x.reshape(x.numel() // group_size, group_size)
75
- amax = x_.abs().max(dim=-1,
76
- keepdim=True)[0].clamp(min=eps).to(torch.float32)
77
- x_s = amax / fp8_max
78
- x_q = (x_ / x_s).clamp(min=fp8_min, max=fp8_max).to(dtype)
79
- x_q = x_q.reshape(x.shape)
80
- x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))
81
-
82
- return x_q, x_s
83
-
84
-
85
-def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
86
- """Fused moe with block-wise quantization using native torch."""
87
- B, D = a.shape
88
- a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
89
- out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
90
- score = torch.softmax(score, dim=-1, dtype=torch.float32)
91
- topk_weight, topk_ids = torch.topk(score, topk)
92
- topk_weight = topk_weight.view(-1)
93
- topk_ids = topk_ids.view(-1)
94
-
95
- _, block_k = block_shape[0], block_shape[1]
96
- a_q, a_s = native_per_token_group_quant_fp8(a, block_k)
97
- a_q = a_q.to(torch.float32)
98
- for i in range(w1.shape[0]):
99
- mask = topk_ids == i
100
- if mask.sum():
101
- inter_out = native_w8a8_block_matmul(a_q[mask],
102
- w1[i],
103
- a_s[mask],
104
- w1_s[i],
105
- block_shape,
106
- output_dtype=a.dtype)
107
- act_out = SiluAndMul().forward_native(inter_out)
108
- act_out_q, act_out_s = native_per_token_group_quant_fp8(
109
- act_out, block_k)
110
- act_out = act_out.to(torch.float32)
111
- out[mask] = native_w8a8_block_matmul(act_out_q,
112
- w2[i],
113
- act_out_s,
114
- w2_s[i],
115
- block_shape,
116
- output_dtype=a.dtype)
117
- return (out.view(B, -1, w2.shape[1]) *
118
- topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
119
-
120
-
47
# Skip all tests if CUDA is not available
48
pytest.importorskip("torch.cuda")
49
@@ -177,111 +103,6 @@ def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
103
assert rel_diff < 0.001
104
105
180
-@pytest.mark.parametrize(
181
- "M,N,K,E,topk,block_size,dtype,seed",
182
- itertools.product(M_moe, N_moe, K_moe, E, TOP_KS, BLOCK_SIZE, DTYPES,
183
- SEEDS))
184
-@torch.inference_mode()
185
-def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
186
- if topk > E:
187
- pytest.skip(f"Skipping test; topk={topk} > E={E}")
188
-
189
- torch.manual_seed(seed)
190
- factor_for_scale = 1e-2
191
- fp8_info = torch.finfo(torch.float8_e4m3fn)
192
- fp8_max, fp8_min = fp8_info.max, fp8_info.min
193
-
194
- a = torch.randn((M, K), dtype=dtype) / 10
195
-
196
- w1_bf16 = (torch.rand(
197
- (E, 2 * N, K), dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
198
- w1 = w1_bf16.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
199
- del w1_bf16
200
-
201
- w2_bf16 = (torch.rand((E, K, N), dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
202
- w2 = w2_bf16.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
203
- del w2_bf16
204
-
205
- block_n, block_k = block_size[0], block_size[1]
206
- n_tiles_w1 = (2 * N + block_n - 1) // block_n
207
- n_tiles_w2 = (K + block_n - 1) // block_n
208
- k_tiles_w1 = (K + block_k - 1) // block_k
209
- k_tiles_w2 = (N + block_k - 1) // block_k
210
-
211
- w1_s = torch.rand(
212
- (E, n_tiles_w1, k_tiles_w1), dtype=torch.float32) * factor_for_scale
213
- w2_s = torch.rand(
214
- (E, n_tiles_w2, k_tiles_w2), dtype=torch.float32) * factor_for_scale
215
-
216
- score = torch.randn((M, E), dtype=dtype)
217
-
218
- m_fused_moe = modular_triton_fused_moe(use_fp8_w8a8=True,
219
- use_int8_w8a8=False,
220
- use_int8_w8a16=False,
221
- use_int4_w4a16=False,
222
- per_channel_quant=False,
223
- block_shape=block_size)
224
-
225
- # Set the context to avoid lots of warning spam.
226
- with set_current_vllm_config(vllm_config):
227
- out = fused_moe(
228
- a,
229
- w1,
230
- w2,
231
- score,
232
- topk,
233
- renormalize=False,
234
- use_fp8_w8a8=True,
235
- w1_scale=w1_s,
236
- w2_scale=w2_s,
237
- block_shape=block_size,
238
- )
239
- ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score, topk,
240
- block_size)
241
-
242
- topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
243
- m_out = m_fused_moe(a,
244
- w1,
245
- w2,
246
- topk_weights,
247
- topk_ids,
248
- global_num_experts=E,
249
- w1_scale=w1_s,
250
- w2_scale=w2_s)
251
-
252
- #print(f"{out.sum()=}")
253
- #print(f"{ref_out.sum()=}")
254
-
255
- rel_diff = (torch.mean(
256
- torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
257
- torch.mean(torch.abs(ref_out.to(torch.float32))))
258
- assert rel_diff < 0.03
259
-
260
- rel_diff = (torch.mean(
261
- torch.abs(m_out.to(torch.float32) - ref_out.to(torch.float32))) /
262
- torch.mean(torch.abs(ref_out.to(torch.float32))))
263
- assert rel_diff < 0.03
264
-
265
-
266
-def per_block_cast_to_fp8(
267
- x: torch.Tensor,
268
- block_size_n: int = 128) -> tuple[torch.Tensor, torch.Tensor]:
269
- assert x.dim() == 2
270
- m, n = x.shape
271
- x_padded = torch.zeros(
272
- (deep_gemm.ceil_div(m, 128) * 128,
273
- deep_gemm.ceil_div(n, block_size_n) * block_size_n),
274
- dtype=x.dtype,
275
- device=x.device)
276
- x_padded[:m, :n] = x
277
- x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, block_size_n)
278
- x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
279
- x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn)
280
- x_scaled_sub = x_scaled.view_as(x_padded)[:m, :n].contiguous()
281
- scales = (x_amax / 448.0).view(x_view.size(0), x_view.size(2))
282
- return x_scaled_sub, scales
283
-
284
-
106
@pytest.mark.parametrize(
107
"M,N,K,block_size,out_dtype,seed",
108
itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))
@@ -324,187 +145,3 @@ def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
145
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
146
torch.mean(torch.abs(ref_out.to(torch.float32))))
147
assert rel_diff < 0.001
327
-
328
-
329
-def fp8_perm(m, idx):
330
- if torch.is_floating_point(m) and torch.finfo(m.dtype).bits == 8:
331
- return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
332
- else:
333
- return m[idx, ...]
334
-
335
-
336
-def _moe_permute(a, a_s, topk_ids, num_groups, topk, block_m):
337
- M, K = a.shape
338
-
339
- sorted_token_ids, m_indices, num_pad = moe_align_block_size(
340
- topk_ids, block_m, num_groups, None, pad_sorted_ids=True)
341
-
342
- num_tokens = topk * M
343
-
344
- sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
345
- m_indices = torch.repeat_interleave(m_indices, block_m, dim=0)
346
- inv_perm = torch.argsort(sorted_token_ids)[:M * topk]
347
-
348
- a = fp8_perm(a, sorted_token_ids // topk)
349
- if a_s is not None:
350
- a_s = a_s[sorted_token_ids // topk]
351
-
352
- return a, a_s, m_indices, inv_perm
353
-
354
-
355
-def _moe_unpermute(out, inv_perm, topk, K, topk_weight):
356
- M = topk_weight.shape[0]
357
- out = out[inv_perm, ...]
358
- tmp_out = out.view(-1, topk, K)
359
- return (tmp_out * topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
360
-
361
-
362
-def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, score, topk,
363
- block_shape):
364
- """Fused moe with block-wise quantization using DeepGemm grouped gemm."""
365
- num_groups = w1.shape[0]
366
- M, K = a.shape
367
- N = w2.shape[-1]
368
-
369
- topk_weight, topk_ids, token_expert_indices = fused_topk(
370
- a, score.float(), topk, False)
371
-
372
- block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
373
-
374
- _, block_k = block_shape[0], block_shape[1]
375
-
376
- a_q, a_s = per_token_group_quant_fp8(a, block_m)
377
-
378
- a_q, a_s, m_indices, inv_perm = _moe_permute(a_q, a_s, topk_ids,
379
- num_groups, topk, block_m)
380
-
381
- inter_out = torch.zeros((a_q.shape[0], N * 2),
382
- dtype=torch.bfloat16,
383
- device=a.device)
384
-
385
- deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous((a_q, a_s), (w1, w1_s),
386
- inter_out, m_indices)
387
-
388
- act_out = SiluAndMul().forward_native(inter_out)
389
- act_out_q, act_out_s = per_token_group_quant_fp8(act_out, block_k)
390
-
391
- out = torch.zeros(a_q.shape[0], K, dtype=torch.bfloat16, device=a.device)
392
-
393
- deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
394
- (act_out_q, act_out_s), (w2, w2_s), out, m_indices)
395
-
396
- final_out = _moe_unpermute(out, inv_perm, topk, K, topk_weight)
397
-
398
- return final_out
399
-
400
-
401
-@pytest.mark.parametrize(
402
- "M,N,K,E,topk,seed",
403
- itertools.product(M_moe_dg, N_moe, K_moe, E, TOP_KS, SEEDS))
404
-@pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
405
-@torch.inference_mode()
406
-def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
407
- monkeypatch):
408
- if topk > E:
409
- pytest.skip(f"Skipping test: topk={topk} > E={E}")
410
-
411
- if not _valid_deep_gemm_shape(M, N, K):
412
- pytest.skip(f"Skipping test: invalid size m={M}, n={N}, k={K}")
413
-
414
- chunk_size = 1024
415
-
416
- torch.manual_seed(seed)
417
-
418
- monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
419
-
420
- block_m = deep_gemm.get_m_alignment_for_contiguous_layout()
421
- block_size = [block_m, block_m]
422
- dtype = torch.bfloat16
423
-
424
- fp8_info = torch.finfo(torch.float8_e4m3fn)
425
- fp8_max, fp8_min = fp8_info.max, fp8_info.min
426
-
427
- a = torch.randn((M, K), dtype=dtype) / 10
428
-
429
- w1_bf16 = ((torch.rand((E, 2 * N, K), dtype=torch.bfloat16) - 0.5) * 2 *
430
- fp8_max).clamp(min=fp8_min, max=fp8_max)
431
-
432
- w2_bf16 = ((torch.rand((E, K, N), dtype=torch.bfloat16) - 0.5) * 2 *
433
- fp8_max).clamp(min=fp8_min, max=fp8_max)
434
-
435
- score = torch.randn((M, E), dtype=dtype)
436
-
437
- block_n, block_k = block_size[0], block_size[1]
438
- n_tiles_w1 = ((2 * N) + block_n - 1) // block_n
439
- k_tiles_w1 = (K + block_k - 1) // block_k
440
- n_tiles_w2 = (K + block_n - 1) // block_n
441
- k_tiles_w2 = (N + block_k - 1) // block_k
442
-
443
- w1 = torch.empty_like(w1_bf16, dtype=torch.float8_e4m3fn)
444
- w2 = torch.empty_like(w2_bf16, dtype=torch.float8_e4m3fn)
445
-
446
- w1_s = torch.empty((E, n_tiles_w1, k_tiles_w1), dtype=torch.float32)
447
- w2_s = torch.empty((E, n_tiles_w2, k_tiles_w2), dtype=torch.float32)
448
-
449
- w1_s = deep_gemm.get_col_major_tma_aligned_tensor(w1_s).contiguous()
450
- w2_s = deep_gemm.get_col_major_tma_aligned_tensor(w2_s).contiguous()
451
-
452
- assert w1_s.shape == (E, (2 * N + 127) // 128, (K + 127) // 128)
453
- assert (w2.shape[-2] + block_n - 1) // block_n == w2_s.shape[-2]
454
-
455
- for i in range(E):
456
- w1[i], w1_s[i] = per_block_cast_to_fp8(w1_bf16[i])
457
- w2[i], w2_s[i] = per_block_cast_to_fp8(w2_bf16[i])
458
-
459
- # Note: for now use_compile will error out if the problem size is
460
- # large enough to trigger chunking. I'm leaving the flag and
461
- # setup code in case we are able to revisit this later.
462
- use_compile = False
463
-
464
- use_cudagraph = (chunk_size < M and N >= 1024 and K >= 1024
465
- and current_platform.is_cuda_alike())
466
-
467
- # Set the context to avoid lots of warning spam.
468
- with set_current_vllm_config(vllm_config):
469
- if M >= 128:
470
- ref_out = deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s,
471
- score, topk, block_size)
472
- else:
473
- ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, score,
474
- topk, block_size)
475
-
476
- topk_weights, topk_ids, token_expert_indices = fused_topk(
477
- a, score.float(), topk, False)
478
-
479
- if use_compile:
480
- deep_gemm_moe_fp8_fn = torch.compile(deep_gemm_moe_fp8,
481
- backend="inductor",
482
- fullgraph=True)
483
- torch._dynamo.mark_dynamic(a, 0)
484
- torch._dynamo.mark_dynamic(topk_weights, 0)
485
- torch._dynamo.mark_dynamic(topk_ids, 0)
486
- else:
487
- deep_gemm_moe_fp8_fn = deep_gemm_moe_fp8
488
-
489
- out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights,
490
- topk_ids)
491
-
492
- if use_cudagraph:
493
- out.fill_(0)
494
- stream = torch.cuda.Stream()
495
- graph = torch.cuda.CUDAGraph()
496
- with torch.cuda.graph(graph, stream=stream):
497
- out = deep_gemm_moe_fp8_fn(a, w1, w2, w1_s, w2_s, topk_weights,
498
- topk_ids)
499
- torch.cuda.synchronize()
500
- graph.replay()
501
- torch.cuda.synchronize()
502
-
503
- #print(f"{out.sum()=}")
504
- #print(f"{ref_out.sum()=}")
505
-
506
- rel_diff = (torch.mean(
507
- torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
508
- torch.mean(torch.abs(ref_out.to(torch.float32))))
509
-
510
- assert rel_diff < 0.03
tests/kernels/quantization/test_block_int8.py
+1
-135
@@ -8,9 +8,7 @@ import pytest
8
import torch
9
10
from tests.kernels.quant_utils import native_w8a8_block_matmul
11
-from vllm.config import VllmConfig, set_current_vllm_config
12
-from vllm.model_executor.layers.activation import SiluAndMul
13
-from vllm.model_executor.layers.fused_moe import fused_moe
11
+from vllm.config import VllmConfig
12
from vllm.model_executor.layers.quantization.utils.int8_utils import (
13
w8a8_block_int8_matmul)
14
from vllm.platforms import current_platform
@@ -23,82 +21,10 @@ vllm_config = VllmConfig()
21
vllm_config.scheduler_config.max_num_seqs = 128
22
vllm_config.scheduler_config.max_model_len = 8192
23
26
-
27
-# For test
28
-def native_per_token_group_quant_int8(x,
29
- group_size,
30
- eps=1e-10,
31
- dtype=torch.int8):
32
- """Function to perform per-token-group quantization on an input tensor
33
- `x` using native torch.
34
-
35
- It converts the tensor values into int8 values and returns the
36
- quantized tensor along with the scaling factor used for quantization.
37
- """
38
- assert (x.shape[-1] % group_size == 0
39
- ), "the last dimension of `x` cannot be divisible by `group_size`"
40
- assert x.is_contiguous(), "`x` is not contiguous"
41
-
42
- iinfo = torch.iinfo(dtype)
43
- int8_min = iinfo.min
44
- int8_max = iinfo.max
45
-
46
- x_ = x.reshape(x.numel() // group_size, group_size)
47
- # Use float32 for scale calculation for stability
48
- amax = x_.abs().max(dim=-1,
49
- keepdim=True)[0].clamp(min=eps).to(torch.float32)
50
- x_s = amax / int8_max
51
- x_q = (x_.to(torch.float32) / x_s).round().clamp(
52
- min=int8_min, max=int8_max).to(dtype) # Round before clamping
53
- x_q = x_q.reshape(x.shape)
54
- x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size, ))
55
-
56
- return x_q, x_s
57
-
58
-
59
-# For test
60
-def torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk, block_shape):
61
- """This function performs fused moe with block-wise quantization using
62
- native torch."""
63
- B, D = a.shape
64
- a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
65
- out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
66
- score = torch.softmax(score, dim=-1, dtype=torch.float32)
67
- topk_weight, topk_ids = torch.topk(score, topk)
68
- topk_weight = topk_weight.view(-1)
69
- topk_ids = topk_ids.view(-1)
70
-
71
- _, block_k = block_shape[0], block_shape[1]
72
- a_q, a_s = native_per_token_group_quant_int8(a, block_k)
73
- for i in range(w1.shape[0]):
74
- mask = topk_ids == i
75
- if mask.sum():
76
- inter_out = native_w8a8_block_matmul(a_q[mask],
77
- w1[i],
78
- a_s[mask],
79
- w1_s[i],
80
- block_shape,
81
- output_dtype=a.dtype)
82
- act_out = SiluAndMul().forward_native(inter_out)
83
- act_out_q, act_out_s = native_per_token_group_quant_int8(
84
- act_out, block_k)
85
- act_out = act_out.to(torch.float32)
86
- out[mask] = native_w8a8_block_matmul(act_out_q,
87
- w2[i],
88
- act_out_s,
89
- w2_s[i],
90
- block_shape,
91
- output_dtype=a.dtype)
92
- return (out.view(B, -1, w2.shape[1]) *
93
- topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
94
-
95
-
24
DTYPES = [torch.half, torch.bfloat16]
25
M = [1, 33, 64, 222]
26
N = [128, 1024]
27
K = [256, 4096]
100
-E = [8, 24]
101
-TOP_KS = [2, 6]
28
# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
29
BLOCK_SIZE = [[128, 128]]
30
SEEDS = [0]
@@ -140,63 +66,3 @@ def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed):
66
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
67
torch.mean(torch.abs(ref_out.to(torch.float32))))
68
assert rel_diff < 0.001
143
-
144
-
145
-@pytest.mark.parametrize(
146
- "M, N, K, E, topk, block_size, dtype, seed",
147
- itertools.product(M, N, K, E, TOP_KS, BLOCK_SIZE, DTYPES, SEEDS))
148
-@torch.inference_mode()
149
-def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
150
- """Tests the fused_moe kernel with W8A8 INT8 block quantization against a
151
- native torch reference."""
152
- torch.manual_seed(seed)
153
- # Use a smaller factor for scale initialization to prevent large
154
- # values/overflow especially when output dtype might be float16
155
- factor_for_scale = 1e-2
156
- int8_info = torch.iinfo(torch.int8)
157
- int8_max, int8_min = int8_info.max, int8_info.min
158
-
159
- a = torch.randn((M, K), dtype=dtype) / 10
160
-
161
- w1_fp32 = (torch.rand(
162
- (E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 * int8_max
163
- w1 = w1_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
164
-
165
- w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 * int8_max
166
- w2 = w2_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
167
-
168
- block_n, block_k = block_size[0], block_size[1]
169
- n_tiles_w1 = (2 * N + block_n - 1) // block_n
170
- n_tiles_w2 = (K + block_n - 1) // block_n
171
- k_tiles_w1 = (K + block_k - 1) // block_k
172
- k_tiles_w2 = (N + block_k - 1) // block_k
173
-
174
- w1_s = (torch.rand(
175
- (E, n_tiles_w1, k_tiles_w1), dtype=torch.float32) * factor_for_scale)
176
- w2_s = (torch.rand(
177
- (E, n_tiles_w2, k_tiles_w2), dtype=torch.float32) * factor_for_scale)
178
-
179
- score = torch.randn((M, E), dtype=dtype)
180
-
181
- # Set the context to avoid lots of warning spam.
182
- with set_current_vllm_config(vllm_config):
183
- out = fused_moe(
184
- a,
185
- w1,
186
- w2,
187
- score,
188
- topk,
189
- renormalize=False,
190
- use_int8_w8a8=True,
191
- w1_scale=w1_s,
192
- w2_scale=w2_s,
193
- block_shape=block_size,
194
- )
195
- ref_out = torch_w8a8_block_int8_moe(a, w1, w2, w1_s, w2_s, score, topk,
196
- block_size)
197
-
198
- # Check results
199
- rel_diff = (torch.mean(
200
- torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))) /
201
- torch.mean(torch.abs(ref_out.to(torch.float32))))
202
- assert rel_diff < 0.06
tests/kernels/utils.py
+64
-16
@@ -13,8 +13,11 @@ import pytest
13
import torch
14
from torch._prims_common import TensorLikeType
15
16
+from tests.kernels.quant_utils import native_w8a8_block_matmul
17
from vllm.attention import AttentionBackend, AttentionMetadata, AttentionType
18
from vllm.model_executor.layers.activation import SiluAndMul
19
+from vllm.model_executor.layers.fused_moe.utils import (
20
+ moe_kernel_quantize_input)
21
from vllm.platforms.interface import _Backend
22
from vllm.utils import (STR_BACKEND_ENV_VAR, STR_FLASH_ATTN_VAL,
23
STR_XFORMERS_ATTN_VAL, make_tensor_with_pad)
@@ -1054,32 +1057,77 @@ def compute_max_diff(output, output_ref):
1057
torch.abs(output_ref))
1058
1059
1057
-def torch_experts(a: torch.Tensor,
1058
- w1: torch.Tensor,
1059
- w2: torch.Tensor,
1060
- topk_weight: torch.Tensor,
1061
- topk_ids: torch.Tensor,
1062
- global_num_experts: int = -1,
1063
- expert_map: Optional[torch.Tensor] = None) -> torch.Tensor:
1060
+def torch_experts(
1061
+ a: torch.Tensor,
1062
+ w1: torch.Tensor,
1063
+ w2: torch.Tensor,
1064
+ topk_weight: torch.Tensor,
1065
+ topk_ids: torch.Tensor,
1066
+ global_num_experts: int = -1,
1067
+ expert_map: Optional[torch.Tensor] = None,
1068
+ w1_scale: Optional[torch.Tensor] = None,
1069
+ w2_scale: Optional[torch.Tensor] = None,
1070
+ a1_scale: Optional[torch.Tensor] = None,
1071
+ a2_scale: Optional[torch.Tensor] = None,
1072
+ quant_dtype: Optional[torch.dtype] = None,
1073
+ per_act_token_quant=False,
1074
+ block_shape: Optional[list[int]] = None,
1075
+) -> torch.Tensor:
1076
assert (global_num_experts == -1
1077
or (global_num_experts == w1.shape[0] and expert_map is None)
1078
or (expert_map is not None
1079
and global_num_experts == expert_map.shape[0]))
1080
+
1081
+ M, K = a.shape
1082
topk = topk_ids.shape[1]
1069
- B, D = a.shape
1070
- a = a.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
1071
- out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
1072
- topk_weight = topk_weight.view(-1)
1083
+
1084
+ a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
1085
+
1086
+ out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)
1087
+
1088
+ a, a_scale = moe_kernel_quantize_input(a, None, quant_dtype,
1089
+ per_act_token_quant, block_shape)
1090
+
1091
+ num_experts = w1.shape[0]
1092
+
1093
topk_ids = topk_ids.view(-1)
1094
if expert_map is not None:
1095
topk_ids = expert_map[topk_ids]
1076
- for i in range(w1.shape[0]):
1096
+
1097
+ for i in range(num_experts):
1098
mask = topk_ids == i
1099
if mask.sum():
1079
- out[mask] = SiluAndMul()(
1080
- a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
1081
- return (out.view(B, -1, w2.shape[1]) *
1082
- topk_weight.view(B, -1, 1).to(out.dtype)).sum(dim=1)
1100
+ if quant_dtype is None:
1101
+ tmp1 = a[mask] @ w1[i].transpose(0, 1)
1102
+ tmp2 = SiluAndMul()(tmp1)
1103
+ out[mask] = tmp2 @ w2[i].transpose(0, 1)
1104
+ elif block_shape is not None:
1105
+ assert (a_scale is not None and w1_scale is not None
1106
+ and w2_scale is not None)
1107
+ tmp1 = native_w8a8_block_matmul(a[mask], w1[i], a_scale[mask],
1108
+ w1_scale[i], block_shape,
1109
+ out.dtype)
1110
+ tmp2 = SiluAndMul()(tmp1)
1111
+ tmp2, b_scale = moe_kernel_quantize_input(
1112
+ tmp2, None, quant_dtype, per_act_token_quant, block_shape)
1113
+
1114
+ out[mask] = native_w8a8_block_matmul(tmp2, w2[i], b_scale,
1115
+ w2_scale[i], block_shape,
1116
+ out.dtype)
1117
+ else:
1118
+ assert (a_scale is not None and w1_scale is not None
1119
+ and w2_scale is not None)
1120
+ f32 = torch.float32
1121
+ scales = a_scale if a_scale.numel() == 1 else a_scale[mask]
1122
+ tmp1 = a[mask].to(f32) * scales
1123
+ w1_dq = (w1[i].to(f32) * w1_scale[i]).transpose(0, 1)
1124
+ tmp1 = tmp1 @ w1_dq
1125
+ tmp2 = SiluAndMul()(tmp1)
1126
+ w2_dq = (w2[i].to(f32) * w2_scale[i]).transpose(0, 1)
1127
+ out[mask] = (tmp2 @ w2_dq).to(out.dtype)
1128
+
1129
+ return (out.view(M, -1, w2.shape[1]) *
1130
+ topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
1131
1132
1133
def torch_moe(a: torch.Tensor,
vllm/_custom_ops.py
+1
-1
@@ -1274,7 +1274,7 @@ def scaled_fp8_quant(
1274
scale = torch.zeros(1, device=input.device, dtype=torch.float32)
1275
torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
1276
else:
1277
- assert scale.numel() == 1
1277
+ assert scale.numel() == 1, f"{scale.shape}"
1278
torch.ops._C.static_scaled_fp8_quant(output, input, scale)
1279
1280
return output, scale
vllm/model_executor/layers/fused_moe/__init__.py
+25
-1
@@ -4,8 +4,12 @@
4
from contextlib import contextmanager
5
from typing import Any, Optional
6
7
+from vllm.model_executor.layers.fused_moe.config import FusedMoEConfig
8
from vllm.model_executor.layers.fused_moe.layer import (
9
FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
10
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
11
+ FusedMoEActivationFormat, FusedMoEPermuteExpertsUnpermute,
12
+ FusedMoEPrepareAndFinalize)
13
from vllm.triton_utils import HAS_TRITON
14
15
_config: Optional[dict[str, Any]] = None
@@ -26,8 +30,12 @@ def get_config() -> Optional[dict[str, Any]]:
30
31
__all__ = [
32
"FusedMoE",
33
+ "FusedMoEConfig",
34
"FusedMoEMethodBase",
35
"FusedMoeWeightScaleSupported",
36
+ "FusedMoEPermuteExpertsUnpermute",
37
+ "FusedMoEActivationFormat",
38
+ "FusedMoEPrepareAndFinalize",
39
"override_config",
40
"get_config",
41
]
@@ -36,11 +44,21 @@ if HAS_TRITON:
44
# import to register the custom ops
45
import vllm.model_executor.layers.fused_moe.fused_marlin_moe # noqa
46
import vllm.model_executor.layers.fused_moe.fused_moe # noqa
47
+ from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
48
+ BatchedDeepGemmExperts)
49
+ from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
50
+ BatchedTritonOrDeepGemmExperts)
51
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
40
- cutlass_moe_fp4, cutlass_moe_fp8)
52
+ CutlassExpertsFp8, cutlass_moe_fp4, cutlass_moe_fp8)
53
+ from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
54
+ DeepGemmExperts)
55
+ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
56
+ BatchedTritonExperts)
57
from vllm.model_executor.layers.fused_moe.fused_moe import (
58
TritonExperts, fused_experts, fused_moe, fused_topk,
59
get_config_file_name, grouped_topk)
60
+ from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
61
+ TritonOrDeepGemmExperts)
62
63
__all__ += [
64
"fused_moe",
@@ -50,5 +68,11 @@ if HAS_TRITON:
68
"grouped_topk",
69
"cutlass_moe_fp8",
70
"cutlass_moe_fp4",
71
+ "CutlassExpertsFp8",
72
"TritonExperts",
73
+ "BatchedTritonExperts",
74
+ "DeepGemmExperts",
75
+ "BatchedDeepGemmExperts",
76
+ "TritonOrDeepGemmExperts",
77
+ "BatchedTritonOrDeepGemmExperts",
78
]
vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py
+26
-8
@@ -5,6 +5,7 @@ import torch
5
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
from vllm.logger import init_logger
8
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
10
from vllm.triton_utils import tl, triton
11
@@ -179,28 +180,44 @@ def silu_mul_fp8_quant_deep_gemm(
180
class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
181
182
# The Deep Gemm kernels only support block size of 128
182
- DEEPGEMM_BLOCK_SHAPE = 128
183
-
184
- def __init__(self, max_num_tokens: int, world_size: int, dp_size: int,
185
- block_shape: list[int]):
183
+ DEEPGEMM_BLOCK_SHAPE: list[int] = [128, 128]
184
+
185
+ def __init__(self,
186
+ max_num_tokens: int,
187
+ world_size: int,
188
+ dp_size: int,
189
+ block_shape: list[int],
190
+ per_act_token_quant=False):
191
"""
192
max_num_tokens: Maximum number of tokens from a DP Rank
193
world_size: Number of EP ranks
194
dp_size: Number of data-parallel ranks
195
block_shape: Block quantization block shape
196
"""
192
- super().__init__()
197
+ super().__init__(
198
+ FusedMoEQuantConfig(
199
+ quant_dtype=torch.float8_e4m3fn,
200
+ per_act_token_quant=per_act_token_quant,
201
+ block_shape=block_shape,
202
+ ))
203
+ assert self.block_shape == self.DEEPGEMM_BLOCK_SHAPE
204
self.max_num_tokens = max_num_tokens
205
self.world_size = world_size
206
self.dp_size = dp_size
196
- self.block_shape = block_shape
207
198
- assert (len(self.block_shape) == 2 and all(
199
- [v == self.DEEPGEMM_BLOCK_SHAPE for v in self.block_shape]))
208
+ @property
209
+ def activation_formats(
210
+ self
211
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
212
+ return (mk.FusedMoEActivationFormat.BatchedExperts,
213
+ mk.FusedMoEActivationFormat.BatchedExperts)
214
215
def supports_chunking(self) -> bool:
216
return False
217
218
+ def supports_expert_map(self) -> bool:
219
+ return False
220
+
221
def workspace_shapes(
222
self,
223
a: torch.Tensor,
@@ -248,6 +265,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
265
):
266
import deep_gemm as dg
267
assert hidden_states.ndim == 3
268
+ assert self.block_shape is not None
269
270
a1q = hidden_states
271
_, N, K = w1.size()
vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py
+44
-21
@@ -6,6 +6,7 @@ import torch
6
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (
8
BatchedDeepGemmExperts)
9
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10
from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
11
BatchedTritonExperts)
12
@@ -20,43 +21,45 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
21
use_int8_w8a8: bool = False,
22
use_int8_w8a16: bool = False,
23
use_int4_w4a16: bool = False,
23
- per_channel_quant: bool = False,
24
block_shape: Optional[list[int]] = None,
25
+ per_act_token_quant: bool = False,
26
allow_deep_gemm: bool = False):
26
- super().__init__()
27
assert not use_int8_w8a8, "NYI"
28
assert not use_int8_w8a16, "NYI"
29
assert not use_int4_w4a16, "NYI"
30
31
+ super().__init__(
32
+ FusedMoEQuantConfig.make(
33
+ use_fp8_w8a8=use_fp8_w8a8,
34
+ use_int8_w8a8=use_int8_w8a8,
35
+ use_int8_w8a16=use_int8_w8a16,
36
+ use_int4_w4a16=use_int4_w4a16,
37
+ block_shape=block_shape,
38
+ per_act_token_quant=per_act_token_quant,
39
+ ))
40
self.max_num_tokens = max_num_tokens
41
self.world_size = world_size
42
self.dp_size = dp_size
34
- self.use_fp8_w8a8 = use_fp8_w8a8
35
- self.use_int8_w8a8 = use_int8_w8a8
36
- self.use_int8_w8a16 = use_int8_w8a16
37
- self.use_int4_w4a16 = use_int4_w4a16
38
- self.per_channel_quant = per_channel_quant
39
- self.block_shape = block_shape
43
self.allow_deep_gemm = allow_deep_gemm
44
45
# BatchedTritonKernel doesn't support block quantization
46
# at the moment.
47
self.batched_triton_experts = BatchedTritonExperts(
48
max_num_tokens=self.max_num_tokens,
46
- use_fp8_w8a8=self.use_fp8_w8a8,
47
- use_int8_w8a8=self.use_int8_w8a8,
48
- use_int8_w8a16=self.use_int8_w8a16,
49
- use_int4_w4a16=self.use_int4_w4a16,
50
- per_channel_quant=self.per_channel_quant,
51
- block_shape=self.block_shape,
49
world_size=self.world_size,
53
- dp_size=self.dp_size) if self.block_shape is None else None
50
+ dp_size=self.dp_size,
51
+ use_fp8_w8a8=use_fp8_w8a8,
52
+ use_int8_w8a8=use_int8_w8a8,
53
+ use_int8_w8a16=use_int8_w8a16,
54
+ use_int4_w4a16=use_int4_w4a16,
55
+ per_act_token_quant=self.per_act_token_quant,
56
+ block_shape=self.block_shape,
57
+ ) if self.block_shape is None else None
58
+
59
+ is_fp8_128_block_quantized = (
60
+ use_fp8_w8a8 and self.block_shape
61
+ == BatchedDeepGemmExperts.DEEPGEMM_BLOCK_SHAPE)
62
55
- is_fp8_128_block_quantized = (self.use_fp8_w8a8
56
- and self.block_shape is not None
57
- and len(self.block_shape) == 2 and all(
58
- [b == 128
59
- for b in self.block_shape]))
63
self.batched_deep_gemm_experts = BatchedDeepGemmExperts(
64
max_num_tokens=self.max_num_tokens,
65
world_size=self.world_size,
@@ -67,12 +70,31 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
70
assert (self.batched_deep_gemm_experts is not None
71
or self.batched_triton_experts is not None)
72
73
+ @property
74
+ def activation_formats(
75
+ self
76
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
77
+ if self.batched_triton_experts is not None:
78
+ assert (self.batched_deep_gemm_experts is None
79
+ or self.batched_deep_gemm_experts.activation_formats
80
+ == self.batched_triton_experts.activation_formats)
81
+ return self.batched_triton_experts.activation_formats
82
+ else:
83
+ assert self.batched_deep_gemm_experts is not None
84
+ return self.batched_deep_gemm_experts.activation_formats
85
+
86
def supports_chunking(self) -> bool:
87
bdge = self.batched_deep_gemm_experts
88
bte = self.batched_triton_experts
89
return ((bdge is None or bdge.supports_chunking())
90
and (bte is None or bte.supports_chunking()))
91
92
+ def supports_expert_map(self) -> bool:
93
+ bdge = self.batched_deep_gemm_experts
94
+ bte = self.batched_triton_experts
95
+ return ((bdge is None or bdge.supports_expert_map())
96
+ and (bte is None or bte.supports_expert_map()))
97
+
98
def workspace_shapes(
99
self,
100
a: torch.Tensor,
@@ -87,7 +109,8 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
109
# Note: the deep gemm workspaces are strictly larger than the triton
110
# workspaces so we can be pessimistic here and allocate for DeepGemm
111
# even if we fall back to triton later, e.g. if expert maps are set.
90
- if self.allow_deep_gemm and self.batched_deep_gemm_experts is not None:
112
+ if self.allow_deep_gemm:
113
+ assert self.batched_deep_gemm_experts is not None
114
return self.batched_deep_gemm_experts.workspace_shapes(
115
a, aq, M, N, K, topk, global_num_experts, local_num_experts)
116
else:
vllm/model_executor/layers/fused_moe/config.py
new
+410
@@ -0,0 +1,410 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+from typing import Optional, Union
5
+
6
+import torch
7
+from compressed_tensors.quantization import (QuantizationArgs,
8
+ QuantizationStrategy,
9
+ QuantizationType)
10
+
11
+import vllm.envs as envs
12
+from vllm.config import ParallelConfig
13
+from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
14
+from vllm.logger import init_logger
15
+from vllm.model_executor.layers.quantization.base_config import (
16
+ QuantizationConfig)
17
+
18
+logger = init_logger(__name__)
19
+
20
+
21
+def _get_quant_config_quantization_args(
22
+ quant_config: Optional[QuantizationConfig],
23
+ prop_name: str,
24
+) -> Optional[QuantizationArgs]:
25
+ if (quant_config is not None and hasattr(quant_config, 'target_scheme_map')
26
+ and "Linear" in quant_config.target_scheme_map and
27
+ "input_activations" in quant_config.target_scheme_map["Linear"]):
28
+ return quant_config.target_scheme_map["Linear"].get(prop_name)
29
+ else:
30
+ return None
31
+
32
+
33
+def get_quant_config_input_quant(
34
+ quant_config: Optional[QuantizationConfig]
35
+) -> Optional[QuantizationArgs]:
36
+ return _get_quant_config_quantization_args(quant_config,
37
+ "input_activations")
38
+
39
+
40
+def get_quant_config_weight_quant(
41
+ quant_config: Optional[QuantizationConfig]
42
+) -> Optional[QuantizationArgs]:
43
+ return _get_quant_config_quantization_args(quant_config, "weights")
44
+
45
+
46
+# TODO (bnell): use scalar_type instead of bools?
47
+def get_config_quant_dtype(
48
+ use_fp8_w8a8: bool,
49
+ use_int8_w8a8: bool,
50
+ use_int8_w8a16: bool,
51
+ use_int4_w4a16: bool,
52
+) -> Optional[torch.dtype]:
53
+ if use_fp8_w8a8:
54
+ return torch.float8_e4m3fn
55
+ elif use_int8_w8a8:
56
+ return torch.int8
57
+ return None
58
+
59
+
60
+@dataclass
61
+class FusedMoEQuantConfig:
62
+ # The post quantization activation type.
63
+ quant_dtype: Optional[torch.dtype] = None
64
+ per_act_token_quant: bool = False
65
+ per_out_ch_quant: bool = False
66
+ block_shape: Optional[list[int]] = None
67
+
68
+ # TODO: add col major flag?
69
+ # add detailed quant info for input, intermediates, weights, etc?
70
+
71
+ @staticmethod
72
+ def make(
73
+ use_fp8_w8a8: bool = False,
74
+ use_int8_w8a8: bool = False,
75
+ use_int8_w8a16: bool = False,
76
+ use_int4_w4a16: bool = False,
77
+ per_act_token_quant: bool = False,
78
+ per_out_ch_quant: bool = False,
79
+ block_shape: Optional[list[int]] = None,
80
+ ) -> "FusedMoEQuantConfig":
81
+ assert sum([
82
+ int(flag) for flag in [
83
+ use_fp8_w8a8,
84
+ use_int8_w8a8,
85
+ use_int8_w8a16,
86
+ use_int4_w4a16,
87
+ ]
88
+ ]) <= 1, "Quantization flags are mutually exclusive."
89
+
90
+ quant_dtype = get_config_quant_dtype(
91
+ use_fp8_w8a8=use_fp8_w8a8,
92
+ use_int8_w8a8=use_int8_w8a8,
93
+ use_int8_w8a16=use_int8_w8a16,
94
+ use_int4_w4a16=use_int4_w4a16,
95
+ )
96
+ return FusedMoEQuantConfig(
97
+ quant_dtype,
98
+ per_act_token_quant,
99
+ per_out_ch_quant,
100
+ block_shape,
101
+ )
102
+
103
+
104
+@dataclass
105
+class FusedMoEParallelConfig:
106
+ tp_size: int
107
+ dp_size: int
108
+ ep_size: int
109
+ tp_rank: int
110
+ dp_rank: int
111
+ ep_rank: int
112
+ world_size: int
113
+
114
+ use_ep: bool # whether to use EP or not
115
+
116
+ @property
117
+ def use_all2all_kernels(self):
118
+ return self.dp_size > 1 and self.use_ep
119
+
120
+ @property
121
+ def use_pplx_kernels(self):
122
+ return (self.use_all2all_kernels
123
+ and envs.VLLM_ALL2ALL_BACKEND == "pplx")
124
+
125
+ @property
126
+ def use_deepep_ht_kernels(self):
127
+ return (self.use_all2all_kernels
128
+ and envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput")
129
+
130
+ @property
131
+ def use_deepep_ll_kernels(self):
132
+ return (self.use_all2all_kernels
133
+ and envs.VLLM_ALL2ALL_BACKEND == "deepep_low_latency")
134
+
135
+ @staticmethod
136
+ def make(tp_size_: int, dp_size_: int, world_size_: int,
137
+ vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
138
+ """
139
+ Determine MoE parallel configuration. Based on the input tp_size_,
140
+ dp_size_, ep_size_ and vllm's parallel config, determine what
141
+ level's of parallelism to use in the fused moe layer.
142
+
143
+ Args:
144
+ tp_size_ (int): tp_size passed into the FusedMoE constructor.
145
+ dp_size_ (int): dp_size passed into the FusedMoE constructor.
146
+ ep_size_ (int): ep_size passed into the FusedMoE constructor.
147
+ world_size_ (int): the world size of the current All2All manager.
148
+ vllm_parallel_config (ParallelConfig): vllm's parallel config
149
+ object.
150
+
151
+ Examples:
152
+ When there is no parallelism requested, i.e. tp_size_ = dp_size_ = 1,
153
+ we simply return the sizes unaltered and the ranks set to 0.
154
+
155
+ Expert Parallelism is considered only when either dp_size_ or tp_size_
156
+ is non trivial.
157
+
158
+ When TP = 2, DP = 1 and EP = False, the configuration on different
159
+ devices,
160
+ - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} //
161
+ legend : {size, rank}
162
+ - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0}
163
+ - Comment : Tensors are sharded across 2 devices.
164
+
165
+ When TP = 1, DP = 2 and EP = False, the configuration on different
166
+ devices,
167
+ - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0}
168
+ - device 1 : TP = {2, 1} DP = {2, 1} EP = {1, 0}
169
+ - Comment: There are 2 engine instances and the tensors are sharded
170
+ across 2 decvices.
171
+
172
+ When TP = 2, DP = 2 and EP = False, the configuration on different
173
+ devices,
174
+ - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0}
175
+ - device 1: TP = {4, 1} DP = {2, 0} EP = {1, 0}
176
+ - device 2: TP = {4, 2} DP = {2, 1} EP = {1, 0}
177
+ - device 3: TP = {4, 3} DP = {2, 1} EP = {1, 0}
178
+ - Comment: There are 2 engine instances and the tensors are sharded
179
+ across 4 devices.
180
+
181
+ When, TP = 2, DP = 1 and EP = True, the configuration on different
182
+ devices,
183
+ - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0}
184
+ - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1}
185
+ - Comment: The experts are split between the 2 devices.
186
+
187
+ When, TP = 1, DP = 2 and EP = True, the configuration on different
188
+ devices,
189
+ - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0}
190
+ - device 1: TP = {1, 0} DP = {2, 1} EP = {2, 1}
191
+ - Comment: There are 2 engine instances and the experts are split
192
+ between the 2 devices.
193
+
194
+ When TP = 2, DP = 2 and EP = True, the configuration on different
195
+ devices,
196
+ - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0}
197
+ - device 1: TP = {1, 0} DP = {2, 0} EP = {4, 1}
198
+ - device 2: TP = {1, 0} DP = {2, 1} EP = {4, 2}
199
+ - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3}
200
+ - Comment: There are 2 engine instances and the experts are split
201
+ between the 4 devices.
202
+ """
203
+
204
+ def flatten_tp_across_dp(dp_rank: int):
205
+ tp_rank = 0 if tp_size_ == 1 else get_tensor_model_parallel_rank()
206
+ # There are actually dp_size_ * tp_size_ devices. Update tp_size
207
+ # and tp_rank so we shard across all devices.
208
+ tp_size = dp_size_ * tp_size_
209
+ tp_rank = dp_rank * tp_size_ + tp_rank
210
+ return tp_size, tp_rank
211
+
212
+ use_ep = (dp_size_ * tp_size_ > 1
213
+ and vllm_parallel_config.enable_expert_parallel)
214
+
215
+ dp_size = dp_size_
216
+ dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
217
+ tp_size, tp_rank = flatten_tp_across_dp(dp_rank)
218
+
219
+ if not use_ep:
220
+ return FusedMoEParallelConfig(tp_size=tp_size,
221
+ tp_rank=tp_rank,
222
+ dp_size=dp_size,
223
+ dp_rank=dp_rank,
224
+ ep_size=1,
225
+ ep_rank=0,
226
+ world_size=world_size_,
227
+ use_ep=False)
228
+ # DP + EP / TP + EP / DP + TP + EP
229
+ assert use_ep
230
+ # In EP, each device owns a set of experts fully. There is no tensor
231
+ # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that.
232
+ ep_size = tp_size
233
+ ep_rank = tp_rank
234
+ return FusedMoEParallelConfig(tp_size=1,
235
+ tp_rank=0,
236
+ dp_size=dp_size,
237
+ dp_rank=dp_rank,
238
+ ep_size=ep_size,
239
+ ep_rank=ep_rank,
240
+ world_size=world_size_,
241
+ use_ep=True)
242
+
243
+
244
+# Adapted from pplx-kernels tests/all_to_all_utils.py
245
+@dataclass
246
+class FusedMoEConfig:
247
+ num_experts: int
248
+ experts_per_token: int
249
+ hidden_dim: int
250
+
251
+ num_local_experts: int
252
+ moe_parallel_config: FusedMoEParallelConfig
253
+
254
+ # The activation type.
255
+ in_dtype: torch.dtype
256
+
257
+ quant_config: Optional[FusedMoEQuantConfig] = None
258
+
259
+ max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
260
+
261
+ def __post_init__(self):
262
+ if self.dp_size > 1:
263
+ logger.debug("Using FusedMoEConfig::max_num_tokens=%d",
264
+ self.max_num_tokens)
265
+
266
+ @property
267
+ def quant_dtype(self) -> Optional[torch.dtype]:
268
+ if self.quant_config is not None:
269
+ return self.quant_config.quant_dtype
270
+ else:
271
+ return None
272
+
273
+ @property
274
+ def block_shape(self) -> Optional[list[int]]:
275
+ if self.quant_config is not None:
276
+ return self.quant_config.block_shape
277
+ else:
278
+ return None
279
+
280
+ @property
281
+ def per_act_token_quant(self) -> bool:
282
+ if self.quant_config is not None:
283
+ return self.quant_config.per_act_token_quant
284
+ else:
285
+ return False
286
+
287
+ @property
288
+ def per_out_ch_quant(self) -> bool:
289
+ if self.quant_config is not None:
290
+ return self.quant_config.per_out_ch_quant
291
+ else:
292
+ return False
293
+
294
+ @property
295
+ def tp_size(self):
296
+ return self.moe_parallel_config.tp_size
297
+
298
+ @property
299
+ def dp_size(self):
300
+ return self.moe_parallel_config.dp_size
301
+
302
+ @property
303
+ def ep_size(self):
304
+ return self.moe_parallel_config.ep_size
305
+
306
+ @property
307
+ def world_size(self):
308
+ return self.moe_parallel_config.world_size
309
+
310
+ @property
311
+ def tp_rank(self):
312
+ return self.moe_parallel_config.tp_rank
313
+
314
+ @property
315
+ def dp_rank(self):
316
+ return self.moe_parallel_config.dp_rank
317
+
318
+ @property
319
+ def ep_rank(self):
320
+ return self.moe_parallel_config.ep_rank
321
+
322
+ @property
323
+ def use_ep(self):
324
+ return self.moe_parallel_config.use_ep
325
+
326
+ @property
327
+ def use_pplx_kernels(self):
328
+ return self.moe_parallel_config.use_pplx_kernels
329
+
330
+ @property
331
+ def use_deepep_ht_kernels(self):
332
+ return self.moe_parallel_config.use_deepep_ht_kernels
333
+
334
+ @property
335
+ def use_deepep_ll_kernels(self):
336
+ return self.moe_parallel_config.use_deepep_ll_kernels
337
+
338
+ @staticmethod
339
+ def make(
340
+ num_experts: int,
341
+ experts_per_token: int,
342
+ hidden_dim: int,
343
+ num_local_experts: int,
344
+ moe_parallel_config: FusedMoEParallelConfig,
345
+ in_dtype: torch.dtype,
346
+ max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE,
347
+ quant_config: Optional[Union[FusedMoEQuantConfig,
348
+ QuantizationConfig]] = None
349
+ ) -> "FusedMoEConfig":
350
+
351
+ _quant_config: Optional[FusedMoEQuantConfig] = None
352
+
353
+ if quant_config is not None and isinstance(quant_config,
354
+ QuantizationConfig):
355
+ if hasattr(quant_config, 'weight_block_size'):
356
+ block_shape = quant_config.weight_block_size
357
+ else:
358
+ block_shape = None
359
+ per_act_token_quant = False
360
+ per_out_ch_quant = False
361
+ quant_dtype: Optional[torch.dtype] = None
362
+
363
+ input_quant = get_quant_config_input_quant(quant_config)
364
+ weight_quant = get_quant_config_weight_quant(quant_config)
365
+
366
+ if input_quant is not None:
367
+ per_act_token_quant = (input_quant.strategy
368
+ == QuantizationStrategy.TOKEN
369
+ if input_quant is not None else False)
370
+
371
+ if input_quant.num_bits == 8:
372
+ if input_quant.type == QuantizationType.FLOAT:
373
+ quant_dtype = torch.float8_e4m3fn
374
+ elif input_quant.type == QuantizationType.INT:
375
+ quant_dtype = torch.int8
376
+
377
+ from vllm.model_executor.layers.quantization.fp8 import Fp8Config
378
+ if quant_dtype is None and isinstance(quant_config, Fp8Config):
379
+ quant_dtype = torch.float8_e4m3fn
380
+
381
+ if weight_quant is not None:
382
+ per_out_ch_quant = (
383
+ weight_quant.strategy == QuantizationStrategy.CHANNEL)
384
+
385
+ if quant_dtype is not None:
386
+ _quant_config = FusedMoEQuantConfig(
387
+ quant_dtype=quant_dtype,
388
+ per_act_token_quant=per_act_token_quant,
389
+ per_out_ch_quant=per_out_ch_quant,
390
+ block_shape=block_shape,
391
+ )
392
+ else:
393
+ _quant_config = FusedMoEQuantConfig()
394
+ logger.warning_once("MoE DP setup unable to determine "
395
+ "quantization scheme or unsupported "
396
+ "quantization type. This model will "
397
+ "not run with DP enabled.")
398
+ else:
399
+ _quant_config = quant_config
400
+
401
+ return FusedMoEConfig(
402
+ num_experts=num_experts,
403
+ experts_per_token=experts_per_token,
404
+ hidden_dim=hidden_dim,
405
+ num_local_experts=num_local_experts,
406
+ moe_parallel_config=moe_parallel_config,
407
+ in_dtype=in_dtype,
408
+ quant_config=_quant_config,
409
+ max_num_tokens=max_num_tokens,
410
+ )
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+47
-26
@@ -7,6 +7,7 @@ import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm import _custom_ops as ops
10
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
12
MoEPrepareAndFinalizeNoEP)
13
from vllm.model_executor.layers.fused_moe.utils import _fp8_perm, _resize_cache
@@ -202,26 +203,47 @@ def run_cutlass_moe_fp8(
203
204
205
# TODO (bnell): split class batched vs. non-batched?
206
+# maybe remove need for passing aq to workspace_shapes
207
class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
208
209
def __init__(
210
self,
211
max_experts_per_worker: int,
210
- out_dtype: torch.dtype,
211
- per_act_token: bool,
212
- per_out_ch: bool,
212
+ out_dtype: Optional[torch.dtype],
213
+ per_act_token_quant: bool,
214
+ per_out_ch_quant: bool,
215
+ block_shape: Optional[list[int]] = None,
216
use_batched_format: bool = False,
217
):
215
- super().__init__()
218
+ super().__init__(
219
+ FusedMoEQuantConfig(
220
+ quant_dtype=torch.float8_e4m3fn,
221
+ per_act_token_quant=per_act_token_quant,
222
+ per_out_ch_quant=per_out_ch_quant,
223
+ block_shape=block_shape,
224
+ ))
225
+ assert max_experts_per_worker > 0
226
self.max_experts_per_worker = max_experts_per_worker
227
self.out_dtype = out_dtype
218
- self.per_act_token = per_act_token
219
- self.per_out_ch = per_out_ch
228
self.use_batched_format = use_batched_format
229
230
+ @property
231
+ def activation_formats(
232
+ self
233
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
234
+ if self.use_batched_format:
235
+ return (mk.FusedMoEActivationFormat.BatchedExperts,
236
+ mk.FusedMoEActivationFormat.BatchedExperts)
237
+ else:
238
+ return (mk.FusedMoEActivationFormat.Standard,
239
+ mk.FusedMoEActivationFormat.Standard)
240
+
241
def supports_chunking(self) -> bool:
242
return not self.use_batched_format
243
244
+ def supports_expert_map(self) -> bool:
245
+ return not self.use_batched_format
246
+
247
def workspace_shapes(
248
self,
249
a: torch.Tensor,
@@ -245,7 +267,8 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
267
workspace1 = (M * topk, max(2 * N, K))
268
workspace2 = (M * topk, N)
269
output = (M * topk, K)
248
- return (workspace1, workspace2, output, self.out_dtype)
270
+ return (workspace1, workspace2, output,
271
+ self.out_dtype if self.out_dtype is not None else a.dtype)
272
273
def apply(
274
self,
@@ -270,13 +293,14 @@ class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
293
assert w1_zp is None, "w1_zp is not supported in CUTLASS MoE"
294
assert w2_zp is None, "w2_zp is not supported in CUTLASS MoE"
295
activation_callable = lambda i, o: self.activation(activation, i, o)
273
- run_cutlass_moe_fp8(output, hidden_states, w1, w2, topk_ids,
274
- activation_callable, global_num_experts,
275
- expert_map, w1_scale, w2_scale, a1q_scale,
276
- a2_scale, workspace13, workspace2,
277
- expert_num_tokens, self.out_dtype,
278
- self.per_act_token, self.per_out_ch,
279
- self.use_batched_format)
296
+ in_dtype = hidden_states.dtype
297
+ run_cutlass_moe_fp8(
298
+ output, hidden_states, w1, w2, topk_ids, activation_callable,
299
+ global_num_experts, expert_map, w1_scale, w2_scale, a1q_scale,
300
+ a2_scale, workspace13, workspace2, expert_num_tokens,
301
+ self.out_dtype if self.out_dtype is not None else in_dtype,
302
+ self.per_act_token_quant, self.per_out_ch_quant,
303
+ self.use_batched_format)
304
305
306
def cutlass_moe_fp8(
@@ -287,6 +311,7 @@ def cutlass_moe_fp8(
311
topk_ids: torch.Tensor,
312
w1_scale: torch.Tensor,
313
w2_scale: torch.Tensor,
314
+ per_act_token: bool,
315
activation: str = "silu",
316
a1_scale: Optional[torch.Tensor] = None,
317
a2_scale: Optional[torch.Tensor] = None,
@@ -330,22 +355,18 @@ def cutlass_moe_fp8(
355
Returns:
356
- torch.Tensor: The fp16 output tensor after applying the MoE layer.
357
"""
333
- per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
334
- a2_scale.numel() != 1 if a2_scale is not None else False)
358
per_out_ch = w1_scale.numel() != w1_q.size(0)
359
337
- out_dtype = a.dtype
360
+ num_experts = global_num_experts if global_num_experts != -1 else w1_q.size(
361
+ 0)
362
363
fn = mk.FusedMoEModularKernel(
340
- MoEPrepareAndFinalizeNoEP(
341
- quant_dtype=torch.float8_e4m3fn,
342
- per_channel_quant=per_act_token,
343
- ),
364
+ MoEPrepareAndFinalizeNoEP(),
365
CutlassExpertsFp8(
345
- max_experts_per_worker=global_num_experts,
346
- out_dtype=out_dtype,
347
- per_act_token=per_act_token,
348
- per_out_ch=per_out_ch,
366
+ max_experts_per_worker=num_experts,
367
+ out_dtype=a.dtype,
368
+ per_act_token_quant=per_act_token,
369
+ per_out_ch_quant=per_out_ch,
370
use_batched_format=False,
371
),
372
)
@@ -358,7 +379,7 @@ def cutlass_moe_fp8(
379
topk_ids,
380
False,
381
activation,
361
- global_num_experts if global_num_experts != -1 else w1_q.size(0),
382
+ num_experts,
383
expert_map,
384
w1_scale,
385
w2_scale,
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+22
-7
@@ -7,13 +7,13 @@ import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
from vllm.logger import init_logger
10
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
11
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)
14
-from vllm.model_executor.layers.fused_moe.utils import _resize_cache
15
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
16
- per_token_group_quant_fp8)
15
+from vllm.model_executor.layers.fused_moe.utils import (
16
+ _resize_cache, per_token_group_quant_fp8)
17
from vllm.utils import has_deep_gemm, round_up
18
19
logger = init_logger(__name__)
@@ -65,16 +65,31 @@ def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
65
class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
66
67
def __init__(self):
68
- super().__init__()
69
- self.block_shape = deep_gemm_block_shape()
68
+ super().__init__(
69
+ FusedMoEQuantConfig(
70
+ quant_dtype=torch.float8_e4m3fn,
71
+ per_act_token_quant=False,
72
+ block_shape=deep_gemm_block_shape(),
73
+ ))
74
+
75
+ @property
76
+ def activation_formats(
77
+ self
78
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
79
+ return (mk.FusedMoEActivationFormat.Standard,
80
+ mk.FusedMoEActivationFormat.Standard)
81
82
def supports_chunking(self) -> bool:
83
return True
84
85
+ def supports_expert_map(self) -> bool:
86
+ return True
87
+
88
def workspace_shapes(
89
self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
90
topk: int, global_num_experts: int, local_num_experts: int
91
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
92
+ assert self.block_shape is not None
93
# We use global_num_experts due to how moe_align_block_size handles
94
# expert_maps.
95
num_experts = global_num_experts
@@ -107,6 +122,7 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
122
expert_num_tokens: Optional[torch.Tensor],
123
):
124
import deep_gemm as dg
125
+ assert self.block_shape is not None
126
127
a1q = hidden_states
128
_, N, K = w1.size()
@@ -213,8 +229,7 @@ def deep_gemm_moe_fp8(
229
- torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
230
"""
231
fn = mk.FusedMoEModularKernel(
216
- MoEPrepareAndFinalizeNoEP(quant_dtype=torch.float8_e4m3fn,
217
- block_shape=deep_gemm_block_shape()),
232
+ MoEPrepareAndFinalizeNoEP(),
233
DeepGemmExperts(),
234
)
235
return fn(
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+36
-35
@@ -6,6 +6,7 @@ import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
from vllm import _custom_ops as ops
9
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10
from vllm.model_executor.layers.fused_moe.utils import (
11
moe_kernel_quantize_input)
12
@@ -15,22 +16,14 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
16
Prepare/Finalize using DeepEP High-Throughput kernels.
17
"""
18
18
- def __init__(self,
19
- buffer: deep_ep.Buffer,
20
- world_size: int,
21
- rank: int,
22
- dp_size: int,
23
- rank_expert_offset: int,
24
- quant_dtype: Optional[torch.dtype] = None,
25
- block_shape: Optional[list[int]] = None):
19
+ def __init__(self, buffer: deep_ep.Buffer, world_size: int, rank: int,
20
+ dp_size: int, rank_expert_offset: int):
21
super().__init__()
22
self.buffer = buffer
23
self.world_size = world_size
24
self.rank = rank
25
self.dp_size = dp_size
26
self.rank_expert_offset = rank_expert_offset
32
- self.quant_dtype = quant_dtype
33
- self.block_shape = block_shape
27
# The dispatch function returns a handle that the combine function
28
# requires. We store the handle here so it is available to the
29
# combine function.
@@ -39,6 +32,10 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
32
# From https://github.com/deepseek-ai/DeepEP/blob/9fe9021f29c9083cd1808ab36b740208524d9f63/deep_ep/buffer.py#L164
33
self.available_rank_configs = [2, 4, 8, 16, 24, 32, 64, 128, 144, 160]
34
35
+ @property
36
+ def activation_format(self) -> mk.FusedMoEActivationFormat:
37
+ return mk.FusedMoEActivationFormat.Standard
38
+
39
def max_num_tokens_per_rank(self) -> Optional[int]:
40
return None
41
@@ -55,13 +52,6 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
52
return None
53
return deep_ep.Buffer.get_combine_config(self.dp_size)
54
58
- def _do_quant(self, tokens: torch.Tensor,
59
- token_scales: Optional[torch.Tensor], per_act_token: bool):
60
- tokens, token_scales = moe_kernel_quantize_input(
61
- tokens, token_scales, self.quant_dtype, per_act_token,
62
- self.block_shape)
63
- return tokens, token_scales
64
-
55
def _do_dispatch(self, tokens: torch.Tensor,
56
token_scales: Optional[torch.Tensor],
57
rank_topk_ids: torch.Tensor,
@@ -130,43 +120,51 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
120
a1: torch.Tensor,
121
a1_scale: Optional[torch.Tensor],
122
a2_scale: Optional[torch.Tensor],
133
- rank_topk_weights: torch.Tensor,
134
- rank_topk_ids: torch.Tensor,
123
+ topk_weights: torch.Tensor,
124
+ topk_ids: torch.Tensor,
125
num_experts: int,
126
expert_map: Optional[torch.Tensor],
127
apply_router_weight_on_input: bool,
128
+ quant_config: FusedMoEQuantConfig,
129
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
130
Optional[torch.Tensor], Optional[torch.Tensor]]:
131
132
if apply_router_weight_on_input:
142
- topk = rank_topk_ids.size(1)
133
+ topk = topk_ids.size(1)
134
# TODO: this only works for topK=1, will need to update for topK>1
135
assert topk == 1, (
136
"apply_router_weight_on_input is only implemented for topk=1")
146
- a1 = a1 * rank_topk_weights.to(a1.dtype)
137
+ a1 = a1 * topk_weights.to(a1.dtype)
138
139
# Check if there is a block_shape / or if we can infer the quantization
140
# schemes from the scales.
141
per_token_quant = None
151
- if all([x is None for x in [self.block_shape, a1_scale, a2_scale]
152
- ]) and self.quant_dtype is not None:
142
+ if all([
143
+ x is None
144
+ for x in [quant_config.block_shape, a1_scale, a2_scale]
145
+ ]) and quant_config.quant_dtype is not None:
146
# Quantization required despite none of the inputs suggesting
147
# quantization. Fallback to per_token_dynamic quant.
148
per_token_quant = True
149
else:
157
- per_token_quant = ((self.block_shape is not None) or
158
- (a1_scale is not None and a1_scale.numel() != 1)
159
- or (a2_scale is not None
160
- and a2_scale.numel() != 1))
150
+ per_token_quant = False
151
152
if per_token_quant:
163
- a1q, a1q_scale = self._do_quant(a1, a1_scale, per_act_token=True)
153
+ a1q, a1q_scale = moe_kernel_quantize_input(
154
+ a1,
155
+ a1_scale,
156
+ quant_dtype=quant_config.quant_dtype,
157
+ per_act_token_quant=True,
158
+ block_shape=quant_config.block_shape,
159
+ )
160
+ if a1q_scale is not None and a1q_scale.numel() == 1:
161
+ a1q_scale = a1q_scale.view(1, 1)
162
(expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,
163
expert_topk_weights) = self._do_dispatch(
164
tokens=a1q,
165
token_scales=a1q_scale,
168
- rank_topk_ids=rank_topk_ids,
169
- rank_topk_weights=rank_topk_weights,
166
+ rank_topk_ids=topk_ids,
167
+ rank_topk_weights=topk_weights,
168
num_experts=num_experts)
169
else:
170
# DeepEP kernels only support dispatching per-token-quant
@@ -175,15 +173,18 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
173
expert_topk_weights) = self._do_dispatch(
174
tokens=a1,
175
token_scales=None,
178
- rank_topk_ids=rank_topk_ids,
179
- rank_topk_weights=rank_topk_weights,
176
+ rank_topk_ids=topk_ids,
177
+ rank_topk_weights=topk_weights,
178
num_experts=num_experts)
179
# quantize now
180
expert_x_scale = None
181
if expert_x.numel() != 0:
184
- expert_x, expert_x_scale = self._do_quant(expert_x,
185
- a1_scale,
186
- per_act_token=False)
182
+ expert_x, expert_x_scale = moe_kernel_quantize_input(
183
+ expert_x,
184
+ a1_scale,
185
+ quant_dtype=quant_config.quant_dtype,
186
+ per_act_token_quant=False,
187
+ block_shape=quant_config.block_shape)
188
189
return (expert_x, expert_x_scale, expert_num_tokens, expert_topk_ids,
190
expert_topk_weights)
vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py
+34
-37
@@ -5,11 +5,13 @@ 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.config import FusedMoEQuantConfig
9
from vllm.model_executor.layers.fused_moe.utils import (
9
- moe_kernel_quantize_input)
10
+ maybe_fix_scales, moe_kernel_quantize_input)
11
12
# DeepEP kernels quantize dispatch inputs in 128 element chunks.
13
DEEPEP_QUANT_BLOCK_SIZE = 128
14
+DEEPEP_QUANT_BLOCK_SHAPE = [DEEPEP_QUANT_BLOCK_SIZE, DEEPEP_QUANT_BLOCK_SIZE]
15
16
17
def dequant_fp8(expert_x_fp8: torch.Tensor,
@@ -35,30 +37,30 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
37
38
# DeepEP low-latency kernels are compiled only for certain
39
# specific hidden sizes.
38
- SUPPORTED_HIDDEN_SIZES = [2560, 4096, 5120, 7168]
40
+ SUPPORTED_HIDDEN_SIZES = [2048, 2560, 4096, 5120, 7168]
41
42
def __init__(self,
43
buffer: deep_ep.Buffer,
44
+ max_tokens_per_rank: int,
45
world_size: int,
46
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.max_tokens_per_rank = max_tokens_per_rank
52
self.world_size = world_size
53
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
54
self.use_fp8_dispatch = use_fp8_dispatch
55
# The dispatch function returns a handle that the combine function
56
# requires. We store the handle here so it is available to the
57
# combine function.
58
self.handle = None
59
60
+ @property
61
+ def activation_format(self) -> mk.FusedMoEActivationFormat:
62
+ return mk.FusedMoEActivationFormat.BatchedExperts
63
+
64
def max_num_tokens_per_rank(self) -> Optional[int]:
65
return self.max_tokens_per_rank
66
@@ -66,12 +68,17 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
68
return torch.int64
69
70
def _do_quant(
69
- self, x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
70
- a1_scale: Optional[torch.Tensor], a2_scale: Optional[torch.Tensor],
71
- a1_dtype: torch.dtype
71
+ self,
72
+ x: Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]],
73
+ a1_scale: Optional[torch.Tensor],
74
+ a2_scale: Optional[torch.Tensor],
75
+ a1_dtype: torch.dtype,
76
+ quant_dtype: Optional[torch.dtype],
77
+ per_act_token_quant: bool,
78
+ block_shape: Optional[list[int]],
79
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
80
74
- block_k = self.block_shape[1] if self.block_shape is not None else None
81
+ block_k = block_shape[1] if block_shape is not None else None
82
if self.use_fp8_dispatch:
83
if block_k == DEEPEP_QUANT_BLOCK_SIZE:
84
# DeepEP kernels did the quantization for us.
@@ -84,32 +91,20 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
91
92
assert isinstance(x, torch.Tensor)
93
87
- # Check if there is a block_shape / or if we can infer the quantization
88
- # schemes from the scales.
89
- per_token_quant = None
90
- if all([v is None for v in [self.block_shape, a1_scale, a2_scale]
91
- ]) and self.quant_dtype is not None:
92
- # Quantization required despite none of the inputs suggesting
93
- # quantization. Fallback to per_token_dynamic quant.
94
- per_token_quant = True
95
- else:
96
- per_token_quant = ((self.block_shape is not None) or
97
- (a1_scale is not None and a1_scale.numel() != 1)
98
- or (a2_scale is not None
99
- and a2_scale.numel() != 1))
94
+ assert not per_act_token_quant
95
96
num_experts, max_tokens, hidden_dim = x.size()
97
98
# TODO (varun): Optimization - Use a batched version of quant
99
x = x.view((-1, hidden_dim))
105
- x, x_scales = moe_kernel_quantize_input(x, a1_scale, self.quant_dtype,
106
- per_token_quant,
107
- self.block_shape)
100
+ x, x_scales = moe_kernel_quantize_input(x, a1_scale, quant_dtype,
101
+ per_act_token_quant,
102
+ block_shape)
103
x = x.view((num_experts, -1, hidden_dim))
104
110
- if per_token_quant:
105
+ if quant_dtype is not None:
106
assert x_scales is not None
112
- x_scales = x_scales.view(num_experts, max_tokens, -1)
107
+ x_scales = maybe_fix_scales(x_scales, num_experts)
108
109
return x, x_scales
110
@@ -118,11 +113,12 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
113
a1: torch.Tensor,
114
a1_scale: Optional[torch.Tensor],
115
a2_scale: Optional[torch.Tensor],
121
- rank_topk_weights: torch.Tensor,
122
- rank_topk_ids: torch.Tensor,
116
+ topk_weights: torch.Tensor,
117
+ topk_ids: torch.Tensor,
118
num_experts: int,
119
expert_map: Optional[torch.Tensor],
120
apply_router_weight_on_input: bool,
121
+ quant_config: FusedMoEQuantConfig,
122
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
123
Optional[torch.Tensor], Optional[torch.Tensor]]:
124
@@ -142,24 +138,25 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
138
"low_latency kernels doesn't support dispatching per-token scales")
139
140
if apply_router_weight_on_input:
145
- topk = rank_topk_ids.size(1)
141
+ topk = 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")
149
- a1 = a1 * rank_topk_weights.to(a1.dtype)
145
+ a1 = a1 * topk_weights.to(a1.dtype)
146
147
# Dispatch
148
expert_x, expert_num_tokens, self.handle, event, hook = \
149
self.buffer.low_latency_dispatch(a1,
154
- rank_topk_ids,
150
+ topk_ids,
151
self.max_tokens_per_rank,
152
num_experts,
153
use_fp8=self.use_fp8_dispatch,
154
async_finish=False,
155
return_recv_hook=False)
156
161
- expert_x, expert_x_scale = self._do_quant(expert_x, a1_scale, a2_scale,
162
- a1.dtype)
157
+ expert_x, expert_x_scale = self._do_quant(
158
+ expert_x, a1_scale, a2_scale, a1.dtype, quant_config.quant_dtype,
159
+ quant_config.per_act_token_quant, quant_config.block_shape)
160
161
return (expert_x, expert_x_scale, expert_num_tokens, None, None)
162
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
+94
-39
@@ -8,6 +8,7 @@ import triton
8
import triton.language as tl
9
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.fused_moe.fused_moe import (
13
get_config_dtype_str, try_get_optimal_moe_config)
14
from vllm.model_executor.layers.fused_moe.utils import (
@@ -317,8 +318,8 @@ def invoke_moe_batched_triton_kernel(
318
expert_num_tokens: torch.Tensor, # [E]
319
compute_type: tl.dtype,
320
# Quantization data
320
- A_scale: torch.Tensor,
321
- B_scale: torch.Tensor,
321
+ A_scale: Optional[torch.Tensor],
322
+ B_scale: Optional[torch.Tensor],
323
B_zp: torch.Tensor,
324
# Quantization schemes
325
use_fp8_w8a8: bool,
@@ -387,14 +388,23 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
388
that the PPLX dispatch/combine kernels use.
389
"""
390
390
- def __init__(self, max_num_tokens: int, world_size: int, dp_size: int,
391
- rank: int):
391
+ def __init__(
392
+ self,
393
+ max_num_tokens: int,
394
+ world_size: int,
395
+ dp_size: int,
396
+ rank: int,
397
+ ):
398
super().__init__()
399
self.world_size = world_size
400
self.dp_size = dp_size
401
self.rank = rank
402
self.max_num_tokens = max_num_tokens
403
404
+ @property
405
+ def activation_format(self) -> mk.FusedMoEActivationFormat:
406
+ return mk.FusedMoEActivationFormat.BatchedExperts
407
+
408
def max_num_tokens_per_rank(self) -> Optional[int]:
409
return self.max_num_tokens
410
@@ -411,6 +421,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
421
num_experts: int,
422
expert_map: Optional[torch.Tensor],
423
apply_router_weight_on_input: bool,
424
+ quant_config: FusedMoEQuantConfig,
425
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
426
Optional[torch.Tensor], Optional[torch.Tensor]]:
427
assert a1.dim() == 2
@@ -435,22 +446,35 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
446
447
num_local_experts = num_experts // self.world_size
448
449
+ if quant_config.quant_dtype is None:
450
+ b_type = a1.dtype
451
+ else:
452
+ b_type = quant_config.quant_dtype
453
+
454
b_a1 = torch.zeros(
455
(num_local_experts, self.max_num_tokens, hidden_dim),
440
- dtype=a1.dtype,
456
+ dtype=b_type,
457
device=a1.device)
458
459
+ b_a1_scale = None
460
+
461
+ assert quant_config.quant_dtype is None, "quantization NYI"
462
+
463
first_expert = num_local_experts * self.rank
464
last_expert = first_expert + num_local_experts
465
466
for expert_id in range(first_expert, last_expert):
467
topks = torch.any(topk_ids == expert_id, dim=1).flatten()
468
rows = torch.count_nonzero(topks.flatten())
449
- b_a1[expert_id -
450
- first_expert, :rows, :] = a1[:topks.numel()][topks]
451
- tokens_per_expert[expert_id - first_expert] = rows
469
+ if rows == 0:
470
+ continue
471
+ idx = expert_id - first_expert
472
+ b_a1[idx, :rows, :] = a1[:topks.numel()][topks]
473
+ tokens_per_expert[idx] = rows
474
453
- return b_a1, a1_scale, tokens_per_expert, None, None
475
+ assert b_a1_scale is None or b_a1_scale.ndim == 3
476
+
477
+ return b_a1, b_a1_scale, tokens_per_expert, None, None
478
479
def finalize(
480
self,
@@ -480,7 +504,7 @@ class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
504
output[topks] = output[topks] + rhs
505
506
483
-class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
507
+class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
508
"""
509
A reference MoE expert class that operates on expert batched format,
510
i.e. E x max_num_tokens x K. This is the format that the pplx
@@ -497,11 +521,17 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
521
use_int8_w8a16: bool = False,
522
use_int4_w4a16: bool = False,
523
block_shape: Optional[list[int]] = None,
500
- block_m: Optional[int] = None,
524
+ per_act_token_quant: bool = False,
525
):
502
- super().__init__()
503
- assert block_shape is None
504
- assert block_m is None
526
+ super().__init__(
527
+ FusedMoEQuantConfig.make(
528
+ use_fp8_w8a8=use_fp8_w8a8,
529
+ use_int8_w8a8=use_int8_w8a8,
530
+ use_int8_w8a16=use_int8_w8a16,
531
+ use_int4_w4a16=use_int4_w4a16,
532
+ per_act_token_quant=per_act_token_quant,
533
+ block_shape=block_shape,
534
+ ))
535
assert not use_fp8_w8a8, "NYI"
536
assert not use_int8_w8a8, "NYI"
537
assert not use_int8_w8a16, "NYI"
@@ -510,9 +540,19 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
540
self.world_size = world_size
541
self.dp_size = dp_size
542
543
+ @property
544
+ def activation_formats(
545
+ self
546
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
547
+ return (mk.FusedMoEActivationFormat.BatchedExperts,
548
+ mk.FusedMoEActivationFormat.BatchedExperts)
549
+
550
def supports_chunking(self) -> bool:
551
return False
552
553
+ def supports_expert_map(self) -> bool:
554
+ return False
555
+
556
def workspace_shapes(
557
self,
558
a: torch.Tensor,
@@ -554,20 +594,12 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
594
assert hidden_states.dim() == 3
595
assert expert_num_tokens is not None
596
557
- max_num_tokens = self.max_num_tokens
558
- num_dp = self.world_size // self.dp_size
597
num_local_experts = w1.size(0)
598
assert num_local_experts == w1.size(0), (
599
f"{num_local_experts} == {w1.size(0)}")
600
601
N = w1.size(1) // 2
602
565
- # Not cudagraph friendly
566
- assert (torch.compiler.is_compiling()
567
- or torch.cuda.is_current_stream_capturing()
568
- or torch.all(expert_num_tokens <= max_num_tokens * num_dp)), (
569
- f"{expert_num_tokens} <= {max_num_tokens * num_dp}")
570
-
603
for expert in range(num_local_experts):
604
# Indexing expert_num_tokens doesn't work w/cudagraphs or inductor
605
if (torch.compiler.is_compiling()
@@ -575,6 +607,10 @@ class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
607
num = hidden_states.shape[1]
608
else:
609
num = int(expert_num_tokens[expert].item())
610
+
611
+ if num == 0:
612
+ continue
613
+
614
tmp = _resize_cache(workspace2, (num, N))
615
input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
616
self.activation(activation, tmp, input)
@@ -590,34 +626,53 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
626
627
def __init__(
628
self,
593
- max_num_tokens: Optional[int] = None,
629
+ max_num_tokens: int,
630
+ world_size: int,
631
+ dp_size: int,
632
use_fp8_w8a8: bool = False,
633
use_int8_w8a8: bool = False,
634
use_int8_w8a16: bool = False,
635
use_int4_w4a16: bool = False,
598
- per_channel_quant: bool = False,
636
+ per_act_token_quant: bool = False,
637
block_shape: Optional[list[int]] = None,
600
- world_size: int = 1,
601
- dp_size: int = 1,
638
):
603
- super().__init__()
639
+ super().__init__(
640
+ FusedMoEQuantConfig.make(
641
+ use_fp8_w8a8=use_fp8_w8a8,
642
+ use_int8_w8a8=use_int8_w8a8,
643
+ use_int8_w8a16=use_int8_w8a16,
644
+ use_int4_w4a16=use_int4_w4a16,
645
+ per_act_token_quant=per_act_token_quant,
646
+ block_shape=block_shape,
647
+ ))
648
+ assert not use_int8_w8a8, "NYI"
649
+ assert not use_int8_w8a16, "NYI"
650
+ assert not use_int4_w4a16, "NYI"
651
self.use_fp8_w8a8 = use_fp8_w8a8
652
self.use_int8_w8a8 = use_int8_w8a8
653
self.use_int4_w4a16 = use_int4_w4a16
654
self.use_int8_w8a16 = use_int8_w8a16
608
- self.block_shape = block_shape
609
- self.per_channel_quant = per_channel_quant
655
self.max_num_tokens = max_num_tokens
656
self.world_size = world_size
657
self.dp_size = dp_size
613
-
614
- assert not use_int8_w8a8, "NYI"
615
- assert not use_int4_w4a16, "NYI"
616
- assert self.block_shape is None, "NYI"
658
+ assert world_size > 0
659
+ assert dp_size > 0
660
+ assert dp_size <= world_size
661
+ assert max_num_tokens > 0
662
+
663
+ @property
664
+ def activation_formats(
665
+ self
666
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
667
+ return (mk.FusedMoEActivationFormat.BatchedExperts,
668
+ mk.FusedMoEActivationFormat.BatchedExperts)
669
670
def supports_chunking(self) -> bool:
671
return False
672
673
+ def supports_expert_map(self) -> bool:
674
+ return False
675
+
676
def workspace_shapes(
677
self,
678
a: torch.Tensor,
@@ -630,10 +685,9 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
685
local_num_experts: int,
686
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
687
assert a.dim() == 2
633
- num_dp = self.world_size // self.dp_size
688
+ num_dp = self.world_size
689
num_experts = local_num_experts
635
- max_num_tokens = a.size(
636
- 0) if self.max_num_tokens is None else self.max_num_tokens
690
+ max_num_tokens = self.max_num_tokens
691
workspace13 = (num_experts, max_num_tokens * num_dp, max(K, N))
692
workspace2 = (num_experts, max_num_tokens * num_dp, (N // 2))
693
output = (num_experts, max_num_tokens * num_dp, K)
@@ -708,7 +762,6 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
762
raise ValueError(
763
f"Unsupported compute_type: {hidden_states.dtype}")
764
711
- #print(f"shape: E={E}, M={num_tokens}, N={N}, K={K}, top_k={top_k_num}")
765
# We can reuse the memory between these because by the time we need
766
# cache3, we're done with cache1
767
intermediate_cache1 = _resize_cache(workspace13,
@@ -734,6 +787,8 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
787
config=config,
788
block_shape=self.block_shape)
789
790
+ intermediate_cache2.fill_(0)
791
+
792
# TODO: would be nice to use expert_num_tokens here to reduce
793
# garbage compute
794
self.activation(activation, intermediate_cache2.view(-1, N // 2),
@@ -745,8 +800,8 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
800
qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
801
A=intermediate_cache2,
802
A_scale=a2_scale,
748
- qtype=torch.float8_e4m3fn if self.use_fp8_w8a8 else None,
749
- per_channel_quant=self.per_channel_quant,
803
+ quant_dtype=self.quant_dtype,
804
+ per_act_token_quant=self.per_act_token_quant,
805
block_shape=self.block_shape)
806
807
qintermediate_cache2 = qintermediate_cache2.view(
vllm/model_executor/layers/fused_moe/fused_moe.py
+44
-53
@@ -12,6 +12,10 @@ import vllm.envs as envs
12
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
from vllm import _custom_ops as ops
14
from vllm.logger import init_logger
15
+# yapf: disable
16
+from vllm.model_executor.layers.fused_moe.config import (
17
+ FusedMoEQuantConfig, get_config_quant_dtype)
18
+# yapf: enable
19
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
20
_valid_deep_gemm, deep_gemm_moe_fp8)
21
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
@@ -980,20 +984,6 @@ def get_config_dtype_str(
984
return None
985
986
983
-# TODO (bnell): use scalar_type instead of bools?
984
-def get_config_qtype(
985
- use_fp8_w8a8: bool,
986
- use_int8_w8a8: bool,
987
- use_int8_w8a16: bool,
988
- use_int4_w4a16: bool,
989
-) -> Optional[torch.dtype]:
990
- if use_fp8_w8a8:
991
- return torch.float8_e4m3fn
992
- elif use_int8_w8a8:
993
- return torch.int8
994
- return None
995
-
996
-
987
def inplace_fused_experts(hidden_states: torch.Tensor,
988
w1: torch.Tensor,
989
w2: torch.Tensor,
@@ -1262,10 +1252,10 @@ def fused_experts_impl(
1252
use_int4_w4a16=use_int4_w4a16,
1253
dtype=hidden_states.dtype)
1254
1265
- qtype = get_config_qtype(use_fp8_w8a8=use_fp8_w8a8,
1266
- use_int8_w8a8=use_int8_w8a8,
1267
- use_int8_w8a16=use_int8_w8a16,
1268
- use_int4_w4a16=use_int4_w4a16)
1255
+ qtype = get_config_quant_dtype(use_fp8_w8a8=use_fp8_w8a8,
1256
+ use_int8_w8a8=use_int8_w8a8,
1257
+ use_int8_w8a16=use_int8_w8a16,
1258
+ use_int4_w4a16=use_int4_w4a16)
1259
1260
get_config_func = functools.partial(
1261
try_get_optimal_moe_config,
@@ -1332,8 +1322,8 @@ def fused_experts_impl(
1322
qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
1323
A=curr_hidden_states,
1324
A_scale=a1_scale,
1335
- qtype=qtype,
1336
- per_channel_quant=per_channel_quant,
1325
+ quant_dtype=qtype,
1326
+ per_act_token_quant=per_channel_quant,
1327
block_shape=block_shape)
1328
1329
sorted_token_ids, expert_ids, num_tokens_post_padded = (
@@ -1373,8 +1363,8 @@ def fused_experts_impl(
1363
qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1364
A=intermediate_cache2,
1365
A_scale=a2_scale,
1376
- qtype=qtype,
1377
- per_channel_quant=per_channel_quant,
1366
+ quant_dtype=qtype,
1367
+ per_act_token_quant=per_channel_quant,
1368
block_shape=block_shape)
1369
1370
invoke_fused_moe_kernel(qintermediate_cache2,
@@ -1521,30 +1511,41 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1511
1512
def __init__(
1513
self,
1524
- use_fp8_w8a8: bool,
1525
- use_int8_w8a8: bool,
1526
- use_int8_w8a16: bool,
1527
- use_int4_w4a16: bool,
1528
- per_channel_quant: bool,
1514
+ use_fp8_w8a8: bool = False,
1515
+ use_int8_w8a8: bool = False,
1516
+ use_int8_w8a16: bool = False,
1517
+ use_int4_w4a16: bool = False,
1518
+ per_act_token_quant: bool = False,
1519
block_shape: Optional[list[int]] = None,
1530
- block_m: Optional[int] = None,
1520
):
1532
- super().__init__()
1521
+ super().__init__(
1522
+ FusedMoEQuantConfig.make(
1523
+ use_fp8_w8a8=use_fp8_w8a8,
1524
+ use_int8_w8a8=use_int8_w8a8,
1525
+ use_int8_w8a16=use_int8_w8a16,
1526
+ use_int4_w4a16=use_int4_w4a16,
1527
+ per_act_token_quant=per_act_token_quant,
1528
+ block_shape=block_shape,
1529
+ ))
1530
+
1531
self.use_fp8_w8a8 = use_fp8_w8a8
1532
self.use_int4_w4a16 = use_int4_w4a16
1533
self.use_int8_w8a8 = use_int8_w8a8
1534
self.use_int8_w8a16 = use_int8_w8a16
1537
- self.block_shape = block_shape
1538
- self.block_m = block_m
1539
- self.qtype = get_config_qtype(use_fp8_w8a8=use_fp8_w8a8,
1540
- use_int8_w8a8=use_int8_w8a8,
1541
- use_int8_w8a16=use_int8_w8a16,
1542
- use_int4_w4a16=use_int4_w4a16)
1543
- self.per_channel_quant = per_channel_quant
1535
+
1536
+ @property
1537
+ def activation_formats(
1538
+ self
1539
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
1540
+ return (mk.FusedMoEActivationFormat.Standard,
1541
+ mk.FusedMoEActivationFormat.Standard)
1542
1543
def supports_chunking(self) -> bool:
1544
return True
1545
1546
+ def supports_expert_map(self) -> bool:
1547
+ return True
1548
+
1549
def workspace_shapes(
1550
self,
1551
a: torch.Tensor,
@@ -1660,7 +1661,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1661
use_int8_w8a8=self.use_int8_w8a8,
1662
use_int8_w8a16=self.use_int8_w8a16,
1663
use_int4_w4a16=self.use_int4_w4a16,
1663
- per_channel_quant=self.per_channel_quant,
1664
+ per_channel_quant=self.per_act_token_quant,
1665
block_shape=self.block_shape)
1666
1667
self.activation(activation, intermediate_cache2,
@@ -1669,8 +1670,8 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1670
a2q_scale: Optional[torch.Tensor] = None
1671
1672
qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1672
- intermediate_cache2, a2_scale, self.qtype, self.per_channel_quant,
1673
- self.block_shape)
1673
+ intermediate_cache2, a2_scale, self.quant_dtype,
1674
+ self.per_act_token_quant, self.block_shape)
1675
1676
invoke_fused_moe_kernel(qintermediate_cache2,
1677
w2,
@@ -1690,7 +1691,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1691
use_int8_w8a8=self.use_int8_w8a8,
1692
use_int8_w8a16=self.use_int8_w8a16,
1693
use_int4_w4a16=self.use_int4_w4a16,
1693
- per_channel_quant=self.per_channel_quant,
1694
+ per_channel_quant=self.per_act_token_quant,
1695
block_shape=self.block_shape)
1696
1697
@@ -1699,27 +1700,17 @@ def modular_triton_fused_moe(
1700
use_int8_w8a8: bool,
1701
use_int8_w8a16: bool,
1702
use_int4_w4a16: bool,
1702
- per_channel_quant: bool,
1703
+ per_act_token_quant: bool,
1704
block_shape: Optional[list[int]] = None,
1705
) -> mk.FusedMoEModularKernel:
1705
- qtype = get_config_qtype(
1706
- use_fp8_w8a8=use_fp8_w8a8,
1707
- use_int8_w8a8=use_int8_w8a8,
1708
- use_int8_w8a16=use_int8_w8a16,
1709
- use_int4_w4a16=use_int4_w4a16,
1710
- )
1706
return mk.FusedMoEModularKernel(
1712
- MoEPrepareAndFinalizeNoEP(
1713
- quant_dtype=qtype,
1714
- per_channel_quant=per_channel_quant,
1715
- block_shape=block_shape,
1716
- ),
1707
+ MoEPrepareAndFinalizeNoEP(),
1708
TritonExperts(
1709
use_fp8_w8a8=use_fp8_w8a8,
1710
use_int8_w8a8=use_int8_w8a8,
1711
use_int8_w8a16=use_int8_w8a16,
1712
use_int4_w4a16=use_int4_w4a16,
1722
- per_channel_quant=per_channel_quant,
1713
+ per_act_token_quant=per_act_token_quant,
1714
block_shape=block_shape,
1715
),
1716
)
vllm/model_executor/layers/fused_moe/layer.py
+77
-311
@@ -3,27 +3,30 @@
3
4
from abc import abstractmethod
5
from collections.abc import Iterable
6
-from dataclasses import dataclass
6
from enum import Enum
8
-from typing import Callable, Literal, Optional, Union, overload
7
+from typing import Callable, Literal, Optional, overload
8
9
import torch
10
import torch.nn.functional as F
12
-from compressed_tensors.quantization import (QuantizationArgs,
13
- QuantizationStrategy,
14
- QuantizationType)
11
from torch.nn.parameter import UninitializedParameter
12
13
import vllm.envs as envs
18
-from vllm.config import ParallelConfig, get_current_vllm_config
14
+from vllm.config import get_current_vllm_config
15
from vllm.distributed import (get_dp_group, get_ep_group,
20
- get_tensor_model_parallel_rank,
16
get_tensor_model_parallel_world_size,
17
+ get_world_group,
18
tensor_model_parallel_all_reduce)
19
from vllm.distributed.eplb.eplb_state import EplbState
20
from vllm.forward_context import ForwardContext, get_forward_context
21
from vllm.logger import init_logger
22
from vllm.model_executor.custom_op import CustomOp
23
+# yapf: disable
24
+from vllm.model_executor.layers.fused_moe.config import (
25
+ FusedMoEConfig, FusedMoEParallelConfig)
26
+# yapf: enable
27
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
28
+ FusedMoEActivationFormat, FusedMoEModularKernel,
29
+ FusedMoEPermuteExpertsUnpermute, FusedMoEPrepareAndFinalize)
30
from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
31
is_rocm_aiter_moe_enabled)
32
from vllm.model_executor.layers.quantization.base_config import (
@@ -36,14 +39,12 @@ from vllm.utils import direct_register_custom_op, has_deep_ep, has_pplx
39
if current_platform.is_cuda_alike():
40
from .fused_batched_moe import BatchedTritonExperts
41
from .fused_moe import TritonExperts, fused_experts
39
- from .modular_kernel import (FusedMoEModularKernel,
40
- FusedMoEPermuteExpertsUnpermute,
41
- FusedMoEPrepareAndFinalize)
42
if has_pplx():
43
- from .pplx_prepare_finalize import PplxPrepareAndFinalize
43
+ from .pplx_prepare_finalize import (PplxPrepareAndFinalize,
44
+ pplx_hidden_dim_scale_bytes)
45
if has_deep_ep():
46
from .deepep_ht_prepare_finalize import DeepEPHTPrepareAndFinalize
46
- from .deepep_ll_prepare_finalize import (DEEPEP_QUANT_BLOCK_SIZE,
47
+ from .deepep_ll_prepare_finalize import (DEEPEP_QUANT_BLOCK_SHAPE,
48
DeepEPLLPrepareAndFinalize)
49
else:
50
fused_experts = None # type: ignore
@@ -60,207 +61,8 @@ if current_platform.is_tpu():
61
from .moe_pallas import fused_moe as fused_moe_pallas
62
else:
63
fused_moe_pallas = None # type: ignore
63
-logger = init_logger(__name__)
64
-
65
-
66
-@dataclass
67
-class FusedMoEParallelConfig:
68
- tp_size: int
69
- dp_size: int
70
- ep_size: int
71
- tp_rank: int
72
- dp_rank: int
73
- ep_rank: int
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):
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,
98
- vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
99
- """
100
- Determine MoE parallel configuration. Based on the input tp_size_,
101
- dp_size_, ep_size_ and vllm's parallel config, determine what
102
- level's of parallelism to use in the fused moe layer.
103
-
104
- Args:
105
- tp_size_ (int): tp_size passed into the FusedMoE constructor.
106
- dp_size_ (int): dp_size passed into the FusedMoE constructor.
107
- ep_size_ (int): ep_size passed into the FusedMoE constructor.
108
- vllm_parallel_config (ParallelConfig): vllm's parallel config
109
- object.
110
-
111
- Examples:
112
- When there is no parallelism requested, i.e. tp_size_ = dp_size_ = 1,
113
- we simply return the sizes unaltered and the ranks set to 0.
114
-
115
- Expert Parallelism is considered only when either dp_size_ or tp_size_
116
- is non trivial.
117
-
118
- When TP = 2, DP = 1 and EP = False, the configuration on different
119
- devices,
120
- - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} //
121
- legend : {size, rank}
122
- - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0}
123
- - Comment : Tensors are sharded across 2 devices.
124
-
125
- When TP = 1, DP = 2 and EP = False, the configuration on different
126
- devices,
127
- - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0}
128
- - device 1 : TP = {2, 1} DP = {2, 1} EP = {1, 0}
129
- - Comment: There are 2 engine instances and the tensors are sharded
130
- across 2 decvices.
131
-
132
- When TP = 2, DP = 2 and EP = False, the configuration on different
133
- devices,
134
- - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0}
135
- - device 1: TP = {4, 1} DP = {2, 0} EP = {1, 0}
136
- - device 2: TP = {4, 2} DP = {2, 1} EP = {1, 0}
137
- - device 3: TP = {4, 3} DP = {2, 1} EP = {1, 0}
138
- - Comment: There are 2 engine instances and the tensors are sharded
139
- across 4 devices.
140
-
141
- When, TP = 2, DP = 1 and EP = True, the configuration on different
142
- devices,
143
- - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0}
144
- - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1}
145
- - Comment: The experts are split between the 2 devices.
146
-
147
- When, TP = 1, DP = 2 and EP = True, the configuration on different
148
- devices,
149
- - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0}
150
- - device 1: TP = {1, 0} DP = {2, 1} EP = {2, 1}
151
- - Comment: There are 2 engine instances and the experts are split
152
- between the 2 devices.
153
-
154
- When TP = 2, DP = 2 and EP = True, the configuration on different
155
- devices,
156
- - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0}
157
- - device 1: TP = {1, 0} DP = {2, 0} EP = {4, 1}
158
- - device 2: TP = {1, 0} DP = {2, 1} EP = {4, 2}
159
- - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3}
160
- - Comment: There are 2 engine instances and the experts are split
161
- between the 4 devices.
162
- """
163
-
164
- def flatten_tp_across_dp(dp_rank: int):
165
- tp_rank = 0 if tp_size_ == 1 else get_tensor_model_parallel_rank()
166
- # There are actually dp_size_ * tp_size_ devices. Update tp_size
167
- # and tp_rank so we shard across all devices.
168
- tp_size = dp_size_ * tp_size_
169
- tp_rank = dp_rank * tp_size_ + tp_rank
170
- return tp_size, tp_rank
171
-
172
- use_ep = (dp_size_ * tp_size_ > 1
173
- and vllm_parallel_config.enable_expert_parallel)
174
-
175
- dp_size = dp_size_
176
- dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
177
- tp_size, tp_rank = flatten_tp_across_dp(dp_rank)
178
-
179
- if not use_ep:
180
- return FusedMoEParallelConfig(tp_size=tp_size,
181
- tp_rank=tp_rank,
182
- dp_size=dp_size,
183
- dp_rank=dp_rank,
184
- ep_size=1,
185
- ep_rank=0,
186
- use_ep=False)
187
- # DP + EP / TP + EP / DP + TP + EP
188
- assert use_ep
189
- # In EP, each device owns a set of experts fully. There is no tensor
190
- # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that.
191
- ep_size = tp_size
192
- ep_rank = tp_rank
193
- return FusedMoEParallelConfig(tp_size=1,
194
- tp_rank=0,
195
- dp_size=dp_size,
196
- dp_rank=dp_rank,
197
- ep_size=ep_size,
198
- ep_rank=ep_rank,
199
- use_ep=True)
200
-
201
-
202
-# Adapted from pplx-kernels tests/all_to_all_utils.py
203
-@dataclass
204
-class MoEConfig:
205
- num_experts: int
206
- experts_per_token: int
207
- hidden_dim: int
208
-
209
- num_local_experts: int
210
- moe_parallel_config: FusedMoEParallelConfig
211
-
212
- in_dtype: torch.dtype # The activation type.
213
- quant_dtype: torch.dtype = None
214
-
215
- # TODO: add more quantization params, blocked, per-token, etc.
216
- block_size: int = 128
217
-
218
- max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE
219
-
220
- def __post_init__(self):
221
- if self.dp_size > 1:
222
- logger.debug("Using MOEConfig::max_num_tokens=%d",
223
- self.max_num_tokens)
224
-
225
- @property
226
- def tp_size(self):
227
- return self.moe_parallel_config.tp_size
228
-
229
- @property
230
- def dp_size(self):
231
- return self.moe_parallel_config.dp_size
232
-
233
- @property
234
- def ep_size(self):
235
- return self.moe_parallel_config.ep_size
236
-
237
- @property
238
- def tp_rank(self):
239
- return self.moe_parallel_config.tp_rank
240
-
241
- @property
242
- def dp_rank(self):
243
- return self.moe_parallel_config.dp_rank
244
-
245
- @property
246
- def ep_rank(self):
247
- return self.moe_parallel_config.ep_rank
248
-
249
- @property
250
- def use_ep(self):
251
- return self.moe_parallel_config.use_ep
64
253
- @property
254
- def use_pplx_kernels(self):
255
- return self.moe_parallel_config.use_pplx_kernels
256
-
257
- @property
258
- def use_deepep_ht_kernels(self):
259
- return self.moe_parallel_config.use_deepep_ht_kernels
260
-
261
- @property
262
- def use_deepep_ll_kernels(self):
263
- return self.moe_parallel_config.use_deepep_ll_kernels
65
+logger = init_logger(__name__)
66
67
68
class FusedMoeWeightScaleSupported(Enum):
@@ -270,21 +72,9 @@ class FusedMoeWeightScaleSupported(Enum):
72
BLOCK = "block"
73
74
273
-def get_quant_config_input_activations(
274
- quant_config: Optional[QuantizationConfig]
275
-) -> Optional[QuantizationArgs]:
276
- if (quant_config is not None and hasattr(quant_config, 'target_scheme_map')
277
- and "Linear" in quant_config.target_scheme_map and
278
- "input_activations" in quant_config.target_scheme_map["Linear"]):
279
- return quant_config.target_scheme_map["Linear"].get(
280
- "input_activations")
281
- else:
282
- return None
283
-
284
-
75
class FusedMoEMethodBase(QuantizeMethodBase):
76
287
- moe: MoEConfig
77
+ moe: FusedMoEConfig
78
79
@abstractmethod
80
def create_weights(self, layer: torch.nn.Module, num_experts: int,
@@ -292,23 +82,25 @@ class FusedMoEMethodBase(QuantizeMethodBase):
82
params_dtype: torch.dtype, **extra_weight_attrs):
83
raise NotImplementedError
84
295
- def init_prepare_finalize(self, moe: MoEConfig,
85
+ def init_prepare_finalize(self, moe: FusedMoEConfig,
86
quant_config: Optional[QuantizationConfig]):
87
all2all_manager = get_ep_group().device_communicator.all2all_manager
88
assert all2all_manager is not None
89
90
self.moe = moe
301
- quant_dtype = None
302
- act_quant_block_size = None
303
- from vllm.model_executor.layers.quantization.fp8 import Fp8Config
304
- if isinstance(quant_config, Fp8Config):
305
- act_quant_block_size = quant_config.weight_block_size
306
- quant_dtype = torch.float8_e4m3fn
307
-
308
- prepare_finalize: Optional[Union[PplxPrepareAndFinalize,
309
- DeepEPHTPrepareAndFinalize,
310
- DeepEPLLPrepareAndFinalize]] = None
91
+
92
+ prepare_finalize: Optional[FusedMoEPrepareAndFinalize] = None
93
+
94
if moe.use_pplx_kernels:
95
+ hidden_dim_bytes, hidden_scale_bytes = pplx_hidden_dim_scale_bytes(
96
+ moe.max_num_tokens,
97
+ moe.hidden_dim,
98
+ moe.in_dtype,
99
+ moe.quant_dtype,
100
+ per_act_token_quant=moe.per_act_token_quant,
101
+ block_shape=moe.block_shape,
102
+ )
103
+
104
all_to_all_args = dict(
105
max_num_tokens=moe.max_num_tokens,
106
num_experts=moe.num_experts,
@@ -318,14 +110,8 @@ class FusedMoEMethodBase(QuantizeMethodBase):
110
# dp_size actually means tp_size, bug in pplx kernels
111
dp_size=all2all_manager.tp_group.world_size,
112
hidden_dim=moe.hidden_dim,
321
- hidden_dim_bytes=moe.hidden_dim * moe.quant_dtype.itemsize,
322
- # For blocked per token: set to
323
- # ceil_div(hidden_dim, block_size) * sizeof(float32)
324
- # For per-token: set to sizeof(float32)
325
- hidden_dim_scale_bytes=(
326
- 0 if moe.quant_dtype.itemsize != 1 else
327
- ((moe.hidden_dim + moe.block_size - 1) // moe.block_size *
328
- torch.float32.itemsize)),
113
+ hidden_dim_bytes=hidden_dim_bytes,
114
+ hidden_dim_scale_bytes=hidden_scale_bytes,
115
)
116
117
# Intranode pplx a2a takes a group name while internode does not.
@@ -335,9 +121,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
121
122
handle = all2all_manager.get_handle(all_to_all_args)
123
338
- input_activations = get_quant_config_input_activations(
339
- quant_config)
340
-
124
prepare_finalize = PplxPrepareAndFinalize(
125
handle,
126
max_num_tokens=moe.max_num_tokens,
@@ -345,10 +128,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
128
rank=all2all_manager.rank,
129
# dp_size actually means tp_size, bug in pplx kernels
130
dp_size=all2all_manager.tp_group.world_size,
348
- quant_dtype=moe.quant_dtype,
349
- per_act_token=(input_activations.strategy
350
- == QuantizationStrategy.TOKEN
351
- if input_activations is not None else False),
131
)
132
elif moe.use_deepep_ht_kernels:
133
assert moe.dp_size == all2all_manager.dp_world_size
@@ -362,8 +141,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):
141
dp_size=all2all_manager.dp_world_size,
142
rank_expert_offset=all2all_manager.rank *
143
moe.num_local_experts,
365
- quant_dtype=quant_dtype,
366
- block_shape=act_quant_block_size,
144
)
145
146
elif moe.use_deepep_ll_kernels:
@@ -380,25 +157,25 @@ class FusedMoEMethodBase(QuantizeMethodBase):
157
158
# Note : We may want to use FP8 dispatch even otherwise just to
159
# reduce datamovement
383
- assert act_quant_block_size is not None
384
- use_fp8_dispatch = (quant_dtype == current_platform.fp8_dtype()
385
- and act_quant_block_size[1]
386
- == DEEPEP_QUANT_BLOCK_SIZE)
160
+ use_fp8_dispatch = (moe.quant_config is not None
161
+ and moe.quant_config.quant_dtype
162
+ == current_platform.fp8_dtype()
163
+ and moe.quant_config.block_shape
164
+ == DEEPEP_QUANT_BLOCK_SHAPE)
165
166
# Note (varun): Whether to use FP8 dispatch or not needs some
167
# profiling. Turning it off for now.
168
prepare_finalize = DeepEPLLPrepareAndFinalize(
169
handle,
170
+ max_tokens_per_rank=moe.max_num_tokens,
171
world_size=all2all_manager.world_size,
172
dp_size=all2all_manager.dp_world_size,
394
- max_tokens_per_rank=moe.max_num_tokens,
395
- quant_dtype=quant_dtype,
396
- block_shape=act_quant_block_size,
173
use_fp8_dispatch=use_fp8_dispatch,
174
)
175
176
self.topk_indices_dtype = None
177
if prepare_finalize is not None:
178
+ logger.debug("%s", prepare_finalize.__class__.__name__)
179
self.topk_indices_dtype = prepare_finalize.topk_indices_dtype()
180
experts = self.select_gemm_impl(prepare_finalize, moe)
181
self.fused_experts = FusedMoEModularKernel(
@@ -407,13 +184,15 @@ class FusedMoEMethodBase(QuantizeMethodBase):
184
)
185
186
def select_gemm_impl(
410
- self, prepare_finalize: FusedMoEPrepareAndFinalize,
411
- moe: Optional[MoEConfig]) -> FusedMoEPermuteExpertsUnpermute:
187
+ self,
188
+ prepare_finalize: FusedMoEPrepareAndFinalize,
189
+ moe: FusedMoEConfig,
190
+ ) -> FusedMoEPermuteExpertsUnpermute:
191
# based on the all2all implementation, select the appropriate
192
# gemm implementation
193
raise NotImplementedError(
415
- "Subclass must select appropriate gemm implementation"
416
- " based on the prepare_finalize")
194
+ f"{self.__class__.__name__} must select appropriate gemm "
195
+ "implementation based on the prepare_finalize")
196
197
@abstractmethod
198
def apply(
@@ -445,7 +224,7 @@ class FusedMoEMethodBase(QuantizeMethodBase):
224
class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
225
"""MoE method without quantization."""
226
448
- def __init__(self, moe: MoEConfig):
227
+ def __init__(self, moe: FusedMoEConfig):
228
super().__init__()
229
self.fused_experts = fused_experts # type: ignore
230
self.topk_indices_dtype = None
@@ -458,44 +237,30 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
237
else:
238
self.rocm_aiter_fused_experts = None # type: ignore
239
461
- def select_gemm_impl(self, prepare_finalize: FusedMoEPrepareAndFinalize,
462
- moe: Optional[MoEConfig]):
240
+ def select_gemm_impl(
241
+ self,
242
+ prepare_finalize: FusedMoEPrepareAndFinalize,
243
+ moe: FusedMoEConfig,
244
+ ) -> FusedMoEPermuteExpertsUnpermute:
245
246
assert self.fused_experts == fused_experts
247
248
all2all_manager = get_ep_group().device_communicator.all2all_manager
249
assert all2all_manager is not None
250
469
- experts: Optional[FusedMoEPermuteExpertsUnpermute] = None
470
-
471
- use_batched_experts = prepare_finalize.max_num_tokens_per_rank(
472
- ) is not None
473
- if use_batched_experts:
251
+ if (prepare_finalize.activation_format ==
252
+ FusedMoEActivationFormat.BatchedExperts):
253
logger.debug("BatchedTritonExperts %s", self.moe)
254
assert self.moe.dp_size == all2all_manager.dp_world_size
476
- experts = BatchedTritonExperts(
255
+ return BatchedTritonExperts(
256
max_num_tokens=self.moe.max_num_tokens,
257
world_size=all2all_manager.world_size,
258
# dp_size actually means tp_size, bug in pplx kernels
259
dp_size=all2all_manager.tp_group.world_size,
481
- use_fp8_w8a8=False,
482
- use_int8_w8a8=False,
483
- use_int8_w8a16=False,
484
- use_int4_w4a16=False,
485
- block_shape=None,
486
- per_channel_quant=False,
260
)
261
else:
262
logger.debug("TritonExperts %s", self.moe)
490
- experts = TritonExperts(
491
- use_fp8_w8a8=False,
492
- use_int8_w8a8=False,
493
- use_int8_w8a16=False,
494
- use_int4_w4a16=False,
495
- block_shape=None,
496
- per_channel_quant=False,
497
- )
498
- return experts
263
+ return TritonExperts()
264
265
def create_weights(self, layer: torch.nn.Module, num_experts: int,
266
hidden_size: int, intermediate_size_per_partition: int,
@@ -883,13 +648,18 @@ class FusedMoE(torch.nn.Module):
648
params_dtype = torch.get_default_dtype()
649
self.params_dtype = params_dtype
650
651
+ tp_size_ = (tp_size if tp_size is not None else
652
+ get_tensor_model_parallel_world_size())
653
+ dp_size_ = (dp_size
654
+ if dp_size is not None else get_dp_group().world_size)
655
+ world_size_ = get_world_group().world_size
656
+
657
vllm_config = get_current_vllm_config()
658
self.moe_parallel_config: FusedMoEParallelConfig = (
659
FusedMoEParallelConfig.make(
889
- tp_size_=(tp_size if tp_size is not None else
890
- get_tensor_model_parallel_world_size()),
891
- dp_size_=(dp_size if dp_size is not None else
892
- get_dp_group().world_size),
660
+ tp_size_=tp_size_,
661
+ dp_size_=dp_size_,
662
+ world_size_=world_size_,
663
vllm_parallel_config=vllm_config.parallel_config))
664
665
self.global_num_experts = num_experts + num_redundant_experts
@@ -948,25 +718,22 @@ class FusedMoE(torch.nn.Module):
718
from vllm_hpu_extension.ops import DynamicFusedMOE
719
self.hpu_fused_moe = DynamicFusedMOE(self.global_num_experts)
720
951
- # Only support float8 for now.
952
- quant_dtype = params_dtype
953
- if quant_config is not None:
954
- input_activations = get_quant_config_input_activations(
955
- quant_config)
956
- if (input_activations is not None
957
- and input_activations.num_bits == 8
958
- and input_activations.type == QuantizationType.FLOAT):
959
- quant_dtype = torch.float8_e4m3fn
960
-
961
- moe = MoEConfig(
721
+ if vllm_config.model_config is not None:
722
+ model_dtype = vllm_config.model_config.dtype
723
+ else:
724
+ # TODO (bnell): This is a hack to get test_mixtral_moe to work
725
+ # since model_config is not set in the pytest test.
726
+ model_dtype = params_dtype
727
+
728
+ moe = FusedMoEConfig.make(
729
num_experts=self.global_num_experts,
730
experts_per_token=top_k,
731
hidden_dim=hidden_size,
732
num_local_experts=self.local_num_experts,
733
moe_parallel_config=self.moe_parallel_config,
967
- in_dtype=params_dtype,
968
- quant_dtype=quant_dtype,
734
+ in_dtype=model_dtype,
735
max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
736
+ quant_config=quant_config,
737
)
738
self.moe_config = moe
739
self.quant_config = quant_config
@@ -1017,16 +784,15 @@ class FusedMoE(torch.nn.Module):
784
self.batched_router_logits: Optional[torch.Tensor] = None
785
if (self.moe_parallel_config.use_pplx_kernels
786
or self.moe_parallel_config.use_deepep_ll_kernels):
1020
- act_dtype = vllm_config.model_config.dtype
787
self.batched_hidden_states = torch.zeros(
1022
- (envs.VLLM_MOE_DP_CHUNK_SIZE, self.hidden_size),
1023
- dtype=act_dtype,
788
+ (moe.max_num_tokens, self.hidden_size),
789
+ dtype=moe.in_dtype,
790
device=torch.cuda.current_device())
791
792
# Note here we use `num_experts` which is logical expert count
793
self.batched_router_logits = torch.zeros(
1028
- (envs.VLLM_MOE_DP_CHUNK_SIZE, num_experts),
1029
- dtype=act_dtype,
794
+ (moe.max_num_tokens, num_experts),
795
+ dtype=moe.in_dtype,
796
device=torch.cuda.current_device())
797
798
@property
@@ -1588,7 +1354,7 @@ class FusedMoE(torch.nn.Module):
1354
1355
assert (self.batched_hidden_states.size(0) # type: ignore
1356
>= chunk_size)
1591
- assert (self.batched_router_logits.size(0) # type: ignore
1357
+ assert (self.batched_router_logits.size(0) # type: ignore
1358
>= chunk_size)
1359
staged_hidden_states = self.batched_hidden_states[:
1360
chunk_size, :] # type: ignore
vllm/model_executor/layers/fused_moe/modular_kernel.py
+84
-3
@@ -1,12 +1,14 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from abc import ABC, abstractmethod
4
+from enum import Enum
5
from math import prod
5
-from typing import Optional
6
+from typing import Optional, final
7
8
import torch
9
10
import vllm.envs as envs
11
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
12
from vllm.model_executor.layers.fused_moe.utils import _resize_cache
13
from vllm.utils import cdiv
14
@@ -82,6 +84,18 @@ def _moe_problem_size(
84
return E, M, N, K, topk
85
86
87
+class FusedMoEActivationFormat(Enum):
88
+ """
89
+ The standard activation format (num_tokens, hidden dim).
90
+ """
91
+ Standard = "standard",
92
+ """
93
+ The batched experts format (num experts, max tokens per expert, hidden dim)
94
+ """
95
+ BatchedExperts = "batched_experts",
96
+
97
+
98
+# TODO: pass FusedMoEParallelConfig in as ctor parameter?
99
class FusedMoEPrepareAndFinalize(ABC):
100
"""
101
An abstract base class for the [Quantize-Prepare] and [Finalize] steps
@@ -99,6 +113,7 @@ class FusedMoEPrepareAndFinalize(ABC):
113
num_experts: int,
114
expert_map: Optional[torch.Tensor],
115
apply_router_weight_on_input: bool,
116
+ quant_config: FusedMoEQuantConfig,
117
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
118
Optional[torch.Tensor], Optional[torch.Tensor]]:
119
"""
@@ -148,6 +163,15 @@ class FusedMoEPrepareAndFinalize(ABC):
163
"""
164
raise NotImplementedError
165
166
+ @property
167
+ @abstractmethod
168
+ def activation_format(self) -> FusedMoEActivationFormat:
169
+ """
170
+ A property indicating the output format of the activations for the
171
+ 'prepare' method.
172
+ """
173
+ raise NotImplementedError
174
+
175
@abstractmethod
176
def topk_indices_dtype(self) -> Optional[torch.dtype]:
177
"""
@@ -176,6 +200,41 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
200
above.
201
"""
202
203
+ def __init__(
204
+ self,
205
+ quant_config: Optional[FusedMoEQuantConfig],
206
+ ):
207
+ if quant_config is not None:
208
+ self.quant_config = quant_config
209
+ else:
210
+ self.quant_config = FusedMoEQuantConfig()
211
+
212
+ @property
213
+ @abstractmethod
214
+ def activation_formats(
215
+ self) -> tuple[FusedMoEActivationFormat, FusedMoEActivationFormat]:
216
+ """
217
+ A property which is a tuple of the input and output activation formats
218
+ for the 'apply' method.
219
+ """
220
+ raise NotImplementedError
221
+
222
+ @property
223
+ def quant_dtype(self) -> Optional[torch.dtype]:
224
+ return self.quant_config.quant_dtype
225
+
226
+ @property
227
+ def block_shape(self) -> Optional[list[int]]:
228
+ return self.quant_config.block_shape
229
+
230
+ @property
231
+ def per_act_token_quant(self) -> bool:
232
+ return self.quant_config.per_act_token_quant
233
+
234
+ @property
235
+ def per_out_ch_quant(self) -> bool:
236
+ return self.quant_config.per_out_ch_quant
237
+
238
# TODO (bnell): make this return a CHUNK_SIZE or None instead?
239
@abstractmethod
240
def supports_chunking(self) -> bool:
@@ -185,6 +244,13 @@ class FusedMoEPermuteExpertsUnpermute(ABC):
244
"""
245
raise NotImplementedError
246
247
+ @abstractmethod
248
+ def supports_expert_map(self) -> bool:
249
+ """
250
+ A flag indicating whether or not this class supports expert maps
251
+ """
252
+ raise NotImplementedError
253
+
254
@abstractmethod
255
def workspace_shapes(
256
self,
@@ -297,6 +363,7 @@ def _chunk_scales(scales: Optional[torch.Tensor], start: int,
363
return None
364
365
366
+@final
367
class FusedMoEModularKernel(torch.nn.Module):
368
"""
369
This class combines a FusedMoEPrepareAndFinalize instance and
@@ -318,6 +385,12 @@ class FusedMoEModularKernel(torch.nn.Module):
385
super().__init__()
386
self.prepare_finalize = prepare_finalize
387
self.fused_experts = fused_experts
388
+ assert prepare_finalize.activation_format == \
389
+ fused_experts.activation_formats[0], (
390
+ f"{prepare_finalize.__class__.__name__}."
391
+ f"{prepare_finalize.activation_format} == "
392
+ f"{fused_experts.__class__.__name__}."
393
+ f"{fused_experts.activation_formats[0]}")
394
395
def forward(
396
self,
@@ -383,8 +456,16 @@ class FusedMoEModularKernel(torch.nn.Module):
456
457
(a1q, a1q_scale, expert_num_tokens, _expert_topk_ids,
458
_expert_topk_weights) = self.prepare_finalize.prepare(
386
- a1, a1_scale, a2_scale, topk_weights, topk_ids,
387
- global_num_experts, expert_map, apply_router_weight_on_input)
459
+ a1,
460
+ a1_scale,
461
+ a2_scale,
462
+ topk_weights,
463
+ topk_ids,
464
+ global_num_experts,
465
+ expert_map,
466
+ apply_router_weight_on_input,
467
+ self.fused_experts.quant_config,
468
+ )
469
470
# Maybe prepare gathered topk_ids and topk_weights from other EP ranks.
471
topk_ids = topk_ids if _expert_topk_ids is None else _expert_topk_ids
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
+77
-30
@@ -6,33 +6,76 @@ import pplx_kernels as pplx
6
import torch
7
8
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
9
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10
from vllm.model_executor.layers.fused_moe.utils import (
11
moe_kernel_quantize_input)
12
+from vllm.utils import cdiv, round_up
13
+
14
+
15
+def pplx_hidden_dim_scale_bytes(
16
+ max_num_tokens: int,
17
+ hidden_dim: int,
18
+ in_dtype: torch.dtype,
19
+ quant_dtype: Optional[torch.dtype],
20
+ per_act_token_quant: bool,
21
+ block_shape: Optional[list[int]],
22
+):
23
+ # All pplx byte sizes must be 16-byte aligned.
24
+ align = 16
25
+
26
+ # For blocked per token: set to
27
+ # ceil_div(hidden_dim, block_size) * sizeof(float32)
28
+ # For per-token: set to 4 * sizeof(float32) (x4 for alignment)
29
+ if quant_dtype is not None:
30
+ assert quant_dtype.itemsize == 1
31
+ hidden_dim_bytes = hidden_dim * quant_dtype.itemsize
32
+ elem_size = torch.float32.itemsize
33
+
34
+ if per_act_token_quant:
35
+ # per-token
36
+ assert block_shape is None
37
+ hidden_scale_bytes = elem_size
38
+ elif block_shape is not None:
39
+ # per-group
40
+ block_size = block_shape[1]
41
+ num_blocks = cdiv(hidden_dim, block_size)
42
+ hidden_scale_bytes = num_blocks * elem_size
43
+ else:
44
+ # per-tensor
45
+ hidden_scale_bytes = elem_size
46
+ else:
47
+ hidden_dim_bytes = hidden_dim * in_dtype.itemsize
48
+ hidden_scale_bytes = 0
49
+
50
+ return (
51
+ round_up(hidden_dim_bytes, align),
52
+ round_up(hidden_scale_bytes, align),
53
+ )
54
55
56
# The max_num_tokens, world_size and dp_size must be the same
57
# as the ones used to create the AllToAll.
58
class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
59
17
- def __init__(self,
18
- a2a: pplx.AllToAll,
19
- max_num_tokens: int,
20
- world_size: int,
21
- rank: int,
22
- dp_size: int,
23
- quant_dtype: Optional[torch.dtype] = None,
24
- block_shape: Optional[list[int]] = None,
25
- per_act_token: bool = False):
60
+ def __init__(
61
+ self,
62
+ a2a: pplx.AllToAll,
63
+ max_num_tokens: int,
64
+ world_size: int,
65
+ rank: int,
66
+ dp_size: int,
67
+ ):
68
super().__init__()
69
assert max_num_tokens > 0
70
self.a2a = a2a
29
- self.block_shape = block_shape
71
self.max_num_tokens = max_num_tokens
72
self.world_size = world_size
73
self.rank = rank
74
self.dp_size = dp_size
34
- self.quant_dtype = quant_dtype
35
- self.per_act_token = per_act_token
75
+
76
+ @property
77
+ def activation_format(self) -> mk.FusedMoEActivationFormat:
78
+ return mk.FusedMoEActivationFormat.BatchedExperts
79
80
def max_num_tokens_per_rank(self) -> Optional[int]:
81
return self.max_num_tokens
@@ -45,36 +88,43 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
88
a1: torch.Tensor,
89
a1_scale: Optional[torch.Tensor],
90
a2_scale: Optional[torch.Tensor],
48
- rank_topk_weights: torch.Tensor,
49
- rank_topk_ids: torch.Tensor,
91
+ topk_weights: torch.Tensor,
92
+ topk_ids: torch.Tensor,
93
num_experts: int,
94
expert_map: Optional[torch.Tensor],
95
apply_router_weight_on_input: bool,
96
+ quant_config: FusedMoEQuantConfig,
97
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
98
Optional[torch.Tensor], Optional[torch.Tensor]]:
99
num_tokens = a1.size(0) # M
100
hidden_dim = a1.size(-1) # K
101
58
- assert rank_topk_ids.size(0) == num_tokens
102
+ assert topk_ids.size(0) == num_tokens
103
# assert expert_map is None, "NYI"
104
105
# Is this always going to be a1.device?
106
device = a1.device
107
108
if apply_router_weight_on_input:
65
- topk = rank_topk_ids.size(1)
109
+ topk = topk_ids.size(1)
110
# TODO: this only works for topK=1, will need to update for topK>1
111
assert topk == 1, (
112
"apply_router_weight_on_input is only implemented for topk=1")
69
- a1 = a1 * rank_topk_weights.to(a1.dtype)
113
+ a1 = a1 * topk_weights.to(a1.dtype)
114
115
repeat_cols = 4
72
- repeat_rows = 1 if self.per_act_token else a1.size(0)
116
+ repeat_rows = 1 if quant_config.per_act_token_quant else a1.size(0)
117
a1q, a1q_scale = moe_kernel_quantize_input(
74
- a1, (None if self.per_act_token else a1_scale), self.quant_dtype,
75
- self.per_act_token, self.block_shape)
118
+ a1, (None if quant_config.per_act_token_quant else a1_scale),
119
+ quant_dtype=quant_config.quant_dtype,
120
+ per_act_token_quant=quant_config.per_act_token_quant,
121
+ block_shape=quant_config.block_shape)
122
123
if a1q_scale is not None:
124
+ if a1q_scale.numel() == 1:
125
+ orig_a_scale_block_shape = 1
126
+ else:
127
+ orig_a_scale_block_shape = a1q_scale.shape[-1]
128
a1q_scale = a1q_scale.repeat(repeat_rows, repeat_cols)
129
130
# rem_experts need to be 0 for pplx to work properly.
@@ -98,15 +148,12 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
148
149
expert_x_scale: Optional[torch.Tensor] = None
150
if a1q.dtype.itemsize == 1:
101
- float32_size = torch.float32.itemsize
102
- block_size = (self.block_shape[0] if self.block_shape is not None
103
- else 1) * float32_size
151
+ block_size = (quant_config.block_shape[1]
152
+ if quant_config.block_shape is not None else 1)
153
expert_x_scale = torch.empty(
105
- (
106
- num_local_experts,
107
- expert_x.size(1),
108
- (expert_x.size(2) + block_size - 1) // block_size,
109
- ),
154
+ (num_local_experts, expert_x.size(1),
155
+ round_up(
156
+ (expert_x.size(2) + block_size - 1) // block_size, 4)),
157
dtype=torch.float32,
158
device=device,
159
)
@@ -121,11 +168,11 @@ class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
168
out_expert_x_scale=expert_x_scale,
169
dp_x=a1q,
170
dp_x_scale=a1q_scale,
124
- indices=rank_topk_ids,
171
+ indices=topk_ids,
172
bound_m=bound_m,
173
)
174
if expert_x_scale is not None:
128
- expert_x_scale = expert_x_scale[:, :, 0:1]
175
+ expert_x_scale = expert_x_scale[:, :, :orig_a_scale_block_shape]
176
177
return expert_x, expert_x_scale, expert_num_tokens, None, None
178
vllm/model_executor/layers/fused_moe/prepare_finalize.py
+9
-15
@@ -5,6 +5,7 @@ from typing import Optional
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9
from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
10
_moe_unpermute_and_reduce)
11
from vllm.model_executor.layers.fused_moe.utils import (
@@ -13,16 +14,9 @@ from vllm.model_executor.layers.fused_moe.utils import (
14
15
class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
16
16
- def __init__(
17
- self,
18
- quant_dtype: Optional[torch.dtype] = None,
19
- per_channel_quant: bool = False,
20
- block_shape: Optional[list[int]] = None,
21
- ):
22
- super().__init__()
23
- self.per_channel_quant = per_channel_quant
24
- self.block_shape = block_shape
25
- self.quant_dtype = quant_dtype
17
+ @property
18
+ def activation_format(self) -> mk.FusedMoEActivationFormat:
19
+ return mk.FusedMoEActivationFormat.Standard
20
21
def max_num_tokens_per_rank(self) -> Optional[int]:
22
return None
@@ -39,7 +33,8 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
33
topk_ids: torch.Tensor,
34
num_experts: int,
35
expert_map: Optional[torch.Tensor],
42
- apply_router_weight_on_input: bool = False,
36
+ apply_router_weight_on_input: bool,
37
+ quant_config: FusedMoEQuantConfig,
38
) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor],
39
Optional[torch.Tensor], Optional[torch.Tensor]]:
40
@@ -50,10 +45,9 @@ class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
45
"apply_router_weight_on_input is only implemented for topk=1"
46
a1.mul_(topk_weights.to(a1.dtype))
47
53
- a1q, a1q_scale = moe_kernel_quantize_input(a1, a1_scale,
54
- self.quant_dtype,
55
- self.per_channel_quant,
56
- self.block_shape)
48
+ a1q, a1q_scale = moe_kernel_quantize_input(
49
+ a1, a1_scale, quant_config.quant_dtype,
50
+ quant_config.per_act_token_quant, quant_config.block_shape)
51
52
return a1q, a1q_scale, None, None, None
53
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
+46
-22
@@ -5,6 +5,7 @@ from typing import Optional
5
import torch
6
7
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
9
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
10
DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape)
11
from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
@@ -12,34 +13,59 @@ from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
13
14
class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
15
15
- def __init__(self,
16
- use_fp8_w8a8: bool = False,
17
- use_int8_w8a8: bool = False,
18
- use_int8_w8a16: bool = False,
19
- use_int4_w4a16: bool = False,
20
- per_channel_quant: bool = False,
21
- block_shape: Optional[list[int]] = None,
22
- block_m: Optional[int] = None,
23
- allow_deep_gemm: bool = False):
24
- super().__init__()
25
- self.triton_expert = TritonExperts(use_fp8_w8a8=use_fp8_w8a8,
26
- use_int8_w8a8=use_int8_w8a8,
27
- use_int4_w4a16=use_int4_w4a16,
28
- use_int8_w8a16=use_int8_w8a16,
29
- per_channel_quant=per_channel_quant,
30
- block_shape=block_shape,
31
- block_m=block_m)
32
- self.allow_deep_gemm = allow_deep_gemm
33
- self.use_fp8_w8a8 = use_fp8_w8a8
16
+ def __init__(
17
+ self,
18
+ use_fp8_w8a8: bool = False,
19
+ use_int8_w8a8: bool = False,
20
+ use_int8_w8a16: bool = False,
21
+ use_int4_w4a16: bool = False,
22
+ per_act_token_quant: bool = False,
23
+ block_shape: Optional[list[int]] = None,
24
+ allow_deep_gemm: bool = False,
25
+ ):
26
+ super().__init__(
27
+ FusedMoEQuantConfig.make(
28
+ use_fp8_w8a8=use_fp8_w8a8,
29
+ use_int8_w8a8=use_int8_w8a8,
30
+ use_int8_w8a16=use_int8_w8a16,
31
+ use_int4_w4a16=use_int4_w4a16,
32
+ per_act_token_quant=per_act_token_quant,
33
+ block_shape=block_shape,
34
+ ))
35
+ self.triton_expert = TritonExperts(
36
+ use_fp8_w8a8=use_fp8_w8a8,
37
+ use_int8_w8a8=use_int8_w8a8,
38
+ use_int4_w4a16=use_int4_w4a16,
39
+ use_int8_w8a16=use_int8_w8a16,
40
+ per_act_token_quant=per_act_token_quant,
41
+ block_shape=block_shape,
42
+ )
43
+ self.allow_deep_gemm = (allow_deep_gemm and not per_act_token_quant
44
+ and use_fp8_w8a8)
45
self.deep_gemm_expert = DeepGemmExperts(
46
) if self.allow_deep_gemm else None
47
48
+ @property
49
+ def activation_formats(
50
+ self
51
+ ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
52
+ assert (self.deep_gemm_expert is None
53
+ or self.triton_expert.activation_formats
54
+ == self.deep_gemm_expert.activation_formats)
55
+ return self.triton_expert.activation_formats
56
+
57
def supports_chunking(self) -> bool:
58
dge = self.deep_gemm_expert
59
te = self.triton_expert
60
return ((dge is None or dge.supports_chunking())
61
and (te is None or te.supports_chunking()))
62
63
+ def supports_expert_map(self) -> bool:
64
+ dge = self.deep_gemm_expert
65
+ te = self.triton_expert
66
+ return ((dge is None or dge.supports_expert_map())
67
+ and (te is None or te.supports_expert_map()))
68
+
69
def workspace_shapes(
70
self,
71
a: torch.Tensor,
@@ -83,9 +109,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
109
workspace2: torch.Tensor,
110
expert_num_tokens: Optional[torch.Tensor],
111
):
86
- N = w1.size(1)
87
-
88
- use_deep_gemm = (self.allow_deep_gemm and self.use_fp8_w8a8 and N > 512
112
+ use_deep_gemm = (self.allow_deep_gemm
113
and _valid_deep_gemm(hidden_states, w1, w2))
114
115
experts = self.deep_gemm_expert if use_deep_gemm else self.triton_expert
vllm/model_executor/layers/fused_moe/utils.py
+22
-7
@@ -37,6 +37,7 @@ def _fp8_quantize(
37
A, A_scale = ops.scaled_fp8_quant(
38
A, A_scale, use_per_token_if_dynamic=per_act_token)
39
else:
40
+ assert not per_act_token
41
assert len(block_shape) == 2
42
_, block_k = block_shape[0], block_shape[1]
43
A, A_scale = per_token_group_quant_fp8(A, block_k)
@@ -64,6 +65,7 @@ def _int8_quantize(
65
"int8 quantization only supports block or channel-wise"
66
A, A_scale = per_token_quant_int8(A)
67
else:
68
+ assert not per_act_token
69
assert len(block_shape) == 2
70
_, block_k = block_shape[0], block_shape[1]
71
A, A_scale = per_token_group_quant_int8(A, block_k)
@@ -75,16 +77,15 @@ def _int8_quantize(
77
def moe_kernel_quantize_input(
78
A: torch.Tensor,
79
A_scale: Optional[torch.Tensor],
78
- qtype: Optional[torch.dtype],
79
- per_channel_quant: bool,
80
+ quant_dtype: Optional[torch.dtype],
81
+ per_act_token_quant: bool,
82
block_shape: Optional[list[int]] = None,
83
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
82
- if qtype == torch.float8_e4m3fn:
83
- return _fp8_quantize(A, A_scale, per_channel_quant, block_shape)
84
- elif qtype == torch.int8:
85
- return _int8_quantize(A, A_scale, per_channel_quant, block_shape)
84
+ if quant_dtype == torch.float8_e4m3fn:
85
+ return _fp8_quantize(A, A_scale, per_act_token_quant, block_shape)
86
+ elif quant_dtype == torch.int8:
87
+ return _int8_quantize(A, A_scale, per_act_token_quant, block_shape)
88
else:
87
- assert A_scale is None
89
return A, A_scale
90
91
@@ -96,3 +97,17 @@ def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
97
return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
98
else:
99
return m[idx, ...]
100
+
101
+
102
+# TODO(bnell): better name
103
+def maybe_fix_scales(scales: Optional[torch.Tensor],
104
+ num_experts: int) -> Optional[torch.Tensor]:
105
+ if scales is not None and scales.ndim < 3:
106
+ if scales.numel() == 1:
107
+ scales = scales.view(1)
108
+ scales = torch.repeat_interleave(scales, num_experts,
109
+ dim=0).view(num_experts, 1, 1)
110
+ else:
111
+ scales = scales.view(num_experts, -1, scales.size(-1))
112
+
113
+ return scales
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+25
-30
@@ -13,8 +13,10 @@ from compressed_tensors.quantization import (ActivationOrdering,
13
import vllm.envs as envs
14
from vllm import _custom_ops as ops
15
from vllm.logger import init_logger
16
-from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
17
- FusedMoeWeightScaleSupported)
16
+from vllm.model_executor.layers.fused_moe import (
17
+ CutlassExpertsFp8, FusedMoE, FusedMoEActivationFormat, FusedMoEConfig,
18
+ FusedMoEMethodBase, FusedMoEPermuteExpertsUnpermute,
19
+ FusedMoEPrepareAndFinalize, FusedMoeWeightScaleSupported, fused_experts)
20
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
21
WNA16_SUPPORTED_BITS, WNA16_SUPPORTED_TYPES_MAP)
22
from vllm.model_executor.layers.quantization.utils import replace_parameter
@@ -32,14 +34,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
34
from vllm.model_executor.utils import set_weight_attrs
35
from vllm.platforms import current_platform
36
from vllm.scalar_type import scalar_types
35
-from vllm.utils import has_pplx
36
-
37
-if current_platform.is_cuda_alike():
38
- from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
39
- BatchedPrepareAndFinalize)
40
- if has_pplx():
41
- from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
42
- PplxPrepareAndFinalize)
37
38
logger = init_logger(__name__)
39
@@ -569,15 +563,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
563
requires_grad=False)
564
565
self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts
572
- else:
573
- from vllm.model_executor.layers.fused_moe import fused_experts
574
- self.fused_experts_func = fused_experts
575
-
576
- if self.use_marlin:
566
+ elif self.use_marlin:
567
prepare_moe_fp8_layer_for_marlin(layer, False)
568
# Activations not quantized for marlin.
569
del layer.w13_input_scale
570
del layer.w2_input_scale
571
+ self.fused_experts_func = None
572
+ else:
573
+ self.fused_experts_func = fused_experts
574
575
def apply(
576
self,
@@ -653,6 +646,8 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
646
global_num_experts=global_num_experts,
647
expert_map=expert_map)
648
649
+ assert self.fused_experts_func is not None
650
+
651
return self.fused_experts_func(
652
hidden_states=x,
653
w1=layer.w13_weight,
@@ -826,28 +821,27 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
821
layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
822
requires_grad=False)
823
829
- def select_gemm_impl(self, prepare_finalize, moe):
830
- from vllm.model_executor.layers.fused_moe.cutlass_moe import (
831
- CutlassExpertsFp8)
824
+ def select_gemm_impl(
825
+ self,
826
+ prepare_finalize: FusedMoEPrepareAndFinalize,
827
+ moe: FusedMoEConfig,
828
+ ) -> FusedMoEPermuteExpertsUnpermute:
829
833
- assert moe is not None
830
+ use_batched_format = (prepare_finalize.activation_format ==
831
+ FusedMoEActivationFormat.BatchedExperts)
832
+
833
+ num_experts = (moe.num_local_experts
834
+ if use_batched_format else moe.num_experts)
835
835
- max_experts_per_worker = (
836
- (moe.num_experts + prepare_finalize.world_size - 1) //
837
- prepare_finalize.world_size)
836
experts = CutlassExpertsFp8(
839
- max_experts_per_worker,
837
+ num_experts,
838
moe.in_dtype,
839
self.input_quant.strategy == QuantizationStrategy.TOKEN,
840
self.weight_quant.strategy == QuantizationStrategy.CHANNEL,
843
- use_batched_format=True,
841
+ use_batched_format=use_batched_format,
842
)
843
846
- if has_pplx() and isinstance(
847
- prepare_finalize,
848
- (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):
849
- # no expert_map support in this case
850
- self.disable_expert_map = True
844
+ self.disable_expert_map = not experts.supports_expert_map()
845
return experts
846
847
def apply(
@@ -888,7 +882,8 @@ class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
882
custom_routing_function=custom_routing_function,
883
scoring_func=scoring_func,
884
e_score_correction_bias=e_score_correction_bias,
891
- indices_type=torch.uint32)
885
+ indices_type=self.topk_indices_dtype,
886
+ )
887
888
return self.fused_experts(
889
x,
vllm/model_executor/layers/quantization/fp8.py
+32
-27
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import functools
5
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
5
+from typing import TYPE_CHECKING, Any, Callable, Optional
6
7
import torch
8
import torch.nn.functional as F
@@ -13,8 +13,11 @@ import vllm.envs as envs
13
from vllm import _custom_ops as ops
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.logger import init_logger
16
-from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEMethodBase,
17
- FusedMoeWeightScaleSupported)
16
+from vllm.model_executor.layers.fused_moe import (
17
+ BatchedTritonOrDeepGemmExperts, FusedMoE, FusedMoEActivationFormat,
18
+ FusedMoEConfig, FusedMoEMethodBase, FusedMoEPermuteExpertsUnpermute,
19
+ FusedMoEPrepareAndFinalize, FusedMoeWeightScaleSupported,
20
+ TritonOrDeepGemmExperts)
21
from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
22
UnquantizedLinearMethod)
23
from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -777,44 +780,46 @@ class Fp8MoEMethod(FusedMoEMethodBase):
780
del layer.w13_input_scale
781
del layer.w2_input_scale
782
780
- def select_gemm_impl(self, prepare_finalize, moe):
781
-
782
- from vllm.model_executor.layers.fused_moe.batched_triton_or_deep_gemm_moe import ( # noqa: E501
783
- BatchedTritonOrDeepGemmExperts)
784
- from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
785
- TritonOrDeepGemmExperts)
786
-
783
+ def select_gemm_impl(
784
+ self,
785
+ prepare_finalize: FusedMoEPrepareAndFinalize,
786
+ moe: FusedMoEConfig,
787
+ ) -> FusedMoEPermuteExpertsUnpermute:
788
assert not self.use_marlin and not self.rocm_aiter_moe_enabled, (
789
"Marlin and ROCm AITER are not supported with all2all yet.")
790
790
- experts: Optional[Union[BatchedTritonOrDeepGemmExperts,
791
- TritonOrDeepGemmExperts]] = None
792
- max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
793
- use_batched_experts = max_num_tokens_per_rank is not None
794
-
795
- if use_batched_experts:
796
- experts = BatchedTritonOrDeepGemmExperts(
791
+ if (prepare_finalize.activation_format ==
792
+ FusedMoEActivationFormat.BatchedExperts):
793
+ max_num_tokens_per_rank = (
794
+ prepare_finalize.max_num_tokens_per_rank())
795
+ assert max_num_tokens_per_rank is not None
796
+ logger.debug(
797
+ "BatchedTritonOrDeepGemmExperts(%s): "
798
+ "max_tokens_per_rank=%s, block_size=%s, per_act_token=%s",
799
+ self.__class__.__name__, max_num_tokens_per_rank,
800
+ self.quant_config.weight_block_size, False)
801
+ return BatchedTritonOrDeepGemmExperts(
802
max_num_tokens=max_num_tokens_per_rank,
798
- world_size=prepare_finalize.world_size,
799
- dp_size=prepare_finalize.dp_size,
803
+ world_size=prepare_finalize.
804
+ world_size, # type: ignore [attr-defined]
805
+ dp_size=prepare_finalize.
806
+ dp_size, # type: ignore [attr-defined]
807
use_fp8_w8a8=True,
801
- use_int8_w8a8=False,
802
- use_int8_w8a16=False,
803
- use_int4_w4a16=False,
804
- per_channel_quant=False,
808
block_shape=self.quant_config.weight_block_size,
809
+ per_act_token_quant=False,
810
allow_deep_gemm=self.allow_deep_gemm,
811
)
812
else:
809
- experts = TritonOrDeepGemmExperts(
813
+ logger.debug(
814
+ "TritonOrDeepGemmExperts(%s): block_size=%s, per_act_token=%s",
815
+ self.__class__.__name__, self.quant_config.weight_block_size,
816
+ False)
817
+ return TritonOrDeepGemmExperts(
818
use_fp8_w8a8=True,
819
block_shape=self.quant_config.weight_block_size,
820
allow_deep_gemm=self.allow_deep_gemm,
821
)
822
815
- assert experts is not None
816
- return experts
817
-
823
def apply(
824
self,
825
layer: torch.nn.Module,