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

[Quantization][Humming] Support MXFP4 weight + block-FP8 activation for MoE (#51332)

Signed-off-by: Elvir Crncevic <elvircrn@gmail.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Signed-off-by: Elvir Crnčević <elvircrn@gmail.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Elvir Crnčević committed Aug 25, 2026 at 16:21 UTC ca29cd48ba05ac9ae1c9fa29271e9f899bc4ce56
17 files changed +1599 -176
csrc/libtorch_stable/activation_kernels.cu
+401 -14
@@ -1,9 +1,17 @@
1 #include <cuda.h>
2 +#include <cuda_runtime.h>
3 #include <torch/csrc/stable/tensor.h>
4 +#include <torch/headeronly/util/Float8_e4m3fn.h>
5 +
6 +#ifndef USE_ROCM
7 + #include <cuda_fp8.h>
8 +#endif
9
10 #include <cmath>
11 +#include <type_traits>
12
13 #include "../cuda_compat.h"
14 +#include "async_util.cuh"
15 #include "cuda_vec_utils.cuh"
16 #include "dispatch_utils.h"
17 #include "torch_utils.h"
@@ -471,15 +479,31 @@ __global__ void swigluoai_and_mul_kernel(
479 // gate_out = beta * tanh(gate / beta) * sigmoid(gate)
480 // up_out = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up
481 // out = gate_out * up_out
474 -// Compute is done in fp32 and written straight to `out` -- no intermediate
475 -// tensors and no full-tensor fp32 upcast (the pure-torch forward_native
476 -// allocated ~8 fp32 temporaries per call, which blows up MoE profiling).
482 +__device__ __forceinline__ float situ_tanh(float x) { return tanhf(x); }
483 +
484 +// Kimi-K3 SITU params; baked into the fused LDG kernel to fold at compile time.
485 +static constexpr float SITU_BETA = 4.0f;
486 +static constexpr float SITU_LINEAR_BETA = 25.0f;
487 +__device__ __forceinline__ float situ_activation(float g, float u, float beta,
488 + float linear_beta,
489 + bool clamp_up, float inv_beta,
490 + float inv_linear_beta) {
491 + // sigmoid(g) == (1 + tanh(g/2)) / 2.
492 + const float gate_out =
493 + (0.5f * beta) * situ_tanh(g * inv_beta) * (1.0f + situ_tanh(g * 0.5f));
494 + const float up_out =
495 + clamp_up ? linear_beta * situ_tanh(u * inv_linear_beta) : u;
496 + return gate_out * up_out;
497 +}
498 +
499 template <typename scalar_t>
500 __global__ void situ_and_mul_kernel(
501 scalar_t* __restrict__ out, // [..., d]
502 const scalar_t* __restrict__ input, // [..., 2, d]
481 - const int d, const float beta, const float linear_beta) {
503 + const int d, const float beta, const float linear_beta,
504 + const int64_t* __restrict__ valid_rows_ptr) {
505 const int64_t row = blockIdx.x;
506 + if (valid_rows_ptr != nullptr && row >= *valid_rows_ptr) return;
507 const scalar_t* gate_ptr = input + row * 2 * d;
508 const scalar_t* up_ptr = gate_ptr + d;
509 scalar_t* out_ptr = out + row * d;
@@ -489,10 +513,258 @@ __global__ void situ_and_mul_kernel(
513 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
514 const float g = (float)VLLM_LDG(&gate_ptr[idx]);
515 const float u = (float)VLLM_LDG(&up_ptr[idx]);
492 - const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
493 - const float up_out =
494 - clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
495 - out_ptr[idx] = (scalar_t)(gate_out * up_out);
516 + out_ptr[idx] = (scalar_t)situ_activation(g, u, beta, linear_beta, clamp_up,
517 + inv_beta, inv_linear_beta);
518 + }
519 +}
520 +
521 +// Match Humming's hardware FP8 conversion. c10::Float8_e4m3fn's software cast
522 +// can round ties differently from `cvt.rn.satfinite.e4m3.f32`.
523 +__device__ __forceinline__ c10::Float8_e4m3fn cvt_fp8_hw(float x) {
524 +#if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
525 + return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3),
526 + c10::Float8_e4m3fn::from_bits());
527 +#else
528 + return static_cast<c10::Float8_e4m3fn>(x);
529 +#endif
530 +}
531 +
532 +template <typename fp8_type>
533 +__device__ __forceinline__ fp8_type quant_to_fp8(float val, float inv_scale,
534 + float fp8_max) {
535 + float x = val * inv_scale;
536 + x = fmaxf(-fp8_max, fminf(x, fp8_max));
537 + if constexpr (std::is_same_v<fp8_type, c10::Float8_e4m3fn>) {
538 + return cvt_fp8_hw(x);
539 + } else {
540 + return static_cast<fp8_type>(x);
541 + }
542 +}
543 +
544 +__device__ __forceinline__ float warp_reduce_max(float v) {
545 +#pragma unroll
546 + for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
547 + v = fmaxf(v, VLLM_SHFL_XOR_SYNC(v, offset));
548 + }
549 + return v;
550 +}
551 +
552 +__device__ __forceinline__ float warp_reduce_absmax(float v) {
553 + return warp_reduce_max(v);
554 +}
555 +
556 +template <typename scalar_t, typename fp8_type, int THREADS, int NUM_STAGES,
557 + int GROUP_SIZE, int GRID_DIM, int D>
558 +__global__ void situ_and_mul_quant_group_pipelined_kernel(
559 + fp8_type* __restrict__ out, float* __restrict__ scale_out,
560 + const scalar_t* __restrict__ input, const int64_t num_rows,
561 + const int32_t* __restrict__ valid_rows_ptr, const int64_t topk) {
562 + const int tid = threadIdx.x;
563 + const int warp_id = tid / WARP_SIZE;
564 + const int lane_id = tid % WARP_SIZE;
565 + static constexpr int NUM_WARPS = THREADS / WARP_SIZE;
566 + static constexpr int NUM_GROUPS = D / GROUP_SIZE;
567 + const int64_t row_bound =
568 + valid_rows_ptr != nullptr
569 + ? max((int64_t)0, min((int64_t)(*valid_rows_ptr) * topk, num_rows))
570 + : num_rows;
571 +
572 + static_assert(NUM_GROUPS % 4 == 0,
573 + "float4 scale fill needs NUM_GROUPS % 4 == 0");
574 + static constexpr int NG4 = NUM_GROUPS / 4;
575 + const int64_t pad_start4 = row_bound * NG4;
576 + const int64_t pad_end4 = num_rows * NG4;
577 + float4* scale4 = reinterpret_cast<float4*>(scale_out);
578 + constexpr float4 ones{1.0f, 1.0f, 1.0f, 1.0f};
579 + for (int64_t i = pad_start4 + (int64_t)blockIdx.x * THREADS + tid;
580 + i < pad_end4; i += (int64_t)GRID_DIM * THREADS) {
581 + scale4[i] = ones;
582 + }
583 +
584 + if constexpr (sizeof(scalar_t) == 2) {
585 + static constexpr int LD_ELTS = 16 / sizeof(scalar_t);
586 + static constexpr int ELTS_PER_LANE = GROUP_SIZE / WARP_SIZE;
587 + static constexpr int STAGE_ELTS = 2 * GROUP_SIZE;
588 +
589 + extern __shared__ __align__(16) unsigned char smem_raw[];
590 + scalar_t* warp_smem = reinterpret_cast<scalar_t*>(smem_raw) +
591 + (size_t)warp_id * NUM_STAGES * STAGE_ELTS;
592 +
593 + static constexpr float beta = SITU_BETA;
594 + static constexpr float linear_beta = SITU_LINEAR_BETA;
595 + static constexpr bool clamp_up = linear_beta > 0.0f;
596 + static constexpr float inv_beta = 1.0f / beta;
597 + static constexpr float inv_linear_beta =
598 + clamp_up ? 1.0f / linear_beta : 0.0f;
599 + static constexpr float FP8_MAX =
600 + std::is_same_v<fp8_type, c10::Float8_e4m3fn> ? 448.0f : 224.0f;
601 + static constexpr float inv_fp8_max = 1.0f / FP8_MAX;
602 +
603 + static_assert(
604 + NUM_GROUPS % NUM_WARPS == 0,
605 + "constexpr num_iters requires groups evenly split across warps");
606 + static constexpr int num_iters = NUM_GROUPS / NUM_WARPS;
607 +
608 + const bool up_half = lane_id >= WARP_SIZE / 2;
609 + const int lane_l = up_half ? lane_id - WARP_SIZE / 2 : lane_id;
610 + const int lane_src_off = (up_half ? D : 0) + lane_l * LD_ELTS;
611 + const int lane_dst_off = (up_half ? GROUP_SIZE : 0) + lane_l * LD_ELTS;
612 + static constexpr int warp_stride = NUM_WARPS * GROUP_SIZE;
613 +
614 + const scalar_t* src_ptr = input + warp_id * GROUP_SIZE + lane_src_off;
615 + static constexpr int src_row_stride = GRID_DIM * 2 * D;
616 + const scalar_t* row_src = src_ptr + blockIdx.x * 2 * D;
617 +
618 + uint32_t* out_ptr =
619 + reinterpret_cast<uint32_t*>(out + warp_id * GROUP_SIZE) + lane_id;
620 + static constexpr int out_row_stride = GRID_DIM * D / 4;
621 + uint32_t* row_out = out_ptr + blockIdx.x * (D / 4);
622 + float* scale_ptr = scale_out + warp_id;
623 + static constexpr int scale_row_stride = GRID_DIM * NUM_GROUPS;
624 + float* row_scale = scale_ptr + blockIdx.x * NUM_GROUPS;
625 +
626 +#pragma unroll 1
627 + for (int64_t row = blockIdx.x; row < row_bound; row += GRID_DIM,
628 + row_src += src_row_stride, row_out += out_row_stride,
629 + row_scale += scale_row_stride) {
630 + const scalar_t* src = row_src;
631 +
632 + auto issue_load = [&](int it, int slot) {
633 + if (it < num_iters) {
634 + cuda_async::cp_async_shared_global_16_cg(
635 + warp_smem + (size_t)slot * STAGE_ELTS + lane_dst_off, src);
636 + src += warp_stride;
637 + }
638 + cuda_async::cp_async_commit_group();
639 + };
640 +
641 + int load_slot = 0;
642 + auto bump = [](int s) { return s + 1 == NUM_STAGES ? 0 : s + 1; };
643 +#pragma unroll
644 + for (int s = 0; s < NUM_STAGES - 1; s++) {
645 + issue_load(s, load_slot);
646 + load_slot = bump(load_slot);
647 + }
648 +
649 + uint32_t* out_st = row_out;
650 + float* scale_st = row_scale;
651 +
652 + int comp_slot = 0;
653 + for (int it = 0; it < num_iters; it++) {
654 + issue_load(it + NUM_STAGES - 1, load_slot);
655 + load_slot = bump(load_slot);
656 + cuda_async::cp_async_wait_group<NUM_STAGES - 1>();
657 +
658 + const scalar_t* stage = warp_smem + (size_t)comp_slot * STAGE_ELTS;
659 + comp_slot = bump(comp_slot);
660 + static_assert(ELTS_PER_LANE == 4,
661 + "expects GROUP_SIZE == 4 * WARP_SIZE");
662 + union V {
663 + float2 f2;
664 + scalar_t s[ELTS_PER_LANE];
665 + };
666 + const float2* gate2 = reinterpret_cast<const float2*>(stage);
667 + const float2* up2 = reinterpret_cast<const float2*>(stage + GROUP_SIZE);
668 + const V gv{gate2[lane_id]};
669 + const V uv{up2[lane_id]};
670 + const scalar_t* gs = gv.s;
671 + const scalar_t* us = uv.s;
672 +
673 + float acts[ELTS_PER_LANE];
674 + float thread_max = 0.0f;
675 +#pragma unroll
676 + for (int e = 0; e < ELTS_PER_LANE; e++) {
677 + acts[e] = (float)(scalar_t)situ_activation(
678 + (float)gs[e], (float)us[e], beta, linear_beta, clamp_up, inv_beta,
679 + inv_linear_beta);
680 + thread_max = fmaxf(thread_max, fabsf(acts[e]));
681 + }
682 + const float absmax = fmaxf(warp_reduce_absmax(thread_max), 1e-30f);
683 + const float scale = absmax * inv_fp8_max;
684 + if (lane_id == 0) *scale_st = scale;
685 + scale_st += NUM_WARPS;
686 + const float inv_scale = __fdividef(1.0f, scale);
687 +
688 + union O {
689 + uint32_t u;
690 + fp8_type f[ELTS_PER_LANE];
691 + };
692 + O o;
693 +#pragma unroll
694 + for (int e = 0; e < ELTS_PER_LANE; e++) {
695 + o.f[e] = quant_to_fp8<fp8_type>(acts[e], inv_scale, FP8_MAX);
696 + }
697 + out_st[0] = o.u;
698 + out_st += NUM_WARPS * GROUP_SIZE / 4;
699 + }
700 + cuda_async::cp_async_wait_group<0>();
701 + }
702 + }
703 +}
704 +
705 +// Scalar fallback for the group path (odd d or the fp32 dispatch branch).
706 +template <typename scalar_t, typename fp8_type, int GROUP_SIZE>
707 +__global__ void situ_and_mul_quant_group_scalar_kernel(
708 + fp8_type* __restrict__ out, // [num_tokens, d]
709 + float* __restrict__ scale_out, // [num_tokens, num_groups]
710 + const scalar_t* __restrict__ input, // [num_tokens, 2, d]
711 + const int d, const int num_groups, const float beta,
712 + const float linear_beta, const int64_t num_rows,
713 + const int32_t* __restrict__ valid_rows_ptr, const int64_t topk) {
714 + static constexpr int ELTS_PER_LANE = GROUP_SIZE / WARP_SIZE;
715 + const int tid = threadIdx.x;
716 + const int warp_id = tid / WARP_SIZE;
717 + const int lane_id = tid % WARP_SIZE;
718 + const int num_warps = blockDim.x / WARP_SIZE;
719 + const int64_t row_bound =
720 + valid_rows_ptr != nullptr
721 + ? max((int64_t)0, min((int64_t)(*valid_rows_ptr) * topk, num_rows))
722 + : num_rows;
723 +
724 + const bool clamp_up = linear_beta > 0.0f;
725 + const float inv_beta = __fdividef(1.0f, beta);
726 + const float inv_linear_beta = clamp_up ? __fdividef(1.0f, linear_beta) : 0.0f;
727 + static constexpr float FP8_MAX =
728 + std::is_same_v<fp8_type, c10::Float8_e4m3fn> ? 448.0f : 224.0f;
729 + static constexpr float inv_fp8_max = 1.0f / FP8_MAX;
730 +
731 + for (int64_t row = blockIdx.x; row < row_bound; row += gridDim.x) {
732 + const scalar_t* gate_ptr = input + row * 2 * (int64_t)d;
733 + const scalar_t* up_ptr = gate_ptr + d;
734 + fp8_type* out_ptr = out + row * (int64_t)d;
735 + float* scale_row = scale_out + row * (int64_t)num_groups;
736 +
737 + for (int g = warp_id; g < num_groups; g += num_warps) {
738 + const int base = g * GROUP_SIZE;
739 + float acts[ELTS_PER_LANE];
740 + float thread_max = 0.0f;
741 +#pragma unroll
742 + for (int e = 0; e < ELTS_PER_LANE; e++) {
743 + const int idx = base + e * WARP_SIZE + lane_id;
744 + const float gv = (float)VLLM_LDG(&gate_ptr[idx]);
745 + const float uv = (float)VLLM_LDG(&up_ptr[idx]);
746 + acts[e] = (float)(scalar_t)situ_activation(
747 + gv, uv, beta, linear_beta, clamp_up, inv_beta, inv_linear_beta);
748 + thread_max = fmaxf(thread_max, fabsf(acts[e]));
749 + }
750 + const float absmax = fmaxf(warp_reduce_max(thread_max), 1e-30f);
751 + const float scale = absmax * inv_fp8_max;
752 + if (lane_id == 0) scale_row[g] = scale;
753 + const float inv_scale = __fdividef(1.0f, scale);
754 +#pragma unroll
755 + for (int e = 0; e < ELTS_PER_LANE; e++) {
756 + const int idx = base + e * WARP_SIZE + lane_id;
757 + out_ptr[idx] = quant_to_fp8<fp8_type>(acts[e], inv_scale, FP8_MAX);
758 + }
759 + }
760 + }
761 +
762 + // Fill skipped padding-row scales with 1 so they don't feed NaN into w2.
763 + const int64_t pad_start = row_bound * (int64_t)num_groups;
764 + const int64_t pad_end = num_rows * (int64_t)num_groups;
765 + for (int64_t i = pad_start + (int64_t)blockIdx.x * blockDim.x + tid;
766 + i < pad_end; i += (int64_t)gridDim.x * blockDim.x) {
767 + scale_out[i] = 1.0f;
768 }
769 }
770
@@ -519,10 +791,8 @@ __global__ void masked_situ_and_mul_kernel(
791 scalar_t* out_ptr = out + row * d;
792 const float g = (float)VLLM_LDG(&gate_ptr[idx]);
793 const float u = (float)VLLM_LDG(&up_ptr[idx]);
522 - const float gate_out = beta * tanhf(g * inv_beta) / (1.0f + expf(-g));
523 - const float up_out =
524 - clamp_up ? linear_beta * tanhf(u * inv_linear_beta) : u;
525 - out_ptr[idx] = (scalar_t)(gate_out * up_out);
794 + out_ptr[idx] = (scalar_t)situ_activation(g, u, beta, linear_beta, clamp_up,
795 + inv_beta, inv_linear_beta);
796 }
797 }
798
@@ -620,7 +890,8 @@ void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d]
890 // through), matching SituAndMul(linear_beta=None) on the Python side.
891 void situ_and_mul(torch::stable::Tensor& out, // [..., d]
892 torch::stable::Tensor& input, // [..., 2 * d]
623 - double beta, double linear_beta) {
893 + double beta, double linear_beta,
894 + std::optional<torch::stable::Tensor> valid_rows) {
895 int d = input.size(-1) / 2;
896 int64_t num_tokens = input.numel() / input.size(-1);
897 if (num_tokens == 0) {
@@ -628,6 +899,14 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d]
899 }
900 dim3 grid(num_tokens);
901 dim3 block(std::min(d, 1024));
902 + const int64_t* valid_rows_ptr = nullptr;
903 + if (valid_rows.has_value()) {
904 + STD_TORCH_CHECK(valid_rows->is_cuda() && valid_rows->get_device_index() ==
905 + input.get_device_index(),
906 + "situ_and_mul: valid_rows must be a CUDA tensor on the "
907 + "same device as input");
908 + valid_rows_ptr = valid_rows->const_data_ptr<int64_t>();
909 + }
910 const torch::stable::accelerator::DeviceGuard device_guard(
911 input.get_device_index());
912 const cudaStream_t stream = get_current_cuda_stream();
@@ -635,7 +914,7 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d]
914 input.scalar_type(), "situ_and_mul_kernel", [&] {
915 vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
916 out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
638 - d, (float)beta, (float)linear_beta);
917 + d, (float)beta, (float)linear_beta, valid_rows_ptr);
918 });
919 }
920
@@ -663,6 +942,114 @@ void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d]
942 (float)beta, (float)linear_beta);
943 });
944 }
945 +
946 +// Fused SITU activation and block-FP8 quantization for Humming w2.
947 +// `valid_rows` is the int32 DeepEP token count (psum[-1]); rows = it * `topk`,
948 +// so padding is excluded on-device and skipped rows receive scale 1.
949 +void situ_and_mul_quant(torch::stable::Tensor& out, // [..., d] (fp8)
950 + torch::stable::Tensor& scale, // [..., 1 or d/128]
951 + torch::stable::Tensor& input, // [..., 2 * d]
952 + double beta, double linear_beta, int64_t group_size,
953 + std::optional<torch::stable::Tensor> valid_rows,
954 + int64_t topk) {
955 + STD_TORCH_CHECK(
956 + out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
957 + out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz,
958 + "situ_and_mul_quant output must be FP8 e4m3");
959 + STD_TORCH_CHECK(
960 + input.scalar_type() == torch::headeronly::ScalarType::Half ||
961 + input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
962 + "situ_and_mul_quant input must be FP16 or BF16");
963 + STD_TORCH_CHECK(scale.scalar_type() == torch::headeronly::ScalarType::Float,
964 + "situ_and_mul_quant scale must be float32");
965 + STD_TORCH_CHECK(input.is_cuda() && out.is_cuda() && scale.is_cuda() &&
966 + out.get_device_index() == input.get_device_index() &&
967 + scale.get_device_index() == input.get_device_index(),
968 + "situ_and_mul_quant: input, out and scale must be CUDA "
969 + "tensors on the same device");
970 + int d = input.size(-1) / 2;
971 + int64_t num_tokens = input.numel() / input.size(-1);
972 + if (num_tokens == 0) {
973 + return;
974 + }
975 + STD_TORCH_CHECK(out.size(-1) == d && out.numel() == num_tokens * (int64_t)d,
976 + "situ_and_mul_quant: out shape must be [num_tokens, d]");
977 + const int32_t* valid_rows_ptr = nullptr;
978 + if (valid_rows.has_value()) {
979 + STD_TORCH_CHECK(valid_rows->is_cuda() && valid_rows->get_device_index() ==
980 + input.get_device_index(),
981 + "situ_and_mul_quant: valid_rows must be a CUDA tensor on "
982 + "the same device as input");
983 + STD_TORCH_CHECK(
984 + valid_rows->scalar_type() == torch::headeronly::ScalarType::Int,
985 + "situ_and_mul_quant: valid_rows must be int32 (psum count)");
986 + valid_rows_ptr = valid_rows->const_data_ptr<int32_t>();
987 + }
988 + const torch::stable::accelerator::DeviceGuard device_guard(
989 + input.get_device_index());
990 + const cudaStream_t stream = get_current_cuda_stream();
991 + static constexpr int BLOCKS_PER_SM = 8; // matches kernel __launch_bounds__
992 + static constexpr int SM_COUNT = 132; // H200 (GH100, 132 SMs)
993 + static constexpr int GRID_DIM = SM_COUNT * BLOCKS_PER_SM;
994 +
995 + STD_TORCH_CHECK(group_size == 128,
996 + "situ_and_mul_quant: only group_size 128 (block-FP8) "
997 + "supported, got ",
998 + group_size);
999 + STD_TORCH_CHECK(d % group_size == 0, "situ_and_mul_quant: d (", d,
1000 + ") must be divisible by group_size ", group_size);
1001 + {
1002 + const int num_groups = d / (int)group_size;
1003 + STD_TORCH_CHECK(scale.size(-1) == num_groups &&
1004 + scale.numel() == num_tokens * (int64_t)num_groups,
1005 + "situ_and_mul_quant: scale shape must be "
1006 + "[num_tokens, d/group_size]");
1007 + dim3 grid(GRID_DIM);
1008 + VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1009 + input.scalar_type(), "situ_and_mul_quant_group_kernel", [&] {
1010 + VLLM_STABLE_DISPATCH_FP8_TYPES(
1011 + out.scalar_type(), "situ_and_mul_quant_group_kernel_fp8", [&] {
1012 +#ifndef USE_ROCM
1013 + // The pipelined kernel's float2-per-lane geometry assumes a
1014 + // 32-lane warp (GROUP_SIZE == 4 * WARP_SIZE); on HIP (64-lane)
1015 + // fall back to the WARP_SIZE-generic scalar kernel.
1016 + constexpr int THREADS = 256;
1017 + constexpr int SITU_D = 3072; // Kimi-K3 fused w2 input dim
1018 + if constexpr (sizeof(scalar_t) == 2) {
1019 + if (d == SITU_D && (float)beta == vllm::SITU_BETA &&
1020 + (float)linear_beta == vllm::SITU_LINEAR_BETA) {
1021 + constexpr int D = SITU_D;
1022 + constexpr int GROUP_STAGES = 4;
1023 + constexpr int NUM_WARPS = THREADS / 32;
1024 + constexpr int STAGE_ELTS = 2 * 128;
1025 + dim3 block(THREADS);
1026 + size_t smem_bytes = (size_t)NUM_WARPS * GROUP_STAGES *
1027 + STAGE_ELTS * sizeof(scalar_t);
1028 + vllm::situ_and_mul_quant_group_pipelined_kernel<
1029 + scalar_t, fp8_t, THREADS, GROUP_STAGES, 128, GRID_DIM,
1030 + D><<<grid, block, smem_bytes, stream>>>(
1031 + out.mutable_data_ptr<fp8_t>(),
1032 + scale.mutable_data_ptr<float>(),
1033 + input.const_data_ptr<scalar_t>(), num_tokens,
1034 + valid_rows_ptr, topk);
1035 + return;
1036 + }
1037 + }
1038 +#endif
1039 + const int num_warps = std::min(num_groups, 1024 / WARP_SIZE);
1040 + dim3 block(num_warps * WARP_SIZE);
1041 + vllm::situ_and_mul_quant_group_scalar_kernel<scalar_t, fp8_t,
1042 + 128>
1043 + <<<grid, block, 0, stream>>>(
1044 + out.mutable_data_ptr<fp8_t>(),
1045 + scale.mutable_data_ptr<float>(),
1046 + input.const_data_ptr<scalar_t>(), d, num_groups,
1047 + (float)beta, (float)linear_beta, num_tokens,
1048 + valid_rows_ptr, topk);
1049 + });
1050 + });
1051 + }
1052 +}
1053 namespace vllm {
1054
1055 // Element-wise activation kernel template.
csrc/libtorch_stable/ops.h
+10 -2
@@ -526,8 +526,16 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
526 double threshold);
527 void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
528 double alpha = 1.702, double limit = 7.0);
529 -void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
530 - double beta = 1.0, double linear_beta = -1.0);
529 +void situ_and_mul(
530 + torch::stable::Tensor& out, torch::stable::Tensor& input, double beta = 1.0,
531 + double linear_beta = -1.0,
532 + std::optional<torch::stable::Tensor> valid_rows = std::nullopt);
533 +void situ_and_mul_quant(
534 + torch::stable::Tensor& out, torch::stable::Tensor& scale,
535 + torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0,
536 + int64_t group_size = 0,
537 + std::optional<torch::stable::Tensor> valid_rows = std::nullopt,
538 + int64_t topk = 1);
539 void masked_situ_and_mul(torch::stable::Tensor& out,
540 torch::stable::Tensor& input,
541 const torch::stable::Tensor& expert_num_tokens,
csrc/libtorch_stable/torch_bindings.cpp
+10 -1
@@ -611,7 +611,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
611 // SituGLU implementation used in Kimi models.
612 ops.def(
613 "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float "
614 - "linear_beta=-1.0) -> ()");
614 + "linear_beta=-1.0, Tensor? valid_rows=None) -> ()");
615 + // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path
616 + // (writes the fp8 down input and its float32 scale). group_size=0 ->
617 + // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [..,
618 + // d/128].
619 + ops.def(
620 + "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, "
621 + "float beta=1.0, float linear_beta=-1.0, int group_size=0, "
622 + "Tensor? valid_rows=None, int topk=1) -> ()");
623 ops.def(
624 "masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
625 "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
@@ -830,6 +838,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
838 ops.impl("fatrelu_and_mul", TORCH_BOX(&fatrelu_and_mul));
839 ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul));
840 ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
841 + ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant));
842 ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
843 ops.impl("gelu_new", TORCH_BOX(&gelu_new));
844 ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
tests/evals/gsm8k/configs/humming/Qwen3-30B-A3B-MXFP4A16-humming-act-fp8-block.yaml new
+11
@@ -0,0 +1,11 @@
1 +model_name: "nm-testing/Qwen3-30B-A3B-MXFP4A16"
2 +accuracy_threshold: 0.86
3 +num_questions: 1319
4 +num_fewshot: 5
5 +server_args: >-
6 + --enforce-eager
7 + --max-model-len 8192
8 + --quantization humming
9 + --linear-backend humming
10 +env:
11 + VLLM_HUMMING_INPUT_QUANT_CONFIG: '{"dtype":"float8e4m3","group_size":128}'
tests/evals/gsm8k/configs/humming/config-act-fp8.txt
+1
@@ -1,5 +1,6 @@
1 gpt-oss-20b-humming-act-fp8.yaml
2 Qwen3-30B-A3B-MXFP4A16-humming-act-fp8.yaml
3 +Qwen3-30B-A3B-MXFP4A16-humming-act-fp8-block.yaml
4 Qwen2-1.5B-Instruct-FP8W8-humming-act-fp8.yaml
5 Qwen3-0.6B-MXFP8-humming-act-fp8.yaml
6 Qwen3-30B-A3B-Fp8-v1-humming-act-fp8.yaml
tests/kernels/moe/test_moe.py
+1
@@ -1408,6 +1408,7 @@ def test_humming_delegates_to_instance_activation():
1408 activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
1409 input=input,
1410 output=output,
1411 + valid_rows=None,
1412 )
1413
1414
tests/kernels/moe/test_moe_fused_mul_sum.py new
+406
@@ -0,0 +1,406 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the fused MoE topk weight-and-reduce kernel.
4 +
5 +Run `pytest tests/kernels/moe/test_moe_fused_mul_sum.py`.
6 +"""
7 +
8 +import pytest
9 +import torch
10 +
11 +from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import (
12 + moe_fused_mul_sum,
13 +)
14 +from vllm.utils.torch_utils import set_random_seed
15 +
16 +NUM_TOKENS = [1, 17, 256]
17 +TOP_KS = [2, 8]
18 +HIDDEN_SIZE = 512
19 +# One BLOCK_K tile, a partial trailing tile, and Kimi-K3's real dims
20 +# (moe_intermediate_size 3072, hidden_size 7168) -- exercises the per-tile loop
21 +# and its tail mask (BLOCK_K caps at 512 for <=2-byte dtypes).
22 +HIDDEN_SIZES = [512, 1000, 3072, 7168]
23 +DTYPES = [torch.bfloat16, torch.float32]
24 +NUM_EXPERTS = 16
25 +
26 +
27 +def _reference(
28 + inputs: torch.Tensor,
29 + topk_weights: torch.Tensor,
30 + topk_ids: torch.Tensor,
31 + expert_map: torch.Tensor | None,
32 +) -> torch.Tensor:
33 + valid = topk_ids >= 0
34 + if expert_map is not None:
35 + valid &= expert_map[topk_ids.clamp(min=0)] >= 0
36 + weights = torch.where(valid, topk_weights, torch.zeros_like(topk_weights))
37 + return (inputs.float() * weights.float().unsqueeze(-1)).sum(dim=1)
38 +
39 +
40 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
41 +@pytest.mark.parametrize("top_k", TOP_KS)
42 +@pytest.mark.parametrize("use_expert_map", [True, False])
43 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
44 +def test_invalid_slots_excluded(num_tokens: int, top_k: int, use_expert_map: bool):
45 + """Slots marked -1 must not contribute, and must not be dereferenced.
46 +
47 + The expert GEMM never writes rows for topk slots whose expert id is -1
48 + (a non-local expert under EP, or an all2all padding row), so those rows
49 + hold stale workspace data. Filling them with NaN here means any failure to
50 + mask them shows up as NaN in the output rather than a small numeric drift.
51 + Indexing `expert_map` with -1 is also an out-of-bounds read.
52 + """
53 + set_random_seed(0)
54 + device = "cuda"
55 +
56 + inputs = torch.randn(
57 + num_tokens, top_k, HIDDEN_SIZE, dtype=torch.bfloat16, device=device
58 + )
59 + topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device)
60 + topk_ids = torch.randint(
61 + 0, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device
62 + )
63 +
64 + # Mark roughly half the slots invalid and poison the rows behind them.
65 + invalid = torch.rand(num_tokens, top_k, device=device) < 0.5
66 + topk_ids[invalid] = -1
67 + inputs[invalid] = float("nan")
68 +
69 + expert_map = None
70 + if use_expert_map:
71 + # Second half of the experts is non-local, so valid ids can still map
72 + # to -1 -- exercising the second masking condition.
73 + expert_map = torch.full((NUM_EXPERTS,), -1, dtype=torch.int32, device=device)
74 + local = NUM_EXPERTS // 2
75 + expert_map[:local] = torch.arange(local, dtype=torch.int32, device=device)
76 + inputs[(topk_ids >= 0) & (expert_map[topk_ids.clamp(min=0)] < 0)] = float("nan")
77 +
78 + # The _reference implementation zeros out outputs. We need to match that here,
79 + # otherwise the test fails.
80 + outputs = torch.zeros(num_tokens, HIDDEN_SIZE, dtype=torch.bfloat16, device=device)
81 + out = moe_fused_mul_sum(
82 + inputs=inputs,
83 + topk_weights=topk_weights,
84 + outputs=outputs,
85 + topk_ids=topk_ids,
86 + expert_map=expert_map,
87 + )
88 +
89 + assert not out.isnan().any(), "invalid slots leaked into the reduction"
90 + ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map)
91 + torch.testing.assert_close(out.float(), ref, atol=1e-2, rtol=1e-2)
92 +
93 +
94 +def _local_expert_map(device: str) -> torch.Tensor:
95 + """First half of the experts is local, second half maps to -1 (non-local)."""
96 + expert_map = torch.full((NUM_EXPERTS,), -1, dtype=torch.int32, device=device)
97 + local = NUM_EXPERTS // 2
98 + expert_map[:local] = torch.arange(local, dtype=torch.int32, device=device)
99 + return expert_map
100 +
101 +
102 +SENTINEL = 42.0 # exact in bf16; marks output the kernel must leave untouched
103 +
104 +
105 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
106 +def test_nonlocal_row_is_zeroed():
107 + """A real row whose positive ids all map to non-local experts is written 0.
108 +
109 + Such a row is not padding (its ids are >= 0, so row_valid > 0 and the block
110 + is not early-returned). Every slot is then masked out by expert_map, so the
111 + accumulator stays 0 and 0 is stored -- the row must not be left holding stale
112 + buffer data, because downstream reduction treats it as a real token.
113 + """
114 + set_random_seed(0)
115 + device = "cuda"
116 + num_tokens, top_k = 17, 8
117 + local = NUM_EXPERTS // 2
118 + expert_map = _local_expert_map(device)
119 +
120 + # Poison inputs so any failure to mask the non-local slots surfaces as NaN.
121 + inputs = torch.full(
122 + (num_tokens, top_k, HIDDEN_SIZE),
123 + float("nan"),
124 + dtype=torch.bfloat16,
125 + device=device,
126 + )
127 + topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device)
128 + # Every id lands in the non-local half [local, NUM_EXPERTS).
129 + topk_ids = torch.randint(
130 + local, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device
131 + )
132 +
133 + outputs = torch.full(
134 + (num_tokens, HIDDEN_SIZE), SENTINEL, dtype=torch.bfloat16, device=device
135 + )
136 + out = moe_fused_mul_sum(
137 + inputs=inputs,
138 + topk_weights=topk_weights,
139 + outputs=outputs,
140 + topk_ids=topk_ids,
141 + expert_map=expert_map,
142 + )
143 +
144 + assert not out.isnan().any(), "non-local row leaked NaN into the output"
145 + torch.testing.assert_close(out.float(), torch.zeros_like(out.float()))
146 +
147 +
148 +@pytest.mark.parametrize("use_expert_map", [True, False])
149 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
150 +def test_padding_row_left_untouched(use_expert_map: bool):
151 + """An all -1 padding row is skipped: its output row is left untouched.
152 +
153 + This is the CUDA-graph decode contract -- padding rows past num_recv are
154 + never read downstream, so the kernel elides them (whole-padding blocks
155 + early-return) rather than zeroing. The skip is gated on topk_ids alone, so
156 + it must hold whether or not expert_map is passed (the humming path passes
157 + None); we pin both.
158 + """
159 + set_random_seed(0)
160 + device = "cuda"
161 + num_tokens, top_k = 17, 8
162 + expert_map = _local_expert_map(device) if use_expert_map else None
163 +
164 + inputs = torch.full(
165 + (num_tokens, top_k, HIDDEN_SIZE),
166 + float("nan"),
167 + dtype=torch.bfloat16,
168 + device=device,
169 + )
170 + topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device)
171 + topk_ids = torch.full((num_tokens, top_k), -1, dtype=torch.int32, device=device)
172 +
173 + outputs = torch.full(
174 + (num_tokens, HIDDEN_SIZE), SENTINEL, dtype=torch.bfloat16, device=device
175 + )
176 + out = moe_fused_mul_sum(
177 + inputs=inputs,
178 + topk_weights=topk_weights,
179 + outputs=outputs,
180 + topk_ids=topk_ids,
181 + expert_map=expert_map,
182 + )
183 +
184 + assert torch.all(out == SENTINEL), "padding row output was modified"
185 +
186 +
187 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
188 +def test_nonlocal_and_padding_adjacent():
189 + """A non-local real row next to an all -1 padding row.
190 +
191 + Each CTA owns one token row: the non-local row must be zeroed (its slots are
192 + all masked by expert_map) while the padding row beside it is left untouched.
193 + Guards against a change that lets one row's fate leak into its neighbour.
194 + """
195 + set_random_seed(0)
196 + device = "cuda"
197 + num_tokens, top_k = 17, 8
198 + local = NUM_EXPERTS // 2
199 + expert_map = _local_expert_map(device)
200 +
201 + inputs = torch.full(
202 + (num_tokens, top_k, HIDDEN_SIZE),
203 + float("nan"),
204 + dtype=torch.bfloat16,
205 + device=device,
206 + )
207 + topk_weights = torch.rand(num_tokens, top_k, dtype=torch.bfloat16, device=device)
208 + # All rows are non-local real rows except row 1, an all -1 padding row.
209 + topk_ids = torch.randint(
210 + local, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device
211 + )
212 + topk_ids[1] = -1
213 +
214 + outputs = torch.full(
215 + (num_tokens, HIDDEN_SIZE), SENTINEL, dtype=torch.bfloat16, device=device
216 + )
217 + out = moe_fused_mul_sum(
218 + inputs=inputs,
219 + topk_weights=topk_weights,
220 + outputs=outputs,
221 + topk_ids=topk_ids,
222 + expert_map=expert_map,
223 + )
224 +
225 + padding = torch.zeros(num_tokens, dtype=torch.bool, device=device)
226 + padding[1] = True
227 + assert torch.all(out[1] == SENTINEL), "padding row beside a real row was modified"
228 + written = out[~padding].float()
229 + assert not written.isnan().any(), "non-local row leaked NaN into the output"
230 + torch.testing.assert_close(written, torch.zeros_like(written))
231 +
232 +
233 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
234 +@pytest.mark.parametrize("top_k", TOP_KS)
235 +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
236 +@pytest.mark.parametrize("dtype", DTYPES)
237 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
238 +def test_weighted_sum_matches_reference(
239 + num_tokens: int, top_k: int, hidden_size: int, dtype: torch.dtype
240 +):
241 + """Weighted sum over live slots matches the reference across hidden tiles.
242 +
243 + Covers the deployed humming path (globalized ids, expert_map=None): a random
244 + -1 mix with NaN-poisoned dead slots, over hidden sizes that span a single
245 + tile, a partial trailing tile, and many tiles, in both bf16 and fp32.
246 + """
247 + set_random_seed(0)
248 + device = "cuda"
249 +
250 + inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device)
251 + topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device)
252 + topk_ids = torch.randint(
253 + 0, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device
254 + )
255 + invalid = torch.rand(num_tokens, top_k, device=device) < 0.5
256 + topk_ids[invalid] = -1
257 + inputs[invalid] = float("nan")
258 +
259 + # All-padding rows early-return untouched, so pre-zero to match the reference.
260 + outputs = torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device)
261 + out = moe_fused_mul_sum(
262 + inputs=inputs,
263 + topk_weights=topk_weights,
264 + outputs=outputs,
265 + topk_ids=topk_ids,
266 + expert_map=None,
267 + )
268 +
269 + assert not out.isnan().any(), "invalid slots leaked into the reduction"
270 + ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map=None)
271 + atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4)
272 + torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol)
273 +
274 +
275 +@pytest.mark.parametrize("dtype", DTYPES)
276 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
277 +def test_no_topk_ids_sums_all_slots(dtype: torch.dtype):
278 + """With topk_ids=None every slot is live: plain weighted sum over top_k."""
279 + set_random_seed(0)
280 + device = "cuda"
281 + num_tokens, top_k, hidden_size = 17, 8, 1000
282 +
283 + inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device)
284 + topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device)
285 + out = moe_fused_mul_sum(inputs=inputs, topk_weights=topk_weights)
286 +
287 + ref = (inputs.float() * topk_weights.float().unsqueeze(-1)).sum(dim=1)
288 + atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4)
289 + torch.testing.assert_close(out.float(), ref, atol=atol, rtol=rtol)
290 +
291 +
292 +@pytest.mark.parametrize("hidden_size", [512, 7168])
293 +@pytest.mark.parametrize("dtype", DTYPES)
294 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
295 +def test_expert_map_is_required_for_nonlocal_ids(hidden_size: int, dtype: torch.dtype):
296 + """expert_map is load-bearing, not redundant: dropping it leaks non-local ids.
297 +
298 + A positive id whose expert is non-local marks a slot the GEMM never wrote, so
299 + its `inputs` row is stale. `topk_ids >= 0` alone cannot tell it apart from a
300 + real slot -- only expert_map can. This pins the contract the humming path
301 + relies on (globalized ids are still routed through expert_map as a safety
302 + net): with the map the stale slot is masked; without it, it corrupts the sum.
303 + Regression guard against re-dropping expert_map as an optimization.
304 + """
305 + set_random_seed(0)
306 + device = "cuda"
307 + num_tokens, top_k = 8, 8
308 + local = NUM_EXPERTS // 2
309 + expert_map = _local_expert_map(device)
310 +
311 + inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device)
312 + topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device)
313 + # Each row keeps at least one local slot (so the mapped output is finite) and
314 + # at least one non-local slot whose stale row is poisoned with NaN.
315 + half = top_k // 2
316 + topk_ids = torch.empty(num_tokens, top_k, dtype=torch.int32, device=device)
317 + topk_ids[:, :half] = torch.randint(
318 + 0, local, (num_tokens, half), dtype=torch.int32, device=device
319 + )
320 + topk_ids[:, half:] = torch.randint(
321 + local, NUM_EXPERTS, (num_tokens, top_k - half), dtype=torch.int32, device=device
322 + )
323 + inputs[:, half:] = float("nan")
324 +
325 + atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4)
326 +
327 + mapped = moe_fused_mul_sum(
328 + inputs=inputs,
329 + topk_weights=topk_weights,
330 + outputs=torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device),
331 + topk_ids=topk_ids,
332 + expert_map=expert_map,
333 + )
334 + assert not mapped.isnan().any(), "expert_map failed to mask stale non-local slot"
335 + ref = _reference(inputs.nan_to_num(), topk_weights, topk_ids, expert_map)
336 + torch.testing.assert_close(mapped.float(), ref, atol=atol, rtol=rtol)
337 +
338 + # Same call without the map: the non-local slot is summed and the NaN leaks.
339 + unmapped = moe_fused_mul_sum(
340 + inputs=inputs,
341 + topk_weights=topk_weights,
342 + outputs=torch.zeros(num_tokens, hidden_size, dtype=dtype, device=device),
343 + topk_ids=topk_ids,
344 + expert_map=None,
345 + )
346 + assert unmapped.isnan().any(), "expert_map is not redundant: dropping it must leak"
347 +
348 +
349 +@pytest.mark.parametrize("dtype", DTYPES)
350 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
351 +def test_num_valid_tokens_bounds_stale_padding(dtype: torch.dtype):
352 + """Rows past num_valid_tokens are left untouched, even with stale ids/NaN.
353 +
354 + Under cudagraph decode the grid is the static padded token count and the
355 + tail rows past num_recv hold stale (non -1) ids and stale NaN in `inputs`
356 + from a prior replay -- the -1 mask alone cannot spot them. Bounding on
357 + num_valid_tokens must skip the tail before any load so its NaN never reaches
358 + the output; without the bound it is summed and a masked finalize can then
359 + propagate the NaN. Regression guard for the decode accuracy fix.
360 + """
361 + set_random_seed(0)
362 + device = "cuda"
363 + num_tokens, num_recv, top_k, hidden_size = 17, 10, 8, 512
364 +
365 + inputs = torch.randn(num_tokens, top_k, hidden_size, dtype=dtype, device=device)
366 + topk_weights = torch.rand(num_tokens, top_k, dtype=dtype, device=device)
367 + # All ids valid (>= 0): the stale tail is indistinguishable by the -1 mask.
368 + topk_ids = torch.randint(
369 + 0, NUM_EXPERTS, (num_tokens, top_k), dtype=torch.int32, device=device
370 + )
371 + inputs[num_recv:] = float("nan") # stale down_output behind the padding tail
372 + num_valid = torch.tensor([num_recv], dtype=torch.int32, device=device)
373 +
374 + atol, rtol = (1e-2, 1e-2) if dtype == torch.bfloat16 else (1e-4, 1e-4)
375 +
376 + bounded = moe_fused_mul_sum(
377 + inputs=inputs,
378 + topk_weights=topk_weights,
379 + outputs=torch.full(
380 + (num_tokens, hidden_size), SENTINEL, dtype=dtype, device=device
381 + ),
382 + topk_ids=topk_ids,
383 + num_valid_tokens=num_valid,
384 + )
385 + assert torch.all(bounded[num_recv:] == SENTINEL), (
386 + "padding tail past num_recv written"
387 + )
388 + assert not bounded[:num_recv].isnan().any(), "real rows leaked NaN"
389 + ref = _reference(
390 + inputs[:num_recv].nan_to_num(),
391 + topk_weights[:num_recv],
392 + topk_ids[:num_recv],
393 + expert_map=None,
394 + )
395 + torch.testing.assert_close(bounded[:num_recv].float(), ref, atol=atol, rtol=rtol)
396 +
397 + # Without the bound the stale tail is summed and leaks NaN into its rows.
398 + unbounded = moe_fused_mul_sum(
399 + inputs=inputs,
400 + topk_weights=topk_weights,
401 + outputs=torch.full(
402 + (num_tokens, hidden_size), SENTINEL, dtype=dtype, device=device
403 + ),
404 + topk_ids=topk_ids,
405 + )
406 + assert unbounded[num_recv:].isnan().any(), "num_valid_tokens bound is load-bearing"
tests/kernels/moe/test_situ_mul_fp8_quant.py new
+224
@@ -0,0 +1,224 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Correctness tests for the persistent cp.async-pipelined SITU + block-FP8
4 +quant kernel (``situ_and_mul_quant``).
5 +
6 +Mirrors ``test_silu_mul_fp8_quant_deep_gemm.py`` in structure: build a reference
7 +activation + block-FP8 quant, drive random per-row scale magnitudes, exercise
8 +the masked (``valid_rows``) contiguous layout, and compare per valid row.
9 +
10 +The kernel computes SITU in fp32 and converts to fp8 with hardware rounding,
11 +which torch cannot reproduce bit-exactly. It is therefore validated against the
12 +*ideal* FP8 quantizer of the exact fp64 activation: the difference must stay
13 +within the unavoidable FP8 block-quant error (within 1 fp8 code for ~all
14 +elements, never off by more than 2).
15 +"""
16 +
17 +import pytest
18 +import torch
19 +
20 +from vllm.model_executor.layers.fused_moe.activation import situ_and_mul_quant
21 +from vllm.platforms import current_platform
22 +from vllm.utils.torch_utils import set_random_seed
23 +
24 +fp8_dtype = current_platform.fp8_dtype()
25 +DEVICE = current_platform.device_type
26 +FP8_MAX = torch.finfo(fp8_dtype).max
27 +
28 +# The kernel uses fp8 e4m3 conversion, so it is CUDA-only and needs SM89+
29 +# (Ada/Hopper). There is no Triton fallback.
30 +pytestmark = pytest.mark.skipif(
31 + not current_platform.is_cuda() or not current_platform.has_device_capability(89),
32 + reason="situ_and_mul_quant persistent kernel requires CUDA SM89+ (fp8 e4m3).",
33 +)
34 +
35 +# Kimi-K3 baked SITU params. Only (d=3072, beta=4.0, linear_beta=25.0,
36 +# group_size=128, fp16/bf16) routes to the persistent cp.async-pipelined kernel;
37 +# anything else falls back to the scalar group kernel.
38 +SITU_BETA = 4.0
39 +SITU_LINEAR_BETA = 25.0
40 +SITU_D = 3072
41 +GROUP_SIZE = 128
42 +
43 +# (num_tokens, valid_rows). valid_rows=None => all rows valid. 2048/4096 rows
44 +# span multiple grid-stride waves of the persistent grid (GRID_DIM = 132*8).
45 +CASES = [
46 + (1, None),
47 + (17, None),
48 + (128, None),
49 + (2048, None),
50 + (4096, None),
51 + (2048, 1500), # DeepEP v2 contiguous padding
52 + (2048, 1),
53 + (4096, 4095),
54 + (512, 0), # whole tensor is padding
55 +]
56 +
57 +
58 +def token_random(num_tokens: int, twod: int, dtype: torch.dtype) -> torch.Tensor:
59 + """Per-row random exponent so groups span many scale magnitudes."""
60 + base = torch.randn(num_tokens, twod, dtype=torch.float32, device=DEVICE)
61 + exps = torch.randint(1, 13, (num_tokens, 1), device=DEVICE).float()
62 + return (base * (2.0**exps)).to(dtype)
63 +
64 +
65 +def ref_situ_block_fp8_quant(
66 + inp: torch.Tensor, beta: float, linear_beta: float, group_size: int
67 +) -> tuple[torch.Tensor, torch.Tensor]:
68 + """Ideal FP8 block-quant of the exact fp64 SITU activation.
69 +
70 + The kernel rounds the activation to ``scalar_t`` before quantizing (it stores
71 + ``(float)(scalar_t)situ(...)``), so the reference does the same before taking
72 + the per-group absmax.
73 + """
74 + d = inp.shape[-1] // 2
75 + gate, up = inp.double().chunk(2, dim=-1)
76 + gate_out = beta * torch.tanh(gate / beta) * torch.sigmoid(gate)
77 + up_out = linear_beta * torch.tanh(up / linear_beta) if linear_beta > 0 else up
78 + act = (gate_out * up_out).to(inp.dtype).float()
79 +
80 + m = act.shape[0]
81 + ng = d // group_size
82 + a = act.view(m, ng, group_size)
83 + scale = a.abs().amax(dim=2).clamp_min(1e-30) / FP8_MAX
84 + q = (a / scale.unsqueeze(2)).clamp(-FP8_MAX, FP8_MAX).to(fp8_dtype)
85 + return q.view(m, d), scale
86 +
87 +
88 +def fp8_code_dist(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
89 + """Signed-magnitude fp8 code distance (valid for e4m3fn and e4m3fnuz)."""
90 +
91 + def ordered(x: torch.Tensor) -> torch.Tensor:
92 + x = x.view(torch.uint8).to(torch.int32)
93 + return torch.where((x >> 7) == 1, -(x & 0x7F), x & 0x7F)
94 +
95 + return (ordered(a) - ordered(b)).abs()
96 +
97 +
98 +@pytest.mark.parametrize("num_tokens,valid", CASES)
99 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half])
100 +@torch.inference_mode()
101 +def test_situ_and_mul_quant_pipelined(
102 + num_tokens: int, valid: int | None, dtype: torch.dtype
103 +) -> None:
104 + set_random_seed(42)
105 + d = SITU_D
106 + inp = token_random(num_tokens, 2 * d, dtype)
107 +
108 + # Sentinels: padding rows must keep their bytes (out) and get scale 1.0.
109 + out = torch.full((num_tokens, d), 17.0, device=DEVICE).to(fp8_dtype)
110 + scale = torch.full(
111 + (num_tokens, d // GROUP_SIZE), -1.0, dtype=torch.float32, device=DEVICE
112 + )
113 + out_sentinel = out.clone()
114 + valid_rows = (
115 + None if valid is None else torch.tensor(valid, dtype=torch.int32, device=DEVICE)
116 + )
117 +
118 + situ_and_mul_quant(
119 + out,
120 + scale,
121 + inp,
122 + beta=SITU_BETA,
123 + linear_beta=SITU_LINEAR_BETA,
124 + group_size=GROUP_SIZE,
125 + valid_rows=valid_rows,
126 + )
127 +
128 + # Padding contract: skipped rows keep their bytes; their scale becomes 1.0.
129 + if valid is not None and valid < num_tokens:
130 + assert torch.equal(scale[valid:], torch.ones_like(scale[valid:]))
131 + assert torch.equal(
132 + out[valid:].view(torch.uint8), out_sentinel[valid:].view(torch.uint8)
133 + )
134 +
135 + rows = num_tokens if valid is None else valid
136 + if rows == 0:
137 + return
138 +
139 + ref_q, ref_s = ref_situ_block_fp8_quant(
140 + inp[:rows], SITU_BETA, SITU_LINEAR_BETA, GROUP_SIZE
141 + )
142 +
143 + assert torch.isfinite(scale[:rows]).all()
144 + assert not torch.isnan(out[:rows].float()).any()
145 +
146 + # Scales: the kernel's fp32 SITU shifts the per-group absmax only
147 + # marginally vs the fp64 reference.
148 + torch.testing.assert_close(scale[:rows], ref_s, rtol=3e-2, atol=1e-5)
149 +
150 + # FP8 codes: must land within 1 code of the ideal quantizer for ~all
151 + # elements, never off by more than 2.
152 + dist = fp8_code_dist(out[:rows], ref_q)
153 + assert dist.max().item() <= 2
154 + assert (dist <= 1).float().mean().item() >= 0.99
155 +
156 +
157 +@pytest.mark.parametrize("tokens,topk", [(300, 8), (1, 4), (256, 6)])
158 +@torch.inference_mode()
159 +def test_situ_and_mul_quant_topk_row_bound(tokens: int, topk: int) -> None:
160 + """valid_rows is a token count; the kernel bounds rows at tokens*topk."""
161 + set_random_seed(7)
162 + d = SITU_D
163 + num_tokens = tokens * topk + 37 # padding tail past the valid rows
164 + inp = token_random(num_tokens, 2 * d, torch.bfloat16)
165 +
166 + out = torch.full((num_tokens, d), 17.0, device=DEVICE).to(fp8_dtype)
167 + scale = torch.full(
168 + (num_tokens, d // GROUP_SIZE), -1.0, dtype=torch.float32, device=DEVICE
169 + )
170 + out_sentinel = out.clone()
171 + valid_rows = torch.tensor(tokens, dtype=torch.int32, device=DEVICE)
172 +
173 + situ_and_mul_quant(
174 + out,
175 + scale,
176 + inp,
177 + beta=SITU_BETA,
178 + linear_beta=SITU_LINEAR_BETA,
179 + group_size=GROUP_SIZE,
180 + valid_rows=valid_rows,
181 + topk=topk,
182 + )
183 +
184 + rows = tokens * topk
185 + assert torch.equal(scale[rows:], torch.ones_like(scale[rows:]))
186 + assert torch.equal(
187 + out[rows:].view(torch.uint8), out_sentinel[rows:].view(torch.uint8)
188 + )
189 +
190 + ref_q, ref_s = ref_situ_block_fp8_quant(
191 + inp[:rows], SITU_BETA, SITU_LINEAR_BETA, GROUP_SIZE
192 + )
193 + torch.testing.assert_close(scale[:rows], ref_s, rtol=3e-2, atol=1e-5)
194 + assert fp8_code_dist(out[:rows], ref_q).max().item() <= 2
195 +
196 +
197 +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half])
198 +@torch.inference_mode()
199 +def test_situ_and_mul_quant_pipelined_deterministic(dtype: torch.dtype) -> None:
200 + """Two runs on the same input must be byte-identical."""
201 + set_random_seed(0)
202 + d = SITU_D
203 + num_tokens = 2048
204 + inp = token_random(num_tokens, 2 * d, dtype)
205 +
206 + def run() -> tuple[torch.Tensor, torch.Tensor]:
207 + out = torch.empty(num_tokens, d, dtype=fp8_dtype, device=DEVICE)
208 + scale = torch.empty(
209 + num_tokens, d // GROUP_SIZE, dtype=torch.float32, device=DEVICE
210 + )
211 + situ_and_mul_quant(
212 + out,
213 + scale,
214 + inp,
215 + beta=SITU_BETA,
216 + linear_beta=SITU_LINEAR_BETA,
217 + group_size=GROUP_SIZE,
218 + )
219 + return out, scale
220 +
221 + out1, s1 = run()
222 + out2, s2 = run()
223 + assert torch.equal(out1.view(torch.uint8), out2.view(torch.uint8))
224 + assert torch.equal(s1, s2)
tests/quantization/test_humming_mxfp4_block_fp8.py new
+122
@@ -0,0 +1,122 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for Humming MXFP4-weight GEMM paired with a block-FP8 activation.
4 +
5 +The Humming kernel already implements MXFP4 weight (group-32 e8m0) with a
6 +grouped FP8 activation via WGMMA software dequant -- see the grouped-fp8
7 +`test_mxfp4.py` cases in the humming repo, which run on SM90 (H200). Enabling
8 +it in vLLM comes down to two things, both exercised here:
9 +
10 + 1. A grouped FP8 activation must classify to a float32-scaled key
11 + (kFp8Dynamic128Sym), not a uint8 microscale key. MX microscale
12 + activations (group size 32) keep their e8m0 (uint8) scale.
13 + 2. (kMxfp4Static, kFp8Dynamic128Sym) must be an allowed MoE
14 + (weight, activation) pair.
15 +"""
16 +
17 +import pytest
18 +
19 +from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
20 + HummingExpertsBase,
21 +)
22 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
23 + kFp8Dynamic128Sym,
24 + kFp8DynamicTokenSym,
25 + kMxfp4Static,
26 + kMxfp8Dynamic,
27 +)
28 +from vllm.utils.import_utils import has_humming
29 +
30 +# make_humming_moe_quant_config stores per-sublayer humming LayerConfigs but the
31 +# config-plumbing asserts below never read them, so a placeholder satisfies the
32 +# required arg without pulling in the humming package.
33 +_PLACEHOLDER_HUMMING_CONFIGS = {"w13": None, "w2": None}
34 +
35 +
36 +def test_mxfp4_weight_with_block_fp8_activation_is_supported():
37 + """MXFP4 weight + block-FP8 (group-128) activation is an allowed pair."""
38 + assert HummingExpertsBase._supports_quant_scheme(kMxfp4Static, kFp8Dynamic128Sym)
39 +
40 +
41 +def test_mxfp4_weight_with_per_token_fp8_activation_still_supported():
42 + """The pre-existing per-token FP8 pairing must keep working."""
43 + assert HummingExpertsBase._supports_quant_scheme(kMxfp4Static, kFp8DynamicTokenSym)
44 +
45 +
46 +@pytest.mark.skipif(not has_humming(), reason="humming is not installed")
47 +@pytest.mark.parametrize(
48 + "group_size,expected_key",
49 + [
50 + # per-token FP8 -> float32 per-token scale
51 + (0, kFp8DynamicTokenSym),
52 + # block FP8 -> float32 group-128 scale (the newly-fixed case)
53 + (128, kFp8Dynamic128Sym),
54 + # MXFP8 microscale -> e8m0 (uint8) group-32 scale (must be preserved)
55 + (32, kMxfp8Dynamic),
56 + ],
57 +)
58 +def test_humming_fp8_input_schema_to_quant_key(group_size, expected_key):
59 + from vllm.model_executor.layers.quantization.utils.humming_utils import (
60 + _humming_input_schema_to_quant_key,
61 + )
62 + from vllm.utils.humming import HummingInputSchema
63 + from vllm.utils.humming import dtypes as humming_dtypes
64 +
65 + schema = HummingInputSchema(
66 + a_dtype=humming_dtypes.float8e4m3,
67 + input_scale_group_size=group_size,
68 + )
69 + assert _humming_input_schema_to_quant_key(schema) == expected_key
70 +
71 +
72 +@pytest.mark.skipif(not has_humming(), reason="humming is not installed")
73 +def test_humming_bf16_input_schema_is_unquantized():
74 + from vllm.model_executor.layers.quantization.utils.humming_utils import (
75 + _humming_input_schema_to_quant_key,
76 + )
77 + from vllm.utils.humming import HummingInputSchema
78 +
79 + # No a_dtype -> unquantized (bf16/fp16) inputs -> None.
80 + assert _humming_input_schema_to_quant_key(HummingInputSchema()) is None
81 +
82 +
83 +def test_block_fp8_activation_quant_config_is_block_quantized():
84 + """A block-FP8 activation group shape yields a block-quantized MoE config
85 + (FP8 dtype, [1, 128] block, not per-act-token). This is what makes the
86 + DeepEP prepare/finalize step quantize activations to block FP8 *before* the
87 + all-to-all dispatch instead of deferring to Humming -- and what makes
88 + HummingExpertsBase.expects_unquantized_inputs return False."""
89 + from vllm.model_executor.layers.quantization.utils.humming_utils import (
90 + make_humming_moe_quant_config,
91 + )
92 + from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
93 + from vllm.platforms import current_platform
94 +
95 + qc = make_humming_moe_quant_config(
96 + quant_dtype=current_platform.fp8_dtype(),
97 + weight_dtype="float4e2m1",
98 + activation_group_shape=GroupShape(row=1, col=128),
99 + humming_configs=_PLACEHOLDER_HUMMING_CONFIGS,
100 + )
101 + assert qc.is_block_quantized
102 + assert qc.block_shape == [1, 128]
103 + assert not qc.per_act_token_quant
104 + assert qc.quant_dtype == current_platform.fp8_dtype()
105 +
106 +
107 +def test_default_activation_quant_config_defers_to_humming():
108 + """Without a block activation shape the config stays per-token (the deferred
109 + path): Humming quantizes internally, preserving pre-existing behavior for
110 + every non-block-FP8 scheme."""
111 + from vllm.model_executor.layers.quantization.utils.humming_utils import (
112 + make_humming_moe_quant_config,
113 + )
114 + from vllm.platforms import current_platform
115 +
116 + qc = make_humming_moe_quant_config(
117 + quant_dtype=current_platform.fp8_dtype(),
118 + weight_dtype="float4e2m1",
119 + humming_configs=_PLACEHOLDER_HUMMING_CONFIGS,
120 + )
121 + assert not qc.is_block_quantized
122 + assert qc.per_act_token_quant
vllm/model_executor/layers/fused_moe/activation.py
+44
@@ -201,6 +201,7 @@ def apply_moe_activation(
201 activation_config: ApplyMoEActivationConfig | None = None,
202 topk_ids: torch.Tensor | None = None,
203 expert_map: torch.Tensor | None = None,
204 + valid_rows: torch.Tensor | None = None,
205 ) -> torch.Tensor:
206 """Apply MoE activation function.
207
@@ -256,6 +257,7 @@ def apply_moe_activation(
257 -1.0
258 if config.activation_situ_linear_beta is None
259 else config.activation_situ_linear_beta,
260 + valid_rows,
261 )
262 elif activation == MoEActivation.SWIGLUOAI:
263 torch.ops._C.swigluoai_and_mul(output, input)
@@ -286,3 +288,45 @@ def apply_moe_activation(
288 raise ValueError(f"Unsupported FusedMoe activation: {activation}")
289
290 return output
291 +
292 +
293 +def situ_and_mul_quant(
294 + output: torch.Tensor,
295 + scale: torch.Tensor,
296 + input: torch.Tensor,
297 + *,
298 + beta: float,
299 + linear_beta: float | None,
300 + group_size: int = 0,
301 + valid_rows: torch.Tensor | None = None,
302 + topk: int = 1,
303 +) -> None:
304 + """Fused Kimi SITU activation + dynamic FP8 quantization.
305 +
306 + Writes the quantized fp8 down-projection input into ``output`` [M, d] and its
307 + float32 scale into ``scale`` (dequant = q * scale), fusing the SITU
308 + activation with the FP8 quant that Humming's w2 GEMM would otherwise perform
309 + in a separate pass -- this avoids materializing (and rounding through) the
310 + bf16 activation_output buffer.
311 +
312 + ``group_size`` selects the quant granularity: ``0`` gives per-token scales
313 + ``scale`` [M, 1]; ``128`` gives k-major block-FP8 group scales ``scale``
314 + [M, d // 128], matching humming ``quant_input(group_size=128, float32)``.
315 +
316 + ``linear_beta`` <= 0 (or None) means "unset" (up passed through), matching
317 + ``SituAndMul(linear_beta=None)``. ``valid_rows`` (int32 scalar tensor) is the
318 + DeepEP v2 contiguous-layout valid *token* count (e.g. ``psum[-1:]``); the
319 + kernel multiplies it by ``topk`` to bound rows, so no host-side cast/multiply
320 + is needed. Padding rows are skipped and receive a benign scale of 1.0. Kept
321 + in this module so that every ``torch.ops._C.situ*`` call lives in one file.
322 + """
323 + torch.ops._C.situ_and_mul_quant(
324 + output,
325 + scale,
326 + input,
327 + beta,
328 + -1.0 if linear_beta is None else linear_beta,
329 + group_size,
330 + valid_rows,
331 + topk,
332 + )
vllm/model_executor/layers/fused_moe/deep_gemm_utils.py
+3 -1
@@ -525,7 +525,9 @@ def deepgemm_moe_permute(
525 expert_num_tokens = None
526 if expert_tokens_meta is not None:
527 expert_num_tokens = expert_tokens_meta.expert_num_tokens
528 - else:
528 + if expert_num_tokens is None:
529 + # A cudagraph-decode carrier may hold only psum_recv_per_rank and leave
530 + # expert_num_tokens None; count locally as in the meta-absent case.
531 expert_num_tokens = count_expert_num_tokens(
532 topk_ids, local_num_experts, expert_map
533 )
vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+1
@@ -668,6 +668,7 @@ class NaiveBatchedExperts(mk.FusedMoEExpertsModular):
668 assert hidden_states.dim() == 3
669 assert expert_tokens_meta is not None
670 expert_num_tokens = expert_tokens_meta.expert_num_tokens
671 + assert expert_num_tokens is not None
672
673 num_local_experts = w1.size(0)
674 assert num_local_experts == w1.size(0), f"{num_local_experts} == {w1.size(0)}"
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+224 -35
@@ -119,6 +119,37 @@ def get_humming_moe_gemm_type() -> str:
119 return gemm_type
120
121
122 +def _fixup_moe_tuning_config(tuning_config: list, max_k_block: int = 128) -> None:
123 + """Fix up each MoE tile in place: cap the K-block and widen warp-N.
124 +
125 + - block_shape[2] (K-block) > ``max_k_block``: the driver rejects the TMA
126 + descriptor at launch (CUDA_ERROR_MISALIGNED_ADDRESS). Cap at 128, which
127 + Humming already uses for larger M.
128 + - warp_shape[1] (warp-N) < 32: block-FP8 (group-128) activations route the
129 + w13 (gate/up) GEMM to a tuning table that pins warp-N to 16, under-filling
130 + the Hopper WGMMA N dimension and corrupting the GEMM output (gsm8k
131 + 0.94 -> 0.89). Widen to 32 whenever block_n % 32 == 0 -- on every tile,
132 + not just the K-capped ones. w2 (down) already uses warp_n=32.
133 + """
134 + logger.info_once("Attempting to override humming GEMM config")
135 + for entry in tuning_config:
136 + config = entry[2]
137 + block_shape = config.get("block_shape")
138 + if not (block_shape and len(block_shape) == 3):
139 + continue
140 + block_m, block_n, block_k = block_shape
141 +
142 + logger.info_once(f"Overriding humming GEMM config. Previous config\n: {config}")
143 + if block_k > max_k_block:
144 + config["block_shape"] = [block_m, block_n, max_k_block]
145 +
146 + warp_shape = config.get("warp_shape")
147 + if warp_shape and warp_shape[1] < 32 and block_n % 32 == 0:
148 + config["warp_shape"] = [warp_shape[0], 32, warp_shape[2]]
149 +
150 + logger.info_once(f"Overridden humming GEMM config. Current config\n: {config}")
151 +
152 +
153 class HummingExpertsBase(mk.FusedMoEExpertsModular):
154 def __init__(
155 self,
@@ -170,6 +201,10 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
201 gemm_type=self.humming_gemm_type(),
202 )
203 self.compute_config_str = json.dumps(self.compute_config)
204 + # Fix up the heuristic-chosen tiles before freezing to JSON -- see
205 + # _fixup_moe_tuning_config.
206 + _fixup_moe_tuning_config(self.w13_tuning_config)
207 + _fixup_moe_tuning_config(self.w2_tuning_config)
208 self.w13_tuning_config_str = json.dumps(self.w13_tuning_config)
209 self.w2_tuning_config_str = json.dumps(self.w2_tuning_config)
210
@@ -178,12 +213,17 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
213 sublayer_name: str,
214 inputs: torch.Tensor,
215 quanted_input: torch.Tensor | None,
216 + input_scale: torch.Tensor | None = None,
217 ) -> tuple[torch.Tensor, torch.Tensor | None]:
218 from vllm.utils.humming import may_quant_input
219
220 + # input_scale is set for block-FP8 (group-128) activations that were
221 + # quantized before the EP dispatch: may_quant_input then skips the
222 + # redundant w13 quantization and forwards the pre-computed scale.
223 return may_quant_input(
224 self.humming_configs[sublayer_name],
225 inputs=inputs,
226 + input_scale=input_scale,
227 quanted_input=quanted_input,
228 )
229
@@ -263,6 +303,10 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
303 (kMxfp4Static, kMxfp4Dynamic),
304 (kMxfp4Static, kMxfp8Dynamic),
305 (kMxfp4Static, kFp8DynamicTokenSym),
306 + # MXFP4 weight (group-32 e8m0) with block-FP8 activation
307 + # (group-128 float32). Runs via WGMMA software dequant, so it
308 + # works on Hopper (SM90/H200) as well as Blackwell.
309 + (kMxfp4Static, kFp8Dynamic128Sym),
310 (kNvfp4Static, None),
311 (kNvfp4Static, kFp8DynamicTokenSym),
312 (kMxfp8Static, None),
@@ -299,21 +343,43 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
343 and _is_supported_wna16_weight_key(weight_key)
344 )
345
346 + def _prequantizes_dispatch_activation(self) -> bool:
347 + """
348 + Whether the prepare/finalize step should quantize activations before
349 + the (EP all-to-all) dispatch instead of leaving it to Humming.
350 +
351 + This is enabled only for block-FP8 (group-128) activations: quantizing
352 + to FP8 before dispatch sends FP8 rather than BF16 over the interconnect,
353 + and Humming then consumes the pre-quantized FP8 + scale directly (see
354 + the apply() methods, which forward the dispatch scale into
355 + HummingMethod.may_quant_input, and may_quant_input itself, which is a
356 + no-op when an input scale is already supplied). The scale layout
357 + produced by vLLM's block-FP8 quantization ([M, K // 128] float32,
358 + row-major) matches what the Humming WGMMA grouped GEMM expects.
359 + """
360 + quant_config = self.quant_config
361 + return (
362 + quant_config.is_block_quantized
363 + and quant_config.quant_dtype == current_platform.fp8_dtype()
364 + )
365 +
366 @property
367 def expects_unquantized_inputs(self) -> bool:
368 """
305 - Humming kernels handle input quantization internally in apply().
306 -
307 - This property tells the prepare/finalize step to skip input
308 - quantization (by setting defer_input_quant=True) and pass
309 - unquantized inputs to the experts. This prevents double
310 - quantization: once in prepare and once in Humming's apply().
311 -
312 - Returns:
313 - True to indicate that this expert expects unquantized inputs
314 - and will handle quantization internally.
369 + Whether the prepare/finalize step should defer input quantization to
370 + the experts (by setting defer_input_quant=True and passing unquantized
371 + inputs).
372 +
373 + Humming normally quantizes inputs internally via
374 + HummingMethod.may_quant_input() in apply(), so we defer quantization
375 + (return True) to avoid quantizing twice -- once in prepare and once in
376 + Humming's apply().
377 +
378 + The exception is block-FP8 (group-128) activations, which are quantized
379 + before the dispatch to save interconnect bandwidth (see
380 + _prequantizes_dispatch_activation): for those we must NOT defer.
381 """
316 - return True
382 + return not self._prequantizes_dispatch_activation()
383
384 @staticmethod
385 def _supports_current_device() -> bool:
@@ -585,6 +651,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
651 activation: MoEActivation,
652 output: torch.Tensor,
653 input: torch.Tensor,
654 + valid_rows: torch.Tensor | None = None,
655 ) -> None:
656 activation_config = self.activation_config
657 if (
@@ -601,8 +668,84 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
668 activation=activation,
669 input=input,
670 output=output,
671 + valid_rows=valid_rows,
672 )
673
674 + def fused_situ_quant_enabled(self, activation: MoEActivation) -> bool:
675 + """Whether the SITU activation + w2 quant can be fused into one kernel.
676 +
677 + Fused only for k-major block-FP8 group-128 e4m3 with float32 scales --
678 + the sole layout situ_and_mul_quant supports. A float32 as_dtype rules out
679 + MXMMA (which uses e8m0 m-major scales), so a group-128 scale is
680 + guaranteed k-major. Everything else (m-major MXFP8, other group sizes,
681 + 16-bit passthrough, non-SITU) falls back to the separate
682 + situ_and_mul + quantize_input path.
683 + """
684 + if activation != MoEActivation.SITU:
685 + return False
686 + w2cfg = self.humming_configs["w2"]
687 + # Report the specific blocker once so a silent fallback to the unfused
688 + # situ_and_mul + quant pair is never a mystery in a trace.
689 + reason: str | None = None
690 + if not (w2cfg.a_dtype.num_bits == 8 and str(w2cfg.a_dtype) == "float8e4m3"):
691 + reason = f"w2 a_dtype is {w2cfg.a_dtype} (need float8e4m3)"
692 + elif w2cfg.input_scale_group_size != 128:
693 + reason = (
694 + f"w2 input_scale_group_size is {w2cfg.input_scale_group_size} "
695 + "(need 128)"
696 + )
697 + elif not (w2cfg.as_dtype is not None and str(w2cfg.as_dtype) == "float32"):
698 + reason = f"w2 as_dtype is {w2cfg.as_dtype} (need float32, k-major)"
699 + if reason is not None:
700 + logger.warning_once(
701 + "Humming fused SITU+FP8 quant disabled, using unfused "
702 + "situ_and_mul + quant: %s",
703 + reason,
704 + )
705 + return False
706 + return True
707 +
708 + def fused_situ_quant(
709 + self,
710 + gate_up_output: torch.Tensor,
711 + quanted_down_input: torch.Tensor,
712 + valid_rows: torch.Tensor | None,
713 + topk: int = 1,
714 + ) -> tuple[torch.Tensor, torch.Tensor]:
715 + """Run the fused SITU activation + FP8 quant for the w2 input.
716 +
717 + Returns ``(quanted_down_input, input_scale)`` in the same layout the
718 + unfused ``apply_activation`` + ``quantize_input("w2")`` pair produces: an
719 + fp8 [M, d] tensor plus a k-major block-FP8 float32 [M, d // 128] group
720 + scale. Only reached when ``fused_situ_quant_enabled`` (group_size 128).
721 + """
722 + from vllm.model_executor.layers.fused_moe.activation import (
723 + situ_and_mul_quant,
724 + )
725 +
726 + cfg = self.activation_config
727 + assert cfg.activation_situ_beta is not None, (
728 + "SITU requires activation_situ_beta from FusedMoEConfig"
729 + )
730 + group_size = self.humming_configs["w2"].input_scale_group_size
731 + m, d = quanted_down_input.size(0), quanted_down_input.size(1)
732 + input_scale = torch.empty(
733 + (m, d // group_size),
734 + dtype=torch.float32,
735 + device=quanted_down_input.device,
736 + )
737 + situ_and_mul_quant(
738 + quanted_down_input,
739 + input_scale,
740 + gate_up_output,
741 + beta=cfg.activation_situ_beta,
742 + linear_beta=cfg.activation_situ_linear_beta,
743 + group_size=group_size,
744 + valid_rows=valid_rows,
745 + topk=topk,
746 + )
747 + return quanted_down_input, input_scale
748 +
749
750 class HummingIndexedExperts(HummingExpertsBase):
751 def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
@@ -684,10 +827,12 @@ class HummingIndexedExperts(HummingExpertsBase):
827 """
828 Standard apply implementation for Humming indexed experts.
829
687 - Humming performs activation quantization internally and consumes the
688 - weights supplied by the modular kernel interface.
689 - Results are written to the output buffer supplied by the modular
690 - kernel, which may alias workspace13.
830 + Note: Humming kernels handle weights internally through the layer
831 + object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
832 + path (Humming quantizes the w13 input itself); for block-FP8 activations
833 + it carries the scale computed before dispatch, which is forwarded to
834 + may_quant_input so Humming skips the redundant w13 quantization. The
835 + output is written into workspace13 via the buffer management.
836 """
837 assert not apply_router_weight_on_input
838
@@ -710,6 +855,7 @@ class HummingIndexedExperts(HummingExpertsBase):
855 inputs, input_scale = self.quantize_input(
856 "w13",
857 inputs=hidden_states,
858 + input_scale=a1q_scale,
859 quanted_input=buffers.get("quanted_gate_up_input", None),
860 )
861
@@ -722,17 +868,45 @@ class HummingIndexedExperts(HummingExpertsBase):
868 **moe_kwargs1,
869 )
870
725 - self.apply_activation(
726 - activation=activation,
727 - input=buffers["gate_up_output"],
728 - output=buffers["activation_output"],
729 - )
871 + # psum[-1:] is the DeepEP valid *token* count as a zero-cost int32 view.
872 + # situ (rows = tokens*topk) multiplies by topk on-device; mul_sum bounds
873 + # on tokens, so both read this pointer directly -- no host cast/multiply.
874 + valid_tokens = None
875 + topk = topk_ids.size(1)
876 + if (
877 + expert_tokens_meta is not None
878 + and expert_tokens_meta.psum_recv_per_rank is not None
879 + ):
880 + valid_tokens = expert_tokens_meta.psum_recv_per_rank[-1:]
881 +
882 + if self.fused_situ_quant_enabled(activation):
883 + # Fused SITU + FP8 quant (per-token or block-FP8 group-128) straight
884 + # into the w2 input, skipping the bf16 activation_output round-trip.
885 + inputs, input_scale = self.fused_situ_quant(
886 + gate_up_output=buffers["gate_up_output"],
887 + quanted_down_input=buffers["quanted_down_input"],
888 + valid_rows=valid_tokens,
889 + topk=topk,
890 + )
891 + else:
892 + # Fallback situ_and_mul takes int64 row counts (tokens*topk).
893 + valid_rows = (
894 + valid_tokens.to(torch.int64) * topk
895 + if valid_tokens is not None
896 + else None
897 + )
898 + self.apply_activation(
899 + activation=activation,
900 + input=buffers["gate_up_output"],
901 + output=buffers["activation_output"],
902 + valid_rows=valid_rows,
903 + )
904
731 - inputs, input_scale = self.quantize_input(
732 - "w2",
733 - inputs=buffers["activation_output"],
734 - quanted_input=buffers.get("quanted_down_input", None),
735 - )
905 + inputs, input_scale = self.quantize_input(
906 + "w2",
907 + inputs=buffers["activation_output"],
908 + quanted_input=buffers.get("quanted_down_input", None),
909 + )
910
911 self.humming_forward(
912 "w2",
@@ -743,12 +917,16 @@ class HummingIndexedExperts(HummingExpertsBase):
917 **moe_kwargs2,
918 )
919
920 + # expert_map masks any non-local id; num_valid_tokens bounds the
921 + # persistent kernel to the real token rows [0, num_recv) so the padding
922 + # tail is never iterated (CUDA-graph-safe device scalar).
923 moe_fused_mul_sum(
924 inputs=buffers["down_output"].view(*topk_ids.shape, -1),
925 topk_weights=topk_weights,
926 topk_ids=topk_ids,
927 expert_map=expert_map,
928 outputs=output,
929 + num_valid_tokens=valid_tokens,
930 )
931
932
@@ -787,10 +965,13 @@ class HummingGroupedExperts(HummingExpertsBase):
965 """
966 Standard apply implementation for Humming grouped experts.
967
790 - Humming performs activation quantization internally and consumes the
791 - weights supplied by the modular kernel interface.
792 - Results are written to the output buffer supplied by the modular
793 - kernel, which may alias workspace13.
968 + Note: Humming kernels handle weights internally through the layer
969 + object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
970 + path (Humming quantizes the w13 input itself); for block-FP8 activations
971 + it carries the scale computed before dispatch. It is permuted alongside
972 + the tokens by moe_permute and forwarded to may_quant_input so Humming
973 + skips the redundant w13 quantization. The output is written into
974 + workspace13 via the buffer management.
975 """
976 assert not apply_router_weight_on_input
977
@@ -805,9 +986,9 @@ class HummingGroupedExperts(HummingExpertsBase):
986 )
987 buffers["output"] = output
988
808 - hidden_states, _, expert_first_token_offset, inv_perm, _ = moe_permute(
989 + hidden_states, a1q_scale, expert_first_token_offset, inv_perm, _ = moe_permute(
990 hidden_states=hidden_states,
810 - a1q_scale=None,
991 + a1q_scale=a1q_scale,
992 topk_ids=topk_ids,
993 n_expert=global_num_experts,
994 n_local_expert=self.num_experts,
@@ -818,6 +999,7 @@ class HummingGroupedExperts(HummingExpertsBase):
999 inputs, input_scale = self.quantize_input(
1000 "w13",
1001 inputs=hidden_states,
1002 + input_scale=a1q_scale,
1003 quanted_input=buffers.get("quanted_gate_up_input", None),
1004 )
1005
@@ -901,15 +1083,21 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1083 """
1084 Standard apply implementation for Humming batched grouped experts.
1085
904 - Humming performs activation quantization internally and consumes the
905 - weights supplied by the modular kernel interface.
906 - Results are written to the output buffer supplied by the modular
907 - kernel, which may alias workspace13.
1086 + Note: Humming kernels handle weights internally through the layer
1087 + object, so w1, w2, a2_scale are unused. a1q_scale is None on the usual
1088 + path (Humming quantizes the w13 input itself); for block-FP8 activations
1089 + it carries the scale computed before dispatch, which is forwarded to
1090 + may_quant_input so Humming skips the redundant w13 quantization. The
1091 + output is written into workspace13 via the buffer management.
1092 """
1093 assert not apply_router_weight_on_input
1094 assert expert_tokens_meta is not None
1095
1096 hidden_states = hidden_states.view(-1, hidden_states.size(-1))
1097 + # Keep the (batched) block-FP8 scale row-aligned with the flattened
1098 + # [num_experts * max_tokens, K] hidden states above.
1099 + if a1q_scale is not None and a1q_scale.dim() == 3:
1100 + a1q_scale = a1q_scale.view(-1, a1q_scale.size(-1))
1101 valid_shape_m = self.estimate_local_valid_shape_m(topk_ids)
1102 expert_num_tokens = expert_tokens_meta.expert_num_tokens
1103
@@ -925,6 +1113,7 @@ class BatchedHummingGroupedExperts(HummingExpertsBase):
1113 inputs, input_scale = self.quantize_input(
1114 "w13",
1115 inputs=hidden_states,
1116 + input_scale=a1q_scale,
1117 quanted_input=buffers.get("quanted_gate_up_input", None),
1118 )
1119
vllm/model_executor/layers/fused_moe/modular_kernel.py
+4 -1
@@ -100,8 +100,9 @@ class ExpertTokensMetadata:
100 Metadata regarding expert-token routing.
101 """
102
103 - expert_num_tokens: torch.Tensor
103 + expert_num_tokens: torch.Tensor | None
104 expert_num_tokens_cpu: torch.Tensor | None
105 + psum_recv_per_rank: torch.Tensor | None = None
106
107 @staticmethod
108 def make_from_list(
@@ -901,6 +902,7 @@ class FusedMoEExpertsModular(FusedMoEExperts):
902 *,
903 topk_ids: torch.Tensor | None = None,
904 expert_map: torch.Tensor | None = None,
905 + valid_rows: torch.Tensor | None = None,
906 ) -> None:
907 apply_moe_activation(
908 activation,
@@ -909,6 +911,7 @@ class FusedMoEExpertsModular(FusedMoEExperts):
911 activation_config=self.activation_config,
912 topk_ids=topk_ids,
913 expert_map=expert_map,
914 + valid_rows=valid_rows,
915 )
916
917 @abstractmethod
vllm/model_executor/layers/fused_moe/moe_fused_mul_sum.py
+90 -120
@@ -3,7 +3,6 @@
3 import torch
4 from torch._subclasses.fake_tensor import FakeTensor
5
6 -from vllm.platforms import current_platform
6 from vllm.triton_utils import tl, triton
7
8
@@ -14,122 +13,75 @@ def moe_fused_mul_sum_kernel(
13 outputs_ptr,
14 top_ids_ptr,
15 expert_map_ptr,
17 - num_tokens,
16 + num_valid_tokens_ptr,
17 stride_m,
18 + has_topk_ids: tl.constexpr,
19 has_expert_map: tl.constexpr,
20 + has_num_valid: tl.constexpr,
21 top_k: tl.constexpr,
21 - size: tl.constexpr,
22 - BLOCK_M: tl.constexpr,
22 + hidden_size: tl.constexpr,
23 BLOCK_K: tl.constexpr,
24 ):
25 - pid_k = tl.program_id(0)
26 - pid_m = tl.program_id(1)
27 -
28 - offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
29 - offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K)
30 -
31 - m_mask = offs_m < num_tokens
32 - k_mask = offs_k < size
33 - mask = m_mask[:, None] & k_mask[None, :]
34 -
35 - a_base = inputs_ptr + (offs_m * stride_m)[:, None] + offs_k[None, :]
36 - b_base = topk_weights_ptr + offs_m * top_k
37 -
38 - acc = tl.zeros((BLOCK_M, BLOCK_K), dtype=tl.float32)
39 -
40 - for n in tl.static_range(top_k):
41 - b_val = tl.load(b_base + n, mask=m_mask, other=0.0).to(tl.float32)
42 - if has_expert_map:
43 - id_val = tl.load(top_ids_ptr + offs_m * top_k + n, mask=m_mask, other=0)
44 - expert_mask = tl.load(expert_map_ptr + id_val) >= 0
45 - a_vec = tl.load(
46 - a_base + n * size,
47 - mask=mask & expert_mask[:, None],
48 - other=0.0,
49 - ).to(tl.float32)
50 - else:
51 - a_vec = tl.load(
52 - a_base + n * size,
53 - mask=mask,
54 - other=0.0,
55 - ).to(tl.float32)
56 - acc += a_vec * b_val[:, None]
57 -
58 - out_ptrs = outputs_ptr + (offs_m * size)[:, None] + offs_k[None, :]
59 - tl.store(
60 - out_ptrs,
61 - acc.to(outputs_ptr.dtype.element_ty),
62 - mask=mask,
63 - )
25 + # One CTA owns one token's output row. Resolve which slots to sum once per
26 + # row (the decision is invariant over hidden tiles); the `if take` branch is
27 + # block-uniform, so padding slots are skipped rather than masked-and-loaded.
28 + pid_m = tl.program_id(0)
29 + # Bound to the real recv rows [0, num_recv). Under cudagraph decode the grid
30 + # is the static padded token count, and padding rows past num_recv carry
31 + # stale ids and stale NaN in `inputs`; returning before any load leaves their
32 + # output rows untouched instead of summing that garbage. The constexpr guard
33 + # short-circuits, so num_valid_tokens_ptr is only read when it is provided.
34 + if has_num_valid and pid_m >= tl.load(num_valid_tokens_ptr).to(tl.int32):
35 + return
36 + w_row = topk_weights_ptr + pid_m * top_k
37 +
38 + takes = ()
39 + weights = ()
40 + if has_topk_ids:
41 + any_valid = 0
42 + for n in tl.static_range(top_k):
43 + id_val = tl.load(top_ids_ptr + pid_m * top_k + n)
44 + valid = id_val >= 0
45 + any_valid += valid.to(tl.int32)
46 + take = valid
47 + if has_expert_map:
48 + local_id = tl.load(expert_map_ptr + tl.where(valid, id_val, 0))
49 + take = valid & (local_id >= 0)
50 + takes += (take,) # type: ignore[assignment]
51 + weights += (tl.load(w_row + n).to(tl.float32),) # type: ignore[assignment]
52 + # All-padding rows early-return untouched (decode contract past num_recv).
53 + if any_valid == 0:
54 + return
55 + else:
56 + for n in tl.static_range(top_k):
57 + takes += (True,) # type: ignore[assignment]
58 + weights += (tl.load(w_row + n).to(tl.float32),) # type: ignore[assignment]
59 +
60 + n_tiles: tl.constexpr = (hidden_size + BLOCK_K - 1) // BLOCK_K
61 + a_row = inputs_ptr + pid_m * stride_m
62 + out_row = outputs_ptr + pid_m * hidden_size
63 +
64 + for t in tl.range(0, n_tiles):
65 + offs_k = t * BLOCK_K + tl.arange(0, BLOCK_K)
66 + kmask = offs_k < hidden_size
67 + acc = tl.zeros((BLOCK_K,), dtype=tl.float32)
68 + for n in tl.static_range(top_k):
69 + if takes[n]:
70 + a = tl.load(a_row + n * hidden_size + offs_k, mask=kmask, other=0.0)
71 + acc += a.to(tl.float32) * weights[n]
72 + tl.store(out_row + offs_k, acc.to(outputs_ptr.dtype.element_ty), mask=kmask)
73
74
75 def _heuristic_config(
67 - num_tokens: int,
68 - top_k: int,
69 - size: int,
76 + hidden_size: int,
77 element_size: int,
78 ):
79 is_fp32 = element_size > 2
73 - is_sm90_plus = current_platform.has_device_capability(90)
74 - is_sm80_before = not current_platform.has_device_capability(80)
75 -
76 - if current_platform.has_device_capability(90):
77 - # SM90/SM100+: prefer small tiles + many CTAs.
78 - if is_fp32:
79 - BLOCK_M = 1 if num_tokens <= 4 else 2
80 - else:
81 - if num_tokens <= 4:
82 - BLOCK_M = 1
83 - elif num_tokens <= 128:
84 - BLOCK_M = 2
85 - else:
86 - BLOCK_M = 4
87 - elif is_fp32:
88 - if num_tokens <= 4:
89 - BLOCK_M = 1
90 - elif num_tokens <= 32:
91 - BLOCK_M = 2
92 - elif num_tokens <= 128:
93 - BLOCK_M = 4
94 - else:
95 - BLOCK_M = 4
96 - else:
97 - if num_tokens <= 4:
98 - BLOCK_M = 1
99 - elif num_tokens <= 32:
100 - BLOCK_M = 2
101 - elif num_tokens <= 128:
102 - BLOCK_M = 4
103 - elif num_tokens <= 1024:
104 - BLOCK_M = 16
105 - else:
106 - BLOCK_M = 8
107 -
108 - if is_fp32:
109 - max_block_k = 256
110 - elif is_sm80_before or is_sm90_plus:
111 - max_block_k = 512
112 - else:
113 - max_block_k = 1024
114 - BLOCK_K = min(triton.next_power_of_2(size), max_block_k)
115 - BLOCK_K = max(BLOCK_K, 256)
116 -
117 - total = BLOCK_M * BLOCK_K
118 - if is_fp32:
119 - num_warps = max(8, min(16, total // 64))
120 - else:
121 - num_warps = max(4, min(16, total // 256))
122 -
123 - if is_sm80_before:
124 - num_warps = min(num_warps, 8)
125 - num_stages = 2
126 - elif is_sm90_plus:
127 - num_warps = min(num_warps, 8)
128 - num_stages = 4 if total <= 2048 else 2
129 - else:
130 - num_stages = 4 if total <= 2048 else 2
131 -
132 - return BLOCK_M, BLOCK_K, num_warps, num_stages
80 + max_block_k = 256 if is_fp32 else 512
81 + BLOCK_K = max(128, min(triton.next_power_of_2(hidden_size), max_block_k))
82 + num_warps = 4 if is_fp32 else 2
83 + num_stages = 3
84 + return BLOCK_K, num_warps, num_stages
85
86
87 def moe_fused_mul_sum(
@@ -138,6 +90,7 @@ def moe_fused_mul_sum(
90 outputs: torch.Tensor | None = None,
91 topk_ids: torch.Tensor | None = None,
92 expert_map: torch.Tensor | None = None,
93 + num_valid_tokens: torch.Tensor | None = None,
94 ) -> torch.Tensor:
95 """
96 Fused kernel for MoE (Mixture of Experts) to perform weighted summation
@@ -150,10 +103,21 @@ def moe_fused_mul_sum(
103 Shape: (num_tokens, top_k).
104 outputs: Optional pre-allocated output tensor.
105 Shape: (num_tokens, hidden_size).
153 - topk_ids: Optional indices of the top-k experts. Used when
154 - `expert_map` is provided. Shape: (num_tokens, top_k).
106 + topk_ids: Optional indices of the top-k experts. Shape:
107 + (num_tokens, top_k). A value of -1 marks a slot the expert GEMM
108 + skipped; those slots are excluded from the sum. When provided, rows
109 + with all top ids < 0 (worst-case padding) are skipped and their
110 + output rows left untouched. Required when `expert_map` is provided.
111 expert_map: Optional mapping for Expert Parallelism. A value < 0
156 - indicates an invalid token/expert pair that will be skipped.
112 + indicates an invalid token/expert pair that will be skipped. Only
113 + needed when `topk_ids` may contain non-local expert ids; if every
114 + non-(-1) id is already a local expert, leave it None to skip the
115 + redundant per-slot lookup.
116 + num_valid_tokens: Optional device scalar (1-element tensor) holding the
117 + number of real token rows (num_recv for a decode dispatch). When
118 + provided, rows past it are left untouched, so the static cudagraph
119 + grid never sums stale padding rows. Pass the token count, not
120 + token*top_k.
121
122 Returns:
123 The fused weighted sum of expert outputs.
@@ -166,34 +130,40 @@ def moe_fused_mul_sum(
130 assert inputs.dtype in (torch.float32, torch.float16, torch.bfloat16)
131 assert topk_weights.dtype in (torch.float32, torch.float16, torch.bfloat16)
132
169 - num_tokens, top_k, size = inputs.shape
170 - output_shape = (num_tokens, size)
133 + num_tokens, top_k, hidden_size = inputs.shape
134 + output_shape = (num_tokens, hidden_size)
135 if outputs is None:
136 outputs = torch.empty(output_shape, dtype=inputs.dtype, device=inputs.device)
137
138 assert outputs.shape == output_shape
139 assert topk_weights.shape == (num_tokens, top_k)
140 + assert expert_map is None or topk_ids is not None, (
141 + "topk_ids is required to interpret expert_map"
142 + )
143 + if topk_ids is not None:
144 + assert topk_ids.shape == (num_tokens, top_k)
145 + assert topk_ids.is_contiguous()
146 + assert topk_ids.dtype in (torch.int32, torch.int64)
147
148 if not isinstance(inputs, FakeTensor):
178 - BLOCK_M, BLOCK_K, num_warps, num_stages = _heuristic_config(
179 - num_tokens,
180 - top_k,
181 - size,
149 + BLOCK_K, num_warps, num_stages = _heuristic_config(
150 + hidden_size,
151 inputs.element_size(),
152 )
184 - grid = (triton.cdiv(size, BLOCK_K), triton.cdiv(num_tokens, BLOCK_M))
153 + grid = (num_tokens,)
154 moe_fused_mul_sum_kernel[grid](
155 inputs,
156 topk_weights,
157 outputs,
158 topk_ids,
159 expert_map,
191 - num_tokens,
192 - top_k * size,
160 + num_valid_tokens,
161 + top_k * hidden_size,
162 + topk_ids is not None,
163 expert_map is not None,
164 + num_valid_tokens is not None,
165 top_k,
195 - size,
196 - BLOCK_M,
166 + hidden_size,
167 BLOCK_K,
168 num_warps=num_warps,
169 num_stages=num_stages,
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+13 -1
@@ -32,7 +32,7 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
32 - Worst-case tensor allocation; padding rows zeroed via
33 handle.psum_num_recv_tokens_per_scaleup_rank
34 - Fully cudagraph-capturable
35 - - Expert kernel sorts internally (expert_tokens_meta=None)
35 + - Expert kernel sorts internally (expert_tokens_meta carries no counts)
36
37 **Prefill mode (use_cudagraph=False):**
38 - do_expand=True, do_cpu_sync=True
@@ -266,6 +266,18 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
266 if recv_topk_weights is not None and recv_topk_weights.ndim == 1:
267 recv_topk_weights = recv_topk_weights.unsqueeze(1)
268
269 + if self.use_cudagraph:
270 + # Carry the per-rank prefix sum so SiTU can skip padding rows.
271 + # expert_num_tokens stays None: count-based consumers (DeepGEMM,
272 + # Triton) must treat a None field as "no counts" and derive their
273 + # own, exactly as in the meta-absent decode case.
274 + if expert_tokens_meta is None:
275 + expert_tokens_meta = mk.ExpertTokensMetadata(
276 + expert_num_tokens=None,
277 + expert_num_tokens_cpu=None,
278 + )
279 + expert_tokens_meta.psum_recv_per_rank = psum_recv_per_rank
280 +
281 if not quant_config.is_block_quantized and not defer_input_quant:
282 expert_x_scale = None
283 if expert_x.numel() != 0:
vllm/model_executor/layers/quantization/utils/humming_utils.py
+34 -1
@@ -307,7 +307,18 @@ def _humming_input_schema_to_quant_key(
307 gs = schema.input_scale_group_size
308 group_shape = GroupShape(row=1, col=gs) if gs > 0 else GroupShape.PER_TOKEN
309
310 - scale_dtype = MXFP_SCALE_DTYPE if gs > 0 else torch.float32
310 + # Pick the scale dtype the Humming kernel actually consumes. An explicit
311 + # input_scale_dtype always wins. Otherwise infer from the grouping: MX
312 + # microscale activations (group size 32, e.g. MXFP8) carry an e8m0 (uint8)
313 + # scale, while block-FP8 (group size 128) and per-token FP8/int8 carry a
314 + # float32 scale. Getting this right lets a grouped FP8 activation match
315 + # kFp8Dynamic128Sym instead of an unmatchable uint8-scaled key.
316 + if schema.input_scale_dtype is not None:
317 + scale_dtype = _HUMMING_TO_SCALE_DTYPE[schema.input_scale_dtype]
318 + elif gs == 32:
319 + scale_dtype = MXFP_SCALE_DTYPE
320 + else:
321 + scale_dtype = torch.float32
322
323 scale = ScaleDesc(dtype=scale_dtype, static=False, group_shape=group_shape)
324
@@ -553,6 +564,7 @@ def make_humming_moe_quant_config(
564 quant_dtype: torch.dtype | str | None,
565 weight_dtype: torch.dtype | str | None,
566 weight_group_shape: GroupShape | None = None,
567 + activation_group_shape: GroupShape | None = None,
568 w1_scale: torch.Tensor | None = None,
569 w2_scale: torch.Tensor | None = None,
570 w1_zp: torch.Tensor | None = None,
@@ -569,7 +581,14 @@ def make_humming_moe_quant_config(
581 assert humming_configs is not None
582 if quant_dtype is None:
583 a_quant_desc = FusedMoEQuantDesc(dtype=None)
584 + elif activation_group_shape is not None:
585 + # Pre-dispatch quantization.
586 + a_quant_desc = FusedMoEQuantDesc(
587 + dtype=quant_dtype, shape=activation_group_shape
588 + )
589 else:
590 + # Deferred path: Humming quantizes the activation internally, so the
591 + # descriptor only needs a non-None dtype to mark it as quantized.
592 shape = GroupShape(row=1, col=-1)
593 a_quant_desc = FusedMoEQuantDesc(dtype=quant_dtype, shape=shape)
594
@@ -621,6 +640,19 @@ def get_humming_moe_quant_config(
640 else:
641 q_dtype = str(input_schema.a_dtype)
642
643 + # Block-FP8 (group-128) activations are quantized *before* the EP all-to-all
644 + # dispatch (so FP8 rather than BF16 crosses the interconnect) and consumed by
645 + # Humming as-is.
646 + activation_group_shape: GroupShape | None = None
647 + input_scale_group_size = getattr(input_schema, "input_scale_group_size", 0) or 0
648 + if (
649 + q_dtype is not None
650 + and q_dtype.startswith("float8")
651 + and input_scale_group_size == 128
652 + ):
653 + q_dtype = _HUMMING_TO_QUANT_DTYPE.get(input_schema.a_dtype, FP8_DTYPE)
654 + activation_group_shape = GroupShape(row=1, col=input_scale_group_size)
655 +
656 weight_scale_group_size = weight_schema.weight_scale_group_size
657 weight_scale_group_size_n = weight_schema.weight_scale_group_size_n
658 weight_group_shape: tuple[int, ...] = ()
@@ -638,6 +670,7 @@ def get_humming_moe_quant_config(
670 quant_dtype=q_dtype,
671 weight_dtype=str(weight_schema.b_dtype),
672 weight_group_shape=weight_group_shape,
673 + activation_group_shape=activation_group_shape,
674 w1_scale=getattr(layer, "w13_weight_scale", None),
675 w1_gscale=getattr(layer, "w13_weight_scale_2", None),
676 w1_zp=getattr(layer, "w13_zero_point", None),