[Kernel][MoE] Optimize batched_moe_align_block_size with cooperative writes (#53280)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Sep 11, 2026 at 15:05 UTC
8a7f98c8c3ae409df6130995e364393ead5dbc24
3 files changed
+154
-27
benchmarks/kernels/benchmark_moe_align_block_size.py
+82
-1
@@ -5,6 +5,7 @@ import itertools
5
6
import torch
7
8
+import vllm._custom_ops as ops
9
from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
10
moe_align_block_size,
11
)
@@ -68,8 +69,87 @@ def benchmark(num_tokens, num_experts, topk, ep_size, provider):
69
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
70
71
72
+batched_configs = list(
73
+ itertools.product(
74
+ [8, 32, 64, 256],
75
+ [16, 256, 512, 2048],
76
+ [8, 32, 48, 64],
77
+ ["uniform", "skewed", "mostly_empty"],
78
+ )
79
+)
80
+
81
+
82
+@triton.testing.perf_report(
83
+ triton.testing.Benchmark(
84
+ x_names=[
85
+ "num_batches",
86
+ "max_tokens_per_batch",
87
+ "block_size",
88
+ "distribution",
89
+ ],
90
+ x_vals=batched_configs,
91
+ line_arg="provider",
92
+ line_vals=["vllm"],
93
+ line_names=["vLLM"],
94
+ plot_name="batched-moe-align-block-size-performance",
95
+ args={},
96
+ )
97
+)
98
+def benchmark_batched(
99
+ num_batches,
100
+ max_tokens_per_batch,
101
+ block_size,
102
+ distribution,
103
+ provider,
104
+):
105
+ set_random_seed(0)
106
+ if distribution == "uniform":
107
+ expert_num_tokens = torch.full(
108
+ (num_batches,),
109
+ max_tokens_per_batch,
110
+ dtype=torch.int32,
111
+ device="cuda",
112
+ )
113
+ elif distribution == "skewed":
114
+ expert_num_tokens = torch.ones((num_batches,), dtype=torch.int32, device="cuda")
115
+ expert_num_tokens[0] = max_tokens_per_batch
116
+ else:
117
+ expert_num_tokens = torch.zeros(
118
+ (num_batches,), dtype=torch.int32, device="cuda"
119
+ )
120
+ expert_num_tokens[::8] = max_tokens_per_batch
121
+
122
+ max_num_tokens_padded = (
123
+ num_batches
124
+ * ((max_tokens_per_batch + block_size - 1) // block_size)
125
+ * block_size
126
+ )
127
+ sorted_ids = torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda")
128
+ expert_ids = torch.empty(
129
+ (max_num_tokens_padded // block_size,), dtype=torch.int32, device="cuda"
130
+ )
131
+ num_tokens_post_pad = torch.empty((1,), dtype=torch.int32, device="cuda")
132
+
133
+ quantiles = [0.5, 0.2, 0.8]
134
+ if provider == "vllm":
135
+ ms, min_ms, max_ms = triton.testing.do_bench(
136
+ lambda: ops.batched_moe_align_block_size(
137
+ max_tokens_per_batch,
138
+ block_size,
139
+ expert_num_tokens,
140
+ sorted_ids,
141
+ expert_ids,
142
+ num_tokens_post_pad,
143
+ ),
144
+ quantiles=quantiles,
145
+ )
146
+
147
+ return 1000 * ms, 1000 * max_ms, 1000 * min_ms
148
+
149
+
150
if __name__ == "__main__":
151
parser = argparse.ArgumentParser()
152
+ parser.add_argument("--batched", action="store_true")
153
parser.add_argument(
154
"--num_experts",
155
type=int,
@@ -85,4 +165,5 @@ if __name__ == "__main__":
165
)
166
args = parser.parse_args()
167
88
- benchmark.run(print_data=True, show_plots=True)
168
+ selected_benchmark = benchmark_batched if args.batched else benchmark
169
+ selected_benchmark.run(print_data=True, show_plots=True)
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+60
-18
@@ -23,20 +23,18 @@ namespace batched_moe_align_block_size {
23
// Note num_threads needs to be 1024 for BlockScan Reduction in the kernel.
24
static constexpr int32_t num_threads = 1024;
25
static constexpr int32_t num_blocks = 1;
26
+template <bool cooperative_writes>
27
__global__ void batched_moe_align_block_size_kernel(
28
int32_t const num_batches, int32_t const max_tokens_per_batch,
29
int32_t const block_size, int32_t const* __restrict__ batch_num_tokens,
30
int32_t* __restrict__ sorted_ids, int32_t* __restrict__ block_ids,
31
int32_t* __restrict__ num_tokens_post_pad) {
31
- // TODO(varun): This is a naive implementation. Could be optimized.
32
-
33
- size_t const batch_id = threadIdx.x;
32
size_t const stride = blockDim.x * gridDim.x;
33
int32_t const num_blocks_per_batch =
34
CEILDIV(max_tokens_per_batch, block_size);
37
- int32_t const sorted_ids_size =
38
- num_blocks_per_batch * num_batches * block_size;
39
- int32_t const block_ids_size = sorted_ids_size / block_size;
35
+ size_t const sorted_ids_size =
36
+ static_cast<size_t>(num_blocks_per_batch) * num_batches * block_size;
37
+ size_t const block_ids_size = sorted_ids_size / block_size;
38
int32_t const SENTINEL =
39
num_batches * max_tokens_per_batch; // To denote invalid entries.
40
// Initialize sorted_ids
@@ -49,8 +47,8 @@ __global__ void batched_moe_align_block_size_kernel(
47
}
48
49
int32_t b_num_tokens = 0;
52
- if (batch_id < num_batches) {
53
- b_num_tokens = batch_num_tokens[batch_id];
50
+ if (threadIdx.x < num_batches) {
51
+ b_num_tokens = batch_num_tokens[threadIdx.x];
52
}
53
int32_t const ceil_b_num_tokens =
54
CEILDIV(b_num_tokens, block_size) * block_size;
@@ -62,20 +60,56 @@ __global__ void batched_moe_align_block_size_kernel(
60
BlockScan(temp_storage).ExclusiveSum(ceil_b_num_tokens, cumsum_val);
61
__syncthreads();
62
65
- bool const is_last_batch = batch_id == (num_batches - 1);
66
- if (is_last_batch) {
63
+ if (threadIdx.x == num_batches - 1) {
64
*num_tokens_post_pad = cumsum_val + ceil_b_num_tokens;
65
}
66
70
- if (batch_id < num_batches) {
71
- int32_t const batch_offset = batch_id * max_tokens_per_batch;
67
+ if constexpr (cooperative_writes) {
68
+ __shared__ int32_t batch_cumsum[num_threads];
69
+ __shared__ int32_t valid_tokens[num_threads];
70
+ __shared__ int32_t batch_num_blocks[num_threads];
71
+ if (threadIdx.x < num_batches) {
72
+ batch_cumsum[threadIdx.x] = cumsum_val;
73
+ valid_tokens[threadIdx.x] = b_num_tokens;
74
+ batch_num_blocks[threadIdx.x] = ceil_b_num_tokens / block_size;
75
+ }
76
+ __syncthreads();
77
+
78
+ int32_t const max_num_groups = blockDim.x / WARP_SIZE;
79
+ int32_t num_groups = 1;
80
+ while (num_groups < num_batches && num_groups < max_num_groups) {
81
+ num_groups *= 2;
82
+ }
83
+ int32_t const threads_per_batch = blockDim.x / num_groups;
84
+ int32_t const group_id = threadIdx.x / threads_per_batch;
85
+ int32_t const group_offset = threadIdx.x % threads_per_batch;
86
+
87
+ // Assign at least one warp to each batch when possible.
88
+ for (int32_t batch_id = group_id; batch_id < num_batches;
89
+ batch_id += num_groups) {
90
+ size_t const batch_offset =
91
+ static_cast<size_t>(batch_id) * max_tokens_per_batch;
92
+ size_t const cumsum = batch_cumsum[batch_id];
93
+ for (size_t i = group_offset; i < valid_tokens[batch_id];
94
+ i += threads_per_batch) {
95
+ sorted_ids[cumsum + i] = static_cast<int32_t>(batch_offset + i);
96
+ }
97
+
98
+ size_t const block_start = cumsum / block_size;
99
+ for (size_t i = group_offset; i < batch_num_blocks[batch_id];
100
+ i += threads_per_batch) {
101
+ block_ids[block_start + i] = batch_id;
102
+ }
103
+ }
104
+ } else if (threadIdx.x < num_batches) {
105
+ size_t const batch_id = threadIdx.x;
106
+ size_t const batch_offset = batch_id * max_tokens_per_batch;
107
for (size_t i = 0; i < b_num_tokens; ++i) {
73
- sorted_ids[cumsum_val + i] = batch_offset + i;
108
+ sorted_ids[cumsum_val + i] = static_cast<int32_t>(batch_offset + i);
109
}
110
76
- int32_t const block_start = cumsum_val / block_size;
77
- int32_t const num_blocks = ceil_b_num_tokens / block_size;
78
- for (size_t i = 0; i < num_blocks; ++i) {
111
+ size_t const block_start = cumsum_val / block_size;
112
+ for (size_t i = 0; i < ceil_b_num_tokens / block_size; ++i) {
113
block_ids[block_start + i] = batch_id;
114
}
115
}
@@ -747,8 +781,16 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
781
STD_TORCH_CHECK(num_tokens_post_pad.size(0) == 1);
782
STD_TORCH_CHECK(B <= batched_kernel::num_threads);
783
750
- batched_kernel::batched_moe_align_block_size_kernel<<<
751
- batched_kernel::num_blocks, batched_kernel::num_threads, 0, stream>>>(
784
+ // Avoid coordination overhead for small capacities or many batches.
785
+ int64_t const cooperative_threshold = std::max<int64_t>(256, 8 * B);
786
+ bool const use_cooperative_writes =
787
+ max_tokens_per_batch >= cooperative_threshold;
788
+ auto kernel =
789
+ use_cooperative_writes
790
+ ? batched_kernel::batched_moe_align_block_size_kernel<true>
791
+ : batched_kernel::batched_moe_align_block_size_kernel<false>;
792
+ kernel<<<batched_kernel::num_blocks, batched_kernel::num_threads, 0,
793
+ stream>>>(
794
B, max_tokens_per_batch, block_size,
795
reinterpret_cast<const int32_t*>(batch_num_tokens.const_data_ptr()),
796
reinterpret_cast<int32_t*>(sorted_ids.mutable_data_ptr()),
tests/kernels/moe/test_moe_align_block_size.py
+12
-8
@@ -329,13 +329,15 @@ def test_moe_align_block_size_deterministic():
329
330
@pytest.mark.parametrize("max_tokens_per_batch", [13, 16, 512])
331
@pytest.mark.parametrize("num_experts", [8, 16, 32, 64])
332
-@pytest.mark.parametrize("block_size", [8, 16, 32, 64])
333
-@pytest.mark.parametrize("simulate_empty_batches", [False, True])
332
+@pytest.mark.parametrize("block_size", [8, 16, 32, 48, 64])
333
+@pytest.mark.parametrize(
334
+ "token_distribution", ["random", "alternating_empty", "skewed"]
335
+)
336
def test_batched_moe_align_block_size(
337
max_tokens_per_batch: int,
338
num_experts: int,
339
block_size: int,
338
- simulate_empty_batches: bool,
340
+ token_distribution: str,
341
):
342
def ref_outputs(
343
expert_num_tokens: torch.Tensor,
@@ -384,15 +386,17 @@ def test_batched_moe_align_block_size(
386
# Compute expert_num_tokens
387
expert_num_tokens = torch.randint(
388
low=0,
387
- high=max_tokens_per_batch,
389
+ high=max_tokens_per_batch + 1,
390
size=(num_experts,),
391
device="cpu",
392
dtype=torch.int32,
393
)
392
- if simulate_empty_batches:
393
- # mark half the batches to have 0 tokens
394
- zero_batches = torch.randperm(num_experts)[: num_experts // 2]
395
- expert_num_tokens[zero_batches] = 0
394
+ if token_distribution == "alternating_empty":
395
+ expert_num_tokens[::2] = 0
396
+ elif token_distribution == "skewed":
397
+ expert_num_tokens.zero_()
398
+ expert_num_tokens[0] = max_tokens_per_batch
399
+ expert_num_tokens[-1] = 1
400
401
# ref outputs
402
ref_sorted_ids, ref_expert_ids, ref_num_tokens_post_pad = ref_outputs(