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, 09:45 UTC
Modularize fused experts and integrate PPLX kernels (#15956)
bnellnm committed
May 14, 2025 at 16:11 UTC
f9c069c85e029830094ff9abb926ffbf37b7c7e7
42 files changed
+3835
-665
csrc/activation_kernels.cu
+3
@@ -70,6 +70,9 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
70
int64_t num_tokens = input.numel() / input.size(-1); \
71
dim3 grid(num_tokens); \
72
dim3 block(std::min(d, 1024)); \
73
+ if (num_tokens == 0) { \
74
+ return; \
75
+ } \
76
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
77
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
78
VLLM_DISPATCH_FLOATING_TYPES( \
csrc/dispatch_utils.h
+14
@@ -65,5 +65,19 @@
65
AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \
66
AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__)
67
68
+#define VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(...) \
69
+ AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__) \
70
+ AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__) \
71
+ AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \
72
+ AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__) \
73
+ AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__) \
74
+ AT_DISPATCH_CASE(at::ScalarType::UInt16, __VA_ARGS__) \
75
+ AT_DISPATCH_CASE(at::ScalarType::UInt32, __VA_ARGS__) \
76
+ AT_DISPATCH_CASE(at::ScalarType::UInt64, __VA_ARGS__)
77
+
78
#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
79
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
80
+
81
+#define VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(TYPE, NAME, ...) \
82
+ AT_DISPATCH_SWITCH( \
83
+ TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_AND_UNSIGNED_TYPES(__VA_ARGS__))
csrc/moe/moe_align_sum_kernels.cu
+4
-4
@@ -326,7 +326,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
326
}
327
328
if (use_global_memory) {
329
- VLLM_DISPATCH_INTEGRAL_TYPES(
329
+ VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
330
topk_ids.scalar_type(), "moe_align_block_size_global_mem_kernel", [&] {
331
// calc needed amount of shared mem for `tokens_cnts` and `cumsum`
332
// tensors
@@ -351,7 +351,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
351
cumsum_buffer.data_ptr<int32_t>());
352
});
353
} else if (use_i16) {
354
- VLLM_DISPATCH_INTEGRAL_TYPES(
354
+ VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
355
topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
356
// set dynamic shared mem
357
auto kernel =
@@ -366,7 +366,7 @@ void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
366
topk_ids.numel());
367
});
368
} else {
369
- VLLM_DISPATCH_INTEGRAL_TYPES(
369
+ VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
370
topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
371
auto kernel =
372
vllm::moe::moe_align_block_size_kernel<scalar_t, int32_t>;
@@ -391,7 +391,7 @@ void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
391
TORCH_CHECK(num_experts == 256,
392
"sgl_moe_align_block_size kernel only supports deepseek v3.");
393
394
- VLLM_DISPATCH_INTEGRAL_TYPES(
394
+ VLLM_DISPATCH_INTEGRAL_AND_UNSIGNED_TYPES(
395
topk_ids.scalar_type(), "sgl_moe_align_block_size_kernel", [&] {
396
// calc needed amount of shared mem for `cumsum` tensors
397
auto options_int =
csrc/moe/topk_softmax_kernels.cu
+45
-18
@@ -108,9 +108,17 @@ __launch_bounds__(TPB) __global__
108
}
109
}
110
111
-template <int TPB>
112
-__launch_bounds__(TPB) __global__ void moeTopK(const float* inputs_after_softmax, const bool* finished, float* output,
113
- int* indices, int* source_rows, const int num_experts, const int k, const int start_expert, const int end_expert)
111
+template <int TPB, typename IndType>
112
+__launch_bounds__(TPB) __global__ void moeTopK(
113
+ const float* inputs_after_softmax,
114
+ const bool* finished,
115
+ float* output,
116
+ IndType* indices,
117
+ int* source_rows,
118
+ const int num_experts,
119
+ const int k,
120
+ const int start_expert,
121
+ const int end_expert)
122
{
123
124
using cub_kvp = cub::KeyValuePair<int, float>;
@@ -182,9 +190,9 @@ __launch_bounds__(TPB) __global__ void moeTopK(const float* inputs_after_softmax
190
2) This implementation assumes k is small, but will work for any k.
191
*/
192
185
-template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG>
193
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, typename IndType>
194
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
187
- void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, int* indices,
195
+ void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, IndType* indices,
196
int* source_rows, const int k, const int start_expert, const int end_expert)
197
{
198
// We begin by enforcing compile time assertions and setting up compile time constants.
@@ -397,8 +405,8 @@ struct TopkConstants
405
};
406
} // namespace detail
407
400
-template <int EXPERTS, int WARPS_PER_TB>
401
-void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, int* indices,
408
+template <int EXPERTS, int WARPS_PER_TB, typename IndType>
409
+void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, IndType* indices,
410
int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream)
411
{
412
static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
@@ -421,10 +429,11 @@ void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, f
429
token_expert_indices, num_tokens, topk, 0, num_experts, \
430
stream);
431
432
+template <typename IndType>
433
void topkGatingSoftmaxKernelLauncher(
434
const float* gating_output,
435
float* topk_weights,
427
- int* topk_indicies,
436
+ IndType* topk_indicies,
437
int* token_expert_indices,
438
float* softmax_workspace,
439
const int num_tokens,
@@ -493,14 +502,32 @@ void topk_softmax(
502
const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
503
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
504
torch::Tensor softmax_workspace = torch::empty({workspace_size}, gating_output.options());
496
- vllm::moe::topkGatingSoftmaxKernelLauncher(
497
- gating_output.data_ptr<float>(),
498
- topk_weights.data_ptr<float>(),
499
- topk_indices.data_ptr<int>(),
500
- token_expert_indices.data_ptr<int>(),
501
- softmax_workspace.data_ptr<float>(),
502
- num_tokens,
503
- num_experts,
504
- topk,
505
- stream);
505
+
506
+ if(topk_indices.scalar_type() == at::ScalarType::Int)
507
+ {
508
+ vllm::moe::topkGatingSoftmaxKernelLauncher(
509
+ gating_output.data_ptr<float>(),
510
+ topk_weights.data_ptr<float>(),
511
+ topk_indices.data_ptr<int>(),
512
+ token_expert_indices.data_ptr<int>(),
513
+ softmax_workspace.data_ptr<float>(),
514
+ num_tokens,
515
+ num_experts,
516
+ topk,
517
+ stream);
518
+ }
519
+ else
520
+ {
521
+ assert(topk_indices.scalar_type() == at::ScalarType::UInt32);
522
+ vllm::moe::topkGatingSoftmaxKernelLauncher(
523
+ gating_output.data_ptr<float>(),
524
+ topk_weights.data_ptr<float>(),
525
+ topk_indices.data_ptr<uint32_t>(),
526
+ token_expert_indices.data_ptr<int>(),
527
+ softmax_workspace.data_ptr<float>(),
528
+ num_tokens,
529
+ num_experts,
530
+ topk,
531
+ stream);
532
+ }
533
}
examples/offline_inference/data_parallel.py
+16
-6
@@ -65,11 +65,17 @@ def parse_args():
65
type=int,
66
default=0,
67
help="Master node port")
68
+ parser.add_argument("--enforce-eager",
69
+ action='store_true',
70
+ help="Enforce eager mode execution.")
71
+ parser.add_argument("--trust-remote-code",
72
+ action='store_true',
73
+ help="Trust remote code.")
74
return parser.parse_args()
75
76
77
def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
72
- dp_master_port, GPUs_per_dp_rank):
78
+ dp_master_port, GPUs_per_dp_rank, enforce_eager, trust_remote_code):
79
os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
80
os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
81
os.environ["VLLM_DP_SIZE"] = str(dp_size)
@@ -109,10 +115,13 @@ def main(model, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
115
max_tokens=[16, 20][global_dp_rank % 2])
116
117
# Create an LLM.
112
- llm = LLM(model=model,
113
- tensor_parallel_size=GPUs_per_dp_rank,
114
- enforce_eager=True,
115
- enable_expert_parallel=True)
118
+ llm = LLM(
119
+ model=model,
120
+ tensor_parallel_size=GPUs_per_dp_rank,
121
+ enforce_eager=enforce_eager,
122
+ enable_expert_parallel=True,
123
+ trust_remote_code=trust_remote_code,
124
+ )
125
outputs = llm.generate(prompts, sampling_params)
126
# Print the outputs.
127
for i, output in enumerate(outputs):
@@ -155,7 +164,8 @@ if __name__ == "__main__":
164
proc = Process(target=main,
165
args=(args.model, dp_size, local_dp_rank,
166
global_dp_rank, dp_master_ip, dp_master_port,
158
- tp_size))
167
+ tp_size, args.enforce_eager,
168
+ args.trust_remote_code))
169
proc.start()
170
procs.append(proc)
171
exit_code = 0
tests/kernels/moe/test_batched_moe.py
new
+114
@@ -0,0 +1,114 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from dataclasses import dataclass
4
+
5
+import pytest
6
+import torch
7
+import triton.language as tl
8
+
9
+from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
10
+ invoke_moe_batched_triton_kernel)
11
+
12
+
13
+@dataclass
14
+class BatchedMMConfig:
15
+ dtype: torch.dtype
16
+ num_experts: int
17
+ max_tokens_per_expert: int
18
+ K: int
19
+ N: int
20
+
21
+
22
+@dataclass
23
+class BatchedMMTensors:
24
+ A: torch.Tensor # [E, max_tokens, K]
25
+ B: torch.Tensor # [E, K, N] - column major
26
+ C: torch.Tensor # [E, max_tokens, N]
27
+ num_expert_tokens: torch.Tensor # [E]
28
+
29
+ @staticmethod
30
+ def make_tensors(config: BatchedMMConfig):
31
+ A = torch.randn(
32
+ (config.num_experts, config.max_tokens_per_expert, config.K),
33
+ device="cuda",
34
+ dtype=config.dtype) / 10
35
+ B = torch.randn((config.num_experts, config.N, config.K),
36
+ device="cuda",
37
+ dtype=config.dtype)
38
+ C = torch.zeros(
39
+ (config.num_experts, config.max_tokens_per_expert, config.N),
40
+ device="cuda",
41
+ dtype=config.dtype)
42
+ num_expert_tokens = torch.randint(low=0,
43
+ high=config.max_tokens_per_expert,
44
+ size=(config.num_experts, ),
45
+ device="cuda",
46
+ dtype=torch.int32)
47
+ return BatchedMMTensors(A, B, C, num_expert_tokens)
48
+
49
+
50
+def ref_impl(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor,
51
+ num_expert_tokens: torch.Tensor) -> torch.Tensor:
52
+
53
+ num_expert_tokens_cpu = num_expert_tokens.clone()
54
+ num_expert_tokens_cpu = num_expert_tokens_cpu.to(device="cpu")
55
+ num_experts = num_expert_tokens.size(0)
56
+
57
+ for e in range(num_experts):
58
+ num_tokens = num_expert_tokens_cpu[e]
59
+ C[e, :num_tokens, :] = A[e, :num_tokens, :] @ B[e].transpose(0, 1)
60
+
61
+ return C
62
+
63
+
64
+@pytest.mark.parametrize("num_experts", [16, 32])
65
+@pytest.mark.parametrize("max_tokens_per_expert",
66
+ [32, 64, 128, 192, 224, 256, 512])
67
+@pytest.mark.parametrize("K", [128, 256, 1024])
68
+@pytest.mark.parametrize("N", [128, 256, 512, 1024])
69
+@pytest.mark.parametrize("dtype",
70
+ [torch.float32, torch.float16, torch.bfloat16])
71
+def test_batched_mm(num_experts: int, max_tokens_per_expert: int, K: int,
72
+ N: int, dtype: torch.dtype):
73
+
74
+ config = BatchedMMConfig(dtype, num_experts, max_tokens_per_expert, K, N)
75
+ tensors = BatchedMMTensors.make_tensors(config)
76
+
77
+ test_output = tensors.C
78
+ ref_output = test_output.clone()
79
+
80
+ compute_tl_dtype = {
81
+ torch.float16: tl.float16,
82
+ torch.bfloat16: tl.bfloat16,
83
+ torch.float32: tl.float32
84
+ }[test_output.dtype]
85
+ invoke_moe_batched_triton_kernel(
86
+ tensors.A,
87
+ tensors.B,
88
+ test_output,
89
+ tensors.num_expert_tokens,
90
+ compute_tl_dtype,
91
+ # Quantization data
92
+ None,
93
+ None,
94
+ None,
95
+ # Quantization schemes
96
+ False,
97
+ False,
98
+ False,
99
+ config={
100
+ "BLOCK_SIZE_M": 16,
101
+ "BLOCK_SIZE_N": 16,
102
+ "BLOCK_SIZE_K": 16
103
+ })
104
+
105
+ ref_output = ref_impl(tensors.A, tensors.B, ref_output,
106
+ tensors.num_expert_tokens)
107
+
108
+ rtol, atol = {
109
+ torch.float16: (6e-2, 6e-2),
110
+ torch.bfloat16: (6e-2, 6e-2),
111
+ torch.float32: (1e-2, 1e-2),
112
+ }[test_output.dtype]
113
+
114
+ torch.testing.assert_close(test_output, ref_output, atol=atol, rtol=rtol)
tests/kernels/moe/test_cutlass_moe.py
+21
-25
@@ -30,6 +30,11 @@ MNK_FACTORS = [
30
(224, 3072, 1536),
31
]
32
33
+vllm_config = VllmConfig(parallel_config=ParallelConfig(
34
+ pipeline_parallel_size=1))
35
+vllm_config.scheduler_config.max_num_seqs = 128
36
+vllm_config.scheduler_config.max_model_len = 8192
37
+
38
39
@dataclasses.dataclass
40
class MOETensors:
@@ -190,7 +195,7 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
195
'w1_q': moe_tensors.w1_q.transpose(1, 2), # type: ignore[union-attr]
196
'w2_q': moe_tensors.w2_q.transpose(1, 2), # type: ignore[union-attr]
197
'topk_weights': topk_weights,
193
- 'topk_ids_': topk_ids,
198
+ 'topk_ids': topk_ids,
199
'ab_strides1': moe_tensors.ab_strides1,
200
'c_strides1': moe_tensors.c_strides1,
201
'ab_strides2': moe_tensors.ab_strides2,
@@ -231,18 +236,15 @@ def test_cutlass_moe_8_bit_no_graph(
236
per_out_ch: bool,
237
):
238
current_platform.seed_everything(7)
234
- with set_current_vllm_config(
235
- VllmConfig(parallel_config=ParallelConfig(
236
- pipeline_parallel_size=1))):
237
-
239
+ with set_current_vllm_config(vllm_config):
240
mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
241
per_out_ch)
242
243
score = torch.randn((m, e), device="cuda", dtype=torch.half)
242
- topk_weights, topk_ids = fused_topk(mt.a,
243
- score,
244
- topk,
245
- renormalize=False)
244
+ topk_weights, topk_ids, _ = fused_topk(mt.a,
245
+ score,
246
+ topk,
247
+ renormalize=False)
248
249
# Note that we are using the dequantized versions of the tensors.
250
# Using a, w1 and w2 directly results in minor output differences.
@@ -276,20 +278,17 @@ def test_cutlass_moe_8_bit_cuda_graph(
278
per_out_ch: bool,
279
):
280
current_platform.seed_everything(7)
279
- with set_current_vllm_config(
280
- VllmConfig(parallel_config=ParallelConfig(
281
- pipeline_parallel_size=1))):
282
-
281
+ with set_current_vllm_config(vllm_config):
282
dtype = torch.half
283
284
mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
285
per_out_ch)
286
287
score = torch.randn((m, e), device="cuda", dtype=dtype)
289
- topk_weights, topk_ids = fused_topk(mt.a,
290
- score,
291
- topk,
292
- renormalize=False)
288
+ topk_weights, topk_ids, _ = fused_topk(mt.a,
289
+ score,
290
+ topk,
291
+ renormalize=False)
292
293
# Note that we are using the dequantized versions of the tensors.
294
# Using a, w1 and w2 directly results in minor output differences.
@@ -334,18 +333,15 @@ def test_cutlass_moe_8_bit_EP(
333
ep_size: int,
334
):
335
current_platform.seed_everything(7)
337
- with set_current_vllm_config(
338
- VllmConfig(parallel_config=ParallelConfig(
339
- pipeline_parallel_size=1))):
340
-
336
+ with set_current_vllm_config(vllm_config):
337
mt = MOETensors8Bit.make_moe_tensors_8bit(m, k, n, e, per_act_token,
338
per_out_channel)
339
340
score = torch.randn((m, e), device="cuda", dtype=torch.half)
345
- topk_weights, topk_ids = fused_topk(mt.a,
346
- score,
347
- topk,
348
- renormalize=False)
341
+ topk_weights, topk_ids, _ = fused_topk(mt.a,
342
+ score,
343
+ topk,
344
+ renormalize=False)
345
346
# Note that we are using the dequantized versions of the tensors.
347
# Using a, w1 and w2 directly results in minor output differences.
tests/kernels/moe/test_moe.py
+51
-42
@@ -12,6 +12,7 @@ from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock
12
13
import vllm.model_executor.layers.fused_moe # noqa
14
from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
15
+from vllm.config import VllmConfig, set_current_vllm_config
16
from vllm.model_executor.layers.fused_moe import fused_moe
17
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
18
from vllm.model_executor.layers.fused_moe.moe_torch_iterative import (
@@ -32,6 +33,10 @@ NUM_EXPERTS = [8, 64]
33
EP_SIZE = [1, 4]
34
TOP_KS = [2, 6]
35
36
+vllm_config = VllmConfig()
37
+vllm_config.scheduler_config.max_num_seqs = 128
38
+vllm_config.scheduler_config.max_model_len = 8192
39
+
40
41
@pytest.mark.parametrize("m", [1, 33, 64, 222, 1024 * 128])
42
@pytest.mark.parametrize("n", [128, 1024, 2048])
@@ -70,31 +75,33 @@ def test_fused_moe(
75
else:
76
e_map = None
77
73
- torch_output = torch_moe(a, w1, w2, score, topk, e_map)
74
- iterative_output = iterative_moe(a,
75
- w1,
76
- w2,
77
- score,
78
- topk,
79
- global_num_experts=e,
80
- expert_map=e_map,
81
- renormalize=False)
78
+ with set_current_vllm_config(vllm_config):
79
+ torch_output = torch_moe(a, w1, w2, score, topk, e_map)
80
+ iterative_output = iterative_moe(a,
81
+ w1,
82
+ w2,
83
+ score,
84
+ topk,
85
+ global_num_experts=e,
86
+ expert_map=e_map,
87
+ renormalize=False)
88
+
89
+ # Pad the weight if moe padding is enabled
90
+ if padding:
91
+ w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128]
92
+ torch.cuda.empty_cache()
93
+ w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128]
94
+ torch.cuda.empty_cache()
95
+
96
+ triton_output = fused_moe(a,
97
+ w1,
98
+ w2,
99
+ score,
100
+ topk,
101
+ global_num_experts=e,
102
+ expert_map=e_map,
103
+ renormalize=False)
104
83
- # Pad the weight if moe padding is enabled
84
- if padding:
85
- w1 = F.pad(w1, (0, 128), "constant", 0)[..., 0:-128]
86
- torch.cuda.empty_cache()
87
- w2 = F.pad(w2, (0, 128), "constant", 0)[..., 0:-128]
88
- torch.cuda.empty_cache()
89
-
90
- triton_output = fused_moe(a,
91
- w1,
92
- w2,
93
- score,
94
- topk,
95
- global_num_experts=e,
96
- expert_map=e_map,
97
- renormalize=False)
105
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
106
torch.testing.assert_close(iterative_output,
107
torch_output,
@@ -115,7 +122,6 @@ def test_fused_moe(
122
def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
123
ep_size: int, dtype: torch.dtype, group_size: int,
124
has_zp: bool, weight_bits: int):
118
- print(m, n, k, e, topk, dtype, group_size, has_zp, weight_bits)
125
a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
126
w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
127
w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
@@ -194,22 +200,24 @@ def test_fused_moe_wn16(m: int, n: int, k: int, e: int, topk: int,
200
else:
201
e_map = None
202
197
- triton_output = fused_moe(a,
198
- w1_qweight,
199
- w2_qweight,
200
- score,
201
- topk,
202
- renormalize=False,
203
- use_int4_w4a16=weight_bits == 4,
204
- use_int8_w8a16=weight_bits == 8,
205
- global_num_experts=e,
206
- expert_map=e_map,
207
- w1_scale=w1_scales,
208
- w2_scale=w2_scales,
209
- w1_zp=w1_qzeros if has_zp else None,
210
- w2_zp=w2_qzeros if has_zp else None,
211
- block_shape=[0, group_size])
212
- torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, e_map)
203
+ with set_current_vllm_config(vllm_config):
204
+ triton_output = fused_moe(a,
205
+ w1_qweight,
206
+ w2_qweight,
207
+ score,
208
+ topk,
209
+ renormalize=False,
210
+ use_int4_w4a16=weight_bits == 4,
211
+ use_int8_w8a16=weight_bits == 8,
212
+ global_num_experts=e,
213
+ expert_map=e_map,
214
+ w1_scale=w1_scales,
215
+ w2_scale=w2_scales,
216
+ w1_zp=w1_qzeros if has_zp else None,
217
+ w2_zp=w2_qzeros if has_zp else None,
218
+ block_shape=[0, group_size])
219
+ torch_output = torch_moe(a, w1_ref, w2_ref, score, topk, e_map)
220
+
221
torch.testing.assert_close(triton_output, torch_output, atol=2e-2, rtol=0)
222
223
@@ -515,7 +523,8 @@ def test_fused_marlin_moe(
523
524
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
525
518
- torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, e_map)
526
+ with set_current_vllm_config(vllm_config):
527
+ torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, e_map)
528
529
marlin_output = torch.ops.vllm.fused_marlin_moe(
530
a,
tests/kernels/moe/test_pplx_moe.py
new
+691
@@ -0,0 +1,691 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Tests for the MOE layers.
3
+
4
+Run `pytest tests/kernels/test_pplx_moe.py`.
5
+"""
6
+import dataclasses
7
+import os
8
+import traceback
9
+from typing import Callable, Optional
10
+
11
+import pytest
12
+import torch
13
+
14
+try:
15
+ from pplx_kernels import AllToAll
16
+ from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
17
+ nvshmem_finalize, nvshmem_get_unique_id,
18
+ nvshmem_init)
19
+ has_pplx = True
20
+except ImportError:
21
+ has_pplx = False
22
+
23
+from torch.multiprocessing import (
24
+ spawn) # pyright: ignore[reportPrivateImportUsage]
25
+from typing_extensions import Concatenate, ParamSpec
26
+
27
+from vllm.config import VllmConfig, set_current_vllm_config
28
+from vllm.model_executor.layers.activation import SiluAndMul
29
+from vllm.model_executor.layers.fused_moe import override_config
30
+from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
31
+ BatchedExperts, BatchedPrepareAndFinalize, BatchedTritonExperts)
32
+from vllm.model_executor.layers.fused_moe.fused_moe import (fused_topk,
33
+ get_default_config)
34
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
35
+ FusedMoEModularKernel)
36
+from vllm.platforms import current_platform
37
+
38
+PPLX_PREPARE_COMBOS = [(4, 128, 128), (32, 1024, 512), (64, 1024, 512),
39
+ (222, 2048, 1024)]
40
+
41
+PPLX_MOE_COMBOS = [
42
+ (1, 128, 128),
43
+ (2, 128, 512),
44
+ (3, 1024, 2048),
45
+ (32, 128, 1024),
46
+ (45, 512, 2048),
47
+ (64, 1024, 1024),
48
+ (222, 1024, 2048),
49
+]
50
+
51
+NUM_EXPERTS = [8, 64]
52
+EP_SIZE = [1, 4]
53
+TOP_KS = [1, 2, 6]
54
+
55
+vllm_config = VllmConfig()
56
+vllm_config.scheduler_config.max_num_seqs = 128
57
+vllm_config.scheduler_config.max_model_len = 8192
58
+
59
+P = ParamSpec("P")
60
+
61
+requires_pplx = pytest.mark.skipif(
62
+ not has_pplx,
63
+ reason="Requires PPLX kernels",
64
+)
65
+
66
+
67
+@dataclasses.dataclass
68
+class ProcessGroupInfo:
69
+ world_size: int
70
+ world_local_size: int
71
+ rank: int
72
+ node_rank: int
73
+ local_rank: int
74
+ device: torch.device
75
+
76
+
77
+def _worker_parallel_launch(
78
+ local_rank: int,
79
+ world_size: int,
80
+ world_local_size: int,
81
+ node_rank: int,
82
+ init_method: str,
83
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
84
+ *args: P.args,
85
+ **kwargs: P.kwargs,
86
+) -> None:
87
+ rank = node_rank * world_local_size + local_rank
88
+ torch.cuda.set_device(local_rank)
89
+ device = torch.device("cuda", local_rank)
90
+ torch.distributed.init_process_group(
91
+ backend="cpu:gloo,cuda:nccl",
92
+ init_method=init_method,
93
+ rank=rank,
94
+ world_size=world_size,
95
+ device_id=device,
96
+ )
97
+ barrier = torch.tensor([rank], device=device)
98
+ torch.distributed.all_reduce(barrier)
99
+
100
+ try:
101
+ worker(
102
+ ProcessGroupInfo(
103
+ world_size=world_size,
104
+ world_local_size=world_local_size,
105
+ rank=rank,
106
+ node_rank=node_rank,
107
+ local_rank=local_rank,
108
+ device=device,
109
+ ),
110
+ *args,
111
+ **kwargs,
112
+ )
113
+ except Exception as ex:
114
+ print(ex)
115
+ traceback.print_exc()
116
+ raise
117
+ finally:
118
+ torch.distributed.destroy_process_group()
119
+
120
+
121
+def parallel_launch(
122
+ world_size: int,
123
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
124
+ *args: P.args,
125
+ **kwargs: P.kwargs,
126
+) -> None:
127
+ assert not kwargs
128
+ spawn(
129
+ _worker_parallel_launch,
130
+ args=(
131
+ world_size,
132
+ world_size,
133
+ 0,
134
+ "tcp://localhost:29500",
135
+ worker,
136
+ ) + args,
137
+ nprocs=world_size,
138
+ join=True,
139
+ )
140
+
141
+
142
+def parallel_launch_from_env(
143
+ worker: Callable[Concatenate[ProcessGroupInfo, P], None],
144
+ *args: P.args,
145
+ **kwargs: P.kwargs,
146
+) -> None:
147
+ """
148
+ Launches a worker function in parallel across all processes in the current
149
+ environment. The environment must have the following variables set:
150
+ - WORLD_SIZE: The total number of processes.
151
+ - WORLD_LOCAL_SIZE: The number of processes on the current node.
152
+ - NODE_RANK: The rank of the current
153
+ - MASTER_ADDR: The address of the master process.
154
+ - MASTER_PORT: The port of the master process.
155
+ """
156
+ assert not kwargs
157
+ world_size = int(os.environ["WORLD_SIZE"])
158
+ world_local_size = int(os.environ["WORLD_LOCAL_SIZE"])
159
+ node_rank = int(os.environ["NODE_RANK"])
160
+ assert "MASTER_ADDR" in os.environ
161
+ assert "MASTER_PORT" in os.environ
162
+ spawn(
163
+ _worker_parallel_launch,
164
+ args=(
165
+ world_size,
166
+ world_local_size,
167
+ node_rank,
168
+ "env://",
169
+ worker,
170
+ ) + args,
171
+ nprocs=world_local_size,
172
+ join=True,
173
+ )
174
+
175
+
176
+def torch_prepare(
177
+ a: torch.Tensor,
178
+ topk_ids: torch.Tensor,
179
+ num_experts: int,
180
+ max_num_tokens: Optional[int] = None,
181
+) -> tuple[torch.Tensor, torch.Tensor]:
182
+ assert topk_ids.dim() == 2
183
+ assert topk_ids.shape[0] == a.shape[0]
184
+
185
+ num_tokens, hidden_dim = a.shape
186
+ topk = topk_ids.shape[1]
187
+
188
+ tokens_per_expert = torch.bincount(topk_ids.view(-1),
189
+ minlength=num_experts)
190
+
191
+ assert tokens_per_expert.numel() == num_experts
192
+
193
+ if max_num_tokens is None:
194
+ max_num_tokens = int(tokens_per_expert.max().item())
195
+
196
+ b_a = torch.zeros((num_experts, max_num_tokens, hidden_dim),
197
+ dtype=a.dtype,
198
+ device=a.device)
199
+
200
+ token_counts = torch.zeros(num_experts, dtype=torch.int, device=a.device)
201
+
202
+ for token in range(num_tokens):
203
+ for j in range(topk):
204
+ expert_id = topk_ids[token, j]
205
+ idx = token_counts[expert_id]
206
+ b_a[expert_id, idx:idx + 1, :] = a[token, :]
207
+ token_counts[expert_id] = token_counts[expert_id] + 1
208
+
209
+ return b_a, tokens_per_expert
210
+
211
+
212
+def torch_finalize(b_out: torch.Tensor, topk_weight: torch.Tensor,
213
+ topk_ids: torch.Tensor) -> torch.Tensor:
214
+ num_tokens = topk_ids.shape[0]
215
+ num_experts = b_out.shape[0]
216
+ K = b_out.shape[-1]
217
+ out = torch.zeros((num_tokens, K), dtype=b_out.dtype, device=b_out.device)
218
+ expert_counts = torch.zeros(num_experts,
219
+ dtype=torch.int,
220
+ device=b_out.device)
221
+ for token in range(num_tokens):
222
+ expert_ids = topk_ids[token]
223
+ for i in range(expert_ids.numel()):
224
+ expert_id = expert_ids[i]
225
+ idx = expert_counts[expert_id]
226
+ out[token, :] = out[token, :] + b_out[expert_id, idx:idx +
227
+ 1, :] * topk_weight[token, i]
228
+ expert_counts[expert_id] = expert_counts[expert_id] + 1
229
+
230
+ return out
231
+
232
+
233
+def torch_batched_moe(
234
+ a: torch.Tensor,
235
+ w1: torch.Tensor,
236
+ w2: torch.Tensor,
237
+ topk_weight: torch.Tensor,
238
+ topk_ids: torch.Tensor,
239
+) -> torch.Tensor:
240
+ num_experts = w1.shape[0]
241
+ b_a, tokens_per_expert = torch_prepare(a, topk_ids, num_experts)
242
+ assert b_a.dim() == 3
243
+ num_tokens, topk = topk_ids.shape
244
+ _, max_num_tokens, K = b_a.shape
245
+ assert num_experts == b_a.shape[0] and w2.shape[1] == K
246
+ out = torch.zeros((num_experts, max_num_tokens, K),
247
+ dtype=b_a.dtype,
248
+ device=b_a.device)
249
+ tmp = torch.empty((max_num_tokens, w1.shape[1] // 2),
250
+ dtype=b_a.dtype,
251
+ device=b_a.device)
252
+ for expert in range(num_experts):
253
+ num = tokens_per_expert[expert]
254
+ if num > 0:
255
+ torch.ops._C.silu_and_mul(
256
+ tmp[:num], b_a[expert, :num, :] @ w1[expert].transpose(0, 1))
257
+ out[expert, :num, :] = tmp[:num] @ w2[expert].transpose(0, 1)
258
+
259
+ return torch_finalize(out, topk_weight, topk_ids)
260
+
261
+
262
+def batched_moe(
263
+ a: torch.Tensor,
264
+ w1: torch.Tensor,
265
+ w2: torch.Tensor,
266
+ topk_weight: torch.Tensor,
267
+ topk_ids: torch.Tensor,
268
+) -> torch.Tensor:
269
+ num_experts = w1.shape[0]
270
+
271
+ fused_experts = FusedMoEModularKernel(
272
+ BatchedPrepareAndFinalize(a.shape[0], world_size=1, dp_size=1, rank=0),
273
+ BatchedExperts(max_num_tokens=a.shape[0], dp_size=1, world_size=1))
274
+
275
+ return fused_experts(a, w1, w2, topk_weight, topk_ids, num_experts)
276
+
277
+
278
+# Note: same as torch_moe but with fused_topk factored out.
279
+def torch_moe2(
280
+ a: torch.Tensor,
281
+ w1: torch.Tensor,
282
+ w2: torch.Tensor,
283
+ topk_weight: torch.Tensor,
284
+ topk_ids: torch.Tensor,
285
+) -> torch.Tensor:
286
+ M, K = a.shape
287
+ topk = topk_ids.shape[1]
288
+ a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)
289
+ out = torch.zeros(M * topk, w2.shape[1], dtype=a.dtype, device=a.device)
290
+ num_experts = w1.shape[0]
291
+ for i in range(num_experts):
292
+ mask = (topk_ids == i).view(-1)
293
+ if mask.sum():
294
+ out[mask] = SiluAndMul()(
295
+ a[mask] @ w1[i].transpose(0, 1)) @ w2[i].transpose(0, 1)
296
+
297
+ return (out.view(M, -1, w2.shape[1]) *
298
+ topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
299
+
300
+
301
+@pytest.mark.parametrize("m", [1, 33, 64, 222])
302
+@pytest.mark.parametrize("n", [128, 1024, 2048])
303
+@pytest.mark.parametrize("k", [128, 512, 1024])
304
+@pytest.mark.parametrize("e", NUM_EXPERTS)
305
+@pytest.mark.parametrize("topk", TOP_KS)
306
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
307
+def test_fused_moe_batched_experts(
308
+ m: int,
309
+ n: int,
310
+ k: int,
311
+ e: int,
312
+ topk: int,
313
+ dtype: torch.dtype,
314
+):
315
+ current_platform.seed_everything(7)
316
+
317
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
318
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
319
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
320
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
321
+
322
+ with set_current_vllm_config(vllm_config):
323
+ topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
324
+ baseline_output = torch_moe2(a, w1, w2, topk_weight, topk_ids)
325
+ torch_output = torch_batched_moe(a, w1, w2, topk_weight, topk_ids)
326
+ batched_output = batched_moe(a, w1, w2, topk_weight, topk_ids)
327
+
328
+ torch.testing.assert_close(baseline_output,
329
+ torch_output,
330
+ atol=2e-2,
331
+ rtol=0)
332
+ torch.testing.assert_close(baseline_output,
333
+ batched_output,
334
+ atol=2e-2,
335
+ rtol=0)
336
+
337
+
338
+def rank_chunk(num: int, r: int, w: int) -> int:
339
+ rem = num % w
340
+ return (num // w) + (1 if r < rem else 0)
341
+
342
+
343
+def chunk_by_rank(t: torch.Tensor, r: int, w: int) -> torch.Tensor:
344
+ chunk = rank_chunk(t.shape[0], r, w)
345
+ return t[(r * chunk):(r + 1) * chunk]
346
+
347
+
348
+def pplx_prepare_finalize(pgi: ProcessGroupInfo, dp_size: int, a: torch.Tensor,
349
+ topk_weight: torch.Tensor, topk_ids: torch.Tensor,
350
+ num_experts: int) -> torch.Tensor:
351
+ from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
352
+ PplxPrepareAndFinalize)
353
+
354
+ assert torch.cuda.current_device() == pgi.local_rank
355
+
356
+ topk = topk_ids.shape[1]
357
+ num_tokens, hidden_dim = a.shape
358
+ block_size = 128
359
+ device = pgi.device
360
+ rank = pgi.rank
361
+ world_size = pgi.world_size
362
+ max_num_tokens = rank_chunk(num_tokens, 0, world_size)
363
+
364
+ ata = AllToAll.internode(
365
+ max_num_tokens=max_num_tokens,
366
+ num_experts=num_experts,
367
+ experts_per_token=topk,
368
+ rank=rank,
369
+ world_size=world_size,
370
+ dp_size=dp_size,
371
+ hidden_dim=hidden_dim,
372
+ hidden_dim_bytes=hidden_dim * a.dtype.itemsize,
373
+ hidden_dim_scale_bytes=(0 if a.dtype.itemsize != 1 else
374
+ ((hidden_dim + block_size - 1) // block_size *
375
+ torch.float32.itemsize)),
376
+ )
377
+
378
+ topk_ids = topk_ids.to(dtype=torch.uint32)
379
+
380
+ prepare_finalize = PplxPrepareAndFinalize(
381
+ ata,
382
+ max_num_tokens,
383
+ world_size,
384
+ rank,
385
+ dp_size,
386
+ a.dtype,
387
+ )
388
+
389
+ a_chunk = chunk_by_rank(a, rank, world_size).to(device)
390
+ chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
391
+ chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
392
+
393
+ b_a, b_a_scale, expert_num_tokens = prepare_finalize.prepare(
394
+ a_chunk,
395
+ None,
396
+ None,
397
+ chunk_topk_weight,
398
+ chunk_topk_ids,
399
+ num_experts,
400
+ None,
401
+ False,
402
+ )
403
+
404
+ b_a = b_a * 1.5
405
+
406
+ out = torch.full(
407
+ (max_num_tokens, hidden_dim),
408
+ torch.nan,
409
+ dtype=a.dtype,
410
+ device=device,
411
+ )
412
+
413
+ prepare_finalize.finalize(
414
+ out,
415
+ b_a,
416
+ chunk_topk_weight,
417
+ chunk_topk_ids,
418
+ False,
419
+ )
420
+
421
+ torch.cuda.synchronize()
422
+
423
+ ata.destroy()
424
+
425
+ num_tokens = a_chunk.shape[0]
426
+
427
+ return out[:num_tokens]
428
+
429
+
430
+def _pplx_prepare_finalize(
431
+ pgi: ProcessGroupInfo,
432
+ dp_size: int,
433
+ a: torch.Tensor,
434
+ score: torch.Tensor,
435
+ topk: torch.Tensor,
436
+ num_experts: int,
437
+):
438
+ uid = nvshmem_get_unique_id(
439
+ ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
440
+ torch.distributed.broadcast(uid, src=0)
441
+ nvshmem_init(uid, pgi.rank, pgi.world_size)
442
+ device = pgi.device
443
+
444
+ topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
445
+ k = a.shape[1]
446
+
447
+ a_rep = torch.repeat_interleave(a, topk, dim=0).to(device)
448
+
449
+ torch_output = (a_rep.view(-1, topk, k) * 1.5 *
450
+ topk_weight.view(-1, topk, 1).to(device)).sum(dim=1).to(
451
+ a.dtype)
452
+
453
+ pplx_output = pplx_prepare_finalize(pgi, dp_size, a, topk_weight, topk_ids,
454
+ num_experts)
455
+
456
+ torch_output = chunk_by_rank(torch_output, pgi.rank,
457
+ pgi.world_size).to(pplx_output.device)
458
+
459
+ torch.testing.assert_close(pplx_output, torch_output, atol=2e-2, rtol=0)
460
+
461
+ nvshmem_finalize()
462
+
463
+
464
+# TODO (bnell): this test point does not work for odd M due to how the test is
465
+# written, not due to limitations of the pplx kernels. The pplx_moe
466
+# test below is able to deal with odd M.
467
+@pytest.mark.parametrize("mnk", PPLX_PREPARE_COMBOS)
468
+@pytest.mark.parametrize("e", NUM_EXPERTS)
469
+@pytest.mark.parametrize("topk", TOP_KS)
470
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
471
+@pytest.mark.parametrize("world_dp_size", [[2, 1]])
472
+@requires_pplx
473
+def test_pplx_prepare_finalize(
474
+ mnk: tuple[int, int, int],
475
+ e: int,
476
+ topk: int,
477
+ dtype: torch.dtype,
478
+ world_dp_size: tuple[int, int],
479
+):
480
+ current_platform.seed_everything(7)
481
+ m, n, k = mnk
482
+ world_size, dp_size = world_dp_size
483
+ device = "cuda"
484
+ a = torch.randn((m, k), device=device, dtype=dtype) / 10
485
+ score = torch.randn((m, e), device=device, dtype=dtype)
486
+
487
+ parallel_launch(world_size, _pplx_prepare_finalize, dp_size, a, score,
488
+ topk, e)
489
+
490
+
491
+def pplx_moe(
492
+ rank: int,
493
+ world_size: int,
494
+ dp_size: int,
495
+ a: torch.Tensor,
496
+ w1: torch.Tensor,
497
+ w2: torch.Tensor,
498
+ topk_weight: torch.Tensor,
499
+ topk_ids: torch.Tensor,
500
+ use_compile: bool = True,
501
+ use_cudagraphs: bool = True,
502
+) -> torch.Tensor:
503
+ from vllm.model_executor.layers.fused_moe.pplx_prepare_finalize import (
504
+ PplxPrepareAndFinalize)
505
+
506
+ device = torch.device("cuda", rank)
507
+ hidden_dim = a.shape[1]
508
+ num_experts = w1.shape[0]
509
+ block_size = 128
510
+ topk = topk_ids.shape[1]
511
+ max_num_tokens = rank_chunk(a.shape[0], 0, world_size)
512
+
513
+ ata = AllToAll.internode(
514
+ max_num_tokens=max_num_tokens,
515
+ num_experts=num_experts,
516
+ experts_per_token=topk,
517
+ rank=rank,
518
+ world_size=world_size,
519
+ dp_size=dp_size,
520
+ hidden_dim=hidden_dim,
521
+ hidden_dim_bytes=hidden_dim * a.dtype.itemsize,
522
+ hidden_dim_scale_bytes=(0 if a.dtype.itemsize != 1 else
523
+ ((hidden_dim + block_size - 1) // block_size *
524
+ torch.float32.itemsize)),
525
+ )
526
+
527
+ topk_ids = topk_ids.to(dtype=torch.uint32)
528
+
529
+ prepare_finalize = PplxPrepareAndFinalize(
530
+ ata,
531
+ max_num_tokens,
532
+ world_size,
533
+ rank,
534
+ dp_size,
535
+ )
536
+
537
+ experts = BatchedTritonExperts(max_num_tokens=a.shape[0],
538
+ world_size=world_size,
539
+ dp_size=dp_size)
540
+
541
+ fused_experts = FusedMoEModularKernel(
542
+ prepare_finalize,
543
+ experts,
544
+ )
545
+
546
+ # Note: workers with the same dp_rank must use the exact same inputs.
547
+ a_chunk = chunk_by_rank(a, rank, world_size).to(device)
548
+ chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
549
+ chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
550
+
551
+ # Chunking weights like this only works for batched format
552
+ w1_chunk = chunk_by_rank(w1, rank, world_size).to(device)
553
+ w2_chunk = chunk_by_rank(w2, rank, world_size).to(device)
554
+
555
+ if use_compile:
556
+ _fused_experts = torch.compile(fused_experts,
557
+ backend='inductor',
558
+ fullgraph=True)
559
+ else:
560
+ _fused_experts = fused_experts
561
+
562
+ out = _fused_experts(a_chunk,
563
+ w1_chunk,
564
+ w2_chunk,
565
+ chunk_topk_weight,
566
+ chunk_topk_ids,
567
+ global_num_experts=num_experts)
568
+
569
+ if use_cudagraphs:
570
+ out.fill_(0)
571
+ stream = torch.cuda.Stream()
572
+ graph = torch.cuda.CUDAGraph()
573
+ with torch.cuda.graph(graph, stream=stream):
574
+ out = _fused_experts(a_chunk,
575
+ w1_chunk,
576
+ w2_chunk,
577
+ chunk_topk_weight,
578
+ chunk_topk_ids,
579
+ global_num_experts=num_experts)
580
+
581
+ torch.cuda.synchronize()
582
+ graph.replay()
583
+
584
+ torch.cuda.synchronize()
585
+
586
+ ata.destroy()
587
+
588
+ return out
589
+
590
+
591
+def _batched_moe(pgi, dp_size, a, w1, w2, topk_weight, topk_ids):
592
+ assert torch.cuda.current_device() == pgi.local_rank
593
+
594
+ num_experts = w1.shape[0]
595
+ device = pgi.device
596
+ rank = pgi.rank
597
+ world_size = pgi.world_size
598
+ max_num_tokens = rank_chunk(a.shape[0], 0, world_size)
599
+
600
+ prepare_finalize = BatchedPrepareAndFinalize(
601
+ max_num_tokens=max_num_tokens,
602
+ world_size=world_size,
603
+ dp_size=dp_size,
604
+ rank=rank,
605
+ )
606
+
607
+ experts = BatchedExperts(max_num_tokens=a.shape[0],
608
+ world_size=1,
609
+ dp_size=1)
610
+
611
+ fused_experts = FusedMoEModularKernel(
612
+ prepare_finalize,
613
+ experts,
614
+ )
615
+
616
+ # Note: workers with the same dp_rank must use the exact same inputs.
617
+ a_chunk = chunk_by_rank(a, rank, world_size).to(device)
618
+ chunk_topk_weight = chunk_by_rank(topk_weight, rank, world_size).to(device)
619
+ chunk_topk_ids = chunk_by_rank(topk_ids, rank, world_size).to(device)
620
+
621
+ out = fused_experts(
622
+ a_chunk,
623
+ # Chunking weights like this only works for batched format
624
+ chunk_by_rank(w1, rank, world_size).to(device),
625
+ chunk_by_rank(w2, rank, world_size).to(device),
626
+ chunk_topk_weight,
627
+ chunk_topk_ids,
628
+ global_num_experts=num_experts)
629
+
630
+ return out
631
+
632
+
633
+def _pplx_moe(
634
+ pgi: ProcessGroupInfo,
635
+ dp_size: int,
636
+ a: torch.Tensor,
637
+ w1: torch.Tensor,
638
+ w2: torch.Tensor,
639
+ score: torch.Tensor,
640
+ topk: int,
641
+):
642
+ uid = nvshmem_get_unique_id(
643
+ ) if pgi.rank == 0 else nvshmem_alloc_empty_unique_id()
644
+ torch.distributed.broadcast(uid, src=0)
645
+ nvshmem_init(uid, pgi.rank, pgi.world_size)
646
+
647
+ m, k = a.shape
648
+ e, _, n = w2.shape
649
+
650
+ moe_config = get_default_config(m, e, n, k, topk, a.dtype, False)
651
+
652
+ with set_current_vllm_config(vllm_config), override_config(moe_config):
653
+ topk_weight, topk_ids, _ = fused_topk(a, score, topk, False)
654
+ torch_output = torch_moe2(a, w1, w2, topk_weight, topk_ids)
655
+ pplx_output = pplx_moe(pgi.rank, pgi.world_size, dp_size, a, w1, w2,
656
+ topk_weight, topk_ids)
657
+ # TODO (bnell): fix + re-enable
658
+ #batched_output = _batched_moe(pgi, dp_size, a, w1, w2, topk_weight,
659
+ # topk_ids)
660
+
661
+ torch_output = chunk_by_rank(torch_output, pgi.rank,
662
+ pgi.world_size).to(pplx_output.device)
663
+
664
+ torch.testing.assert_close(pplx_output, torch_output, atol=2e-2, rtol=0)
665
+ #torch.testing.assert_close(batched_output, torch_output, atol=2e-2, rtol=0)
666
+
667
+ nvshmem_finalize()
668
+
669
+
670
+@pytest.mark.parametrize("mnk", PPLX_MOE_COMBOS)
671
+@pytest.mark.parametrize("e", NUM_EXPERTS)
672
+@pytest.mark.parametrize("topk", TOP_KS)
673
+@pytest.mark.parametrize("dtype", [torch.bfloat16])
674
+@pytest.mark.parametrize("world_dp_size", [[2, 1]])
675
+@requires_pplx
676
+def test_pplx_moe(
677
+ mnk: tuple[int, int, int],
678
+ e: int,
679
+ topk: int,
680
+ dtype: torch.dtype,
681
+ world_dp_size: tuple[int, int],
682
+):
683
+ current_platform.seed_everything(7)
684
+ m, n, k = mnk
685
+ world_size, dp_size = world_dp_size
686
+ a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
687
+ w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
688
+ w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
689
+ score = torch.randn((m, e), device="cuda", dtype=dtype)
690
+
691
+ parallel_launch(world_size, _pplx_moe, dp_size, a, w1, w2, score, topk)
tests/kernels/moe/test_triton_moe_ptpc_fp8.py
+20
-14
@@ -7,6 +7,7 @@ import pytest
7
import torch
8
9
from vllm import _custom_ops as ops
10
+from vllm.config import VllmConfig, set_current_vllm_config
11
from vllm.model_executor.layers.activation import SiluAndMul
12
from vllm.model_executor.layers.fused_moe import fused_moe
13
from vllm.platforms import current_platform
@@ -15,6 +16,10 @@ if current_platform.get_device_capability() < (9, 0):
16
pytest.skip("FP8 Triton requires CUDA 9.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
24
def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
25
"""Matrix multiplication function that supports per-token input
@@ -137,20 +142,21 @@ def test_w8a8_fp8_fused_moe(M, N, K, E, topk, dtype, seed):
142
w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
143
score = torch.randn((M, E), dtype=dtype)
144
140
- ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk)
141
- out = fused_moe(
142
- a,
143
- w1,
144
- w2,
145
- score,
146
- topk,
147
- renormalize=False,
148
- use_fp8_w8a8=True, # using fp8
149
- per_channel_quant=True,
150
- w1_scale=w1_s,
151
- w2_scale=w2_s,
152
- block_shape=None, # Not using block quantization
153
- )
145
+ with set_current_vllm_config(vllm_config):
146
+ ref_out = torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, score, topk)
147
+ out = fused_moe(
148
+ a,
149
+ w1,
150
+ w2,
151
+ score,
152
+ topk,
153
+ renormalize=False,
154
+ use_fp8_w8a8=True, # using fp8
155
+ per_channel_quant=True,
156
+ w1_scale=w1_s,
157
+ w2_scale=w2_s,
158
+ block_shape=None, # Not using block quantization
159
+ )
160
161
# Check results
162
rel_diff = (torch.mean(
tests/kernels/quantization/test_block_fp8.py
+10
-10
@@ -11,7 +11,7 @@ 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.model_executor.layers.fused_moe.deep_gemm_moe import (
14
- deep_gemm_moe_fp8)
14
+ _valid_deep_gemm_shape, deep_gemm_moe_fp8)
15
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
16
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
17
moe_align_block_size)
@@ -30,6 +30,10 @@ if current_platform.get_device_capability() < (9, 0):
30
pytest.skip("FP8 Triton requires CUDA 9.0 or higher",
31
allow_module_level=True)
32
33
+vllm_config = VllmConfig()
34
+vllm_config.scheduler_config.max_num_seqs = 128
35
+vllm_config.scheduler_config.max_model_len = 8192
36
+
37
# Test configurations
38
DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
39
NUM_TOKENS = [7, 83, 2048]
@@ -210,7 +214,6 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
214
score = torch.randn((M, E), dtype=dtype)
215
216
# Set the context to avoid lots of warning spam.
213
- vllm_config = VllmConfig()
217
with set_current_vllm_config(vllm_config):
218
out = fused_moe(
219
a,
@@ -258,6 +261,7 @@ def per_block_cast_to_fp8(
261
@pytest.mark.parametrize(
262
"M,N,K,block_size,out_dtype,seed",
263
itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS))
264
+@pytest.mark.skipif(not dg_available, reason="DeepGemm kernels not available.")
265
@torch.inference_mode()
266
def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
267
# only aligned sizes
@@ -381,15 +385,11 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed):
385
block_size = [block_m, block_m]
386
dtype = torch.bfloat16
387
384
- # only aligned sizes
385
- if (N % block_m != 0 or K % block_m != 0 or topk > E):
386
- pytest.skip(
387
- f"Skipping test; bad size m={M}, n={N}, k={K}, topk={topk}, E={E}")
388
-
389
- if N <= 512:
390
- pytest.skip("Skipping N <= 512 until performance issues solved.")
388
+ if topk > E:
389
+ pytest.skip(f"Skipping test: topk={topk} > E={E}")
390
392
- vllm_config = VllmConfig()
391
+ if not _valid_deep_gemm_shape(M, N, K):
392
+ pytest.skip(f"Skipping test: invalid size m={M}, n={N}, k={K}")
393
394
torch.manual_seed(seed)
395
fp8_info = torch.finfo(torch.float8_e4m3fn)
tests/kernels/quantization/test_block_int8.py
+4
-1
@@ -18,6 +18,10 @@ if current_platform.get_device_capability() < (7, 0):
18
pytest.skip("INT8 Triton requires CUDA 7.0 or higher",
19
allow_module_level=True)
20
21
+vllm_config = VllmConfig()
22
+vllm_config.scheduler_config.max_num_seqs = 128
23
+vllm_config.scheduler_config.max_model_len = 8192
24
+
25
26
# For test
27
def native_per_token_group_quant_int8(x,
@@ -174,7 +178,6 @@ def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
178
score = torch.randn((M, E), dtype=dtype)
179
180
# Set the context to avoid lots of warning spam.
177
- vllm_config = VllmConfig()
181
with set_current_vllm_config(vllm_config):
182
out = fused_moe(
183
a,
vllm/distributed/parallel_state.py
+51
-2
@@ -23,6 +23,7 @@ If you only need to use the distributed environment without model/pipeline
23
"""
24
import contextlib
25
import gc
26
+import importlib.util
27
import pickle
28
import weakref
29
from collections import namedtuple
@@ -42,7 +43,7 @@ from vllm.distributed.device_communicators.base_device_communicator import (
43
from vllm.distributed.utils import StatelessProcessGroup
44
from vllm.logger import init_logger
45
from vllm.utils import (direct_register_custom_op, resolve_obj_by_qualname,
45
- supports_custom_op)
46
+ run_once, supports_custom_op)
47
48
49
@dataclass
@@ -936,9 +937,49 @@ def init_distributed_environment(
937
"world group already initialized with a different world size")
938
939
940
+PPLX_DID_INIT: bool = False
941
+
942
+
943
+@run_once
944
+def pplx_init(rank, world_size):
945
+ has_pplx = importlib.util.find_spec("pplx_kernels") is not None
946
+
947
+ if has_pplx and world_size > 1:
948
+ from pplx_kernels.nvshmem import (nvshmem_alloc_empty_unique_id,
949
+ nvshmem_get_unique_id, nvshmem_init)
950
+ try:
951
+ global PPLX_DID_INIT
952
+ logger.debug(
953
+ "Initialize NVSHMEM for PPLX kernels: rank=%d, "
954
+ "world size=%d", rank, world_size)
955
+ uid = nvshmem_get_unique_id(
956
+ ) if rank == 0 else nvshmem_alloc_empty_unique_id()
957
+ uid_gpu = uid.cuda()
958
+ get_world_group().broadcast(uid_gpu, src=0)
959
+ uid = uid_gpu.to(device='cpu')
960
+ logger.debug("PPLX NVSHMEM UID = %s", uid)
961
+ nvshmem_init(uid, rank, world_size)
962
+ PPLX_DID_INIT = True
963
+ except Exception as ex:
964
+ logger.error("Failed to initialize NVSHMEM for PPLX: %s", ex)
965
+
966
+
967
+@run_once
968
+def pplx_finalize():
969
+ global PPLX_DID_INIT
970
+ if PPLX_DID_INIT:
971
+ from pplx_kernels.nvshmem import nvshmem_finalize
972
+ logger.debug("PPLX NVSHMEM finalize")
973
+ from vllm.model_executor.layers.fused_moe.layer import (
974
+ _all_to_all_cache)
975
+ _all_to_all_cache.destroy()
976
+ nvshmem_finalize()
977
+
978
+
979
def initialize_model_parallel(
980
tensor_model_parallel_size: int = 1,
981
pipeline_model_parallel_size: int = 1,
982
+ enable_expert_parallel: bool = False,
983
backend: Optional[str] = None,
984
) -> None:
985
"""
@@ -1041,10 +1082,14 @@ def initialize_model_parallel(
1082
_DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group,
1083
_EP.rank_in_group)
1084
1085
+ if enable_expert_parallel:
1086
+ pplx_init(rank, world_size)
1087
+
1088
1089
def ensure_model_parallel_initialized(
1090
tensor_model_parallel_size: int,
1091
pipeline_model_parallel_size: int,
1092
+ enable_expert_parallel: bool = False,
1093
backend: Optional[str] = None,
1094
) -> None:
1095
"""Helper to initialize model parallel groups if they are not initialized,
@@ -1055,7 +1100,8 @@ def ensure_model_parallel_initialized(
1100
get_world_group().device_group)
1101
if not model_parallel_is_initialized():
1102
initialize_model_parallel(tensor_model_parallel_size,
1058
- pipeline_model_parallel_size, backend)
1103
+ pipeline_model_parallel_size,
1104
+ enable_expert_parallel, backend)
1105
return
1106
1107
assert (
@@ -1133,6 +1179,9 @@ def get_tensor_model_parallel_rank():
1179
def destroy_model_parallel():
1180
"""Set the groups to none and destroy them."""
1181
global _TP
1182
+
1183
+ pplx_finalize()
1184
+
1185
if _TP:
1186
_TP.destroy()
1187
_TP = None
vllm/distributed/utils.py
+6
-7
@@ -23,7 +23,7 @@ from torch.distributed.rendezvous import rendezvous
23
24
import vllm.envs as envs
25
from vllm.logger import init_logger
26
-from vllm.utils import get_tcp_uri
26
+from vllm.utils import get_tcp_uri, is_torch_equal_or_newer
27
28
logger = init_logger(__name__)
29
@@ -362,12 +362,11 @@ def stateless_destroy_torch_distributed_process_group(
362
Destroy ProcessGroup returned by
363
stateless_init_torch_distributed_process_group().
364
"""
365
- # Lazy import for non-CUDA backends.
366
- try:
367
- # pytorch <= 2.6
365
+ if is_torch_equal_or_newer("2.7"):
366
+ pg.shutdown()
367
+ else:
368
+ # Lazy import for non-CUDA backends.
369
from torch.distributed.distributed_c10d import _shutdown_backend
370
_shutdown_backend(pg)
370
- except ImportError:
371
- # pytorch >= 2.7
372
- pg.shutdown()
371
+
372
_unregister_process_group(pg.group_name)
vllm/forward_context.py
+4
-1
@@ -27,6 +27,7 @@ batchsize_forward_time: defaultdict = defaultdict(list)
27
28
@dataclass
29
class DPMetadata:
30
+ max_tokens_across_dp_cpu: torch.Tensor
31
cu_tokens_across_dp_cpu: torch.Tensor
32
33
@@ -90,8 +91,10 @@ def set_forward_context(attn_metadata: Any,
91
dtype=torch.int32)
92
from vllm.distributed.parallel_state import get_dp_group
93
dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
94
+ max_tokens_across_dp_cpu = torch.max(num_tokens_tensor)
95
cu_tokens_across_dp_cpu = torch.cumsum(num_tokens_tensor, dim=0)
94
- dp_metadata = DPMetadata(cu_tokens_across_dp_cpu)
96
+ dp_metadata = DPMetadata(max_tokens_across_dp_cpu,
97
+ cu_tokens_across_dp_cpu)
98
99
global _forward_context
100
prev_context = _forward_context
vllm/model_executor/layers/fused_moe/__init__.py
+3
-2
@@ -38,8 +38,8 @@ if HAS_TRITON:
38
from vllm.model_executor.layers.fused_moe.cutlass_moe import (
39
cutlass_moe_fp4, cutlass_moe_fp8)
40
from vllm.model_executor.layers.fused_moe.fused_moe import (
41
- fused_experts, fused_moe, fused_topk, get_config_file_name,
42
- grouped_topk)
41
+ TritonExperts, fused_experts, fused_moe, fused_topk,
42
+ get_config_file_name, grouped_topk)
43
44
__all__ += [
45
"fused_moe",
@@ -49,4 +49,5 @@ if HAS_TRITON:
49
"grouped_topk",
50
"cutlass_moe_fp8",
51
"cutlass_moe_fp4",
52
+ "TritonExperts",
53
]
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+195
-108
@@ -5,10 +5,176 @@ from typing import Optional
5
6
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.prepare_finalize import (
11
+ MoEPrepareAndFinalizeNoEP)
12
+from vllm.model_executor.layers.fused_moe.utils import _fp8_perm, _resize_cache
13
from vllm.scalar_type import scalar_types
14
15
16
+class CutlassExpertsFp8(mk.FusedMoEPermuteExpertsUnpermute):
17
+
18
+ def __init__(
19
+ self,
20
+ ab_strides1: torch.Tensor,
21
+ c_strides1: torch.Tensor,
22
+ ab_strides2: torch.Tensor,
23
+ c_strides2: torch.Tensor,
24
+ out_dtype: torch.dtype,
25
+ ):
26
+ super().__init__()
27
+ self.ab_strides1 = ab_strides1
28
+ self.c_strides1 = c_strides1
29
+ self.ab_strides2 = ab_strides2
30
+ self.c_strides2 = c_strides2
31
+ self.out_dtype = out_dtype
32
+
33
+ def workspace_shapes(
34
+ self,
35
+ a: torch.Tensor,
36
+ M: int,
37
+ N: int,
38
+ K: int,
39
+ topk: int,
40
+ num_experts: int,
41
+ ) -> tuple[int, int, torch.dtype]:
42
+ # Note that K, N are transposed
43
+ N, K = K, N
44
+ workspace1 = M * topk * max(2 * N, K)
45
+ workspace2 = M * topk * N
46
+ return (workspace1, workspace2, self.out_dtype)
47
+
48
+ def apply(
49
+ self,
50
+ hidden_states: torch.Tensor,
51
+ w1: torch.Tensor,
52
+ w2: torch.Tensor,
53
+ topk_ids: torch.Tensor,
54
+ activation: str,
55
+ global_num_experts: int,
56
+ expert_map: Optional[torch.Tensor],
57
+ w1_scale: Optional[torch.Tensor],
58
+ w2_scale: Optional[torch.Tensor],
59
+ w1_zp: Optional[torch.Tensor],
60
+ w2_zp: Optional[torch.Tensor],
61
+ a1q_scale: Optional[torch.Tensor],
62
+ a2_scale: Optional[torch.Tensor],
63
+ workspace13: torch.Tensor,
64
+ workspace2: torch.Tensor,
65
+ expert_num_tokens: Optional[torch.Tensor],
66
+ ) -> torch.Tensor:
67
+ a1q = hidden_states
68
+
69
+ assert w1_scale is not None
70
+ assert w2_scale is not None
71
+ assert w1.dtype == torch.float8_e4m3fn
72
+ assert w2.dtype == torch.float8_e4m3fn
73
+ assert a1q.shape[1] == w1.shape[1], "Hidden size mismatch w1"
74
+ assert w1.shape[2] == w2.shape[1] * 2, "Hidden size mismatch w2"
75
+ assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
76
+ assert a1q_scale is None or a1q_scale.dim(
77
+ ) == 0 or a1q_scale.shape[0] == 1 or a1q_scale.shape[0] == a1q.shape[
78
+ 0], "Input scale shape mismatch"
79
+ assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
80
+ 1] == w1.shape[2], "W1 scale shape mismatch"
81
+ assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
82
+ 1] == w2.shape[2], "W2 scale shape mismatch"
83
+ assert w1.shape[0] == w2.shape[0], "Weights expert number mismatch"
84
+ assert w1.shape[0] == w1_scale.shape[
85
+ 0], "w1 scales expert number mismatch"
86
+ assert w1.shape[0] == w2_scale.shape[
87
+ 0], "w2 scales expert number mismatch"
88
+ assert a2_scale is None or a1q_scale is None or a2_scale.shape == a1q_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
89
+ assert self.ab_strides1.shape[0] == w1.shape[
90
+ 0], "AB Strides 1 expert number mismatch"
91
+ assert self.c_strides1.shape[0] == w1.shape[
92
+ 0], "C Strides 1 expert number mismatch"
93
+ assert self.ab_strides2.shape[0] == w2.shape[
94
+ 0], "AB Strides 2 expert number mismatch"
95
+ assert self.c_strides2.shape[0] == w2.shape[
96
+ 0], "C Strides 2 expert number mismatch"
97
+ assert self.out_dtype in [torch.half,
98
+ torch.bfloat16], "Invalid output dtype"
99
+
100
+ M = a1q.shape[0]
101
+ _, N, K = w2.shape # because w1 + w2 are transposed
102
+ device = a1q.device
103
+
104
+ assert w1.shape[1] == K
105
+ assert global_num_experts != -1
106
+ assert a1q_scale is not None
107
+
108
+ if expert_map is not None:
109
+ "Translate info from expert_map to topk_ids"
110
+ local_topk_ids = torch.where(expert_map[topk_ids] != -1,
111
+ expert_map[topk_ids], -1)
112
+ else:
113
+ local_topk_ids = topk_ids
114
+
115
+ topk = local_topk_ids.shape[1]
116
+
117
+ per_act_token = a1q_scale.numel() != 1 if a1q_scale is not None else (
118
+ a2_scale.numel() != 1 if a2_scale is not None else False)
119
+
120
+ expert_offsets = torch.empty((global_num_experts + 1),
121
+ dtype=torch.int32,
122
+ device=device)
123
+ problem_sizes1 = torch.empty((global_num_experts, 3),
124
+ dtype=torch.int32,
125
+ device=device)
126
+ problem_sizes2 = torch.empty((global_num_experts, 3),
127
+ dtype=torch.int32,
128
+ device=device)
129
+
130
+ # With expert_map each Rank processes only a subset of experts. As
131
+ # a result not all of a_map and c2 tensors are filled. We fill it
132
+ # zeros for correctness.
133
+ if expert_map is not None:
134
+ a_map = torch.zeros((local_topk_ids.numel()),
135
+ dtype=torch.int32,
136
+ device=device)
137
+ else:
138
+ a_map = torch.empty((local_topk_ids.numel()),
139
+ dtype=torch.int32,
140
+ device=device)
141
+
142
+ c_map = torch.empty((local_topk_ids.numel()),
143
+ dtype=torch.int32,
144
+ device=device)
145
+
146
+ ops.get_cutlass_moe_mm_data(local_topk_ids, expert_offsets,
147
+ problem_sizes1, problem_sizes2, a_map,
148
+ c_map, global_num_experts, N, K)
149
+
150
+ a1q = _fp8_perm(a1q, a_map)
151
+ a1q_scale = a1q_scale[a_map] if per_act_token else a1q_scale
152
+
153
+ c1 = _resize_cache(workspace13, (M * topk, N * 2))
154
+ c2 = _resize_cache(workspace2, (M * topk, N))
155
+ c3 = _resize_cache(workspace13, (M * topk, K))
156
+
157
+ ops.cutlass_moe_mm(c1, a1q, w1, a1q_scale, w1_scale,
158
+ expert_offsets[:-1], problem_sizes1,
159
+ self.ab_strides1, self.ab_strides1, self.c_strides1)
160
+
161
+ self.activation(activation, c2, c1)
162
+
163
+ a2q, a2q_scale = ops.scaled_fp8_quant(
164
+ c2, a2_scale, use_per_token_if_dynamic=per_act_token)
165
+
166
+ if expert_map is not None:
167
+ c3.fill_(0)
168
+
169
+ ops.cutlass_moe_mm(c3, a2q, w2, a2q_scale, w2_scale,
170
+ expert_offsets[:-1], problem_sizes2,
171
+ self.ab_strides2, self.ab_strides2, self.c_strides2)
172
+
173
+ c3 = c3[c_map]
174
+
175
+ return c3
176
+
177
+
178
#TODO make the grouped gemm kernel consistent with scaled gemm kernel
179
def cutlass_moe_fp8(
180
a: torch.Tensor,
@@ -17,7 +183,7 @@ def cutlass_moe_fp8(
183
w1_scale: torch.Tensor,
184
w2_scale: torch.Tensor,
185
topk_weights: torch.Tensor,
20
- topk_ids_: torch.Tensor,
186
+ topk_ids: torch.Tensor,
187
ab_strides1: torch.Tensor,
188
c_strides1: torch.Tensor,
189
ab_strides2: torch.Tensor,
@@ -59,7 +225,7 @@ def cutlass_moe_fp8(
225
- a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
226
quantize the intermediate result between the gemms.
227
Shape: scalar or [M]
62
- - out_dtype (torch.Tensor): The output tensor type.
228
+ - out_dtype (torch.dtype): The output tensor type.
229
- expert_map (Optional[torch.Tensor]): In the case of Expert parallel,
230
every Rank is responsible for a subset of experts. expert_map is a
231
mapping from global expert-id to local expert-id. When expert_map[i]
@@ -71,115 +237,36 @@ def cutlass_moe_fp8(
237
Returns:
238
- torch.Tensor: The fp16 output tensor after applying the MoE layer.
239
"""
74
-
75
- assert topk_weights.shape == topk_ids_.shape, "topk shape mismatch"
76
- assert w1_q.dtype == torch.float8_e4m3fn
77
- assert w2_q.dtype == torch.float8_e4m3fn
78
- assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
79
- assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2"
80
- assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch"
81
- assert a1_scale is None or a1_scale.dim(
82
- ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[0] == a.shape[
83
- 0], "Input scale shape mismatch"
84
- assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
85
- 1] == w1_q.shape[2], "W1 scale shape mismatch"
86
- assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
87
- 1] == w2_q.shape[2], "W2 scale shape mismatch"
88
- assert w1_q.shape[0] == w2_q.shape[0], "Weights expert number mismatch"
89
- assert w1_q.shape[0] == w1_scale.shape[
90
- 0], "w1 scales expert number mismatch"
91
- assert w1_q.shape[0] == w2_scale.shape[
92
- 0], "w2 scales expert number mismatch"
93
- assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
94
- assert ab_strides1.shape[0] == w1_q.shape[
95
- 0], "AB Strides 1 expert number mismatch"
96
- assert c_strides1.shape[0] == w1_q.shape[
97
- 0], "C Strides 1 expert number mismatch"
98
- assert ab_strides2.shape[0] == w2_q.shape[
99
- 0], "AB Strides 2 expert number mismatch"
100
- assert c_strides2.shape[0] == w2_q.shape[
101
- 0], "C Strides 2 expert number mismatch"
102
- assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
103
-
104
- num_experts = w1_q.size(0)
105
- m = a.size(0)
106
- k = w1_q.size(1)
107
- n = w2_q.size(1)
108
-
109
- local_topk_ids = topk_ids_
110
- if expert_map is not None:
111
- "Translate info from expert_map to topk_ids"
112
- local_topk_ids = torch.where(expert_map[topk_ids_] != -1,
113
- expert_map[topk_ids_], -1)
114
-
115
- topk = local_topk_ids.size(1)
116
-
240
per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
241
a2_scale.numel() != 1 if a2_scale is not None else False)
119
- if apply_router_weight_on_input:
120
- assert topk == 1, \
121
- "apply_router_weight_on_input is only implemented for topk=1"
122
- # TODO: this only works for topK=1, will need to update for topK>1
123
- a = a * topk_weights.to(out_dtype)
124
-
125
- a_q, a1_scale = ops.scaled_fp8_quant(
126
- a, a1_scale, use_per_token_if_dynamic=per_act_token)
127
- device = a_q.device
128
-
129
- expert_offsets = torch.empty((num_experts + 1),
130
- dtype=torch.int32,
131
- device=device)
132
- problem_sizes1 = torch.empty((num_experts, 3),
133
- dtype=torch.int32,
134
- device=device)
135
- problem_sizes2 = torch.empty((num_experts, 3),
136
- dtype=torch.int32,
137
- device=device)
138
-
139
- a_map_initializer = torch.empty
140
- c2_initializer = torch.empty
141
- if expert_map is not None:
142
- # With expert_map each Rank processes only a subset of experts. As
143
- # a result not all of a_map and c2 tensors are filled. We fill it
144
- # zeros for correctness.
145
- a_map_initializer = torch.zeros
146
- c2_initializer = torch.zeros
147
-
148
- a_map = a_map_initializer((local_topk_ids.numel()),
149
- dtype=torch.int32,
150
- device=device)
151
- c_map = torch.empty((local_topk_ids.numel()),
152
- dtype=torch.int32,
153
- device=device)
242
155
- ops.get_cutlass_moe_mm_data(local_topk_ids, expert_offsets, problem_sizes1,
156
- problem_sizes2, a_map, c_map, num_experts, n,
157
- k)
158
-
159
- rep_a_q = a_q.view(dtype=torch.uint8)[a_map].view(dtype=a_q.dtype)
160
- rep_a1_scales = a1_scale[a_map] if per_act_token else a1_scale
161
-
162
- c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
163
- c2 = c2_initializer((m * topk, k), device=device, dtype=out_dtype)
164
-
165
- ops.cutlass_moe_mm(c1, rep_a_q, w1_q, rep_a1_scales, w1_scale,
166
- expert_offsets[:-1], problem_sizes1, ab_strides1,
167
- ab_strides1, c_strides1)
168
-
169
- intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
170
- torch.ops._C.silu_and_mul(intermediate, c1)
171
-
172
- intemediate_q, a2_scale = ops.scaled_fp8_quant(
173
- intermediate, a2_scale, use_per_token_if_dynamic=per_act_token)
174
-
175
- ops.cutlass_moe_mm(c2, intemediate_q, w2_q, a2_scale, w2_scale,
176
- expert_offsets[:-1], problem_sizes2, ab_strides2,
177
- ab_strides2, c_strides2)
178
- # Gather tokens
179
- c2 = c2[c_map].view(m, topk, k)
180
- if not apply_router_weight_on_input:
181
- c2 = c2 * topk_weights.view(m, topk, 1).to(out_dtype)
182
- return c2.sum(dim=1)
243
+ fn = mk.FusedMoEModularKernel(
244
+ MoEPrepareAndFinalizeNoEP(
245
+ per_channel_quant=per_act_token,
246
+ quant_dtype=torch.float8_e4m3fn,
247
+ ),
248
+ CutlassExpertsFp8(
249
+ ab_strides1,
250
+ c_strides1,
251
+ ab_strides2,
252
+ c_strides2,
253
+ out_dtype,
254
+ ),
255
+ )
256
+
257
+ return fn(
258
+ a,
259
+ w1_q,
260
+ w2_q,
261
+ topk_weights,
262
+ topk_ids,
263
+ expert_map=expert_map,
264
+ w1_scale=w1_scale,
265
+ w2_scale=w2_scale,
266
+ a1_scale=a1_scale,
267
+ a2_scale=a2_scale,
268
+ apply_router_weight_on_input=apply_router_weight_on_input,
269
+ )
270
271
272
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+131
-198
@@ -1,16 +1,17 @@
1
# SPDX-License-Identifier: Apache-2.0
2
+import functools
3
import importlib.util
4
from typing import Optional
5
6
import torch
7
7
-import vllm.envs as envs
8
-from vllm import _custom_ops as ops
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.moe_align_block_size import (
11
- moe_align_block_size)
12
-from vllm.model_executor.layers.fused_moe.utils import (_fp8_perm,
13
- _fp8_quantize,
10
+from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
11
+ _moe_permute)
12
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
13
+ MoEPrepareAndFinalizeNoEP)
14
+from vllm.model_executor.layers.fused_moe.utils import (_fp8_quantize,
15
_resize_cache)
16
from vllm.utils import round_up
17
@@ -19,6 +20,19 @@ logger = init_logger(__name__)
20
has_deep_gemm = importlib.util.find_spec("deep_gemm") is not None
21
22
23
+@functools.cache
24
+def deep_gemm_block_shape() -> list[int]:
25
+ # Lazy import to avoid CUDA initialization problems.
26
+ import deep_gemm as dg
27
+ block = dg.get_m_alignment_for_contiguous_layout()
28
+ return [block, block]
29
+
30
+
31
+def _valid_deep_gemm_shape(M: int, N: int, K: int):
32
+ align = deep_gemm_block_shape()[0]
33
+ return align <= M and N % align == 0 and K % align == 0
34
+
35
+
36
def _valid_deep_gemm(hidden_states: torch.Tensor,
37
w1: torch.Tensor,
38
w2: torch.Tensor,
@@ -29,89 +43,112 @@ def _valid_deep_gemm(hidden_states: torch.Tensor,
43
aligned by `dg.get_m_alignment_for_contiguous_layout()`.
44
"""
45
if not has_deep_gemm:
46
+ logger.debug("DeepGemm disabled: deep_gemm not available.")
47
return False
48
34
- # Lazy import to avoid CUDA initialization problems.
35
- import deep_gemm as dg
36
-
37
- # Expert maps not supported yet.
49
if expert_map is not None:
50
+ logger.debug("DeepGemm disabled: expert map NYI.")
51
return False
52
41
- align = dg.get_m_alignment_for_contiguous_layout()
42
- M = hidden_states.shape[0]
43
- _, K, N = w2.shape
44
-
45
- # For now, disable DeepGemm for small N until better permute/unpermute
46
- # ops are available.
47
- if N <= 512:
53
+ M = hidden_states.size(0)
54
+ _, K, N = w2.size()
55
+ if not _valid_deep_gemm_shape(M, N, K):
56
+ logger.debug("DeepGemm disabled: unalinged problem size.")
57
return False
58
50
- if align > M or N % align != 0 or K % align != 0:
59
+ if (w1.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn):
60
+ logger.debug("DeepGemm disabled: invalid weight dtype(s).")
61
return False
62
53
- return (hidden_states.is_contiguous() and w1.is_contiguous()
54
- and w2.is_contiguous())
55
-
56
-
57
-def _moe_permute(
58
- curr_hidden_states: torch.Tensor,
59
- a1q_scale: Optional[torch.Tensor],
60
- curr_topk_ids: torch.Tensor,
61
- global_num_experts: int,
62
- expert_map: Optional[torch.Tensor],
63
- block_m: int,
64
-) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,
65
- Optional[torch.Tensor]]:
66
- """
67
- Determine the sorted_token_ids, expert_ids for the given problem size.
68
- Permute the hidden states and scales according to `sorted_token_ids`.
69
- """
70
- top_k_num = curr_topk_ids.shape[1]
71
-
72
- tokens_in_chunk, _ = curr_hidden_states.shape
63
+ if (not hidden_states.is_contiguous() or not w1.is_contiguous()
64
+ or not w2.is_contiguous()):
65
+ logger.debug(
66
+ "DeepGemm disabled: weights or activations not contiguous.")
67
+ return False
68
74
- sorted_token_ids, expert_ids, num_tokens_post_padded = (
75
- moe_align_block_size(curr_topk_ids,
76
- block_m,
77
- global_num_experts,
78
- expert_map,
79
- pad_sorted_ids=True))
69
+ return True
70
+
71
+
72
+class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
73
+
74
+ def __init__(self):
75
+ super().__init__()
76
+ self.block_shape = deep_gemm_block_shape()
77
+
78
+ def workspace_shapes(
79
+ self,
80
+ a: torch.Tensor,
81
+ M: int,
82
+ N: int,
83
+ K: int,
84
+ topk: int,
85
+ num_experts: int,
86
+ ) -> tuple[int, int, torch.dtype]:
87
+ block_m = self.block_shape[0]
88
+ M_sum = (M * topk) + num_experts * (block_m - 1)
89
+ M_sum = round_up(M_sum, block_m)
90
+ workspace1 = M_sum * max(N * 2, K)
91
+ workspace2 = M_sum * N
92
+ return (workspace1, workspace2, a.dtype)
93
+
94
+ def apply(
95
+ self,
96
+ hidden_states: torch.Tensor,
97
+ w1: torch.Tensor,
98
+ w2: torch.Tensor,
99
+ topk_ids: torch.Tensor,
100
+ activation: str,
101
+ global_num_experts: int,
102
+ expert_map: Optional[torch.Tensor],
103
+ w1_scale: Optional[torch.Tensor],
104
+ w2_scale: Optional[torch.Tensor],
105
+ w1_zp: Optional[torch.Tensor],
106
+ w2_zp: Optional[torch.Tensor],
107
+ a1q_scale: Optional[torch.Tensor],
108
+ a2_scale: Optional[torch.Tensor],
109
+ workspace13: torch.Tensor,
110
+ workspace2: torch.Tensor,
111
+ expert_num_tokens: Optional[torch.Tensor],
112
+ ) -> torch.Tensor:
113
+ import deep_gemm as dg
114
+
115
+ a1q = hidden_states
116
+ _, N, K = w1.size()
117
+
118
+ assert global_num_experts != -1
119
+ assert w2.size(1) == K
120
+
121
+ a1q, a1q_scale, _, expert_ids, inv_perm = _moe_permute(
122
+ a1q,
123
+ a1q_scale,
124
+ topk_ids,
125
+ global_num_experts,
126
+ expert_map,
127
+ self.block_shape[0],
128
+ )
129
+
130
+ # Note: M_sum is different than the pre-permuted shape of a1q.
131
+ M_sum = a1q.size(0)
132
+ workspace1 = _resize_cache(workspace13, (M_sum, N))
133
+ workspace2 = _resize_cache(workspace2, (M_sum, N // 2))
134
+ workspace3 = _resize_cache(workspace13, (M_sum, K))
135
81
- inv_perm: Optional[torch.Tensor] = None
136
+ dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
137
+ (a1q, a1q_scale), (w1, w1_scale), workspace1, expert_ids)
138
83
- num_tokens = top_k_num * tokens_in_chunk
84
- sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
85
- expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
86
- inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
139
+ self.activation(activation, workspace2, workspace1.view(-1, N))
140
88
- # Permute according to sorted token ids.
89
- curr_hidden_states = _fp8_perm(curr_hidden_states,
90
- sorted_token_ids // top_k_num)
141
+ a2q_scale: Optional[torch.Tensor] = None
142
92
- if a1q_scale is not None:
93
- a1q_scale = a1q_scale[sorted_token_ids // top_k_num]
143
+ a2q, a2q_scale = _fp8_quantize(workspace2, a2_scale, False,
144
+ self.block_shape)
145
95
- return (curr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
96
- inv_perm)
146
+ dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
147
+ (a2q, a2q_scale), (w2, w2_scale), workspace3, expert_ids)
148
149
+ workspace3 = workspace3[inv_perm, ...]
150
99
-def _moe_unpermute_and_reduce(
100
- out: torch.Tensor,
101
- curr_hidden: torch.Tensor,
102
- inv_perm: Optional[torch.Tensor],
103
- topk_weight: torch.Tensor,
104
-) -> None:
105
- """
106
- Unpermute the final result and apply topk_weights, then perform the final
107
- reduction on the hidden states.
108
- """
109
- M, topk = topk_weight.shape
110
- K = curr_hidden.shape[1]
111
- curr_hidden = curr_hidden[inv_perm, ...]
112
- curr_hidden = curr_hidden.view(-1, topk, K)
113
- curr_hidden.mul_(topk_weight.view(M, -1, 1))
114
- ops.moe_sum(curr_hidden, out)
151
+ return workspace3
152
153
154
def deep_gemm_moe_fp8(
@@ -128,6 +165,7 @@ def deep_gemm_moe_fp8(
165
expert_map: Optional[torch.Tensor] = None,
166
a1_scale: Optional[torch.Tensor] = None,
167
a2_scale: Optional[torch.Tensor] = None,
168
+ apply_router_weight_on_input=False,
169
) -> torch.Tensor:
170
"""
171
This function computes a a8w8-quantized Mixture of Experts (MoE) layer
@@ -166,129 +204,24 @@ def deep_gemm_moe_fp8(
204
Returns:
205
- torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
206
"""
169
- # Lazy import to avoid CUDA initialization problems.
170
- import deep_gemm as dg
171
-
172
- assert expert_map is None, "Expert maps not supported yet"
173
-
174
- assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
175
-
176
- assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
177
- assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
178
- assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
179
- assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
180
- assert hidden_states.dtype in [
181
- torch.float32, torch.float16, torch.bfloat16
182
- ]
183
- assert w1.dtype == torch.float8_e4m3fn
184
- assert w2.dtype == torch.float8_e4m3fn
185
- assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
186
- assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
187
- assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
188
- assert a1_scale is None or a1_scale.dim(
189
- ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[
190
- 0] == hidden_states.shape[0], "Input scale shape mismatch"
191
- assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
192
-
193
- num_tokens, _ = hidden_states.shape
194
- E, N, _ = w1.shape
195
- K = w2.shape[1]
196
- if global_num_experts == -1:
197
- global_num_experts = E
198
-
199
- # We execute the fused_moe kernel in chunks to circumvent this issue:
200
- # https://github.com/vllm-project/vllm/issues/5938
201
- CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
202
-
203
- assert _valid_deep_gemm(hidden_states, w1, w2, expert_map)
204
-
205
- if inplace:
206
- out_hidden_states = hidden_states
207
- else:
208
- out_hidden_states = torch.empty_like(hidden_states)
209
-
210
- block_m = dg.get_m_alignment_for_contiguous_layout()
211
- block_shape = [block_m, block_m]
212
-
213
- assert w1_scale is not None
214
- assert w2_scale is not None
215
-
216
- # We attempt to transpose and align offline in Fp8MoEMethod, in which
217
- # case these calls will be nops. Otherwise, they'll be performed every
218
- # time the layer is executed.
219
- w1_scale = dg.get_col_major_tma_aligned_tensor(w1_scale).contiguous()
220
- w2_scale = dg.get_col_major_tma_aligned_tensor(w2_scale).contiguous()
221
-
222
- M_sum = topk_ids.numel() + global_num_experts * (block_m - 1)
223
- M_sum = round_up(M_sum, block_m)
224
-
225
- num_chunks = (num_tokens // CHUNK_SIZE) + 1
226
-
227
- # We can reuse the memory between cache1 and cache3 because by the time
228
- # we need cache3, we're done with cache1
229
- workspace13 = torch.empty(M_sum * max(N, K),
230
- device=hidden_states.device,
231
- dtype=hidden_states.dtype)
232
-
233
- workspace1 = workspace13[:M_sum * N].view(M_sum, N)
234
- workspace2 = torch.empty((M_sum, N // 2),
235
- device=hidden_states.device,
236
- dtype=hidden_states.dtype)
237
- workspace3 = workspace13[:M_sum * K].view(M_sum, K)
238
-
239
- for chunk in range(num_chunks):
240
- begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
241
- min((chunk + 1) * CHUNK_SIZE,
242
- num_tokens))
243
- curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
244
- tokens_in_chunk, _ = curr_hidden_states.shape
245
-
246
- if tokens_in_chunk == 0:
247
- break
248
-
249
- curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
250
- curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
251
-
252
- a1q_scale: Optional[torch.Tensor] = None
253
-
254
- qcurr_hidden_states, a1q_scale = _fp8_quantize(curr_hidden_states,
255
- a1_scale, block_shape)
256
-
257
- (qcurr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
258
- inv_perm) = _moe_permute(qcurr_hidden_states, a1q_scale,
259
- curr_topk_ids, global_num_experts,
260
- expert_map, block_m)
261
-
262
- # Adjust the intermediate cache size and config for the last chunk.
263
- # Note that in most cases we only have one chunk so the cache size
264
- # and config are already set correctly and do not need to be adjusted.
265
- if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
266
- curr_M = sorted_token_ids.numel()
267
- workspace1 = _resize_cache(workspace1, (curr_M, N))
268
- workspace2 = _resize_cache(workspace2, (curr_M, N // 2))
269
- workspace3 = _resize_cache(workspace3, (curr_M, K))
270
-
271
- dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
272
- (qcurr_hidden_states, a1q_scale), (w1, w1_scale), workspace1,
273
- expert_ids)
274
-
275
- if activation == "silu":
276
- torch.ops._C.silu_and_mul(workspace2, workspace1.view(-1, N))
277
- elif activation == "gelu":
278
- torch.ops._C.gelu_and_mul(workspace2, workspace1.view(-1, N))
279
- else:
280
- raise ValueError(f"Unsupported FusedMoe activation: {activation}")
281
-
282
- a2q_scale: Optional[torch.Tensor] = None
283
-
284
- qworkspace2, a2q_scale = _fp8_quantize(workspace2, a2_scale,
285
- block_shape)
286
-
287
- dg.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
288
- (qworkspace2, a2q_scale), (w2, w2_scale), workspace3, expert_ids)
289
-
290
- _moe_unpermute_and_reduce(
291
- out_hidden_states[begin_chunk_idx:end_chunk_idx],
292
- workspace3.view(*workspace3.shape), inv_perm, curr_topk_weights)
293
-
294
- return out_hidden_states
207
+ fn = mk.FusedMoEModularKernel(
208
+ MoEPrepareAndFinalizeNoEP(quant_dtype=torch.float8_e4m3fn,
209
+ block_shape=deep_gemm_block_shape()),
210
+ DeepGemmExperts(),
211
+ )
212
+ return fn(
213
+ hidden_states,
214
+ w1,
215
+ w2,
216
+ topk_weights,
217
+ topk_ids,
218
+ inplace,
219
+ activation,
220
+ global_num_experts,
221
+ expert_map,
222
+ w1_scale=w1_scale,
223
+ w2_scale=w2_scale,
224
+ a1_scale=a1_scale,
225
+ a2_scale=a2_scale,
226
+ apply_router_weight_on_input=apply_router_weight_on_input,
227
+ )
vllm/model_executor/layers/fused_moe/fused_batched_moe.py
new
+755
@@ -0,0 +1,755 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+"""Fused batched MoE kernel."""
3
+from typing import Optional
4
+
5
+import torch
6
+import triton
7
+import triton.language as tl
8
+
9
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
+from vllm.model_executor.layers.fused_moe.fused_moe import (
11
+ get_config_dtype_str, try_get_optimal_moe_config)
12
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache
13
+
14
+
15
+@triton.jit
16
+def moe_mmk(
17
+ a_ptrs,
18
+ b_ptrs,
19
+ K,
20
+ expert_id,
21
+ a_scale_ptr,
22
+ b_scale_ptr,
23
+ # The stride variables represent how much to increase the ptr by when
24
+ # moving by 1 element in a particular dimension. E.g. `stride_am` is
25
+ # how much to increase `a_ptr` by to get the element one row down
26
+ # (A has M rows).
27
+ stride_ak,
28
+ stride_bk,
29
+ stride_asm,
30
+ stride_ask,
31
+ stride_bse,
32
+ stride_bsk,
33
+ stride_bsn,
34
+ # Offsets and masks
35
+ offs_m,
36
+ offs_n,
37
+ mask_m,
38
+ # Block size for block-wise quantization
39
+ group_n: tl.constexpr,
40
+ group_k: tl.constexpr,
41
+ # Meta-parameters
42
+ BLOCK_M: tl.constexpr,
43
+ BLOCK_N: tl.constexpr,
44
+ BLOCK_K: tl.constexpr,
45
+ compute_type: tl.constexpr,
46
+ use_w8a8: tl.constexpr,
47
+ use_w8a16: tl.constexpr):
48
+
49
+ offs_k = tl.arange(0, BLOCK_K)
50
+
51
+ if use_w8a16:
52
+ b_scale_ptrs = b_scale_ptr + expert_id * stride_bse + offs_n[
53
+ None, :] * stride_bsn
54
+ b_scale = tl.load(b_scale_ptrs)
55
+
56
+ if use_w8a8:
57
+ # block-wise
58
+ if group_k > 0 and group_n > 0:
59
+ a_scale_ptrs = a_scale_ptr + offs_m * stride_asm
60
+ offs_bsn = offs_n // group_n
61
+ b_scale_ptrs = (b_scale_ptr + expert_id * stride_bse +
62
+ offs_bsn * stride_bsn)
63
+ # tensor-wise
64
+ else:
65
+ a_scale = tl.load(a_scale_ptr)
66
+ b_scale = tl.load(b_scale_ptr + expert_id)
67
+
68
+ # -----------------------------------------------------------
69
+ # Iterate to compute a block of the C matrix.
70
+ # We accumulate into a `[BLOCK_SIZE_M, BLOCK_SIZE_N]` block
71
+ # of fp32 values for higher accuracy.
72
+ # `accumulator` will be converted back to fp16 after the loop.
73
+ accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
74
+ for k in range(0, tl.cdiv(K, BLOCK_K)):
75
+ # Load the next block of A and B, generate a mask by checking the
76
+ # K dimension.
77
+ a = tl.load(a_ptrs,
78
+ mask=mask_m[:, None] & (offs_k[None, :] < K - k * BLOCK_K),
79
+ other=0.0)
80
+ b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_K, other=0.0)
81
+ # We accumulate along the K dimension.
82
+ if use_w8a16:
83
+ accumulator = tl.dot(a, b.to(compute_type), acc=accumulator)
84
+ elif use_w8a8:
85
+ if group_k > 0 and group_n > 0:
86
+ k_start = k * BLOCK_K
87
+ offs_ks = k_start // group_k
88
+ a_scale = tl.load(a_scale_ptrs + offs_ks * stride_ask,
89
+ mask=mask_m,
90
+ other=0.0)
91
+ b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk)
92
+
93
+ accumulator += tl.dot(a, b) * a_scale[:,
94
+ None] * b_scale[None, :]
95
+ else:
96
+ if use_w8a8:
97
+ # acc used to enable fp8_fast_accum
98
+ accumulator = tl.dot(a, b, acc=accumulator)
99
+ else:
100
+ accumulator += tl.dot(a, b)
101
+ else:
102
+ accumulator += tl.dot(a, b)
103
+ # Advance the ptrs to the next K block.
104
+ a_ptrs += BLOCK_K * stride_ak
105
+ b_ptrs += BLOCK_K * stride_bk
106
+
107
+ if use_w8a16:
108
+ accumulator = (accumulator * b_scale).to(compute_type)
109
+ elif use_w8a8:
110
+ if group_k > 0 and group_n > 0:
111
+ accumulator = accumulator.to(compute_type)
112
+ else:
113
+ accumulator = (accumulator * a_scale * b_scale).to(compute_type)
114
+ else:
115
+ accumulator = accumulator.to(compute_type)
116
+
117
+ return accumulator
118
+
119
+
120
+@triton.jit
121
+def expert_triton_kernel(
122
+ a_ptr, #[max_tokens, K]
123
+ b_ptr, #[K, N]
124
+ c_ptr, #[max_tokens, N]
125
+ expert_id,
126
+ compute_type: tl.constexpr,
127
+ # Dimensions
128
+ M,
129
+ N,
130
+ K,
131
+ # Quantization data
132
+ a_scale_ptr,
133
+ b_scale_ptr,
134
+ b_zp_ptr,
135
+ # strides
136
+ stride_am,
137
+ stride_ak,
138
+ stride_bk,
139
+ stride_bn,
140
+ stride_cm,
141
+ stride_cn,
142
+ stride_asm,
143
+ stride_ask,
144
+ stride_bse,
145
+ stride_bsk,
146
+ stride_bsn,
147
+ # Blockwise quantization data
148
+ group_n,
149
+ group_k,
150
+ # Quantization schemes
151
+ use_fp8_w8a8: tl.constexpr,
152
+ use_int8_w8a16: tl.constexpr,
153
+ # Kernel config
154
+ BLOCK_M: tl.constexpr,
155
+ BLOCK_N: tl.constexpr,
156
+ BLOCK_K: tl.constexpr):
157
+
158
+ offs_m = tl.arange(0, BLOCK_M)
159
+ offs_n = tl.arange(0, BLOCK_N) % N
160
+ offs_k = tl.arange(0, BLOCK_K)
161
+ mask_m = offs_m < M
162
+
163
+ a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
164
+ b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
165
+
166
+ accumulator = moe_mmk(
167
+ a_ptrs,
168
+ b_ptrs,
169
+ K,
170
+ expert_id,
171
+ a_scale_ptr,
172
+ b_scale_ptr,
173
+ # The stride variables represent how much to increase the ptr by when
174
+ # moving by 1 element in a particular dimension. E.g. `stride_am` is
175
+ # how much to increase `a_ptr` by to get the element one row down
176
+ # (A has M rows).
177
+ stride_ak,
178
+ stride_bk,
179
+ stride_asm,
180
+ stride_ask,
181
+ stride_bse,
182
+ stride_bsk,
183
+ stride_bsn,
184
+ # Offsets and masks
185
+ offs_m,
186
+ offs_n,
187
+ mask_m,
188
+ # Block size for block-wise quantization
189
+ group_n,
190
+ group_k,
191
+ # Meta-parameters
192
+ BLOCK_M,
193
+ BLOCK_N,
194
+ BLOCK_K,
195
+ compute_type,
196
+ use_fp8_w8a8,
197
+ use_int8_w8a16)
198
+
199
+ # store in C
200
+ offs_cn = tl.arange(0, BLOCK_N)
201
+ c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_cn[None, :] * stride_cn
202
+ c_mask = mask_m[:, None] & (offs_cn[None, :] < N)
203
+ tl.store(c_ptrs, accumulator, mask=c_mask)
204
+
205
+
206
+@triton.jit
207
+def batched_triton_kernel(
208
+ a_ptr, # [E, max_num_tokens, K]
209
+ b_ptr, # [E, K, N]
210
+ c_ptr, # [E, max_num_tokens, N]
211
+ expert_num_tokens, # [E]
212
+ compute_type: tl.constexpr,
213
+ # Dimensions
214
+ max_num_tokens,
215
+ K,
216
+ N,
217
+ # Quantization data
218
+ a_scale_ptr,
219
+ b_scale_ptr,
220
+ b_zp_ptr,
221
+ # The stride variables represent how much to increase the ptr by when
222
+ # moving by 1 element in a particular dimension. E.g. `stride_am` is
223
+ # how much to increase `a_ptr` by to get the element one row down
224
+ # (A has M rows).
225
+ stride_ae,
226
+ stride_am,
227
+ stride_ak,
228
+ stride_be,
229
+ stride_bk,
230
+ stride_bn,
231
+ stride_ce,
232
+ stride_cm,
233
+ stride_cn,
234
+ stride_asm,
235
+ stride_ask,
236
+ stride_bse,
237
+ stride_bsk,
238
+ stride_bsn,
239
+ # Blockwise quantization data
240
+ group_n: tl.constexpr,
241
+ group_k: tl.constexpr,
242
+ # Quantization schemes
243
+ use_fp8_w8a8: tl.constexpr,
244
+ use_int8_w8a16: tl.constexpr,
245
+ # Kernel config
246
+ BLOCK_M: tl.constexpr,
247
+ BLOCK_N: tl.constexpr,
248
+ BLOCK_K: tl.constexpr):
249
+ expert_id = tl.program_id(axis=0)
250
+ e_num_tokens = tl.load(expert_num_tokens + expert_id)
251
+ if e_num_tokens == 0:
252
+ # Early exit
253
+ return
254
+
255
+ pid_mn = tl.program_id(axis=1)
256
+ #num_pid_m = tl.cdiv(max_num_tokens, BLOCK_M)
257
+ num_pid_n = tl.cdiv(N, BLOCK_N)
258
+ pid_m = pid_mn // num_pid_n
259
+ pid_n = pid_mn % num_pid_n
260
+
261
+ cta_m_start = pid_m * BLOCK_M
262
+ cta_n_start = pid_n * BLOCK_N
263
+ if cta_m_start >= e_num_tokens:
264
+ # Early exit
265
+ return
266
+
267
+ cta_m_size = min(BLOCK_M, e_num_tokens - cta_m_start)
268
+ cta_n_size = min(BLOCK_N, N - cta_n_start)
269
+
270
+ a_ptr = a_ptr + expert_id * stride_ae + cta_m_start * stride_am
271
+ b_ptr = b_ptr + expert_id * stride_be + cta_n_start * stride_bn
272
+ c_ptr = (c_ptr + expert_id * stride_ce + cta_m_start * stride_cm +
273
+ cta_n_start * stride_cn)
274
+
275
+ expert_triton_kernel(
276
+ a_ptr,
277
+ b_ptr,
278
+ c_ptr,
279
+ expert_id,
280
+ compute_type,
281
+ cta_m_size, # M
282
+ cta_n_size, # N
283
+ K, # K
284
+ a_scale_ptr,
285
+ b_scale_ptr,
286
+ b_zp_ptr,
287
+ # Strides
288
+ stride_am,
289
+ stride_ak,
290
+ stride_bk,
291
+ stride_bn,
292
+ stride_cm,
293
+ stride_cn,
294
+ stride_asm,
295
+ stride_ask,
296
+ stride_bse,
297
+ stride_bsk,
298
+ stride_bsn,
299
+ # Blockwise quantization data
300
+ group_n,
301
+ group_k,
302
+ # Quantization schemes
303
+ use_fp8_w8a8,
304
+ use_int8_w8a16,
305
+ # Kernel config
306
+ BLOCK_M,
307
+ BLOCK_N,
308
+ BLOCK_K)
309
+
310
+
311
+def invoke_moe_batched_triton_kernel(
312
+ A: torch.Tensor, # [E, max_tokens, K]
313
+ B: torch.Tensor, # [E, K, N]
314
+ C: torch.Tensor, # [E, max_tokens, N]
315
+ expert_num_tokens: torch.Tensor, # [E]
316
+ compute_type: tl.dtype,
317
+ # Quantization data
318
+ A_scale: torch.Tensor,
319
+ B_scale: torch.Tensor,
320
+ B_zp: torch.Tensor,
321
+ # Quantization schemes
322
+ use_fp8_w8a8: bool,
323
+ use_int8_w8a16: bool,
324
+ use_int4_w4a16: bool,
325
+ config: dict[str, int],
326
+ block_shape: Optional[list[int]] = None):
327
+
328
+ assert not use_int4_w4a16
329
+ max_num_tokens = A.size(1)
330
+ K = A.size(2)
331
+ N = C.size(2)
332
+
333
+ BLOCK_M = config['BLOCK_SIZE_M']
334
+ BLOCK_N = config['BLOCK_SIZE_N']
335
+ BLOCK_K = config['BLOCK_SIZE_K']
336
+ assert (torch.compiler.is_compiling()
337
+ or torch.cuda.is_current_stream_capturing()
338
+ or max_num_tokens % BLOCK_M == 0)
339
+
340
+ grid = (expert_num_tokens.size(0), triton.cdiv(max_num_tokens, BLOCK_M) *
341
+ triton.cdiv(B.size(1), BLOCK_N))
342
+
343
+ batched_triton_kernel[grid](
344
+ A,
345
+ B,
346
+ C,
347
+ expert_num_tokens,
348
+ compute_type,
349
+ # Dimensions
350
+ max_num_tokens,
351
+ K,
352
+ N,
353
+ # Quantization data
354
+ A_scale,
355
+ B_scale,
356
+ B_zp,
357
+ # Strides
358
+ A.stride(0),
359
+ A.stride(1),
360
+ A.stride(2),
361
+ B.stride(0),
362
+ B.stride(2),
363
+ B.stride(1),
364
+ C.stride(0),
365
+ C.stride(1),
366
+ C.stride(2),
367
+ A_scale.stride(0) if A_scale is not None and A_scale.ndim == 2 else 0,
368
+ A_scale.stride(1) if A_scale is not None and A_scale.ndim == 2 else 0,
369
+ B_scale.stride(0) if B_scale is not None and B_scale.ndim >= 2 else 0,
370
+ B_scale.stride(2) if B_scale is not None and B_scale.ndim == 3 else 0,
371
+ B_scale.stride(1) if B_scale is not None and B_scale.ndim >= 2 else 0,
372
+ # Blockwise quantization data
373
+ 0 if block_shape is None else block_shape[0],
374
+ 0 if block_shape is None else block_shape[1],
375
+ # Quantization schemes
376
+ use_fp8_w8a8,
377
+ use_int8_w8a16,
378
+ # Kernel config
379
+ BLOCK_M=BLOCK_M,
380
+ BLOCK_N=BLOCK_N,
381
+ BLOCK_K=BLOCK_K)
382
+
383
+
384
+class BatchedPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
385
+ """
386
+ A reference prepare/finalize class that reorganizes the tokens into
387
+ expert batched format, i.e. E x max_num_tokens x K. This is the format
388
+ that the PPLX dispatch/combine kernels use.
389
+ """
390
+
391
+ def __init__(self, max_num_tokens: Optional[int], world_size: int,
392
+ dp_size: int, rank: int):
393
+ super().__init__()
394
+ self.world_size = world_size
395
+ self.dp_size = dp_size
396
+ self.rank = rank
397
+ self.max_num_tokens = max_num_tokens
398
+
399
+ def prepare(
400
+ self,
401
+ a1: torch.Tensor,
402
+ a1_scale: Optional[torch.Tensor],
403
+ a2_scale: Optional[torch.Tensor],
404
+ topk_weights: torch.Tensor,
405
+ topk_ids: torch.Tensor,
406
+ num_experts: int,
407
+ expert_map: Optional[torch.Tensor],
408
+ apply_router_weight_on_input: bool,
409
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
410
+ assert a1.dim() == 2
411
+ assert topk_ids.dim() == 2
412
+ assert topk_ids.size(0) == a1.size(0)
413
+
414
+ if apply_router_weight_on_input:
415
+ topk = topk_ids.size(1)
416
+ # TODO: this only works for topK=1, will need to update for topK>1
417
+ assert topk == 1, \
418
+ "apply_router_weight_on_input is only implemented for topk=1"
419
+ a1.mul_(topk_weights.to(a1.dtype))
420
+
421
+ num_tokens, hidden_dim = a1.size()
422
+ topk = topk_ids.size(1)
423
+
424
+ if self.max_num_tokens is None:
425
+ tokens_per_expert = torch.bincount(topk_ids.view(-1),
426
+ minlength=num_experts)
427
+ self.max_num_tokens = int(tokens_per_expert.max().item())
428
+ else:
429
+ tokens_per_expert = torch.zeros(num_experts,
430
+ dtype=torch.int,
431
+ device=a1.device)
432
+
433
+ assert num_experts % self.world_size == 0
434
+
435
+ num_local_experts = num_experts // self.world_size
436
+
437
+ b_a1 = torch.zeros(
438
+ (num_local_experts, self.max_num_tokens, hidden_dim),
439
+ dtype=a1.dtype,
440
+ device=a1.device)
441
+
442
+ first_expert = num_local_experts * self.rank
443
+ last_expert = first_expert + num_local_experts
444
+
445
+ for expert_id in range(first_expert, last_expert):
446
+ topks = torch.any(topk_ids == expert_id, dim=1).flatten()
447
+ rows = torch.count_nonzero(topks.flatten())
448
+ b_a1[expert_id -
449
+ first_expert, :rows, :] = a1[:topks.numel()][topks]
450
+ tokens_per_expert[expert_id - first_expert] = rows
451
+
452
+ return b_a1, a1_scale, tokens_per_expert
453
+
454
+ def finalize(
455
+ self,
456
+ output: torch.Tensor,
457
+ fused_expert_output: torch.Tensor,
458
+ topk_weights: torch.Tensor,
459
+ topk_ids: torch.Tensor,
460
+ apply_router_weight_on_input: bool,
461
+ ) -> None:
462
+ num_tokens = topk_ids.size(0)
463
+ num_local_experts = fused_expert_output.size(0)
464
+ K = fused_expert_output.size(-1)
465
+ assert output.size(0) == num_tokens and output.size(1) == K
466
+
467
+ output.fill_(0)
468
+
469
+ first_expert = num_local_experts * self.rank
470
+ last_expert = first_expert + num_local_experts
471
+
472
+ for expert_id in range(first_expert, last_expert):
473
+ matching_tokens = topk_ids == expert_id
474
+ topks = torch.any(matching_tokens, dim=1).flatten()
475
+ rows = torch.count_nonzero(topks)
476
+ rhs = fused_expert_output[expert_id - first_expert, :rows, :]
477
+ if not apply_router_weight_on_input:
478
+ rhs.mul_(topk_weights[matching_tokens].view(rhs.size(0), 1))
479
+ output[topks] = output[topks] + rhs
480
+
481
+
482
+class BatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):
483
+ """
484
+ A reference MoE expert class that operates on expert batched format,
485
+ i.e. E x max_num_tokens x K. This is the format that the pplx
486
+ dispatch/combine kernels use.
487
+ """
488
+
489
+ def __init__(
490
+ self,
491
+ world_size: int,
492
+ dp_size: int,
493
+ max_num_tokens: Optional[int] = None,
494
+ use_fp8_w8a8: bool = False,
495
+ use_int8_w8a8: bool = False,
496
+ use_int8_w8a16: bool = False,
497
+ use_int4_w4a16: bool = False,
498
+ block_shape: Optional[list[int]] = None,
499
+ block_m: Optional[int] = None,
500
+ ):
501
+ super().__init__()
502
+ assert block_shape is None
503
+ assert block_m is None
504
+ assert not use_fp8_w8a8, "NYI"
505
+ assert not use_int8_w8a8, "NYI"
506
+ assert not use_int8_w8a16, "NYI"
507
+ assert not use_int4_w4a16, "NYI"
508
+ self.max_num_tokens = max_num_tokens
509
+ self.world_size = world_size
510
+ self.dp_size = dp_size
511
+
512
+ def workspace_shapes(
513
+ self,
514
+ a: torch.Tensor,
515
+ M: int,
516
+ N: int,
517
+ K: int,
518
+ topk: int,
519
+ num_experts: int,
520
+ ) -> tuple[int, int, torch.dtype]:
521
+ assert a.dim() == 2
522
+ num_dp = self.world_size // self.dp_size
523
+ max_num_tokens = a.size(
524
+ 0) if self.max_num_tokens is None else self.max_num_tokens
525
+ #print(f"WORKSPACE {max_num_tokens} {num_dp}")
526
+ workspace13 = num_experts * max_num_tokens * num_dp * K
527
+ workspace2 = max_num_tokens * num_dp * N
528
+ return (workspace13, workspace2, a.dtype)
529
+
530
+ def apply(
531
+ self,
532
+ hidden_states: torch.Tensor,
533
+ w1: torch.Tensor,
534
+ w2: torch.Tensor,
535
+ topk_ids: torch.Tensor,
536
+ activation: str,
537
+ global_num_experts: int,
538
+ expert_map: Optional[torch.Tensor],
539
+ w1_scale: Optional[torch.Tensor],
540
+ w2_scale: Optional[torch.Tensor],
541
+ w1_zp: Optional[torch.Tensor],
542
+ w2_zp: Optional[torch.Tensor],
543
+ a1q_scale: Optional[torch.Tensor],
544
+ a2_scale: Optional[torch.Tensor],
545
+ workspace13: torch.Tensor,
546
+ workspace2: torch.Tensor,
547
+ expert_num_tokens: Optional[torch.Tensor],
548
+ ) -> torch.Tensor:
549
+ assert hidden_states.dim() == 3
550
+ assert expert_num_tokens is not None
551
+ hidden_dim = hidden_states.size(-1)
552
+
553
+ if self.max_num_tokens is None:
554
+ max_num_tokens = hidden_states.size(1)
555
+ else:
556
+ max_num_tokens = self.max_num_tokens
557
+
558
+ num_dp = self.world_size // self.dp_size
559
+ num_experts = global_num_experts
560
+ out = _resize_cache(workspace13,
561
+ (num_experts, max_num_tokens * num_dp, hidden_dim))
562
+ num_local_experts = w1.size(0)
563
+ assert num_local_experts == w1.size(0), (
564
+ f"{num_local_experts} == {w1.size(0)}")
565
+
566
+ N = w1.size(1) // 2
567
+
568
+ # Not cudagraph friendly
569
+ assert (torch.compiler.is_compiling()
570
+ or torch.cuda.is_current_stream_capturing()
571
+ or torch.all(expert_num_tokens <= max_num_tokens * num_dp)), (
572
+ f"{expert_num_tokens} <= {max_num_tokens * num_dp}")
573
+
574
+ for expert in range(num_local_experts):
575
+ # Indexing expert_num_tokens doesn't work w/cudagraphs or inductor
576
+ if (torch.compiler.is_compiling()
577
+ or torch.cuda.is_current_stream_capturing()):
578
+ num = max_num_tokens * num_dp
579
+ else:
580
+ num = int(expert_num_tokens[expert].item())
581
+ tmp = _resize_cache(workspace2, (num, N))
582
+ input = hidden_states[expert, :num, :] @ w1[expert].transpose(0, 1)
583
+ self.activation(activation, tmp, input)
584
+ out[expert, :num, :] = tmp @ w2[expert].transpose(0, 1)
585
+
586
+ return out
587
+
588
+
589
+class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
590
+ """
591
+ A Triton based MoE expert class that operates on expert batched format,
592
+ i.e. E x max_num_tokens x K. This is the format that the pplx
593
+ dispatch/combine kernels use.
594
+ """
595
+
596
+ def __init__(
597
+ self,
598
+ max_num_tokens: Optional[int] = None,
599
+ use_fp8_w8a8: bool = False,
600
+ use_int8_w8a8: bool = False,
601
+ use_int8_w8a16: bool = False,
602
+ use_int4_w4a16: bool = False,
603
+ block_shape: Optional[list[int]] = None,
604
+ world_size: int = 1,
605
+ dp_size: int = 1,
606
+ ):
607
+ super().__init__()
608
+ self.use_fp8_w8a8 = use_fp8_w8a8
609
+ self.use_int8_w8a8 = use_int8_w8a8
610
+ self.use_int4_w4a16 = use_int4_w4a16
611
+ self.use_int8_w8a16 = use_int8_w8a16
612
+ self.block_shape = block_shape
613
+ self.max_num_tokens = max_num_tokens
614
+ assert not use_int8_w8a8, "NYI"
615
+ assert not use_int4_w4a16, "NYI"
616
+ self.world_size = world_size
617
+ self.dp_size = dp_size
618
+
619
+ def workspace_shapes(
620
+ self,
621
+ a: torch.Tensor,
622
+ M: int,
623
+ N: int,
624
+ K: int,
625
+ topk: int,
626
+ num_experts: int,
627
+ ) -> tuple[int, int, torch.dtype]:
628
+ assert a.dim() == 2
629
+ num_dp = self.world_size // self.dp_size
630
+ max_num_tokens = a.size(
631
+ 0) if self.max_num_tokens is None else self.max_num_tokens
632
+ workspace13 = num_experts * max_num_tokens * num_dp * max(K, N)
633
+ workspace2 = num_experts * max_num_tokens * num_dp * (N // 2)
634
+ return (workspace13, workspace2, a.dtype)
635
+
636
+ def apply(
637
+ self,
638
+ hidden_states: torch.Tensor,
639
+ w1: torch.Tensor,
640
+ w2: torch.Tensor,
641
+ topk_ids: torch.Tensor,
642
+ activation: str,
643
+ global_num_experts: int,
644
+ expert_map: Optional[torch.Tensor],
645
+ w1_scale: Optional[torch.Tensor],
646
+ w2_scale: Optional[torch.Tensor],
647
+ w1_zp: Optional[torch.Tensor],
648
+ w2_zp: Optional[torch.Tensor],
649
+ a1q_scale: Optional[torch.Tensor],
650
+ a2_scale: Optional[torch.Tensor],
651
+ workspace13: torch.Tensor,
652
+ workspace2: torch.Tensor,
653
+ expert_num_tokens: Optional[torch.Tensor],
654
+ ) -> torch.Tensor:
655
+ # Check constraints.
656
+ if self.use_int4_w4a16:
657
+ assert hidden_states.size(-1) // 2 == w1.size(2), (
658
+ "Hidden size mismatch")
659
+ else:
660
+ assert hidden_states.size(-1) == w1.size(2), (
661
+ f"Hidden size mismatch {hidden_states.size(-1)} "
662
+ f"!= {w1.size(2)}")
663
+
664
+ assert hidden_states.is_contiguous(
665
+ ), "Hidden_states must be contiguous"
666
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
667
+ assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
668
+ assert hidden_states.dtype in [
669
+ torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
670
+ ]
671
+
672
+ # TODO: num_tokens -> max_num_tokens?
673
+ E, num_tokens, N, K, top_k_num = mk._moe_problem_size(
674
+ hidden_states, w1, w2, topk_ids)
675
+
676
+ assert w1.size(0) == E
677
+ assert w2.size(0) == E
678
+
679
+ config_dtype = get_config_dtype_str(use_fp8_w8a8=self.use_fp8_w8a8,
680
+ use_int8_w8a16=self.use_int8_w8a16,
681
+ use_int4_w4a16=self.use_int4_w4a16,
682
+ dtype=hidden_states.dtype)
683
+
684
+ config = try_get_optimal_moe_config(
685
+ w1.size(),
686
+ w2.size(),
687
+ top_k_num,
688
+ config_dtype,
689
+ num_tokens,
690
+ block_shape=self.block_shape,
691
+ )
692
+
693
+ if hidden_states.dtype == torch.bfloat16:
694
+ compute_type = tl.bfloat16
695
+ elif hidden_states.dtype == torch.float16:
696
+ compute_type = tl.float16
697
+ elif hidden_states.dtype == torch.float32:
698
+ compute_type = tl.float32
699
+ elif hidden_states.dtype == torch.float8_e4m3fn:
700
+ compute_type = tl.bfloat16
701
+ else:
702
+ raise ValueError(
703
+ f"Unsupported compute_type: {hidden_states.dtype}")
704
+
705
+ #print(f"shape: E={E}, M={num_tokens}, N={N}, K={K}, top_k={top_k_num}")
706
+ # We can reuse the memory between these because by the time we need
707
+ # cache3, we're done with cache1
708
+ intermediate_cache1 = _resize_cache(workspace13, (E, num_tokens, N))
709
+ intermediate_cache2 = _resize_cache(workspace2,
710
+ (E, num_tokens, N // 2))
711
+ intermediate_cache3 = _resize_cache(workspace13, (E, num_tokens, K))
712
+
713
+ # MM1
714
+ invoke_moe_batched_triton_kernel(A=hidden_states,
715
+ B=w1,
716
+ C=intermediate_cache1,
717
+ expert_num_tokens=expert_num_tokens,
718
+ compute_type=compute_type,
719
+ A_scale=a1q_scale,
720
+ B_scale=w1_scale,
721
+ B_zp=w1_zp,
722
+ use_fp8_w8a8=self.use_fp8_w8a8,
723
+ use_int8_w8a16=self.use_int8_w8a16,
724
+ use_int4_w4a16=self.use_int4_w4a16,
725
+ config=config,
726
+ block_shape=self.block_shape)
727
+
728
+ # TODO: would be nice to use expert_num_tokens here to reduce
729
+ # garbage compute
730
+ self.activation(activation, intermediate_cache2.view(-1, N // 2),
731
+ intermediate_cache1.view(-1, N))
732
+
733
+ #qintermediate_cache2 = intermediate_cache2
734
+ a2q_scale = a2_scale
735
+ # TODO (varun) : support w8a8
736
+ assert not self.use_fp8_w8a8
737
+ #if self.use_fp8_w8a8:
738
+ # qintermediate_cache2, a2q_scale = _fp8_quantize(
739
+ # intermediate_cache2, a2_scale, self.block_shape)
740
+
741
+ invoke_moe_batched_triton_kernel(A=intermediate_cache2,
742
+ B=w2,
743
+ C=intermediate_cache3,
744
+ expert_num_tokens=expert_num_tokens,
745
+ compute_type=compute_type,
746
+ A_scale=a2q_scale,
747
+ B_scale=w2_scale,
748
+ B_zp=w2_zp,
749
+ use_fp8_w8a8=self.use_fp8_w8a8,
750
+ use_int8_w8a16=self.use_int8_w8a16,
751
+ use_int4_w4a16=self.use_int4_w4a16,
752
+ config=config,
753
+ block_shape=self.block_shape)
754
+
755
+ return intermediate_cache3
vllm/model_executor/layers/fused_moe/fused_moe.py
+287
-101
@@ -8,16 +8,17 @@ from typing import Any, Callable, Optional
8
import torch
9
10
import vllm.envs as envs
11
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm import _custom_ops as ops
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
15
_valid_deep_gemm, deep_gemm_moe_fp8)
16
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
17
moe_align_block_size)
17
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18
- per_token_group_quant_fp8)
19
-from vllm.model_executor.layers.quantization.utils.int8_utils import (
20
- per_token_group_quant_int8, per_token_quant_int8)
18
+from vllm.model_executor.layers.fused_moe.prepare_finalize import (
19
+ MoEPrepareAndFinalizeNoEP)
20
+from vllm.model_executor.layers.fused_moe.utils import (
21
+ _resize_cache, moe_kernel_quantize_input)
22
from vllm.platforms import current_platform
23
from vllm.triton_utils import tl, triton
24
from vllm.utils import direct_register_custom_op
@@ -484,6 +485,20 @@ def invoke_fused_moe_kernel(A: torch.Tensor,
485
assert topk_weights is None or topk_weights.stride(1) == 1
486
assert sorted_token_ids.stride(0) == 1
487
488
+ if use_fp8_w8a8 or use_int8_w8a8:
489
+ assert B_scale is not None
490
+ assert (block_shape is None or triton.cdiv(B.shape[-2], block_shape[0])
491
+ == B_scale.shape[-2])
492
+ assert (block_shape is None or triton.cdiv(B.shape[-1], block_shape[1])
493
+ == B_scale.shape[-1])
494
+
495
+ elif use_int8_w8a16 or use_int4_w4a16:
496
+ assert B_scale is not None
497
+ assert block_shape is None or block_shape[0] == 0
498
+ else:
499
+ assert A_scale is None
500
+ assert B_scale is None
501
+
502
M = A.shape[0]
503
num_tokens = M * top_k
504
@@ -855,6 +870,7 @@ def fused_topk(
870
gating_output: torch.Tensor,
871
topk: int,
872
renormalize: bool,
873
+ indices_type: Optional[torch.dtype] = None,
874
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
875
assert hidden_states.shape[0] == gating_output.shape[0], (
876
"Number of tokens mismatch")
@@ -865,10 +881,11 @@ def fused_topk(
881
topk,
882
dtype=torch.float32,
883
device=hidden_states.device)
868
- topk_ids = torch.empty(M,
869
- topk,
870
- dtype=torch.int32,
871
- device=hidden_states.device)
884
+ topk_ids = torch.empty(
885
+ M,
886
+ topk,
887
+ dtype=torch.int32 if indices_type is None else indices_type,
888
+ device=hidden_states.device)
889
token_expert_indices = torch.empty(M,
890
topk,
891
dtype=torch.int32,
@@ -962,6 +979,20 @@ def get_config_dtype_str(
979
return None
980
981
982
+# TODO (bnell): use scalar_type instead of bools?
983
+def get_config_qtype(
984
+ use_fp8_w8a8: bool,
985
+ use_int8_w8a8: bool,
986
+ use_int8_w8a16: bool,
987
+ use_int4_w4a16: bool,
988
+) -> Optional[torch.dtype]:
989
+ if use_fp8_w8a8:
990
+ return torch.float8_e4m3fn
991
+ elif use_int8_w8a8:
992
+ return torch.int8
993
+ return None
994
+
995
+
996
def inplace_fused_experts(hidden_states: torch.Tensor,
997
w1: torch.Tensor,
998
w2: torch.Tensor,
@@ -1128,7 +1159,10 @@ def fused_experts(hidden_states: torch.Tensor,
1159
a2_scale: Optional[torch.Tensor] = None,
1160
block_shape: Optional[list[int]] = None,
1161
allow_deep_gemm: bool = False) -> torch.Tensor:
1131
- if (allow_deep_gemm and use_fp8_w8a8
1162
+ # For now, disable DeepGemm for small N (<= 512) until better
1163
+ # permute/unpermute ops are available.
1164
+ N = w1.shape[1]
1165
+ if (allow_deep_gemm and use_fp8_w8a8 and N > 512
1166
and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):
1167
assert apply_router_weight_on_input is False
1168
return deep_gemm_moe_fp8(
@@ -1145,6 +1179,7 @@ def fused_experts(hidden_states: torch.Tensor,
1179
w2_scale=w2_scale,
1180
a1_scale=a1_scale,
1181
a2_scale=a2_scale,
1182
+ apply_router_weight_on_input=apply_router_weight_on_input,
1183
)
1184
else:
1185
return dispatch_fused_experts_func(inplace)(
@@ -1171,87 +1206,37 @@ def fused_experts(hidden_states: torch.Tensor,
1206
block_shape=block_shape)
1207
1208
1174
-def moe_kernel_prepare_input(
1175
- A: torch.Tensor,
1176
- B: torch.Tensor,
1177
- A_scale: Optional[torch.Tensor],
1178
- B_scale: Optional[torch.Tensor],
1179
- use_fp8_w8a8: bool,
1180
- use_int8_w8a8: bool,
1181
- use_int8_w8a16: bool,
1182
- use_int4_w4a16: bool,
1183
- per_channel_quant: bool,
1209
+def fused_experts_impl(
1210
+ hidden_states: torch.Tensor,
1211
+ w1: torch.Tensor,
1212
+ w2: torch.Tensor,
1213
+ topk_weights: torch.Tensor,
1214
+ topk_ids: torch.Tensor,
1215
+ inplace: bool = False,
1216
+ activation: str = "silu",
1217
+ apply_router_weight_on_input: bool = False,
1218
+ use_fp8_w8a8: bool = False,
1219
+ use_int8_w8a8: bool = False,
1220
+ use_int8_w8a16: bool = False,
1221
+ use_int4_w4a16: bool = False,
1222
+ per_channel_quant: bool = False,
1223
+ global_num_experts: int = -1,
1224
+ expert_map: Optional[torch.Tensor] = None,
1225
+ w1_scale: Optional[torch.Tensor] = None,
1226
+ w2_scale: Optional[torch.Tensor] = None,
1227
+ w1_zp: Optional[torch.Tensor] = None,
1228
+ w2_zp: Optional[torch.Tensor] = None,
1229
+ a1_scale: Optional[torch.Tensor] = None,
1230
+ a2_scale: Optional[torch.Tensor] = None,
1231
block_shape: Optional[list[int]] = None,
1185
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
1186
- if use_fp8_w8a8:
1187
- assert B_scale is not None
1188
- if block_shape is None:
1189
- # If weights are per-channel (per_channel_quant=True), then
1190
- # activations apply per-token quantization. Otherwise, assume
1191
- # activation tensor-wise fp8 quantization, dynamic or static
1192
- A, A_scale = ops.scaled_fp8_quant(
1193
- A, A_scale, use_per_token_if_dynamic=per_channel_quant)
1194
- else:
1195
- # activation block-wise fp8 quantization
1196
- assert len(block_shape) == 2
1197
- _, block_k = block_shape[0], block_shape[1]
1198
- A, A_scale = per_token_group_quant_fp8(A, block_k)
1199
- assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
1200
- # assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
1201
- # assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
1202
- elif use_int8_w8a8:
1203
- assert B_scale is not None
1204
- if block_shape is None:
1205
- # activation channel-wise int8 quantization
1206
- assert (per_channel_quant
1207
- ), "int8 quantization only supports block or channel-wise"
1208
- A, A_scale = per_token_quant_int8(A)
1209
- else:
1210
- # activation block-wise int8 quantization
1211
- assert len(block_shape) == 2
1212
- _, block_k = block_shape[0], block_shape[1]
1213
- A, A_scale = per_token_group_quant_int8(A, block_k)
1214
- assert triton.cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
1215
- # assert triton.cdiv(B.shape[-2], block_n) == B_scale.shape[-2]
1216
- # assert triton.cdiv(B.shape[-1], block_k) == B_scale.shape[-1]
1217
- elif use_int8_w8a16 or use_int4_w4a16:
1218
- assert B_scale is not None
1219
- assert block_shape is None or block_shape[0] == 0
1220
- else:
1221
- assert A_scale is None
1222
- assert B_scale is None
1223
-
1224
- return A, A_scale
1225
-
1226
-
1227
-def fused_experts_impl(hidden_states: torch.Tensor,
1228
- w1: torch.Tensor,
1229
- w2: torch.Tensor,
1230
- topk_weights: torch.Tensor,
1231
- topk_ids: torch.Tensor,
1232
- inplace: bool = False,
1233
- activation: str = "silu",
1234
- apply_router_weight_on_input: bool = False,
1235
- use_fp8_w8a8: bool = False,
1236
- use_int8_w8a8: bool = False,
1237
- use_int8_w8a16: bool = False,
1238
- use_int4_w4a16: bool = False,
1239
- per_channel_quant: bool = False,
1240
- global_num_experts: int = -1,
1241
- expert_map: Optional[torch.Tensor] = None,
1242
- w1_scale: Optional[torch.Tensor] = None,
1243
- w2_scale: Optional[torch.Tensor] = None,
1244
- w1_zp: Optional[torch.Tensor] = None,
1245
- w2_zp: Optional[torch.Tensor] = None,
1246
- a1_scale: Optional[torch.Tensor] = None,
1247
- a2_scale: Optional[torch.Tensor] = None,
1248
- block_shape: Optional[list[int]] = None):
1232
+) -> torch.Tensor:
1233
# Check constraints.
1234
if use_int4_w4a16:
1235
assert hidden_states.shape[1] // 2 == w1.shape[
1236
2], "Hidden size mismatch"
1237
else:
1254
- assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
1238
+ assert hidden_states.shape[1] == w1.shape[2], (
1239
+ f"Hidden size mismatch {hidden_states.shape[1]} != {w1.shape[2]}")
1240
1241
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
1242
assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
@@ -1261,7 +1246,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1246
torch.float32, torch.float16, torch.bfloat16
1247
]
1248
1264
- num_tokens, _ = hidden_states.shape
1249
+ num_tokens = hidden_states.shape[0]
1250
E, N, _ = w1.shape
1251
K = w2.shape[1]
1252
if global_num_experts == -1:
@@ -1276,6 +1261,11 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1261
use_int4_w4a16=use_int4_w4a16,
1262
dtype=hidden_states.dtype)
1263
1264
+ qtype = get_config_qtype(use_fp8_w8a8=use_fp8_w8a8,
1265
+ use_int8_w8a8=use_int8_w8a8,
1266
+ use_int8_w8a16=use_int8_w8a16,
1267
+ use_int4_w4a16=use_int4_w4a16)
1268
+
1269
get_config_func = functools.partial(
1270
try_get_optimal_moe_config,
1271
w1.shape,
@@ -1338,15 +1328,10 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1328
curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
1329
curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
1330
1341
- qcurr_hidden_states, qa1_scale = moe_kernel_prepare_input(
1331
+ qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
1332
A=curr_hidden_states,
1343
- B=w1,
1333
A_scale=a1_scale,
1345
- B_scale=w1_scale,
1346
- use_fp8_w8a8=use_fp8_w8a8,
1347
- use_int8_w8a8=use_int8_w8a8,
1348
- use_int8_w8a16=use_int8_w8a16,
1349
- use_int4_w4a16=use_int4_w4a16,
1334
+ qtype=qtype,
1335
per_channel_quant=per_channel_quant,
1336
block_shape=block_shape)
1337
@@ -1357,7 +1342,7 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1342
invoke_fused_moe_kernel(qcurr_hidden_states,
1343
w1,
1344
intermediate_cache1,
1360
- qa1_scale,
1345
+ a1q_scale,
1346
w1_scale,
1347
w1_zp,
1348
curr_topk_weights,
@@ -1384,22 +1369,17 @@ def fused_experts_impl(hidden_states: torch.Tensor,
1369
else:
1370
raise ValueError(f"Unsupported FusedMoe activation: {activation}")
1371
1387
- qintermediate_cache2, qa2_scale = moe_kernel_prepare_input(
1372
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1373
A=intermediate_cache2,
1389
- B=w2,
1374
A_scale=a2_scale,
1391
- B_scale=w2_scale,
1392
- use_fp8_w8a8=use_fp8_w8a8,
1393
- use_int8_w8a8=use_int8_w8a8,
1394
- use_int8_w8a16=use_int8_w8a16,
1395
- use_int4_w4a16=use_int4_w4a16,
1375
+ qtype=qtype,
1376
per_channel_quant=per_channel_quant,
1377
block_shape=block_shape)
1378
1379
invoke_fused_moe_kernel(qintermediate_cache2,
1380
w2,
1381
intermediate_cache3,
1402
- qa2_scale,
1382
+ a2q_scale,
1383
w2_scale,
1384
w2_zp,
1385
curr_topk_weights,
@@ -1534,3 +1514,209 @@ def fused_moe(
1514
a1_scale=a1_scale,
1515
a2_scale=a2_scale,
1516
block_shape=block_shape)
1517
+
1518
+
1519
+class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
1520
+
1521
+ def __init__(
1522
+ self,
1523
+ use_fp8_w8a8: bool,
1524
+ use_int8_w8a8: bool,
1525
+ use_int8_w8a16: bool,
1526
+ use_int4_w4a16: bool,
1527
+ per_channel_quant: bool,
1528
+ block_shape: Optional[list[int]] = None,
1529
+ block_m: Optional[int] = None,
1530
+ ):
1531
+ super().__init__()
1532
+ self.use_fp8_w8a8 = use_fp8_w8a8
1533
+ self.use_int4_w4a16 = use_int4_w4a16
1534
+ self.use_int8_w8a8 = use_int8_w8a8
1535
+ self.use_int8_w8a16 = use_int8_w8a16
1536
+ self.block_shape = block_shape
1537
+ self.block_m = block_m
1538
+ self.qtype = get_config_qtype(use_fp8_w8a8=use_fp8_w8a8,
1539
+ use_int8_w8a8=use_int8_w8a8,
1540
+ use_int8_w8a16=use_int8_w8a16,
1541
+ use_int4_w4a16=use_int4_w4a16)
1542
+ self.per_channel_quant = per_channel_quant
1543
+
1544
+ def workspace_shapes(
1545
+ self,
1546
+ a: torch.Tensor,
1547
+ M: int,
1548
+ N: int,
1549
+ K: int,
1550
+ topk: int,
1551
+ num_experts: int,
1552
+ ) -> tuple[int, int, torch.dtype]:
1553
+ factor = num_experts if a.dim() == 3 else 1
1554
+ workspace1 = M * topk * max(N * 2, K) * factor
1555
+ workspace2 = M * topk * N * factor
1556
+ return (workspace1, workspace2, a.dtype)
1557
+
1558
+ def apply(
1559
+ self,
1560
+ hidden_states: torch.Tensor,
1561
+ w1: torch.Tensor,
1562
+ w2: torch.Tensor,
1563
+ topk_ids: torch.Tensor,
1564
+ activation: str,
1565
+ global_num_experts: int,
1566
+ expert_map: Optional[torch.Tensor],
1567
+ w1_scale: Optional[torch.Tensor],
1568
+ w2_scale: Optional[torch.Tensor],
1569
+ w1_zp: Optional[torch.Tensor],
1570
+ w2_zp: Optional[torch.Tensor],
1571
+ a1q_scale: Optional[torch.Tensor],
1572
+ a2_scale: Optional[torch.Tensor],
1573
+ workspace13: torch.Tensor,
1574
+ workspace2: torch.Tensor,
1575
+ expert_num_tokens: Optional[torch.Tensor],
1576
+ ) -> torch.Tensor:
1577
+ # Check constraints.
1578
+ if self.use_int4_w4a16:
1579
+ assert hidden_states.size(-1) // 2 == w1.size(2), (
1580
+ "Hidden size mismatch")
1581
+ else:
1582
+ assert hidden_states.size(-1) == w1.size(2), \
1583
+ (f"Hidden size mismatch {hidden_states.size(-1)} "
1584
+ f"!= {w1.size(2)}")
1585
+
1586
+ assert hidden_states.is_contiguous(
1587
+ ), "Hidden_states must be contiguous"
1588
+ assert hidden_states.dim() == 2
1589
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
1590
+ assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
1591
+ assert hidden_states.dtype in [
1592
+ torch.float32, torch.float16, torch.bfloat16, torch.float8_e4m3fn
1593
+ ]
1594
+
1595
+ E, num_tokens, N, K, top_k_num = mk._moe_problem_size(
1596
+ hidden_states, w1, w2, topk_ids)
1597
+
1598
+ if global_num_experts == -1:
1599
+ global_num_experts = E
1600
+
1601
+ config_dtype = get_config_dtype_str(use_fp8_w8a8=self.use_fp8_w8a8,
1602
+ use_int8_w8a16=self.use_int8_w8a16,
1603
+ use_int4_w4a16=self.use_int4_w4a16,
1604
+ dtype=hidden_states.dtype)
1605
+
1606
+ config = try_get_optimal_moe_config(
1607
+ w1.shape,
1608
+ w2.shape,
1609
+ top_k_num,
1610
+ config_dtype,
1611
+ num_tokens,
1612
+ block_shape=self.block_shape,
1613
+ )
1614
+
1615
+ if hidden_states.dtype == torch.bfloat16:
1616
+ compute_type = tl.bfloat16
1617
+ elif hidden_states.dtype == torch.float16:
1618
+ compute_type = tl.float16
1619
+ elif hidden_states.dtype == torch.float32:
1620
+ compute_type = tl.float32
1621
+ elif hidden_states.dtype == torch.float8_e4m3fn:
1622
+ compute_type = tl.bfloat16
1623
+ else:
1624
+ raise ValueError(
1625
+ f"Unsupported compute_type: {hidden_states.dtype}")
1626
+
1627
+ # We can reuse the memory between these because by the time we need
1628
+ # cache3, we're done with cache1
1629
+ intermediate_cache1 = _resize_cache(workspace13,
1630
+ (num_tokens, top_k_num, N))
1631
+ intermediate_cache2 = _resize_cache(workspace2,
1632
+ (num_tokens * top_k_num, N // 2))
1633
+ intermediate_cache3 = _resize_cache(workspace13,
1634
+ (num_tokens, top_k_num, K))
1635
+
1636
+ sorted_token_ids, expert_ids, num_tokens_post_padded = (
1637
+ moe_align_block_size(topk_ids, config['BLOCK_SIZE_M'],
1638
+ global_num_experts, expert_map))
1639
+
1640
+ invoke_fused_moe_kernel(hidden_states,
1641
+ w1,
1642
+ intermediate_cache1,
1643
+ a1q_scale,
1644
+ w1_scale,
1645
+ w1_zp,
1646
+ None,
1647
+ sorted_token_ids,
1648
+ expert_ids,
1649
+ num_tokens_post_padded,
1650
+ False,
1651
+ top_k_num,
1652
+ config,
1653
+ compute_type=compute_type,
1654
+ use_fp8_w8a8=self.use_fp8_w8a8,
1655
+ use_int8_w8a8=self.use_int8_w8a8,
1656
+ use_int8_w8a16=self.use_int8_w8a16,
1657
+ use_int4_w4a16=self.use_int4_w4a16,
1658
+ per_channel_quant=self.per_channel_quant,
1659
+ block_shape=self.block_shape)
1660
+
1661
+ self.activation(activation, intermediate_cache2,
1662
+ intermediate_cache1.view(-1, N))
1663
+
1664
+ a2q_scale: Optional[torch.Tensor] = None
1665
+
1666
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
1667
+ intermediate_cache2, a2_scale, self.qtype, self.per_channel_quant,
1668
+ self.block_shape)
1669
+
1670
+ invoke_fused_moe_kernel(qintermediate_cache2,
1671
+ w2,
1672
+ intermediate_cache3,
1673
+ a2q_scale,
1674
+ w2_scale,
1675
+ w2_zp,
1676
+ None,
1677
+ sorted_token_ids,
1678
+ expert_ids,
1679
+ num_tokens_post_padded,
1680
+ False,
1681
+ 1,
1682
+ config,
1683
+ compute_type=compute_type,
1684
+ use_fp8_w8a8=self.use_fp8_w8a8,
1685
+ use_int8_w8a8=self.use_int8_w8a8,
1686
+ use_int8_w8a16=self.use_int8_w8a16,
1687
+ use_int4_w4a16=self.use_int4_w4a16,
1688
+ per_channel_quant=self.per_channel_quant,
1689
+ block_shape=self.block_shape)
1690
+
1691
+ return intermediate_cache3
1692
+
1693
+
1694
+def modular_triton_fused_moe(
1695
+ use_fp8_w8a8: bool,
1696
+ use_int8_w8a8: bool,
1697
+ use_int8_w8a16: bool,
1698
+ use_int4_w4a16: bool,
1699
+ per_channel_quant: bool,
1700
+ block_shape: Optional[list[int]] = None,
1701
+) -> mk.FusedMoEModularKernel:
1702
+ qtype = get_config_qtype(
1703
+ use_fp8_w8a8=use_fp8_w8a8,
1704
+ use_int8_w8a8=use_int8_w8a8,
1705
+ use_int8_w8a16=use_int8_w8a16,
1706
+ use_int4_w4a16=use_int4_w4a16,
1707
+ )
1708
+ return mk.FusedMoEModularKernel(
1709
+ MoEPrepareAndFinalizeNoEP(
1710
+ quant_dtype=qtype,
1711
+ per_channel_quant=per_channel_quant,
1712
+ block_shape=block_shape,
1713
+ ),
1714
+ TritonExperts(
1715
+ use_fp8_w8a8=use_fp8_w8a8,
1716
+ use_int8_w8a8=use_int8_w8a8,
1717
+ use_int8_w8a16=use_int8_w8a16,
1718
+ use_int4_w4a16=use_int4_w4a16,
1719
+ per_channel_quant=per_channel_quant,
1720
+ block_shape=block_shape,
1721
+ ),
1722
+ )
vllm/model_executor/layers/fused_moe/layer.py
+499
-52
@@ -1,15 +1,19 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import importlib
4
+import threading
5
from abc import abstractmethod
6
+from dataclasses import dataclass
7
from enum import Enum
8
from typing import Callable, Optional
9
+from weakref import WeakValueDictionary
10
11
import torch
12
import torch.nn.functional as F
13
from torch.nn.parameter import UninitializedParameter
14
15
import vllm.envs as envs
12
-from vllm.config import get_current_vllm_config
16
+from vllm.config import ParallelConfig, get_current_vllm_config
17
from vllm.distributed import (get_dp_group, get_ep_group,
18
get_tensor_model_parallel_rank,
19
get_tensor_model_parallel_world_size,
@@ -26,8 +30,17 @@ from vllm.platforms import current_platform
30
from vllm.platforms.interface import CpuArchEnum
31
from vllm.utils import direct_register_custom_op
32
33
+has_pplx = importlib.util.find_spec("pplx_kernels") is not None
34
+
35
if current_platform.is_cuda_alike():
30
- from .fused_moe import fused_experts
36
+ from .fused_batched_moe import (BatchedPrepareAndFinalize,
37
+ BatchedTritonExperts)
38
+ 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
44
else:
45
fused_experts = None # type: ignore
46
if is_rocm_aiter_moe_enabled():
@@ -42,6 +55,179 @@ else:
55
fused_moe_pallas = None # type: ignore
56
logger = init_logger(__name__)
57
58
+# Note: this limit is somewhat arbitrary and might be changed later.
59
+# The size of the activations will be E x MOE_DP_CHUNK_SIZE x hidden_dim.
60
+MOE_DP_CHUNK_SIZE = 256
61
+
62
+
63
+@dataclass
64
+class FusedMoEParallelConfig:
65
+ tp_size: int
66
+ dp_size: int
67
+ ep_size: int
68
+ tp_rank: int
69
+ dp_rank: int
70
+ ep_rank: int
71
+
72
+ use_ep: bool # whether to use EP or not
73
+
74
+ @property
75
+ def use_pplx_kernels(self):
76
+ return self.dp_size > 1 and self.use_ep and has_pplx
77
+
78
+ @staticmethod
79
+ def make(tp_size_: int, dp_size_: int,
80
+ vllm_parallel_config: ParallelConfig) -> "FusedMoEParallelConfig":
81
+ """
82
+ Determine MoE parallel configuration. Based on the input tp_size_,
83
+ dp_size_, ep_size_ and vllm's parallel config, determine what
84
+ level's of parallelism to use in the fused moe layer.
85
+
86
+ Args:
87
+ tp_size_ (int): tp_size passed into the FusedMoE constructor.
88
+ dp_size_ (int): dp_size passed into the FusedMoE constructor.
89
+ ep_size_ (int): ep_size passed into the FusedMoE constructor.
90
+ vllm_parallel_config (ParallelConfig): vllm's parallel config
91
+ object.
92
+
93
+ Examples:
94
+ When there is no parallelism requested, i.e. tp_size_ = dp_size_ = 1,
95
+ we simply return the sizes unaltered and the ranks set to 0.
96
+
97
+ Expert Parallelism is considered only when either dp_size_ or tp_size_
98
+ is non trivial.
99
+
100
+ When TP = 2, DP = 1 and EP = False, the configuration on different
101
+ devices,
102
+ - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} //
103
+ legend : {size, rank}
104
+ - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0}
105
+ - Comment : Tensors are sharded across 2 devices.
106
+
107
+ When TP = 1, DP = 2 and EP = False, the configuration on different
108
+ devices,
109
+ - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0}
110
+ - device 1 : TP = {2, 1} DP = {2, 1} EP = {1, 0}
111
+ - Comment: There are 2 engine instances and the tensors are sharded
112
+ across 2 decvices.
113
+
114
+ When TP = 2, DP = 2 and EP = False, the configuration on different
115
+ devices,
116
+ - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0}
117
+ - device 1: TP = {4, 1} DP = {2, 0} EP = {1, 0}
118
+ - device 2: TP = {4, 2} DP = {2, 1} EP = {1, 0}
119
+ - device 3: TP = {4, 3} DP = {2, 1} EP = {1, 0}
120
+ - Comment: There are 2 engine instances and the tensors are sharded
121
+ across 4 devices.
122
+
123
+ When, TP = 2, DP = 1 and EP = True, the configuration on different
124
+ devices,
125
+ - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0}
126
+ - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1}
127
+ - Comment: The experts are split between the 2 devices.
128
+
129
+ When, TP = 1, DP = 2 and EP = True, the configuration on different
130
+ devices,
131
+ - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0}
132
+ - device 1: TP = {1, 0} DP = {2, 1} EP = {2, 1}
133
+ - Comment: There are 2 engine instances and the experts are split
134
+ between the 2 devices.
135
+
136
+ When TP = 2, DP = 2 and EP = True, the configuration on different
137
+ devices,
138
+ - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0}
139
+ - device 1: TP = {1, 0} DP = {2, 0} EP = {4, 1}
140
+ - device 2: TP = {1, 0} DP = {2, 1} EP = {4, 2}
141
+ - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3}
142
+ - Comment: There are 2 engine instances and the experts are split
143
+ between the 4 devices.
144
+ """
145
+
146
+ def flatten_tp_across_dp(dp_rank: int):
147
+ tp_rank = 0 if tp_size_ == 1 else get_tensor_model_parallel_rank()
148
+ # There are actually dp_size_ * tp_size_ devices. Update tp_size
149
+ # and tp_rank so we shard across all devices.
150
+ tp_size = dp_size_ * tp_size_
151
+ tp_rank = dp_rank * tp_size_ + tp_rank
152
+ return tp_size, tp_rank
153
+
154
+ use_ep = (dp_size_ * tp_size_ > 1
155
+ and vllm_parallel_config.enable_expert_parallel)
156
+
157
+ dp_size = dp_size_
158
+ dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
159
+ tp_size, tp_rank = flatten_tp_across_dp(dp_rank)
160
+
161
+ if not use_ep:
162
+ return FusedMoEParallelConfig(tp_size=tp_size,
163
+ tp_rank=tp_rank,
164
+ dp_size=dp_size,
165
+ dp_rank=dp_rank,
166
+ ep_size=1,
167
+ ep_rank=0,
168
+ use_ep=False)
169
+ # DP + EP / TP + EP / DP + TP + EP
170
+ assert use_ep
171
+ # In EP, each device owns a set of experts fully. There is no tensor
172
+ # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that.
173
+ ep_size = tp_size
174
+ ep_rank = tp_rank
175
+ return FusedMoEParallelConfig(tp_size=1,
176
+ tp_rank=0,
177
+ dp_size=dp_size,
178
+ dp_rank=dp_rank,
179
+ ep_size=ep_size,
180
+ ep_rank=ep_rank,
181
+ use_ep=True)
182
+
183
+
184
+# Adapted from pplx-kernels tests/all_to_all_utils.py
185
+@dataclass
186
+class MoEConfig:
187
+ num_experts: int
188
+ experts_per_token: int
189
+ hidden_dim: int
190
+
191
+ num_local_experts: int
192
+ moe_parallel_config: FusedMoEParallelConfig
193
+
194
+ in_dtype: torch.dtype # The activation type.
195
+
196
+ # TODO: add more quantization params, blocked, per-token, etc.
197
+ block_size: int = 128
198
+
199
+ @property
200
+ def tp_size(self):
201
+ return self.moe_parallel_config.tp_size
202
+
203
+ @property
204
+ def dp_size(self):
205
+ return self.moe_parallel_config.dp_size
206
+
207
+ @property
208
+ def ep_size(self):
209
+ return self.moe_parallel_config.ep_size
210
+
211
+ @property
212
+ def tp_rank(self):
213
+ return self.moe_parallel_config.tp_rank
214
+
215
+ @property
216
+ def dp_rank(self):
217
+ return self.moe_parallel_config.dp_rank
218
+
219
+ @property
220
+ def ep_rank(self):
221
+ return self.moe_parallel_config.ep_rank
222
+
223
+ @property
224
+ def use_ep(self):
225
+ return self.moe_parallel_config.use_ep
226
+
227
+ @property
228
+ def use_pplx_kernels(self):
229
+ return self.moe_parallel_config.use_pplx_kernels
230
+
231
232
class FusedMoeWeightScaleSupported(Enum):
233
TENSOR = "tensor"
@@ -58,6 +244,14 @@ class FusedMoEMethodBase(QuantizeMethodBase):
244
params_dtype: torch.dtype, **extra_weight_attrs):
245
raise NotImplementedError
246
247
+ def set_prepare_finalize(
248
+ self,
249
+ dp_size: int,
250
+ world_size: int,
251
+ prepare_finalize: FusedMoEPrepareAndFinalize,
252
+ ) -> bool:
253
+ return False
254
+
255
@abstractmethod
256
def apply(
257
self,
@@ -80,12 +274,54 @@ class FusedMoEMethodBase(QuantizeMethodBase):
274
raise NotImplementedError
275
276
277
+class AllToAllCache:
278
+
279
+ def __init__(self):
280
+ self._cache: WeakValueDictionary = WeakValueDictionary()
281
+ self._lock = threading.RLock() # Reentrant lock for thread safety
282
+
283
+ def destroy(self):
284
+ with self._lock:
285
+ # TODO: can we do del self._cache?
286
+ for _, a2a in self._cache.items():
287
+ a2a.destroy()
288
+
289
+ def get_or_create(self, **kwargs):
290
+ assert has_pplx
291
+ import pplx_kernels as pplx
292
+
293
+ # Create a hashable key from the kwargs
294
+ key = tuple(sorted((k, v) for k, v in kwargs.items()))
295
+
296
+ with self._lock:
297
+ instance = self._cache.get(key)
298
+ if instance is None:
299
+ # TODO (varun): Add support to switch to intranode
300
+ # when all communications are within the same
301
+ # node.
302
+ logger.debug("Create AllToAll %s", kwargs)
303
+ instance = pplx.AllToAll.internode(**kwargs)
304
+ self._cache[key] = instance
305
+ return instance
306
+
307
+
308
+# Global singleton
309
+_all_to_all_cache = AllToAllCache()
310
+
311
+
312
+# Factory function as a cleaner interface
313
+def get_all_to_all(**kwargs):
314
+ return _all_to_all_cache.get_or_create(**kwargs)
315
+
316
+
317
@CustomOp.register("unquantized_fused_moe")
318
class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
319
"""MoE method without quantization."""
320
87
- def __init__(self):
321
+ def __init__(self, moe: MoEConfig):
322
super().__init__()
323
+ self.fused_experts = fused_experts
324
+ self.moe = moe
325
326
self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
327
if self.rocm_aiter_moe_enabled:
@@ -193,6 +429,47 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
429
activation=activation,
430
apply_router_weight_on_input=apply_router_weight_on_input)
431
432
+ def set_prepare_finalize(
433
+ self,
434
+ dp_size: int,
435
+ world_size: int,
436
+ prepare_finalize: FusedMoEPrepareAndFinalize,
437
+ ) -> bool:
438
+ assert self.fused_experts == fused_experts
439
+
440
+ experts: Optional[FusedMoEPermuteExpertsUnpermute] = None
441
+
442
+ if isinstance(prepare_finalize,
443
+ (BatchedPrepareAndFinalize, PplxPrepareAndFinalize)):
444
+ logger.debug("BatchedTritonExperts %s", self.moe)
445
+ experts = BatchedTritonExperts(
446
+ max_num_tokens=MOE_DP_CHUNK_SIZE,
447
+ world_size=world_size,
448
+ dp_size=dp_size,
449
+ use_fp8_w8a8=False,
450
+ use_int8_w8a8=False,
451
+ use_int8_w8a16=False,
452
+ use_int4_w4a16=False,
453
+ block_shape=None,
454
+ )
455
+ else:
456
+ logger.debug("TritonExperts %s", self.moe)
457
+ experts = TritonExperts(
458
+ use_fp8_w8a8=False,
459
+ use_int8_w8a8=False,
460
+ use_int8_w8a16=False,
461
+ use_int4_w4a16=False,
462
+ block_shape=None,
463
+ per_channel_quant=False,
464
+ )
465
+
466
+ self.fused_experts = FusedMoEModularKernel(
467
+ prepare_finalize,
468
+ experts,
469
+ )
470
+
471
+ return True
472
+
473
def forward_cuda(
474
self,
475
layer: torch.nn.Module,
@@ -221,9 +498,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
498
num_expert_group=num_expert_group,
499
custom_routing_function=custom_routing_function,
500
scoring_func=scoring_func,
224
- e_score_correction_bias=e_score_correction_bias)
501
+ e_score_correction_bias=e_score_correction_bias,
502
+ indices_type=torch.uint32 if self.moe.use_pplx_kernels else None)
503
504
if self.rocm_aiter_moe_enabled:
505
+ assert not apply_router_weight_on_input
506
+ assert expert_map is None
507
return self.rocm_aiter_fused_experts(
508
hidden_states=x,
509
w1=layer.w13_weight,
@@ -232,18 +512,19 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
512
topk_ids=topk_ids,
513
activation=activation,
514
apply_router_weight_on_input=apply_router_weight_on_input)
235
-
236
- return fused_experts(
237
- hidden_states=x,
238
- w1=layer.w13_weight,
239
- w2=layer.w2_weight,
240
- topk_weights=topk_weights,
241
- topk_ids=topk_ids,
242
- inplace=True,
243
- activation=activation,
244
- apply_router_weight_on_input=apply_router_weight_on_input,
245
- global_num_experts=global_num_experts,
246
- expert_map=expert_map)
515
+ else:
516
+ return self.fused_experts(
517
+ hidden_states=x,
518
+ w1=layer.w13_weight,
519
+ w2=layer.w2_weight,
520
+ topk_weights=topk_weights,
521
+ topk_ids=topk_ids,
522
+ inplace=True,
523
+ activation=activation,
524
+ apply_router_weight_on_input=apply_router_weight_on_input,
525
+ global_num_experts=global_num_experts,
526
+ expert_map=expert_map,
527
+ )
528
529
def forward_cpu(
530
self,
@@ -399,6 +680,45 @@ def determine_expert_map(
680
return (local_num_experts, expert_map)
681
682
683
+def _construct_prepare_finalize(
684
+ moe: MoEConfig, quant_config: Optional[QuantizationConfig]
685
+) -> Optional[FusedMoEPrepareAndFinalize]:
686
+ max_num_tokens = MOE_DP_CHUNK_SIZE
687
+ world_size = moe.ep_size
688
+ dp_size = moe.ep_size // moe.dp_size # dp_size actually means TP.
689
+ rank = moe.ep_rank
690
+
691
+ if moe.use_pplx_kernels:
692
+ logger.debug("using PplxPrepareAndFinalize")
693
+
694
+ all_to_all = get_all_to_all(
695
+ max_num_tokens=max_num_tokens,
696
+ num_experts=moe.num_experts,
697
+ experts_per_token=moe.experts_per_token, # topk
698
+ rank=rank,
699
+ world_size=world_size,
700
+ dp_size=dp_size,
701
+ hidden_dim=moe.hidden_dim,
702
+ hidden_dim_bytes=moe.hidden_dim * moe.in_dtype.itemsize,
703
+ # For blocked per token: set to
704
+ # ceil_div(hidden_dim, block_size) * sizeof(float32)
705
+ # For per-token: set to sizeof(float32)
706
+ hidden_dim_scale_bytes=(0 if moe.in_dtype.itemsize != 1 else
707
+ ((moe.hidden_dim + moe.block_size - 1) //
708
+ moe.block_size * torch.float32.itemsize)))
709
+
710
+ return PplxPrepareAndFinalize(
711
+ all_to_all,
712
+ max_num_tokens=max_num_tokens,
713
+ world_size=world_size,
714
+ rank=rank,
715
+ dp_size=dp_size,
716
+ quant_dtype=moe.in_dtype,
717
+ )
718
+
719
+ return None
720
+
721
+
722
class FusedMoE(torch.nn.Module):
723
"""FusedMoE layer for MoE models.
724
@@ -449,21 +769,16 @@ class FusedMoE(torch.nn.Module):
769
params_dtype = torch.get_default_dtype()
770
self.params_dtype = params_dtype
771
452
- # Note: here we guard against accessing the TP and DP groups when
453
- # uninitialized (this happens when testing)
454
- self.tp_size = (tp_size if tp_size is not None else
455
- get_tensor_model_parallel_world_size())
456
- tp_rank = 0 if self.tp_size == 1 else get_tensor_model_parallel_rank()
457
- self.dp_size = (dp_size
458
- if dp_size is not None else get_dp_group().world_size)
459
- self.dp_rank = (0
460
- if self.dp_size == 1 else get_dp_group().rank_in_group)
461
- self.global_num_experts = num_experts
462
-
463
- # Use expert parallelism instead of tensor parallelism?
772
vllm_config = get_current_vllm_config()
465
- use_ep = (vllm_config.parallel_config.enable_expert_parallel
466
- and self.tp_size * self.dp_size > 1)
773
+ self.moe_parallel_config: FusedMoEParallelConfig = (
774
+ FusedMoEParallelConfig.make(
775
+ tp_size_=(tp_size if tp_size is not None else
776
+ get_tensor_model_parallel_world_size()),
777
+ dp_size_=(dp_size if dp_size is not None else
778
+ get_dp_group().world_size),
779
+ vllm_parallel_config=vllm_config.parallel_config))
780
+
781
+ self.global_num_experts = num_experts
782
783
# For smuggling this layer into the fused moe custom op
784
self.use_direct_call = self.dp_size == 1
@@ -474,28 +789,17 @@ class FusedMoE(torch.nn.Module):
789
compilation_config.static_forward_context[prefix] = self
790
self.layer_name = prefix
791
477
- if use_ep:
478
- # Set TP size to 1 to adjust for EP and adjust EP size and rank
479
- # for DP attention.
480
- self.ep_rank = tp_rank + self.tp_size * self.dp_rank
481
- self.tp_rank = 0
482
- self.ep_size = self.tp_size * self.dp_size
483
- self.tp_size = 1
484
-
792
+ # Determine expert maps
793
+ if self.use_ep:
794
self.local_num_experts, self.expert_map = determine_expert_map(
795
ep_size=self.ep_size,
796
ep_rank=self.ep_rank,
797
global_num_experts=self.global_num_experts)
798
else:
490
- # Adjust TP size for DP attention
491
- self.tp_rank = tp_rank + self.tp_size * self.dp_rank
492
- self.ep_rank = 0
493
- self.tp_size = self.tp_size * self.dp_size
494
- self.ep_size = 1
495
- self.local_num_experts = self.global_num_experts
496
- self.expert_map = None
799
+ self.local_num_experts, self.expert_map = (self.global_num_experts,
800
+ None)
801
+
802
self.top_k = top_k
498
- self.global_num_experts = num_experts
803
804
assert intermediate_size % self.tp_size == 0
805
self.hidden_size = hidden_size
@@ -520,14 +824,40 @@ class FusedMoE(torch.nn.Module):
824
from vllm_hpu_extension.ops import DynamicFusedMOE
825
self.hpu_fused_moe = DynamicFusedMOE(self.global_num_experts)
826
827
+ moe = MoEConfig(
828
+ num_experts=self.global_num_experts,
829
+ experts_per_token=top_k,
830
+ hidden_dim=hidden_size,
831
+ num_local_experts=self.local_num_experts,
832
+ moe_parallel_config=self.moe_parallel_config,
833
+ # TODO (bnell): this needs to be fixed for quantized types.
834
+ in_dtype=params_dtype,
835
+ )
836
+
837
# Note: get_quant_method will look at the layer's local_num_experts
838
# for heuristic purposes, so it must be initialized first.
839
+ quant_method: Optional[QuantizeMethodBase] = None
840
+
841
if quant_config is None:
526
- self.quant_method: Optional[QuantizeMethodBase] = (
527
- UnquantizedFusedMoEMethod())
842
+ quant_method = UnquantizedFusedMoEMethod(moe)
843
+ prepare_finalize = _construct_prepare_finalize(moe, quant_config)
844
else:
529
- self.quant_method = quant_config.get_quant_method(self, prefix)
530
- assert self.quant_method is not None
845
+ quant_method = quant_config.get_quant_method(self, prefix)
846
+ # No pplx for quantized types yet.
847
+ prepare_finalize = None
848
+
849
+ assert quant_method is not None
850
+ assert isinstance(quant_method, FusedMoEMethodBase)
851
+ self.quant_method = quant_method
852
+
853
+ if prepare_finalize is not None:
854
+ world_size = moe.ep_size
855
+ dp_size = int(moe.ep_size // moe.dp_size)
856
+ success = self.quant_method.set_prepare_finalize(
857
+ dp_size, world_size, prepare_finalize)
858
+ if not success:
859
+ logger.warning("DP+EP not supported for %s.",
860
+ type(self.quant_method))
861
862
moe_quant_params = {
863
"num_experts": self.local_num_experts,
@@ -546,6 +876,38 @@ class FusedMoE(torch.nn.Module):
876
877
self.quant_method.create_weights(layer=self, **moe_quant_params)
878
879
+ @property
880
+ def tp_size(self):
881
+ return self.moe_parallel_config.tp_size
882
+
883
+ @property
884
+ def dp_size(self):
885
+ return self.moe_parallel_config.dp_size
886
+
887
+ @property
888
+ def ep_size(self):
889
+ return self.moe_parallel_config.ep_size
890
+
891
+ @property
892
+ def tp_rank(self):
893
+ return self.moe_parallel_config.tp_rank
894
+
895
+ @property
896
+ def dp_rank(self):
897
+ return self.moe_parallel_config.dp_rank
898
+
899
+ @property
900
+ def ep_rank(self):
901
+ return self.moe_parallel_config.ep_rank
902
+
903
+ @property
904
+ def use_ep(self):
905
+ return self.moe_parallel_config.use_ep
906
+
907
+ @property
908
+ def use_pplx_kernels(self):
909
+ return self.moe_parallel_config.use_pplx_kernels
910
+
911
def _load_per_tensor_weight_scale(self, shard_id: str,
912
param: torch.nn.Parameter,
913
loaded_weight: torch.Tensor,
@@ -830,7 +1192,8 @@ class FusedMoE(torch.nn.Module):
1192
num_expert_group: Optional[int] = None,
1193
custom_routing_function: Optional[Callable] = None,
1194
scoring_func: str = "softmax",
833
- e_score_correction_bias: Optional[torch.Tensor] = None):
1195
+ e_score_correction_bias: Optional[torch.Tensor] = None,
1196
+ indices_type: Optional[torch.dtype] = None):
1197
from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
1198
1199
# DeekSeekv2 uses grouped_top_k
@@ -846,21 +1209,52 @@ class FusedMoE(torch.nn.Module):
1209
topk_group=topk_group,
1210
scoring_func=scoring_func,
1211
e_score_correction_bias=e_score_correction_bias)
1212
+ if indices_type is not None:
1213
+ topk_ids = topk_ids.to(dtype=indices_type)
1214
elif custom_routing_function is None:
1215
topk_weights, topk_ids, token_expert_indices = fused_topk(
1216
hidden_states=hidden_states,
1217
gating_output=router_logits,
1218
topk=top_k,
854
- renormalize=renormalize)
1219
+ renormalize=renormalize,
1220
+ indices_type=indices_type,
1221
+ )
1222
else:
1223
topk_weights, topk_ids = custom_routing_function(
1224
hidden_states=hidden_states,
1225
gating_output=router_logits,
1226
topk=top_k,
1227
renormalize=renormalize)
1228
+ if indices_type is not None:
1229
+ topk_ids = topk_ids.to(dtype=indices_type)
1230
1231
return topk_weights, topk_ids
1232
1233
+ def must_reduce_shared_expert_outputs(self) -> bool:
1234
+ """
1235
+ The shared_experts are typically computed using the RowParallelLinear
1236
+ layer. The result of this function is typically used as
1237
+ the reduce_results argument to the module.
1238
+ When just tensor-parallel is used, it is not required to reduce
1239
+ the shared_experts results immediately. Instead we reduce at the
1240
+ once at the end of the MoE op. (Refer to DeepSeekV2MoE module)
1241
+ With EP and the pplx kernels - this is no longer viable as all
1242
+ GPU ranks in DP, produce the complete set of hidden_states.
1243
+ Therefore it is required that we reduce the shared_experts output
1244
+ early.
1245
+ """
1246
+ return self.use_pplx_kernels
1247
+
1248
+ def maybe_all_reduce_tensor_model_parallel(
1249
+ self, final_hidden_states: torch.Tensor):
1250
+ """
1251
+ The pplx combine kernel reduces across GPU ranks by default.
1252
+ """
1253
+ if self.use_pplx_kernels:
1254
+ return final_hidden_states
1255
+ else:
1256
+ return tensor_model_parallel_all_reduce(final_hidden_states)
1257
+
1258
def forward(self, hidden_states: torch.Tensor,
1259
router_logits: torch.Tensor):
1260
if self.use_direct_call:
@@ -869,9 +1263,62 @@ class FusedMoE(torch.nn.Module):
1263
return torch.ops.vllm.moe_forward(hidden_states, router_logits,
1264
self.layer_name)
1265
1266
+ def forward_impl_chunked(self, full_hidden_states: torch.Tensor,
1267
+ full_router_logits: torch.Tensor):
1268
+
1269
+ full_final_hidden_states = torch.empty_like(full_hidden_states)
1270
+
1271
+ def process_chunk(chunk_start, chunk_end, skip_result_store=False):
1272
+ hidden_states = full_hidden_states[chunk_start:chunk_end, :]
1273
+ router_logits = full_router_logits[chunk_start:chunk_end, :]
1274
+
1275
+ # Matrix multiply.
1276
+ final_hidden_states = self.quant_method.apply(
1277
+ layer=self,
1278
+ x=hidden_states,
1279
+ router_logits=router_logits,
1280
+ top_k=self.top_k,
1281
+ renormalize=self.renormalize,
1282
+ use_grouped_topk=self.use_grouped_topk,
1283
+ global_num_experts=self.global_num_experts,
1284
+ expert_map=self.expert_map,
1285
+ topk_group=self.topk_group,
1286
+ num_expert_group=self.num_expert_group,
1287
+ custom_routing_function=self.custom_routing_function,
1288
+ scoring_func=self.scoring_func,
1289
+ e_score_correction_bias=self.e_score_correction_bias,
1290
+ activation=self.activation,
1291
+ )
1292
+
1293
+ if not skip_result_store:
1294
+ full_final_hidden_states[chunk_start:chunk_end, :].copy_(
1295
+ final_hidden_states)
1296
+
1297
+ ctx = get_forward_context()
1298
+ max_tokens_across_dp = ctx.dp_metadata.max_tokens_across_dp_cpu
1299
+ moe_dp_chunk_size_per_rank = MOE_DP_CHUNK_SIZE
1300
+
1301
+ num_tokens = full_hidden_states.size(0)
1302
+ for chunk_start_ in range(0, max_tokens_across_dp,
1303
+ moe_dp_chunk_size_per_rank):
1304
+ chunk_start = chunk_start_
1305
+ chunk_end = min(chunk_start + moe_dp_chunk_size_per_rank,
1306
+ max_tokens_across_dp)
1307
+ # clamp start and end
1308
+ chunk_start = min(chunk_start, num_tokens - 1)
1309
+ chunk_end = min(chunk_end, num_tokens)
1310
+
1311
+ process_chunk(chunk_start,
1312
+ chunk_end,
1313
+ skip_result_store=chunk_start_ >= num_tokens)
1314
+
1315
+ return full_final_hidden_states
1316
+
1317
def forward_impl(self, hidden_states: torch.Tensor,
1318
router_logits: torch.Tensor):
1319
assert self.quant_method is not None
1320
+ if self.moe_parallel_config.use_pplx_kernels:
1321
+ return self.forward_impl_chunked(hidden_states, router_logits)
1322
1323
if self.dp_size > 1:
1324
hidden_states, router_logits = get_ep_group().dispatch(
vllm/model_executor/layers/fused_moe/modular_kernel.py
new
+364
@@ -0,0 +1,364 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from abc import ABC, abstractmethod
3
+from typing import Optional
4
+
5
+import torch
6
+
7
+#
8
+# This file defines a set of base classes used to make MoE kernels more modular.
9
+# The goal is to be able to utilize different communication mechanisms with
10
+# any fused MoE kernel without needing to have combinatoric implementations.
11
+#
12
+# The fused moe kernels are broken down into the following components:
13
+#
14
+# [Router] → [Quantize-Dispatch] → [Permute-Experts-Unpermute] → [Combine]
15
+#
16
+# Each component will be independent of the others except for
17
+# [Quantize-Dispatch] and `[Combine] (see below). The components can then be
18
+# mixed and matched with so that DP+EP can be supported easily for multiple
19
+# MoE kernel implementations.
20
+#
21
+# The following main classes are defined:
22
+# * FusedMoEPrepareAndFinalize - an abstract base class for preparation of MoE
23
+# inputs (e.g. quantization, distribution) and finalization of Moe outputs.
24
+# The prepare method must take care of any needed quantization and the
25
+# finalize method must apply weights and do the final reduction of the output.
26
+# * FusedMoEPermuteExpertsUnpermute - an abstract base class for the main fused
27
+# MoE operation. One important feature to note is that this class does not
28
+# apply topk weights or reduce the final output.
29
+# * FusedMoEModularKernel - an interface class that combines a
30
+# FusedMoEPrepareAndFinalize and a FusedMoEPermuteExpertsUnpermute to
31
+# provide the standard fused MoE kernel interface.
32
+#
33
+# [Quantize-Prepare] and [Finalize] functionality are bundled into a single
34
+# class `FusedMoEPrepareAndFinalize` since they could use collective
35
+# communication mechanisms that need to be consistent.
36
+#
37
+
38
+
39
+def _moe_problem_size(
40
+ a1: torch.Tensor,
41
+ w1: torch.Tensor,
42
+ w2: torch.Tensor,
43
+ topk_ids: torch.Tensor,
44
+) -> tuple[int, int, int, int, int]:
45
+ """
46
+ Extract the MoE problem size from the given tensor arguments:
47
+ - a: The hidden states, input to the MoE layer.
48
+ - w1: The first set of expert weights.
49
+ - w2: The second set of expert weights.
50
+ - topk_ids: The topk ids.
51
+
52
+ Note: extracting the problem shape from the weight and activation tensors is
53
+ not obvious. It needs to be done this way specifically due to subtle issues
54
+ with particular kernels, e.g. the int4 kernels divide the trailing dimension
55
+ by two, so it's not "correct" to extract N or K from the trailing dimension
56
+ of w1 or w2. Similarly, some kernels transpose the weights, so this needs
57
+ to be kept in mind.
58
+ """
59
+ assert w1.dim() == 3 and w2.dim() == 3
60
+ E, N, _ = w1.size()
61
+ K = w2.size(1)
62
+
63
+ if a1.dim() == 2:
64
+ # Make sure we are using the correct a1 (pre-permute).
65
+ assert topk_ids.size(0) == a1.size(0), \
66
+ f"{topk_ids.size(0)} != {a1.size(0)}"
67
+ M = a1.size(0)
68
+ else:
69
+ assert a1.dim() == 3
70
+ assert a1.size(0) == E, f"{a1.size(0)} == {E}"
71
+ M = a1.size(1) # This is max_num_tokens
72
+
73
+ assert topk_ids.dim() == 2
74
+ topk = topk_ids.size(1)
75
+
76
+ return E, M, N, K, topk
77
+
78
+
79
+class FusedMoEPrepareAndFinalize(ABC):
80
+ """
81
+ An abstract base class for the [Quantize-Prepare] and [Finalize] steps
82
+ described above.
83
+ """
84
+
85
+ @abstractmethod
86
+ def prepare(
87
+ self,
88
+ a1: torch.Tensor,
89
+ a1_scale: Optional[torch.Tensor],
90
+ a2_scale: Optional[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
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
97
+ """
98
+ Perform any quantization (and/or) dispatching needed
99
+ for this kernel.
100
+ - a1: The (unquantized) input to the MoE layer.
101
+ - a1_scale: Optional scales for a1
102
+ - a2_scale: Optional scales for the second MoE gemm. Required to make
103
+ sure the quantization is consistent for both gemms.
104
+ - topk_ids: The topk ids.
105
+ - topk_weights: The topk weights.
106
+ - num_experts: The total number of experts in the global expert space.
107
+ - expert_map: A tensor mapping expert indices from the global expert
108
+ space to the local expert space of the expert parallel shard.
109
+ - apply_router_weight_on_input: When True, apply the weights to the
110
+ activations, before quantization + dispatching.
111
+
112
+ Returns a tuple of:
113
+ - quantized + dispatched a.
114
+ - quantized + dispatched a1_scales.
115
+ """
116
+ raise NotImplementedError
117
+
118
+ @abstractmethod
119
+ def finalize(
120
+ self,
121
+ output: torch.Tensor,
122
+ fused_expert_output: torch.Tensor,
123
+ topk_weights: torch.Tensor,
124
+ topk_ids: torch.Tensor,
125
+ apply_router_weight_on_input: bool,
126
+ ) -> None:
127
+ """
128
+ Perform any combine plus apply weights and perform a reduction on the
129
+ fused experts output.
130
+ - output: The output tensor, written in place. Must be (M, K) shape.
131
+ - fused_expert_output: The unweighted, unreduced output of the fused
132
+ experts, it will have (M, topk, K) shape.
133
+ - topk_weights: The weights to be applied to the fused_experts_output.
134
+ - topk_ids: The topk_ids.
135
+ - apply_router_weight_on_input: When False, apply the weights to
136
+ fused_expert_output.
137
+ """
138
+ raise NotImplementedError
139
+
140
+
141
+class FusedMoEPermuteExpertsUnpermute(ABC):
142
+ """
143
+ An abstract base class for the [Permute-Experts-Unpermute] step described
144
+ above.
145
+ """
146
+
147
+ @abstractmethod
148
+ def workspace_shapes(
149
+ self,
150
+ a: torch.Tensor,
151
+ M: int,
152
+ N: int,
153
+ K: int,
154
+ topk: int,
155
+ num_experts: int,
156
+ ) -> tuple[int, int, torch.dtype]:
157
+ """
158
+ Compute the number of elements for the temporary outputs of the two
159
+ gemms and activation in the fused expert function. Since the
160
+ gemms are independent, the workspace for the first gemm can be shared
161
+ with the workspace for the last gemm.
162
+
163
+ Returns a tuple of:
164
+ - Number of workspace13 elements: must be large enough to hold the
165
+ result of either expert gemm.
166
+ - Number of workspace2 elements: must be large enough to hold the
167
+ result of the activation function.
168
+ - Workspace type: The dtype to use for the workspace tensors.
169
+ """
170
+ raise NotImplementedError
171
+
172
+ def activation(self, activation: str, output: torch.Tensor,
173
+ input: torch.Tensor) -> None:
174
+ assert output.size(-1) * 2 == input.size(-1)
175
+ if activation == "silu":
176
+ torch.ops._C.silu_and_mul(output, input)
177
+ elif activation == "gelu":
178
+ torch.ops._C.gelu_and_mul(output, input)
179
+ else:
180
+ raise ValueError(f"Unsupported FusedMoe activation: {activation}")
181
+
182
+ @abstractmethod
183
+ def apply(
184
+ self,
185
+ hidden_states: torch.Tensor,
186
+ w1: torch.Tensor,
187
+ w2: torch.Tensor,
188
+ topk_ids: torch.Tensor,
189
+ activation: str,
190
+ global_num_experts: int,
191
+ expert_map: Optional[torch.Tensor],
192
+ w1_scale: Optional[torch.Tensor],
193
+ w2_scale: Optional[torch.Tensor],
194
+ w1_zp: Optional[torch.Tensor],
195
+ w2_zp: Optional[torch.Tensor],
196
+ a1q_scale: Optional[torch.Tensor],
197
+ a2_scale: Optional[torch.Tensor],
198
+ workspace13: torch.Tensor,
199
+ workspace2: torch.Tensor,
200
+ expert_num_tokens: Optional[torch.Tensor],
201
+ ) -> torch.Tensor:
202
+ """
203
+ This function computes the intermediate result of a Mixture of Experts
204
+ (MoE) layer using two sets of weights, w1 and w2.
205
+
206
+ Parameters:
207
+ - hidden_states: (torch.Tensor): The (quantized) input tensor to the MoE
208
+ layer.
209
+ - w1 (torch.Tensor): The first set of expert weights.
210
+ - w2 (torch.Tensor): The second set of expert weights.
211
+ - topk_ids (torch.Tensor): A map of row to expert id.
212
+ - activation (str): The activation function to apply after the first
213
+ MoE layer.
214
+ - global_num_experts (int): The total number of experts in the global
215
+ expert space.
216
+ - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
217
+ from the global expert space to the local expert space of the expert
218
+ parallel shard.
219
+ - w1_scale (Optional[torch.Tensor]): Optional scale to be used for w1.
220
+ - w2_scale (Optional[torch.Tensor]): Optional scale to be used for w2.
221
+ - w1_zp (Optional[torch.Tensor]): Optional zero points to be used for
222
+ w1.
223
+ - w2_zp (Optional[torch.Tensor]): Optional zero points to be used for
224
+ w2.
225
+ - a1q_scale (Optional[torch.Tensor]): Optional quantized scale to be
226
+ used for a1.
227
+ - a2_scale (Optional[torch.Tensor]): Optional scale to be used for a2.
228
+ - workspace13 (torch.Tensor): A scratch tensor used for gemm outputs
229
+ must be large enough to hold output of either MoE gemm.
230
+ - workspace2 (torch.Tensor): A scratch tensor used for the activation
231
+ function.
232
+ - expert_num_tokens: An optional tensor containing the number of tokens
233
+ assigned to each expert when using batched experts format input.
234
+
235
+ Returns:
236
+ - torch.Tensor: The unweighted, unreduced output tensor
237
+ """
238
+ raise NotImplementedError
239
+
240
+
241
+class FusedMoEModularKernel(torch.nn.Module):
242
+ """
243
+ This class combines a FusedMoEPrepareAndFinalize instance and
244
+ a FusedMoEPermuteExpertsUnpermute to provide an interface that
245
+ is compatible with the `fused_experts` function in fused_moe.py.
246
+
247
+ It takes care of managing any required scratch space.
248
+
249
+ Note: Instances of this class should only be used for a single model
250
+ layer due to any layer specific state that may be used by the component
251
+ objects.
252
+ """
253
+
254
+ def __init__(
255
+ self,
256
+ prepare_finalize: FusedMoEPrepareAndFinalize,
257
+ fused_experts: FusedMoEPermuteExpertsUnpermute,
258
+ ):
259
+ super().__init__()
260
+ self.prepare_finalize = prepare_finalize
261
+ self.fused_experts = fused_experts
262
+
263
+ def forward(
264
+ self,
265
+ hidden_states: torch.Tensor,
266
+ w1: torch.Tensor,
267
+ w2: torch.Tensor,
268
+ topk_weights: torch.Tensor,
269
+ topk_ids: torch.Tensor,
270
+ inplace: bool = False,
271
+ activation: str = "silu",
272
+ global_num_experts: int = -1,
273
+ expert_map: Optional[torch.Tensor] = None,
274
+ w1_scale: Optional[torch.Tensor] = None,
275
+ w2_scale: Optional[torch.Tensor] = None,
276
+ w1_zp: Optional[torch.Tensor] = None,
277
+ w2_zp: Optional[torch.Tensor] = None,
278
+ a1_scale: Optional[torch.Tensor] = None,
279
+ a2_scale: Optional[torch.Tensor] = None,
280
+ apply_router_weight_on_input: bool = False,
281
+ ) -> torch.Tensor:
282
+ """
283
+ This function computes a Mixture of Experts (MoE) layer using two sets
284
+ of weights, w1 and w2, and top-k gating mechanism.
285
+
286
+ Parameters:
287
+ - hidden_states: (torch.Tensor): The input tensor to the MoE layer.
288
+ - w1 (torch.Tensor): The first set of expert weights.
289
+ - w2 (torch.Tensor): The second set of expert weights.
290
+ - topk_weights (torch.Tensor): The topk weights applied at the end of
291
+ the layer.
292
+ - topk_ids (torch.Tensor): A map of row to expert id.
293
+ - inplace (bool): If True, perform the operation in-place.
294
+ Defaults to False.
295
+ - activation (str): The activation function to apply after the first
296
+ MoE layer.
297
+ - global_num_experts (int): The total number of experts in the global
298
+ expert space.
299
+ - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
300
+ from the global expert space to the local expert space of the expert
301
+ parallel shard.
302
+ - w1_scale (Optional[torch.Tensor]): Optional scale to be used for w1.
303
+ - w2_scale (Optional[torch.Tensor]): Optional scale to be used for w2.
304
+ - w1_zp (Optional[torch.Tensor]): Optional zero points to be used for
305
+ w1.
306
+ - w2_zp (Optional[torch.Tensor]): Optional zero points to be used for
307
+ w2.
308
+ - a1_scale (Optional[torch.Tensor]): Optional scale to be used for a1.
309
+ - a2_scale (Optional[torch.Tensor]): Optional scale to be used for a2.
310
+ - apply_router_weight_on_input (bool): When true, the topk weights are
311
+ applied directly on the inputs. This is only applicable when topk is
312
+ 1.
313
+
314
+ Returns:
315
+ - torch.Tensor: The output tensor after applying the MoE layer.
316
+ """
317
+ a1 = hidden_states
318
+ E, M, N, K, top_k = _moe_problem_size(a1, w1, w2, topk_ids)
319
+
320
+ if global_num_experts == -1:
321
+ global_num_experts = E
322
+
323
+ output = a1 if inplace else torch.zeros_like(a1)
324
+
325
+ workspace13_shape, workspace2_shape, workspace_dtype = (
326
+ self.fused_experts.workspace_shapes(a1, M, N, K, top_k,
327
+ global_num_experts))
328
+
329
+ # We can reuse the memory between cache1 and cache3 because by the time
330
+ # we need cache3, we're done with cache1
331
+ workspace13 = torch.zeros(workspace13_shape,
332
+ device=a1.device,
333
+ dtype=workspace_dtype)
334
+ workspace2 = torch.zeros(workspace2_shape,
335
+ device=a1.device,
336
+ dtype=workspace_dtype)
337
+
338
+ a1q, a1q_scale, expert_num_tokens = self.prepare_finalize.prepare(
339
+ a1, a1_scale, a2_scale, topk_weights, topk_ids, global_num_experts,
340
+ expert_map, apply_router_weight_on_input)
341
+
342
+ fused_out = self.fused_experts.apply(
343
+ a1q,
344
+ w1,
345
+ w2,
346
+ topk_ids,
347
+ activation=activation,
348
+ global_num_experts=global_num_experts,
349
+ expert_map=expert_map,
350
+ w1_scale=w1_scale,
351
+ w2_scale=w2_scale,
352
+ w1_zp=w1_zp,
353
+ w2_zp=w2_zp,
354
+ a1q_scale=a1q_scale,
355
+ a2_scale=a2_scale,
356
+ workspace13=workspace13,
357
+ workspace2=workspace2,
358
+ expert_num_tokens=expert_num_tokens,
359
+ )
360
+
361
+ self.prepare_finalize.finalize(output, fused_out, topk_weights,
362
+ topk_ids, apply_router_weight_on_input)
363
+
364
+ return output
vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py
+79
-11
@@ -3,6 +3,74 @@ from typing import Optional
3
4
import torch
5
6
+from vllm import _custom_ops as ops
7
+from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
8
+ moe_align_block_size)
9
+from vllm.model_executor.layers.fused_moe.utils import _fp8_perm
10
+
11
+
12
+def _moe_permute(
13
+ curr_hidden_states: torch.Tensor,
14
+ a1q_scale: Optional[torch.Tensor],
15
+ curr_topk_ids: torch.Tensor,
16
+ global_num_experts: int,
17
+ expert_map: Optional[torch.Tensor],
18
+ block_m: int,
19
+) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, torch.Tensor,
20
+ Optional[torch.Tensor]]:
21
+ """
22
+ Determine the sorted_token_ids, expert_ids for the given problem size.
23
+ Permute the hidden states and scales according to `sorted_token_ids`.
24
+ """
25
+ top_k_num = curr_topk_ids.size(1)
26
+
27
+ tokens_in_chunk = curr_hidden_states.sizze(0)
28
+
29
+ sorted_token_ids, expert_ids, num_tokens_post_padded = (
30
+ moe_align_block_size(curr_topk_ids,
31
+ block_m,
32
+ global_num_experts,
33
+ expert_map,
34
+ pad_sorted_ids=True))
35
+
36
+ inv_perm: Optional[torch.Tensor] = None
37
+
38
+ num_tokens = top_k_num * tokens_in_chunk
39
+ sorted_token_ids = sorted_token_ids.clamp(max=num_tokens - 1)
40
+ expert_ids = torch.repeat_interleave(expert_ids, block_m, dim=0)
41
+ inv_perm = torch.argsort(sorted_token_ids)[:num_tokens]
42
+
43
+ # Permute according to sorted token ids.
44
+ curr_hidden_states = _fp8_perm(curr_hidden_states,
45
+ sorted_token_ids // top_k_num)
46
+
47
+ if a1q_scale is not None:
48
+ a1q_scale = a1q_scale[sorted_token_ids // top_k_num]
49
+
50
+ return (curr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
51
+ inv_perm)
52
+
53
+
54
+def _moe_unpermute_and_reduce(
55
+ out: torch.Tensor,
56
+ curr_hidden: torch.Tensor,
57
+ inv_perm: Optional[torch.Tensor],
58
+ topk_weight: torch.Tensor,
59
+ apply_router_weight_on_input: bool,
60
+) -> None:
61
+ """
62
+ Unpermute the final result and apply topk_weights, then perform the final
63
+ reduction on the hidden states.
64
+ """
65
+ M, topk = topk_weight.size()
66
+ K = curr_hidden.size(-1)
67
+ if inv_perm is not None:
68
+ curr_hidden = curr_hidden[inv_perm, ...]
69
+ curr_hidden = curr_hidden.view(-1, topk, K)
70
+ if not apply_router_weight_on_input:
71
+ curr_hidden.mul_(topk_weight.view(M, -1, 1))
72
+ ops.moe_sum(curr_hidden, out)
73
+
74
75
def moe_permute(
76
hidden_states: torch.Tensor,
@@ -17,21 +85,21 @@ def moe_permute(
85
fill_invalid_expert: int = -1
86
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
87
"""
20
- This function expands and permutes activation to gather uncontinuous tokens
88
+ This function expands and permutes activation to gather uncontinuous tokens
89
for each expert.
90
Parameters:
23
- - hidden_states (torch.Tensor): The input tensor to the MoE layer.
91
+ - hidden_states (torch.Tensor): The input tensor to the MoE layer.
92
- topk_weights (torch.Tensor): topk expert route weight for each token.
93
- topk_ids (torch.Tensor): topk expert route id for each token.
94
- token_expert_indices (torch.Tensor): indice for expanded hidden.
95
- topk (int): The number of top-k experts to select.
96
- n_expert (int): The number of expert.
97
- n_local_expert (int): The number of expert in current EP rank.
30
- - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
31
- from the global expert space to the local expert space of the expert
98
+ - expert_map (Optional[torch.Tensor]): A tensor mapping expert indices
99
+ from the global expert space to the local expert space of the expert
100
parallel shard.
101
- align_block_size (Optional[int]): align group gemm block size for deepgemm
34
- - fill_invalid_expert(int): fill expert id in m_indices for invalid expert
102
+ - fill_invalid_expert(int): fill expert id in m_indices for invalid expert
103
to workaround DeepGemm unsupported -1 in m_indices
104
Returns:
105
- permuted_hidden_states (torch.Tensor): permuted activation.
@@ -39,10 +107,10 @@ def moe_permute(
107
of each expert for standard grouped gemm. if enable 'align_block_size'
108
expert_first_token_offset will align up to 'align_block_size'.
109
- src_row_id2dst_row_id_map (torch.Tensor): idx map for moe_unpermute.
42
- - m_indices: m_indices for grouped gemm in deepgemm,`m_indices[i]` records
110
+ - m_indices: m_indices for grouped gemm in deepgemm,`m_indices[i]` records
111
the group which the j-th row of the LHS belong to.`
112
"""
45
- n_token, n_hidden = hidden_states.shape
113
+ n_token, n_hidden = hidden_states.size()
114
assert (n_hidden * hidden_states.element_size()
115
) % 16 == 0, "permue kernel need hidden dim align to 16B"
116
permuted_row_size = n_token * topk
@@ -87,7 +155,7 @@ def moe_unpermute(
155
n_local_expert: int,
156
) -> torch.Tensor:
157
"""
90
- This function expands and permutes activation to gathering uncontinuous
158
+ This function expands and permutes activation to gathering uncontinuous
159
tokens for each expert.
160
Parameters:
161
- permuted_hidden_states (torch.Tensor): permuted activation.
@@ -99,10 +167,10 @@ def moe_unpermute(
167
- n_expert (int): The number of expert.
168
- n_local_expert (int): The number of expert in current EP rank.
169
Returns:
102
- - hidden_states (torch.Tensor): The reduced and unpermuted activation
103
- tensor.
170
+ - hidden_states (torch.Tensor): The reduced and unpermuted activation
171
+ tensor.
172
"""
105
- n_token, n_hidden = topk_weights.shape[0], permuted_hidden_states.shape[-1]
173
+ n_token, n_hidden = topk_weights.size(0), permuted_hidden_states.size(-1)
174
assert (n_hidden * permuted_hidden_states.element_size()
175
) % 16 == 0, "unpermue kernel need hidden dim align to 16B"
176
hidden_states = torch.empty((n_token, n_hidden),
vllm/model_executor/layers/fused_moe/pplx_prepare_finalize.py
new
+147
@@ -0,0 +1,147 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import pplx_kernels as pplx
5
+import torch
6
+
7
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8
+from vllm.model_executor.layers.fused_moe.utils import (
9
+ moe_kernel_quantize_input)
10
+
11
+
12
+# Note use: layer.get_all_to_all() to get an AllToAll instance
13
+# The max_num_tokens, world_size and dp_size must be the same
14
+# as the ones used to create the AllToAll.
15
+class PplxPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
16
+
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
+ super().__init__()
26
+ assert max_num_tokens > 0
27
+ self.a2a = a2a
28
+ self.block_shape = block_shape
29
+ self.max_num_tokens = max_num_tokens
30
+ self.world_size = world_size
31
+ self.rank = rank
32
+ self.dp_size = dp_size
33
+ self.quant_dtype = quant_dtype
34
+
35
+ def prepare(
36
+ self,
37
+ a1: torch.Tensor,
38
+ a1_scale: Optional[torch.Tensor],
39
+ a2_scale: Optional[torch.Tensor],
40
+ rank_topk_weights: torch.Tensor,
41
+ rank_topk_ids: torch.Tensor,
42
+ num_experts: int,
43
+ expert_map: Optional[torch.Tensor],
44
+ apply_router_weight_on_input: bool,
45
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
46
+ num_tokens = a1.size(0) # M
47
+ hidden_dim = a1.size(-1) # K
48
+
49
+ assert rank_topk_ids.size(0) == num_tokens
50
+ # assert expert_map is None, "NYI"
51
+
52
+ # Is this always going to be a1.device?
53
+ device = a1.device
54
+
55
+ if apply_router_weight_on_input:
56
+ topk = rank_topk_ids.size(1)
57
+ # TODO: this only works for topK=1, will need to update for topK>1
58
+ assert topk == 1, (
59
+ "apply_router_weight_on_input is only implemented for topk=1")
60
+ a1 = a1 * rank_topk_weights.to(a1.dtype)
61
+
62
+ per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
63
+ a2_scale.numel() != 1 if a2_scale is not None else False)
64
+
65
+ a1q, a1q_scale = moe_kernel_quantize_input(a1, a1_scale,
66
+ self.quant_dtype,
67
+ per_act_token,
68
+ self.block_shape)
69
+
70
+ # rem_experts need to be 0 for pplx to work properly.
71
+ rem_experts = num_experts % self.world_size
72
+ assert rem_experts == 0
73
+ num_local_experts = ((num_experts // self.world_size) +
74
+ (1 if self.rank < rem_experts else 0))
75
+
76
+ expert_num_tokens = torch.empty(
77
+ num_local_experts,
78
+ dtype=torch.int32,
79
+ device=device,
80
+ )
81
+
82
+ num_dp = self.world_size // self.dp_size
83
+ expert_x = torch.empty(
84
+ (num_local_experts, self.max_num_tokens * num_dp, hidden_dim),
85
+ dtype=a1q.dtype,
86
+ device=device,
87
+ )
88
+
89
+ expert_x_scale: Optional[torch.Tensor] = None
90
+ if a1q.dtype.itemsize == 1:
91
+ float32_size = torch.float32.itemsize
92
+ block_size = (self.block_shape[0] if self.block_shape is not None
93
+ else 1) * float32_size
94
+ expert_x_scale = torch.empty(
95
+ (
96
+ num_experts,
97
+ expert_x.size(1),
98
+ (expert_x.size(2) + block_size - 1) // block_size,
99
+ ),
100
+ dtype=torch.float32,
101
+ device=device,
102
+ )
103
+
104
+ # This argument is optional, defaults to indices.size(0)
105
+ # There's not much point setting this unless it is != indices.size(0)
106
+ bound_m: Optional[torch.Tensor] = None
107
+
108
+ self.a2a.dispatch(
109
+ out_expert_num_tokens=expert_num_tokens,
110
+ out_expert_x=expert_x,
111
+ out_expert_x_scale=expert_x_scale,
112
+ dp_x=a1q,
113
+ dp_x_scale=a1q_scale,
114
+ indices=rank_topk_ids,
115
+ bound_m=bound_m,
116
+ )
117
+
118
+ return expert_x, expert_x_scale, expert_num_tokens
119
+
120
+ def finalize(
121
+ self,
122
+ output: torch.Tensor,
123
+ fused_expert_output: torch.Tensor,
124
+ topk_weights: torch.Tensor,
125
+ topk_ids: torch.Tensor,
126
+ apply_router_weight_on_input: bool,
127
+ ) -> None:
128
+ num_tokens = output.size(0) # M
129
+ # This argument is optional
130
+ # There's not much point setting this unless it is != topk_ids.size(0)
131
+ bound_m: Optional[torch.Tensor] = None
132
+
133
+ assert topk_ids.size(0) == num_tokens, (
134
+ f"{topk_ids.size(0)} == {num_tokens}")
135
+ assert output.size(0) <= self.max_num_tokens, (
136
+ f"{output.size(0)} <= {self.max_num_tokens}")
137
+ assert output.size(1) == fused_expert_output.size(-1)
138
+
139
+ # Set weights to 1 if we did them in dispatch. This is hacky.
140
+ if apply_router_weight_on_input:
141
+ topk_weights = torch.ones_like(topk_weights)
142
+
143
+ self.a2a.combine(out_tokens=output,
144
+ indices=topk_ids,
145
+ weights=topk_weights,
146
+ expert_y=fused_expert_output,
147
+ bound_m=bound_m)
vllm/model_executor/layers/fused_moe/prepare_finalize.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (
8
+ _moe_unpermute_and_reduce)
9
+from vllm.model_executor.layers.fused_moe.utils import (
10
+ moe_kernel_quantize_input)
11
+
12
+
13
+class MoEPrepareAndFinalizeNoEP(mk.FusedMoEPrepareAndFinalize):
14
+
15
+ def __init__(
16
+ self,
17
+ quant_dtype: Optional[torch.dtype] = None,
18
+ per_channel_quant: bool = False,
19
+ block_shape: Optional[list[int]] = None,
20
+ ):
21
+ super().__init__()
22
+ self.per_channel_quant = per_channel_quant
23
+ self.block_shape = block_shape
24
+ self.quant_dtype = quant_dtype
25
+
26
+ def prepare(
27
+ self,
28
+ a1: torch.Tensor,
29
+ a1_scale: Optional[torch.Tensor],
30
+ a2_scale: Optional[torch.Tensor],
31
+ topk_weights: torch.Tensor,
32
+ topk_ids: torch.Tensor,
33
+ num_experts: int,
34
+ expert_map: Optional[torch.Tensor],
35
+ apply_router_weight_on_input: bool = False,
36
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
37
+ if apply_router_weight_on_input:
38
+ topk = topk_ids.size(1)
39
+ # TODO: this only works for topK=1, will need to update for topK>1
40
+ assert topk == 1, \
41
+ "apply_router_weight_on_input is only implemented for topk=1"
42
+ a1.mul_(topk_weights.to(a1.dtype))
43
+
44
+ a1q, a1q_scale = moe_kernel_quantize_input(a1, a1_scale,
45
+ self.quant_dtype,
46
+ self.per_channel_quant,
47
+ self.block_shape)
48
+
49
+ return a1q, a1q_scale, None
50
+
51
+ def finalize(
52
+ self,
53
+ output: torch.Tensor,
54
+ fused_expert_output: torch.Tensor,
55
+ topk_weights: torch.Tensor,
56
+ topk_ids: torch.Tensor,
57
+ apply_router_weight_on_input: bool,
58
+ ) -> None:
59
+ _moe_unpermute_and_reduce(output, fused_expert_output, None,
60
+ topk_weights, apply_router_weight_on_input)
vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from typing import Optional
3
+
4
+import torch
5
+
6
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7
+from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
8
+ DeepGemmExperts, _valid_deep_gemm, _valid_deep_gemm_shape)
9
+from vllm.model_executor.layers.fused_moe.fused_moe import TritonExperts
10
+
11
+
12
+class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
13
+
14
+ def __init__(self,
15
+ use_fp8_w8a8: bool = False,
16
+ use_int8_w8a8: bool = False,
17
+ use_int8_w8a16: bool = False,
18
+ use_int4_w4a16: bool = False,
19
+ per_channel_quant: bool = False,
20
+ block_shape: Optional[list[int]] = None,
21
+ block_m: Optional[int] = None,
22
+ allow_deep_gemm: bool = False):
23
+ super().__init__()
24
+ self.triton_expert = TritonExperts(use_fp8_w8a8=use_fp8_w8a8,
25
+ use_int8_w8a8=use_int8_w8a8,
26
+ use_int4_w4a16=use_int4_w4a16,
27
+ use_int8_w8a16=use_int8_w8a16,
28
+ per_channel_quant=per_channel_quant,
29
+ block_shape=block_shape,
30
+ block_m=block_m)
31
+ self.deep_gemm_expert = DeepGemmExperts()
32
+ self.allow_deep_gemm = allow_deep_gemm
33
+ self.use_fp8_w8a8 = use_fp8_w8a8
34
+
35
+ def workspace_shapes(
36
+ self,
37
+ a: torch.Tensor,
38
+ M: int,
39
+ N: int,
40
+ K: int,
41
+ topk: int,
42
+ num_experts: int,
43
+ ) -> tuple[int, int, torch.dtype]:
44
+ # Note: the deep gemm workspaces are strictly larger than the triton
45
+ # workspaces so we can be pessimistic here and allocate for DeepGemm
46
+ # even if we fall back to triton later, e.g. if expert maps are set.
47
+ if self.allow_deep_gemm and _valid_deep_gemm_shape(M, N, K):
48
+ return self.deep_gemm_expert.workspace_shapes(
49
+ a, M, N, K, topk, num_experts)
50
+ else:
51
+ return self.triton_expert.workspace_shapes(a, M, N, K, topk,
52
+ num_experts)
53
+
54
+ def apply(
55
+ self,
56
+ hidden_states: torch.Tensor,
57
+ w1: torch.Tensor,
58
+ w2: torch.Tensor,
59
+ topk_ids: torch.Tensor,
60
+ activation: str,
61
+ global_num_experts: int,
62
+ expert_map: Optional[torch.Tensor],
63
+ w1_scale: Optional[torch.Tensor],
64
+ w2_scale: Optional[torch.Tensor],
65
+ w1_zp: Optional[torch.Tensor],
66
+ w2_zp: Optional[torch.Tensor],
67
+ a1q_scale: Optional[torch.Tensor],
68
+ a2_scale: Optional[torch.Tensor],
69
+ workspace13: torch.Tensor,
70
+ workspace2: torch.Tensor,
71
+ expert_num_tokens: Optional[torch.Tensor],
72
+ ) -> torch.Tensor:
73
+ N = w1.size(1)
74
+ if (self.allow_deep_gemm and self.use_fp8_w8a8 and N > 512
75
+ and _valid_deep_gemm(hidden_states, w1, w2, expert_map)):
76
+ return self.deep_gemm_expert.apply(
77
+ hidden_states,
78
+ w1,
79
+ w2,
80
+ topk_ids,
81
+ activation,
82
+ global_num_experts,
83
+ expert_map,
84
+ w1_scale,
85
+ w2_scale,
86
+ w1_zp,
87
+ w2_zp,
88
+ a1q_scale,
89
+ a2_scale,
90
+ workspace13,
91
+ workspace2,
92
+ expert_num_tokens,
93
+ )
94
+ else:
95
+ return self.triton_expert.apply(
96
+ hidden_states,
97
+ w1,
98
+ w2,
99
+ topk_ids,
100
+ activation,
101
+ global_num_experts,
102
+ expert_map,
103
+ w1_scale,
104
+ w2_scale,
105
+ w1_zp,
106
+ w2_zp,
107
+ a1q_scale,
108
+ a2_scale,
109
+ workspace13,
110
+ workspace2,
111
+ expert_num_tokens,
112
+ )
vllm/model_executor/layers/fused_moe/utils.py
+54
-5
@@ -7,6 +7,8 @@ import torch
7
from vllm import _custom_ops as ops
8
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
9
per_token_group_quant_fp8)
10
+from vllm.model_executor.layers.quantization.utils.int8_utils import (
11
+ per_token_group_quant_int8, per_token_quant_int8)
12
from vllm.utils import cdiv
13
14
@@ -15,34 +17,81 @@ def _resize_cache(x: torch.Tensor, v: tuple[int, ...]) -> torch.Tensor:
17
Shrink the given tensor and apply the given view to it. This is
18
used to resize the intermediate fused_moe caches.
19
"""
18
- assert prod(v) <= x.numel()
20
+ assert prod(
21
+ v) <= x.numel(), f"{prod(v)} <= {x.numel()}" # CUDAGRAPH unfriendly?
22
return x.flatten()[:prod(v)].view(*v)
23
24
25
def _fp8_quantize(
26
A: torch.Tensor,
27
A_scale: Optional[torch.Tensor],
25
- block_shape: Optional[list[int]],
28
+ per_act_token: bool,
29
+ block_shape: Optional[list[int]] = None,
30
) -> tuple[torch.Tensor, torch.Tensor]:
31
"""
32
Perform fp8 quantization on the inputs. If a block_shape
33
is provided, the output will be blocked.
34
"""
35
if block_shape is None:
32
- A, A_scale = ops.scaled_fp8_quant(A, A_scale)
36
+ A, A_scale = ops.scaled_fp8_quant(
37
+ A, A_scale, use_per_token_if_dynamic=per_act_token)
38
else:
39
assert len(block_shape) == 2
40
_, block_k = block_shape[0], block_shape[1]
41
A, A_scale = per_token_group_quant_fp8(A, block_k)
37
- assert cdiv(A.shape[-1], block_k) == A_scale.shape[-1]
42
+ assert cdiv(A.size(-1), block_k) == A_scale.size(-1)
43
+
44
return A, A_scale
45
46
47
+def _int8_quantize(
48
+ A: torch.Tensor,
49
+ A_scale: Optional[torch.Tensor],
50
+ per_act_token: bool,
51
+ block_shape: Optional[list[int]] = None,
52
+) -> tuple[torch.Tensor, torch.Tensor]:
53
+ """
54
+ Perform int8 quantization on the inputs. If a block_shape
55
+ is provided, the output will be blocked.
56
+ """
57
+
58
+ # If weights are per-channel (per_channel_quant=True), then
59
+ # activations apply per-token quantization. Otherwise, assume
60
+ # activation tensor-wise fp8/int8 quantization, dynamic or static
61
+ if block_shape is None:
62
+ assert per_act_token, \
63
+ "int8 quantization only supports block or channel-wise"
64
+ A, A_scale = per_token_quant_int8(A)
65
+ else:
66
+ assert len(block_shape) == 2
67
+ _, block_k = block_shape[0], block_shape[1]
68
+ A, A_scale = per_token_group_quant_int8(A, block_k)
69
+ assert cdiv(A.size(-1), block_k) == A_scale.size(-1)
70
+
71
+ return A, A_scale
72
+
73
+
74
+def moe_kernel_quantize_input(
75
+ A: torch.Tensor,
76
+ A_scale: Optional[torch.Tensor],
77
+ qtype: Optional[torch.dtype],
78
+ per_channel_quant: bool,
79
+ block_shape: Optional[list[int]] = None,
80
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
81
+ if qtype == torch.float8_e4m3fn:
82
+ return _fp8_quantize(A, A_scale, per_channel_quant, block_shape)
83
+ elif qtype == torch.int8:
84
+ return _int8_quantize(A, A_scale, per_channel_quant, block_shape)
85
+ else:
86
+ assert A_scale is None
87
+ return A, A_scale
88
+
89
+
90
def _fp8_perm(m: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
91
"""
92
A permutation routine that works on fp8 types.
93
"""
45
- if torch.is_floating_point(m) and torch.finfo(m.dtype).bits == 8:
94
+ if torch.is_floating_point(m) and m.dtype.itemsize == 1:
95
return m.view(dtype=torch.uint8)[idx, ...].view(dtype=m.dtype)
96
else:
97
return m[idx, ...]
vllm/model_executor/layers/quantization/fp8.py
+56
-28
@@ -1,5 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
+import functools
4
import importlib.util
5
from typing import Any, Callable, Optional
6
@@ -9,6 +10,7 @@ from torch.nn import Module
10
from torch.nn.parameter import Parameter
11
12
import vllm.envs as envs
13
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14
from vllm import _custom_ops as ops
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.logger import init_logger
@@ -434,6 +436,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
436
"""
437
438
def __init__(self, quant_config: Fp8Config):
439
+ from vllm.model_executor.layers.fused_moe import fused_experts
440
self.quant_config = quant_config
441
self.block_quant = self.quant_config.weight_block_size is not None
442
@@ -458,6 +461,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):
461
logger.warning_once(
462
"DeepGemm not supported on the current platform.")
463
464
+ self.fused_experts = functools.partial(
465
+ fused_experts,
466
+ block_shape=self.quant_config.weight_block_size,
467
+ allow_deep_gemm=self.allow_deep_gemm)
468
+
469
def create_weights(self, layer: Module, num_experts: int, hidden_size: int,
470
intermediate_size_per_partition: int,
471
params_dtype: torch.dtype, **extra_weight_attrs):
@@ -783,6 +791,31 @@ class Fp8MoEMethod(FusedMoEMethodBase):
791
del layer.w13_input_scale
792
del layer.w2_input_scale
793
794
+ def set_prepare_finalize(
795
+ self,
796
+ dp_size: int,
797
+ world_size: int,
798
+ prepare_finalize: mk.FusedMoEPrepareAndFinalize,
799
+ ) -> bool:
800
+ from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
801
+ TritonOrDeepGemmExperts)
802
+
803
+ if self.use_marlin or self.rocm_aiter_moe_enabled:
804
+ return False
805
+
806
+ experts = TritonOrDeepGemmExperts(
807
+ use_fp8_w8a8=True,
808
+ block_shape=self.quant_config.weight_block_size,
809
+ allow_deep_gemm=self.allow_deep_gemm,
810
+ )
811
+
812
+ self.fused_experts = mk.FusedMoEModularKernel(
813
+ prepare_finalize,
814
+ experts,
815
+ )
816
+
817
+ return True
818
+
819
def apply(
820
self,
821
layer: torch.nn.Module,
@@ -801,10 +834,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
834
apply_router_weight_on_input: bool = False,
835
activation: str = "silu",
836
) -> torch.Tensor:
804
- from vllm.model_executor.layers.fused_moe import fused_experts
805
- from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
806
- rocm_aiter_fused_experts)
807
-
837
topk_weights, topk_ids = FusedMoE.select_experts(
838
hidden_states=x,
839
router_logits=router_logits,
@@ -819,6 +848,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
848
)
849
850
if self.rocm_aiter_moe_enabled:
851
+ from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import ( # noqa: E501
852
+ rocm_aiter_fused_experts)
853
return rocm_aiter_fused_experts(
854
x,
855
layer.w13_weight,
@@ -835,8 +866,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
866
a1_scale=layer.w13_input_scale,
867
a2_scale=layer.w2_input_scale,
868
block_shape=self.quant_config.weight_block_size)
838
-
839
- if self.use_marlin:
869
+ elif self.use_marlin:
870
assert activation == "silu", (
871
f"{activation} not supported for Marlin MoE.")
872
assert not apply_router_weight_on_input, (
@@ -853,28 +883,26 @@ class Fp8MoEMethod(FusedMoEMethodBase):
883
quant_type_id=scalar_types.float8_e4m3fn.id,
884
global_num_experts=global_num_experts,
885
expert_map=expert_map)
856
-
857
- return fused_experts(
858
- x,
859
- layer.w13_weight,
860
- layer.w2_weight,
861
- topk_weights=topk_weights,
862
- topk_ids=topk_ids,
863
- inplace=True,
864
- activation=activation,
865
- use_fp8_w8a8=True,
866
- global_num_experts=global_num_experts,
867
- apply_router_weight_on_input=apply_router_weight_on_input,
868
- expert_map=expert_map,
869
- w1_scale=(layer.w13_weight_scale_inv
870
- if self.block_quant else layer.w13_weight_scale),
871
- w2_scale=(layer.w2_weight_scale_inv
872
- if self.block_quant else layer.w2_weight_scale),
873
- a1_scale=layer.w13_input_scale,
874
- a2_scale=layer.w2_input_scale,
875
- block_shape=self.quant_config.weight_block_size,
876
- allow_deep_gemm=self.allow_deep_gemm,
877
- )
886
+ else:
887
+ return self.fused_experts(
888
+ hidden_states=x,
889
+ w1=layer.w13_weight,
890
+ w2=layer.w2_weight,
891
+ topk_weights=topk_weights,
892
+ topk_ids=topk_ids,
893
+ inplace=True,
894
+ activation=activation,
895
+ use_fp8_w8a8=True,
896
+ global_num_experts=global_num_experts,
897
+ apply_router_weight_on_input=apply_router_weight_on_input,
898
+ expert_map=expert_map,
899
+ w1_scale=(layer.w13_weight_scale_inv
900
+ if self.block_quant else layer.w13_weight_scale),
901
+ w2_scale=(layer.w2_weight_scale_inv
902
+ if self.block_quant else layer.w2_weight_scale),
903
+ a1_scale=layer.w13_input_scale,
904
+ a2_scale=layer.w2_input_scale,
905
+ )
906
907
908
class Fp8KVCacheMethod(BaseKVCacheMethod):
vllm/model_executor/models/dbrx.py
-1
@@ -79,7 +79,6 @@ class DbrxExperts(FusedMoE):
79
prefix=prefix,
80
)
81
self.config = config
82
- self.tp_size = get_tensor_model_parallel_world_size()
82
self.d_model = config.d_model
83
self.intermediate_size = (self.config.ffn_config.ffn_hidden_size //
84
self.tp_size)
vllm/model_executor/models/deepseek_v2.py
+8
-6
@@ -31,9 +31,7 @@ from transformers import PretrainedConfig
31
from vllm.attention import Attention
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, ModelConfig, VllmConfig
34
-from vllm.distributed import (get_pp_group,
35
- get_tensor_model_parallel_world_size,
36
- tensor_model_parallel_all_reduce)
34
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
35
from vllm.model_executor.layers.activation import SiluAndMul
36
from vllm.model_executor.layers.fused_moe import FusedMoE
37
from vllm.model_executor.layers.layernorm import RMSNorm
@@ -143,7 +141,8 @@ class DeepseekV2MoE(nn.Module):
141
intermediate_size=intermediate_size,
142
hidden_act=config.hidden_act,
143
quant_config=quant_config,
146
- reduce_results=False,
144
+ reduce_results=self.experts.must_reduce_shared_expert_outputs(
145
+ ),
146
prefix=f"{prefix}.shared_experts",
147
)
148
@@ -154,6 +153,7 @@ class DeepseekV2MoE(nn.Module):
153
shared_output = self.shared_experts(hidden_states)
154
# router_logits: (num_tokens, n_experts)
155
router_logits, _ = self.gate(hidden_states)
156
+
157
if hidden_states.dtype != torch.float16:
158
final_hidden_states = self.experts(
159
hidden_states=hidden_states,
@@ -171,9 +171,11 @@ class DeepseekV2MoE(nn.Module):
171
# See DeepseekV2DecoderLayer for more details.
172
final_hidden_states = final_hidden_states + shared_output \
173
* (1. / self.routed_scaling_factor)
174
+
175
if self.tp_size > 1:
175
- final_hidden_states = tensor_model_parallel_all_reduce(
176
- final_hidden_states)
176
+ final_hidden_states = (
177
+ self.experts.maybe_all_reduce_tensor_model_parallel(
178
+ final_hidden_states))
179
180
return final_hidden_states.view(num_tokens, hidden_dim)
181
vllm/model_executor/models/llama4.py
+4
-4
@@ -25,8 +25,7 @@ from transformers import Llama4TextConfig
25
from vllm.attention import Attention
26
from vllm.compilation.decorators import support_torch_compile
27
from vllm.config import CacheConfig, VllmConfig
28
-from vllm.distributed import (get_tensor_model_parallel_world_size,
29
- tensor_model_parallel_all_reduce)
28
+from vllm.distributed import get_tensor_model_parallel_world_size
29
from vllm.model_executor.layers.fused_moe import FusedMoE
30
from vllm.model_executor.layers.layernorm import RMSNorm
31
from vllm.model_executor.layers.linear import (QKVParallelLinear,
@@ -89,7 +88,7 @@ class Llama4MoE(nn.Module):
88
quant_config=quant_config,
89
bias=False,
90
prefix=f"{prefix}.shared_expert",
92
- reduce_results=False, # We need to do scatter before reduce
91
+ reduce_results=self.experts.must_reduce_shared_expert_outputs(),
92
)
93
94
def forward(self, hidden_states):
@@ -102,7 +101,8 @@ class Llama4MoE(nn.Module):
101
experts_out = routed_out + shared_out
102
103
if self.tp_size > 1:
105
- experts_out = tensor_model_parallel_all_reduce(experts_out)
104
+ experts_out = self.experts.maybe_all_reduce_tensor_model_parallel(
105
+ experts_out)
106
107
return experts_out
108
vllm/model_executor/models/qwen2_moe.py
+4
-5
@@ -33,9 +33,7 @@ from transformers import PretrainedConfig
33
from vllm.attention import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
-from vllm.distributed import (get_pp_group,
37
- get_tensor_model_parallel_world_size,
38
- tensor_model_parallel_all_reduce)
36
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.logger import init_logger
38
from vllm.model_executor.layers.activation import SiluAndMul
39
from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -129,7 +127,8 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
127
intermediate_size=config.shared_expert_intermediate_size,
128
hidden_act=config.hidden_act,
129
quant_config=quant_config,
132
- reduce_results=False,
130
+ reduce_results=self.experts.must_reduce_shared_expert_outputs(
131
+ ),
132
)
133
else:
134
self.shared_expert = None
@@ -156,7 +155,7 @@ class Qwen2MoeSparseMoeBlock(nn.Module):
155
if shared_output is not None:
156
final_hidden_states = final_hidden_states + shared_output
157
if self.tp_size > 1:
159
- final_hidden_states = tensor_model_parallel_all_reduce(
158
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
159
final_hidden_states)
160
161
return final_hidden_states.view(orig_shape)
vllm/model_executor/models/qwen3_moe.py
+2
-4
@@ -30,9 +30,7 @@ from transformers import PretrainedConfig
30
from vllm.attention import Attention
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
-from vllm.distributed import (get_pp_group,
34
- get_tensor_model_parallel_world_size,
35
- tensor_model_parallel_all_reduce)
33
+from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
from vllm.logger import init_logger
35
from vllm.model_executor.layers.activation import SiluAndMul
36
from vllm.model_executor.layers.fused_moe import FusedMoE
@@ -137,7 +135,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
135
router_logits=router_logits)
136
final_hidden_states = final_hidden_states
137
if self.tp_size > 1:
140
- final_hidden_states = tensor_model_parallel_all_reduce(
138
+ final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel( # noqa E501
139
final_hidden_states)
140
141
return final_hidden_states.view(orig_shape)
vllm/platforms/cuda.py
+1
@@ -158,6 +158,7 @@ class CudaPlatformBase(Platform):
158
"currently not supported with CUDA Graphs.")
159
vllm_config.model_config.enforce_eager = True
160
compilation_config.use_cudagraph = False
161
+ compilation_config.use_inductor = False
162
163
@classmethod
164
def get_current_memory_usage(cls,
vllm/v1/attention/backends/mla/common.py
+4
-2
@@ -865,8 +865,10 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
865
assert output is not None, "Output tensor must be provided."
866
867
if attn_metadata is None:
868
- # Profiling run.
869
- return output
868
+ # The zero fill is required when used with DP + EP
869
+ # to ensure all ranks within a DP group compute the
870
+ # same expert outputs.
871
+ return output.fill_(0)
872
873
num_actual_toks = attn_metadata.num_actual_tokens
874
vllm/v1/worker/gpu_worker.py
+2
-1
@@ -341,7 +341,8 @@ def init_worker_distributed_environment(
341
distributed_init_method, local_rank)
342
343
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
344
- parallel_config.pipeline_parallel_size)
344
+ parallel_config.pipeline_parallel_size,
345
+ parallel_config.enable_expert_parallel)
346
347
ensure_kv_transfer_initialized(vllm_config)
348
vllm/v1/worker/tpu_worker.py
+2
-1
@@ -265,4 +265,5 @@ def init_tpu_worker_distributed_environment(
265
backend="gloo",
266
)
267
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
268
- parallel_config.pipeline_parallel_size)
268
+ parallel_config.pipeline_parallel_size,
269
+ parallel_config.enable_expert_parallel)
vllm/worker/cpu_worker.py
+2
-1
@@ -390,7 +390,8 @@ class CPUWorker(LocalOrDistributedWorkerBase):
390
391
ensure_model_parallel_initialized(
392
parallel_config.tensor_parallel_size,
393
- parallel_config.pipeline_parallel_size)
393
+ parallel_config.pipeline_parallel_size,
394
+ parallel_config.enable_expert_parallel)
395
396
def get_cache_block_size_bytes(self) -> int:
397
"""Return the size in bytes of a single KV cache block.
vllm/worker/hpu_worker.py
+4
-2
@@ -416,7 +416,8 @@ def init_worker_distributed_environment(
416
backend='hccl')
417
418
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
419
- parallel_config.pipeline_parallel_size)
419
+ parallel_config.pipeline_parallel_size,
420
+ parallel_config.enable_expert_parallel)
421
422
if torch.distributed.is_initialized():
423
torch_world_size = torch.distributed.get_world_size()
@@ -442,7 +443,8 @@ def init_worker_distributed_environment(
443
torch.distributed.all_reduce(dummy_tensor_hpu)
444
assert dummy_tensor_hpu.item() == parallel_config.world_size
445
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
445
- parallel_config.pipeline_parallel_size)
446
+ parallel_config.pipeline_parallel_size,
447
+ parallel_config.enable_expert_parallel)
448
449
450
def raise_if_cache_size_invalid(num_gpu_blocks, block_size, max_model_len,
vllm/worker/tpu_worker.py
+2
-1
@@ -76,7 +76,8 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
76
)
77
ensure_model_parallel_initialized(
78
self.parallel_config.tensor_parallel_size,
79
- self.parallel_config.pipeline_parallel_size)
79
+ self.parallel_config.pipeline_parallel_size,
80
+ self.parallel_config.enable_expert_parallel)
81
82
# Device initialization should happen after initializing the distributed
83
# runtime.
vllm/worker/worker.py
+2
-1
@@ -530,7 +530,8 @@ def init_worker_distributed_environment(
530
init_distributed_environment(parallel_config.world_size, rank,
531
distributed_init_method, local_rank)
532
ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
533
- parallel_config.pipeline_parallel_size)
533
+ parallel_config.pipeline_parallel_size,
534
+ parallel_config.enable_expert_parallel)
535
536
ensure_kv_transfer_initialized(vllm_config)
537
vllm/worker/xpu_worker.py
+2
-1
@@ -176,7 +176,8 @@ class XPUWorker(LoRANotSupportedWorkerBase, Worker):
176
177
ensure_model_parallel_initialized(
178
parallel_config.tensor_parallel_size,
179
- parallel_config.pipeline_parallel_size)
179
+ parallel_config.pipeline_parallel_size,
180
+ parallel_config.enable_expert_parallel)
181
# global all_reduce needed for overall oneccl warm up
182
torch.distributed.all_reduce(torch.zeros(1).xpu())
183