add pad-aware reduce path (#48385)
Signed-off-by: gnovack <novackgm@gmail.com>
gnovack committed
Jul 14, 2026 at 18:05 UTC
f7aadae5e565c9a6cf85b8aa74280393fb8b6f28
7 files changed
+239
-34
csrc/libtorch_stable/dispatch_utils.h
+15
@@ -21,6 +21,21 @@
21
#define VLLM_STABLE_DISPATCH_FP8_CASE(enum_type, ...) \
22
THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
23
24
+// Same idea, for dispatching on an int32/int64 index tensor (e.g. topk_ids)
25
+// nested inside a value-type dispatch. Named 'idx_t' instead of 'scalar_t'.
26
+#define VLLM_STABLE_DISPATCH_IDX_CASE(enum_type, ...) \
27
+ THO_PRIVATE_CASE_TYPE_USING_HINT(enum_type, idx_t, __VA_ARGS__)
28
+
29
+#define VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(...) \
30
+ VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Int, \
31
+ __VA_ARGS__) \
32
+ VLLM_STABLE_DISPATCH_IDX_CASE(torch::headeronly::ScalarType::Long, \
33
+ __VA_ARGS__)
34
+
35
+#define VLLM_STABLE_DISPATCH_IDX_TYPES(TYPE, NAME, ...) \
36
+ THO_DISPATCH_SWITCH(TYPE, NAME, \
37
+ VLLM_STABLE_DISPATCH_CASE_IDX_TYPES(__VA_ARGS__))
38
+
39
#define VLLM_STABLE_DISPATCH_CASE_FLOATING_TYPES(...) \
40
THO_DISPATCH_CASE(torch::headeronly::ScalarType::Float, __VA_ARGS__) \
41
THO_DISPATCH_CASE(torch::headeronly::ScalarType::Half, __VA_ARGS__) \
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+158
-20
@@ -360,12 +360,24 @@ __global__ void count_and_sort_expert_tokens_kernel(
360
template <typename scalar_t>
361
constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t);
362
363
-template <typename scalar_t, int TOPK>
363
+template <typename idx_t>
364
+__device__ __forceinline__ bool moe_sum_pad_aware_skip(
365
+ const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map,
366
+ int64_t idx) {
367
+ int64_t expert_id = static_cast<int64_t>(topk_ids[idx]);
368
+ if (expert_id < 0) return true;
369
+ if (expert_map != nullptr && expert_map[expert_id] < 0) return true;
370
+ return false;
371
+}
372
+
373
+template <typename scalar_t, typename idx_t, int TOPK, bool PAD_AWARE>
374
__global__ void moe_sum_vec_kernel(
375
scalar_t* __restrict__ out, // [num_tokens, d], contiguous
376
const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous
377
const int64_t num_tokens, const int d, const int64_t stride_token,
368
- const int64_t stride_topk) {
378
+ const int64_t stride_topk, const idx_t* __restrict__ topk_ids,
379
+ const int32_t* __restrict__ expert_map, const int64_t stride_tk_token,
380
+ const int64_t stride_tk_k) {
381
using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>; // 16-byte pack
382
constexpr int VEC = MOE_SUM_VEC<scalar_t>;
383
const int64_t n_vec = d / VEC;
@@ -375,6 +387,10 @@ __global__ void moe_sum_vec_kernel(
387
const int64_t token = i / n_vec;
388
const int64_t v = i % n_vec;
389
const scalar_t* in_tok = input + token * stride_token + v * VEC;
390
+ const idx_t* tk_tok = nullptr;
391
+ if constexpr (PAD_AWARE) {
392
+ tk_tok = topk_ids + token * stride_tk_token;
393
+ }
394
395
float acc[VEC];
396
#pragma unroll
@@ -382,6 +398,11 @@ __global__ void moe_sum_vec_kernel(
398
399
#pragma unroll
400
for (int k = 0; k < TOPK; ++k) {
401
+ if constexpr (PAD_AWARE) {
402
+ if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) {
403
+ continue;
404
+ }
405
+ }
406
vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
407
#pragma unroll
408
for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
@@ -394,13 +415,16 @@ __global__ void moe_sum_vec_kernel(
415
}
416
}
417
397
-// Runtime-topk variant of the above.
398
-template <typename scalar_t>
418
+// Runtime-topk variant of the above, for topk values outside the templated
419
+// set.
420
+template <typename scalar_t, typename idx_t, bool PAD_AWARE>
421
__global__ void moe_sum_vec_dynamic_kernel(
422
scalar_t* __restrict__ out, // [num_tokens, d], contiguous
423
const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous
424
const int64_t num_tokens, const int d, const int topk,
403
- const int64_t stride_token, const int64_t stride_topk) {
425
+ const int64_t stride_token, const int64_t stride_topk,
426
+ const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map,
427
+ const int64_t stride_tk_token, const int64_t stride_tk_k) {
428
using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>;
429
constexpr int VEC = MOE_SUM_VEC<scalar_t>;
430
const int64_t n_vec = d / VEC;
@@ -410,12 +434,21 @@ __global__ void moe_sum_vec_dynamic_kernel(
434
const int64_t token = i / n_vec;
435
const int64_t v = i % n_vec;
436
const scalar_t* in_tok = input + token * stride_token + v * VEC;
437
+ const idx_t* tk_tok = nullptr;
438
+ if constexpr (PAD_AWARE) {
439
+ tk_tok = topk_ids + token * stride_tk_token;
440
+ }
441
442
float acc[VEC];
443
#pragma unroll
444
for (int j = 0; j < VEC; ++j) acc[j] = 0.f;
445
446
for (int k = 0; k < topk; ++k) {
447
+ if constexpr (PAD_AWARE) {
448
+ if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) {
449
+ continue;
450
+ }
451
+ }
452
vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
453
#pragma unroll
454
for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
@@ -430,17 +463,28 @@ __global__ void moe_sum_vec_dynamic_kernel(
463
464
// Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims
465
// (including a non-contiguous hidden stride) via per-element strided reads.
433
-template <typename scalar_t>
466
+template <typename scalar_t, typename idx_t, bool PAD_AWARE>
467
__global__ void moe_sum_scalar_kernel(
468
scalar_t* __restrict__ out, // [num_tokens, d], contiguous
469
const scalar_t* __restrict__ input, // [num_tokens, topk, d]
470
const int d, const int topk, const int64_t stride_token,
438
- const int64_t stride_topk, const int64_t stride_hidden) {
471
+ const int64_t stride_topk, const int64_t stride_hidden,
472
+ const idx_t* __restrict__ topk_ids, const int32_t* __restrict__ expert_map,
473
+ const int64_t stride_tk_token, const int64_t stride_tk_k) {
474
const int64_t token_idx = blockIdx.x;
475
const scalar_t* in_tok = input + token_idx * stride_token;
476
+ const idx_t* tk_tok = nullptr;
477
+ if constexpr (PAD_AWARE) {
478
+ tk_tok = topk_ids + token_idx * stride_tk_token;
479
+ }
480
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
481
float x = 0.f;
482
for (int k = 0; k < topk; ++k) {
483
+ if constexpr (PAD_AWARE) {
484
+ if (moe_sum_pad_aware_skip(tk_tok, expert_map, k * stride_tk_k)) {
485
+ continue;
486
+ }
487
+ }
488
x += static_cast<float>(
489
VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden]));
490
}
@@ -711,8 +755,9 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
755
}
756
757
void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
714
- torch::stable::Tensor& output) // [num_tokens, hidden_size]
715
-{
758
+ torch::stable::Tensor& output, // [num_tokens, hidden_size]
759
+ std::optional<torch::stable::Tensor> topk_ids,
760
+ std::optional<torch::stable::Tensor> expert_map) {
761
// Output is dense and written in place, so it must be contiguous. The input
762
// is read by its strides (no copy); only the hidden dim needs to be
763
// contiguous to take the vectorized path.
@@ -731,10 +776,102 @@ void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
776
const cudaStream_t stream =
777
get_current_cuda_stream(output.get_device_index());
778
734
-#define LAUNCH_MOE_SUM_VEC(TOPK) \
735
- vllm::moe::moe_sum_vec_kernel<scalar_t, TOPK> \
736
- <<<grid, dim3(block), 0, stream>>>( \
737
- out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk)
779
+ if (topk_ids.has_value()) {
780
+ // Pad-aware reduce path
781
+ const torch::stable::Tensor& tk = topk_ids.value();
782
+ STD_TORCH_CHECK(tk.size(0) == num_tokens && tk.size(1) == topk,
783
+ "moe_sum: topk_ids must have shape [num_tokens, topk]");
784
+ const int64_t stride_tk_token = tk.stride(0);
785
+ const int64_t stride_tk_k = tk.stride(1);
786
+
787
+ const int32_t* expert_map_ptr = nullptr;
788
+ if (expert_map.has_value()) {
789
+ STD_TORCH_CHECK(
790
+ expert_map->scalar_type() == torch::headeronly::ScalarType::Int,
791
+ "moe_sum: expert_map must be int32");
792
+ expert_map_ptr =
793
+ reinterpret_cast<const int32_t*>(expert_map->const_data_ptr());
794
+ }
795
+
796
+#define LAUNCH_MOE_SUM_PAD_AWARE_VEC(TOPK) \
797
+ vllm::moe::moe_sum_vec_kernel<scalar_t, idx_t, TOPK, true> \
798
+ <<<grid, dim3(block), 0, stream>>>( \
799
+ out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk, \
800
+ topk_ids_ptr, expert_map_ptr, stride_tk_token, stride_tk_k)
801
+
802
+ VLLM_STABLE_DISPATCH_FLOATING_TYPES(
803
+ input.scalar_type(), "moe_sum_pad_aware", [&] {
804
+ constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>;
805
+ constexpr int WIDTH = VEC * sizeof(scalar_t);
806
+ auto* out_ptr =
807
+ reinterpret_cast<scalar_t*>(output.mutable_data_ptr());
808
+ auto* in_ptr =
809
+ reinterpret_cast<const scalar_t*>(input.const_data_ptr());
810
+
811
+ const bool can_vec =
812
+ (stride_hidden == 1) && (hidden_size % VEC == 0) &&
813
+ (stride_token % VEC == 0) && (stride_topk % VEC == 0) &&
814
+ (reinterpret_cast<uintptr_t>(in_ptr) % WIDTH == 0) &&
815
+ (reinterpret_cast<uintptr_t>(out_ptr) % WIDTH == 0);
816
+
817
+ VLLM_STABLE_DISPATCH_IDX_TYPES(
818
+ tk.scalar_type(), "moe_sum_pad_aware_idx", [&] {
819
+ auto* topk_ids_ptr =
820
+ reinterpret_cast<const idx_t*>(tk.const_data_ptr());
821
+ if (can_vec) {
822
+ const int64_t n_vec = hidden_size / VEC;
823
+ const int64_t total = num_tokens * n_vec;
824
+ const int block = 256;
825
+ const dim3 grid(
826
+ std::min<int64_t>((total + block - 1) / block, 65535));
827
+ switch (topk) {
828
+ case 1:
829
+ LAUNCH_MOE_SUM_PAD_AWARE_VEC(1);
830
+ break;
831
+ case 2:
832
+ LAUNCH_MOE_SUM_PAD_AWARE_VEC(2);
833
+ break;
834
+ case 4:
835
+ LAUNCH_MOE_SUM_PAD_AWARE_VEC(4);
836
+ break;
837
+ case 6:
838
+ LAUNCH_MOE_SUM_PAD_AWARE_VEC(6);
839
+ break;
840
+ case 8:
841
+ LAUNCH_MOE_SUM_PAD_AWARE_VEC(8);
842
+ break;
843
+ case 9:
844
+ LAUNCH_MOE_SUM_PAD_AWARE_VEC(9);
845
+ break;
846
+ default:
847
+ vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t, idx_t,
848
+ true>
849
+ <<<grid, dim3(block), 0, stream>>>(
850
+ out_ptr, in_ptr, num_tokens, hidden_size, topk,
851
+ stride_token, stride_topk, topk_ids_ptr,
852
+ expert_map_ptr, stride_tk_token, stride_tk_k);
853
+ break;
854
+ }
855
+ } else {
856
+ dim3 grid(num_tokens);
857
+ dim3 block(std::min(hidden_size, 1024));
858
+ vllm::moe::moe_sum_scalar_kernel<scalar_t, idx_t, true>
859
+ <<<grid, block, 0, stream>>>(
860
+ out_ptr, in_ptr, hidden_size, topk, stride_token,
861
+ stride_topk, stride_hidden, topk_ids_ptr,
862
+ expert_map_ptr, stride_tk_token, stride_tk_k);
863
+ }
864
+ });
865
+ });
866
+#undef LAUNCH_MOE_SUM_PAD_AWARE_VEC
867
+ return;
868
+ }
869
+
870
+#define LAUNCH_MOE_SUM_VEC(TOPK) \
871
+ vllm::moe::moe_sum_vec_kernel<scalar_t, int32_t, TOPK, false> \
872
+ <<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens, \
873
+ hidden_size, stride_token, \
874
+ stride_topk, nullptr, nullptr, 0, 0)
875
876
VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] {
877
constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>;
@@ -774,18 +911,19 @@ void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
911
LAUNCH_MOE_SUM_VEC(9);
912
break;
913
default:
777
- vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t>
778
- <<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens,
779
- hidden_size, topk,
780
- stride_token, stride_topk);
914
+ vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t, int32_t, false>
915
+ <<<grid, dim3(block), 0, stream>>>(
916
+ out_ptr, in_ptr, num_tokens, hidden_size, topk, stride_token,
917
+ stride_topk, nullptr, nullptr, 0, 0);
918
break;
919
}
920
} else {
921
dim3 grid(num_tokens);
922
dim3 block(std::min(hidden_size, 1024));
786
- vllm::moe::moe_sum_scalar_kernel<scalar_t><<<grid, block, 0, stream>>>(
787
- out_ptr, in_ptr, hidden_size, topk, stride_token, stride_topk,
788
- stride_hidden);
923
+ vllm::moe::moe_sum_scalar_kernel<scalar_t, int32_t, false>
924
+ <<<grid, block, 0, stream>>>(out_ptr, in_ptr, hidden_size, topk,
925
+ stride_token, stride_topk, stride_hidden,
926
+ nullptr, nullptr, 0, 0);
927
}
928
});
929
#undef LAUNCH_MOE_SUM_VEC
csrc/libtorch_stable/moe/moe_ops.h
+3
-1
@@ -27,7 +27,9 @@ void topk_softplus_sqrt(
27
const std::optional<torch::stable::Tensor>& input_ids,
28
const std::optional<torch::stable::Tensor>& tid2eid);
29
30
-void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output);
30
+void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
31
+ std::optional<torch::stable::Tensor> topk_ids,
32
+ std::optional<torch::stable::Tensor> expert_map);
33
34
void moe_align_block_size(
35
torch::stable::Tensor topk_ids, int64_t num_experts, int64_t block_size,
csrc/libtorch_stable/moe/torch_bindings.cpp
+7
-2
@@ -23,8 +23,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) {
23
"bias, Tensor? input_ids, Tensor? tid2eid) -> ()");
24
25
// Calculate the result of moe by summing up the partial results
26
- // from all selected experts.
27
- m.def("moe_sum(Tensor input, Tensor! output) -> ()");
26
+ // from all selected experts. topk_ids/expert_map are optional and, when
27
+ // both given, enable pad-aware reduce that skips (token, expert)
28
+ // slots that were never actually computed (unrouted, or routed to an
29
+ // expert not owned by this rank under expert parallelism).
30
+ m.def(
31
+ "moe_sum(Tensor input, Tensor! output, Tensor? topk_ids=None, "
32
+ "Tensor? expert_map=None) -> ()");
33
34
// Aligning the number of tokens to be processed by each expert such
35
// that it is divisible by the block size.
tests/kernels/moe/test_moe.py
+26
@@ -1271,6 +1271,32 @@ def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype, layout: str):
1271
opcheck(torch.ops._moe_C.moe_sum, (input, actual))
1272
1273
1274
+@pytest.mark.parametrize("topk", [2, 8])
1275
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
1276
+@pytest.mark.parametrize("topk_ids_dtype", [torch.int32, torch.int64])
1277
+def test_moe_sum_pad_aware(topk: int, dtype: torch.dtype, topk_ids_dtype: torch.dtype):
1278
+ m, k = 17, 128
1279
+ input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
1280
+
1281
+ topk_ids = torch.randint(0, 4, (m, topk), device="cuda", dtype=topk_ids_dtype)
1282
+ # Force some slots unrouted, independent of expert parallelism.
1283
+ topk_ids[0, :] = -1
1284
+ topk_ids[1, 0] = -1
1285
+
1286
+ expert_map = torch.tensor([0, -1, 1, -1], device="cuda", dtype=torch.int32)
1287
+
1288
+ actual = torch.empty((m, k), device="cuda", dtype=dtype)
1289
+ torch.ops._moe_C.moe_sum(input, actual, topk_ids, expert_map)
1290
+
1291
+ routed = expert_map[topk_ids.clamp(min=0).long()] >= 0
1292
+ routed &= topk_ids >= 0
1293
+ expected = (input.float() * routed.unsqueeze(-1)).sum(dim=1).to(dtype)
1294
+
1295
+ torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)
1296
+
1297
+ opcheck(torch.ops._moe_C.moe_sum, (input, actual, topk_ids, expert_map))
1298
+
1299
+
1300
@pytest.mark.usefixtures("default_vllm_config")
1301
@pytest.mark.parametrize("m", [1, 33])
1302
@pytest.mark.parametrize("n,k", [(128, 128)])
vllm/_custom_ops.py
+7
-2
@@ -2138,8 +2138,13 @@ def wvSplitKQ(
2138
2139
2140
# moe
2141
-def moe_sum(input: torch.Tensor, output: torch.Tensor):
2142
- torch.ops._moe_C.moe_sum(input, output)
2141
+def moe_sum(
2142
+ input: torch.Tensor,
2143
+ output: torch.Tensor,
2144
+ topk_ids: torch.Tensor | None = None,
2145
+ expert_map: torch.Tensor | None = None,
2146
+):
2147
+ torch.ops._moe_C.moe_sum(input, output, topk_ids, expert_map)
2148
2149
2150
def moe_align_block_size(
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+23
-9
@@ -176,9 +176,6 @@ def _fused_marlin_moe(
176
if output is None:
177
output = intermediate_cache3
178
179
- if expert_map is not None:
180
- output.zero_()
181
-
179
a_scales2 = None
180
if input_dtype == torch.int8:
181
intermediate_cache2, a_scales2 = marlin_quant_input(
@@ -238,7 +235,7 @@ def fused_marlin_moe(
235
global_num_experts: int = -1,
236
activation: MoEActivation = MoEActivation.SILU,
237
activation_func: Callable[..., None] = apply_moe_activation,
241
- moe_sum: Callable[[torch.Tensor, torch.Tensor], None] | None = None,
238
+ moe_sum: Callable[..., torch.Tensor | None] | None = None,
239
expert_map: torch.Tensor | None = None,
240
input_global_scale1: torch.Tensor | None = None,
241
input_global_scale2: torch.Tensor | None = None,
@@ -385,9 +382,12 @@ def fused_marlin_moe(
382
output = torch.empty_like(hidden_states)
383
384
if moe_sum is None:
385
+ if expert_map is not None:
386
+ ops.moe_sum(moe_output, output, topk_ids, expert_map)
387
+ return output
388
return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output)
389
else:
390
- return moe_sum(moe_output, output)
390
+ return moe_sum(moe_output, output, topk_ids, expert_map)
391
392
393
def batched_fused_marlin_moe(
@@ -874,7 +874,12 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
874
)
875
lora_state["cache2"] = act_output
876
877
- def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None:
877
+ def moe_sum_with_lora(
878
+ moe_out: torch.Tensor,
879
+ out: torch.Tensor,
880
+ topk_ids: torch.Tensor,
881
+ expert_map: torch.Tensor | None,
882
+ ) -> None:
883
# moe_out shape: (M, topk, K)
884
self.apply_w2_lora(
885
ctx,
@@ -890,7 +895,7 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
895
w2=w2,
896
top_k_num=top_k_num,
897
)
893
- self.moe_sum(moe_out, out)
898
+ self.moe_sum(moe_out, out, topk_ids, expert_map)
899
900
return fused_marlin_moe(
901
hidden_states=hidden_states,
@@ -929,8 +934,17 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):
934
gemm1_beta=self.gemm1_beta,
935
)
936
932
- def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None:
933
- ops.moe_sum(input, output)
937
+ def moe_sum(
938
+ self,
939
+ input: torch.Tensor,
940
+ output: torch.Tensor,
941
+ topk_ids: torch.Tensor,
942
+ expert_map: torch.Tensor | None,
943
+ ) -> None:
944
+ if expert_map is not None:
945
+ ops.moe_sum(input, output, topk_ids, expert_map)
946
+ else:
947
+ ops.moe_sum(input, output)
948
949
950
class BatchedMarlinExperts(MarlinExpertsBase):