@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e8c24a7695

[Kernel] Vectorized fp32 `moe_sum` reduction and support any topk (#46643)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>

Michael Goin committed Jun 25, 2026 at 15:02 UTC e8c24a769576fa318ca93fbd927128246a534325
2 files changed +177 -51
csrc/libtorch_stable/moe/moe_align_sum_kernels.cu
+161 -47
@@ -11,6 +11,7 @@
11 #include "../../cuda_compat.h"
12 #include "libtorch_stable/core/math.hpp"
13 #include "libtorch_stable/dispatch_utils.h"
14 +#include "libtorch_stable/quantization/vectorization.cuh"
15 #include "libtorch_stable/torch_utils.h"
16
17 #define CEILDIV(x, y) (((x) + (y) - 1) / (y))
@@ -349,19 +350,102 @@ __global__ void count_and_sort_expert_tokens_kernel(
350 max_num_tokens_padded, nullptr, 0, topk_num, has_expert_map);
351 }
352
353 +// Reduce the topk expert outputs per token (summed in fp32). The output is
354 +// dense [num_tokens, d]; the input is addressed by its strides so non-
355 +// contiguous inputs work without a copy. A 16B-vectorized path is used when
356 +// the hidden dim is contiguous (innermost stride 1) and aligned; otherwise a
357 +// scalar kernel reads via arbitrary strides. topk is a compile-time constant
358 +// for common values and runtime otherwise.
359 +
360 +// Elements per 16-byte vector (8 for bf16/fp16, 4 for fp32).
361 +template <typename scalar_t>
362 +constexpr int MOE_SUM_VEC = 16 / sizeof(scalar_t);
363 +
364 template <typename scalar_t, int TOPK>
353 -__global__ void moe_sum_kernel(
354 - scalar_t* __restrict__ out, // [..., d]
355 - const scalar_t* __restrict__ input, // [..., topk, d]
356 - const int d) {
357 - const int64_t token_idx = blockIdx.x;
358 - for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
359 - scalar_t x = 0.0;
365 +__global__ void moe_sum_vec_kernel(
366 + scalar_t* __restrict__ out, // [num_tokens, d], contiguous
367 + const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous
368 + const int64_t num_tokens, const int d, const int64_t stride_token,
369 + const int64_t stride_topk) {
370 + using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>; // 16-byte pack
371 + constexpr int VEC = MOE_SUM_VEC<scalar_t>;
372 + const int64_t n_vec = d / VEC;
373 + const int64_t total = num_tokens * n_vec;
374 + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total;
375 + i += (int64_t)gridDim.x * blockDim.x) {
376 + const int64_t token = i / n_vec;
377 + const int64_t v = i % n_vec;
378 + const scalar_t* in_tok = input + token * stride_token + v * VEC;
379 +
380 + float acc[VEC];
381 +#pragma unroll
382 + for (int j = 0; j < VEC; ++j) acc[j] = 0.f;
383 +
384 #pragma unroll
385 for (int k = 0; k < TOPK; ++k) {
362 - x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
386 + vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
387 +#pragma unroll
388 + for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
389 + }
390 +
391 + vec_t outp;
392 +#pragma unroll
393 + for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast<scalar_t>(acc[j]);
394 + *reinterpret_cast<vec_t*>(out + token * d + v * VEC) = outp;
395 + }
396 +}
397 +
398 +// Runtime-topk variant of the above.
399 +template <typename scalar_t>
400 +__global__ void moe_sum_vec_dynamic_kernel(
401 + scalar_t* __restrict__ out, // [num_tokens, d], contiguous
402 + const scalar_t* __restrict__ input, // [num_tokens, topk, d], d contiguous
403 + const int64_t num_tokens, const int d, const int topk,
404 + const int64_t stride_token, const int64_t stride_topk) {
405 + using vec_t = vllm::vec_n_t<scalar_t, MOE_SUM_VEC<scalar_t>>;
406 + constexpr int VEC = MOE_SUM_VEC<scalar_t>;
407 + const int64_t n_vec = d / VEC;
408 + const int64_t total = num_tokens * n_vec;
409 + for (int64_t i = blockIdx.x * blockDim.x + threadIdx.x; i < total;
410 + i += (int64_t)gridDim.x * blockDim.x) {
411 + const int64_t token = i / n_vec;
412 + const int64_t v = i % n_vec;
413 + const scalar_t* in_tok = input + token * stride_token + v * VEC;
414 +
415 + float acc[VEC];
416 +#pragma unroll
417 + for (int j = 0; j < VEC; ++j) acc[j] = 0.f;
418 +
419 + for (int k = 0; k < topk; ++k) {
420 + vec_t packed = *reinterpret_cast<const vec_t*>(in_tok + k * stride_topk);
421 +#pragma unroll
422 + for (int j = 0; j < VEC; ++j) acc[j] += static_cast<float>(packed.val[j]);
423 }
364 - out[token_idx * d + idx] = x;
424 +
425 + vec_t outp;
426 +#pragma unroll
427 + for (int j = 0; j < VEC; ++j) outp.val[j] = static_cast<scalar_t>(acc[j]);
428 + *reinterpret_cast<vec_t*>(out + token * d + v * VEC) = outp;
429 + }
430 +}
431 +
432 +// Stride-aware scalar fallback: handles unaligned/non-vectorizable hidden dims
433 +// (including a non-contiguous hidden stride) via per-element strided reads.
434 +template <typename scalar_t>
435 +__global__ void moe_sum_scalar_kernel(
436 + scalar_t* __restrict__ out, // [num_tokens, d], contiguous
437 + const scalar_t* __restrict__ input, // [num_tokens, topk, d]
438 + const int d, const int topk, const int64_t stride_token,
439 + const int64_t stride_topk, const int64_t stride_hidden) {
440 + const int64_t token_idx = blockIdx.x;
441 + const scalar_t* in_tok = input + token_idx * stride_token;
442 + for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
443 + float x = 0.f;
444 + for (int k = 0; k < topk; ++k) {
445 + x += static_cast<float>(
446 + VLLM_LDG(&in_tok[k * stride_topk + idx * stride_hidden]));
447 + }
448 + out[token_idx * d + idx] = static_cast<scalar_t>(x);
449 }
450 }
451
@@ -626,52 +710,82 @@ void batched_moe_align_block_size(int64_t max_tokens_per_batch,
710 void moe_sum(torch::stable::Tensor& input, // [num_tokens, topk, hidden_size]
711 torch::stable::Tensor& output) // [num_tokens, hidden_size]
712 {
713 + // Output is dense and written in place, so it must be contiguous. The input
714 + // is read by its strides (no copy); only the hidden dim needs to be
715 + // contiguous to take the vectorized path.
716 + STD_TORCH_CHECK(output.is_contiguous(),
717 + "moe_sum expects a contiguous output");
718 +
719 const int hidden_size = input.size(-1);
630 - const auto num_tokens = output.numel() / hidden_size;
720 + const int64_t num_tokens = output.numel() / hidden_size;
721 const int topk = input.size(1);
722 + const int64_t stride_token = input.stride(0);
723 + const int64_t stride_topk = input.stride(1);
724 + const int64_t stride_hidden = input.stride(2);
725
633 - dim3 grid(num_tokens);
634 - dim3 block(std::min(hidden_size, 1024));
726 const torch::stable::accelerator::DeviceGuard device_guard(
727 output.get_device_index());
728 const cudaStream_t stream =
729 get_current_cuda_stream(output.get_device_index());
730
640 - switch (topk) {
641 - case 2:
642 - VLLM_STABLE_DISPATCH_FLOATING_TYPES(
643 - input.scalar_type(), "moe_sum_kernel", [&] {
644 - vllm::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
645 - reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
646 - reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
647 - hidden_size);
648 - });
649 - break;
650 -
651 - case 3:
652 - VLLM_STABLE_DISPATCH_FLOATING_TYPES(
653 - input.scalar_type(), "moe_sum_kernel", [&] {
654 - vllm::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
655 - reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
656 - reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
657 - hidden_size);
658 - });
659 - break;
660 -
661 - case 4:
662 - VLLM_STABLE_DISPATCH_FLOATING_TYPES(
663 - input.scalar_type(), "moe_sum_kernel", [&] {
664 - vllm::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
665 - reinterpret_cast<scalar_t*>(output.mutable_data_ptr()),
666 - reinterpret_cast<const scalar_t*>(input.const_data_ptr()),
667 - hidden_size);
668 - });
669 - break;
670 -
671 - default:
672 - torch::stable::sum_out(output, input, std::array<int64_t, 1>{1});
673 - break;
674 - }
731 +#define LAUNCH_MOE_SUM_VEC(TOPK) \
732 + vllm::moe::moe_sum_vec_kernel<scalar_t, TOPK> \
733 + <<<grid, dim3(block), 0, stream>>>( \
734 + out_ptr, in_ptr, num_tokens, hidden_size, stride_token, stride_topk)
735 +
736 + VLLM_STABLE_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum", [&] {
737 + constexpr int VEC = vllm::moe::MOE_SUM_VEC<scalar_t>;
738 + constexpr int WIDTH = VEC * sizeof(scalar_t); // 16 bytes
739 + auto* out_ptr = reinterpret_cast<scalar_t*>(output.mutable_data_ptr());
740 + auto* in_ptr = reinterpret_cast<const scalar_t*>(input.const_data_ptr());
741 +
742 + // Vectorize along hidden only when it is contiguous (innermost stride 1),
743 + // a whole number of vectors, and every row offset stays 16B-aligned.
744 + const bool can_vec = (stride_hidden == 1) && (hidden_size % VEC == 0) &&
745 + (stride_token % VEC == 0) &&
746 + (stride_topk % VEC == 0) &&
747 + (reinterpret_cast<uintptr_t>(in_ptr) % WIDTH == 0) &&
748 + (reinterpret_cast<uintptr_t>(out_ptr) % WIDTH == 0);
749 + if (can_vec) {
750 + const int64_t n_vec = hidden_size / VEC;
751 + const int64_t total = num_tokens * n_vec;
752 + const int block = 256;
753 + const dim3 grid(std::min<int64_t>((total + block - 1) / block, 65535));
754 + switch (topk) {
755 + case 1:
756 + LAUNCH_MOE_SUM_VEC(1);
757 + break;
758 + case 2:
759 + LAUNCH_MOE_SUM_VEC(2);
760 + break;
761 + case 4:
762 + LAUNCH_MOE_SUM_VEC(4);
763 + break;
764 + case 6:
765 + LAUNCH_MOE_SUM_VEC(6);
766 + break;
767 + case 8:
768 + LAUNCH_MOE_SUM_VEC(8);
769 + break;
770 + case 9:
771 + LAUNCH_MOE_SUM_VEC(9);
772 + break;
773 + default:
774 + vllm::moe::moe_sum_vec_dynamic_kernel<scalar_t>
775 + <<<grid, dim3(block), 0, stream>>>(out_ptr, in_ptr, num_tokens,
776 + hidden_size, topk,
777 + stride_token, stride_topk);
778 + break;
779 + }
780 + } else {
781 + dim3 grid(num_tokens);
782 + dim3 block(std::min(hidden_size, 1024));
783 + vllm::moe::moe_sum_scalar_kernel<scalar_t><<<grid, block, 0, stream>>>(
784 + out_ptr, in_ptr, hidden_size, topk, stride_token, stride_topk,
785 + stride_hidden);
786 + }
787 + });
788 +#undef LAUNCH_MOE_SUM_VEC
789 }
790
791 void moe_lora_align_block_size(
tests/kernels/moe/test_moe.py
+16 -4
@@ -1243,15 +1243,27 @@ def test_batched_moe_align_block_size_opcheck():
1243 )
1244
1245
1246 +# topk=8 covers topk > 4; k=511 covers the non-vectorized scalar path. The
1247 +# layouts exercise contiguous input plus the two non-contiguous cases: a
1248 +# transpose (strided hidden -> scalar gather) and a topk-slice (hidden still
1249 +# contiguous -> vectorized).
1250 @pytest.mark.parametrize("m", [1, 33, 222])
1247 -@pytest.mark.parametrize("topk", TOP_KS)
1251 +@pytest.mark.parametrize("topk", [*TOP_KS, 8])
1252 @pytest.mark.parametrize("k", [128, 511, 1024])
1253 @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
1250 -def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype):
1251 - input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
1254 +@pytest.mark.parametrize("layout", ["contig", "transpose", "slice"])
1255 +def test_moe_sum(m: int, topk: int, k: int, dtype: torch.dtype, layout: str):
1256 + if layout == "transpose":
1257 + input = torch.randn((m, k, topk), device="cuda", dtype=dtype).transpose(1, 2)
1258 + elif layout == "slice":
1259 + input = torch.randn((m, 2 * topk, k), device="cuda", dtype=dtype)[:, ::2, :]
1260 + else:
1261 + input = torch.randn((m, topk, k), device="cuda", dtype=dtype)
1262 + assert input.is_contiguous() == (layout == "contig")
1263 actual = torch.empty((m, k), device="cuda", dtype=dtype)
1264
1254 - expected = input.sum(dim=1)
1265 + # Reduction accumulates in fp32.
1266 + expected = input.float().sum(dim=1).to(dtype)
1267 torch.ops._moe_C.moe_sum(input, actual)
1268
1269 torch.testing.assert_close(actual, expected, atol=2e-2, rtol=0)