@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 82b7d49a6e

[MoE] Generalize masked activation for padded layouts (#51217)

Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Misha Goin committed Sep 1, 2026 at 11:34 UTC 82b7d49a6e30cf6d1d53838667ee20f09ac1f08b
15 files changed +1217 -196
csrc/libtorch_stable/activation_kernels.cu
+297 -48
@@ -500,10 +500,8 @@ 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]
503 - const int d, const float beta, const float linear_beta,
504 - const int64_t* __restrict__ valid_rows_ptr) {
503 + const int d, const float beta, const float linear_beta) {
504 const int64_t row = blockIdx.x;
506 - if (valid_rows_ptr != nullptr && row >= *valid_rows_ptr) return;
505 const scalar_t* gate_ptr = input + row * 2 * d;
506 const scalar_t* up_ptr = gate_ptr + d;
507 scalar_t* out_ptr = out + row * d;
@@ -558,15 +556,16 @@ template <typename scalar_t, typename fp8_type, int THREADS, int NUM_STAGES,
556 __global__ void situ_and_mul_quant_group_pipelined_kernel(
557 fp8_type* __restrict__ out, float* __restrict__ scale_out,
558 const scalar_t* __restrict__ input, const int64_t num_rows,
561 - const int32_t* __restrict__ valid_rows_ptr, const int64_t topk) {
559 + const int32_t* __restrict__ num_valid_tokens_ptr, const int64_t topk) {
560 const int tid = threadIdx.x;
561 const int warp_id = tid / WARP_SIZE;
562 const int lane_id = tid % WARP_SIZE;
563 static constexpr int NUM_WARPS = THREADS / WARP_SIZE;
564 static constexpr int NUM_GROUPS = D / GROUP_SIZE;
565 const int64_t row_bound =
568 - valid_rows_ptr != nullptr
569 - ? max((int64_t)0, min((int64_t)(*valid_rows_ptr) * topk, num_rows))
566 + num_valid_tokens_ptr != nullptr
567 + ? max((int64_t)0,
568 + min((int64_t)(*num_valid_tokens_ptr) * topk, num_rows))
569 : num_rows;
570
571 static_assert(NUM_GROUPS % 4 == 0,
@@ -710,15 +709,16 @@ __global__ void situ_and_mul_quant_group_scalar_kernel(
709 const scalar_t* __restrict__ input, // [num_tokens, 2, d]
710 const int d, const int num_groups, const float beta,
711 const float linear_beta, const int64_t num_rows,
713 - const int32_t* __restrict__ valid_rows_ptr, const int64_t topk) {
712 + const int32_t* __restrict__ num_valid_tokens_ptr, const int64_t topk) {
713 static constexpr int ELTS_PER_LANE = GROUP_SIZE / WARP_SIZE;
714 const int tid = threadIdx.x;
715 const int warp_id = tid / WARP_SIZE;
716 const int lane_id = tid % WARP_SIZE;
717 const int num_warps = blockDim.x / WARP_SIZE;
718 const int64_t row_bound =
720 - valid_rows_ptr != nullptr
721 - ? max((int64_t)0, min((int64_t)(*valid_rows_ptr) * topk, num_rows))
719 + num_valid_tokens_ptr != nullptr
720 + ? max((int64_t)0,
721 + min((int64_t)(*num_valid_tokens_ptr) * topk, num_rows))
722 : num_rows;
723
724 const bool clamp_up = linear_beta > 0.0f;
@@ -768,24 +768,61 @@ __global__ void situ_and_mul_quant_group_scalar_kernel(
768 }
769 }
770
771 -template <typename scalar_t>
771 +constexpr int kMaxMaskedTokenBlocks = 32;
772 +
773 +template <bool BATCHED_EXPERTS>
774 +__device__ __forceinline__ bool get_masked_row_range(
775 + const int* __restrict__ valid_token_counts, const int max_num_tokens,
776 + int64_t& first_row, int64_t& end_row, int64_t& row_stride) {
777 + if constexpr (BATCHED_EXPERTS) {
778 + // [E, T, *]: z lanes grid-stride over one expert's valid token prefix.
779 + const int expert = blockIdx.y;
780 + const int num_tokens =
781 + max(0, min(valid_token_counts[expert], max_num_tokens));
782 + const int token_block = blockIdx.z;
783 + if (token_block >= num_tokens) {
784 + return false;
785 + }
786 + const int64_t expert_first_row =
787 + static_cast<int64_t>(expert) * max_num_tokens;
788 + first_row = expert_first_row + token_block;
789 + end_row = expert_first_row + num_tokens;
790 + row_stride = gridDim.z;
791 + } else {
792 + // [T, *]: y lanes grid-stride over the single valid token prefix.
793 + const int num_tokens = max(0, min(valid_token_counts[0], max_num_tokens));
794 + const int token_block = blockIdx.y;
795 + if (token_block >= num_tokens) {
796 + return false;
797 + }
798 + first_row = token_block;
799 + end_row = num_tokens;
800 + row_stride = gridDim.y;
801 + }
802 + return first_row < end_row;
803 +}
804 +
805 +__device__ __forceinline__ int get_masked_feature_index() {
806 + return blockIdx.x * blockDim.x + threadIdx.x;
807 +}
808 +
809 +template <typename scalar_t, bool BATCHED_EXPERTS>
810 __global__ void masked_situ_and_mul_kernel(
811 scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
774 - const int* __restrict__ expert_num_tokens, const int max_num_tokens,
812 + const int* __restrict__ valid_token_counts, const int max_num_tokens,
813 const int d, const float beta, const float linear_beta) {
776 - const int expert = blockIdx.y;
777 - const int num_tokens = expert_num_tokens[expert];
778 - const int idx = blockIdx.x * blockDim.x + threadIdx.x;
779 - if (idx >= d || num_tokens == 0) {
814 + int64_t first_row, end_row, row_stride;
815 + const int idx = get_masked_feature_index();
816 + if (idx >= d ||
817 + !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
818 + first_row, end_row, row_stride)) {
819 return;
820 }
821
822 const bool clamp_up = linear_beta > 0.0f;
823 const float inv_beta = 1.0f / beta;
824 const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
786 - const int64_t expert_row = static_cast<int64_t>(expert) * max_num_tokens;
787 - for (int token = 0; token < num_tokens; ++token) {
788 - const int64_t row = expert_row + token;
825 + for (int64_t row = first_row; row < end_row; row += row_stride) {
826 const scalar_t* gate_ptr = input + row * 2 * d;
827 const scalar_t* up_ptr = gate_ptr + d;
828 scalar_t* out_ptr = out + row * d;
@@ -796,6 +833,119 @@ __global__ void masked_situ_and_mul_kernel(
833 }
834 }
835
836 +template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
837 + bool act_first, bool HAS_CLAMP, bool BATCHED_EXPERTS>
838 +__global__ void masked_act_and_mul_kernel(
839 + scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
840 + const int* __restrict__ valid_token_counts, const int max_num_tokens,
841 + const int d, const float limit, const float alpha, const float beta) {
842 + int64_t first_row, end_row, row_stride;
843 + const int idx = get_masked_feature_index();
844 + if (idx >= d ||
845 + !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
846 + first_row, end_row, row_stride)) {
847 + return;
848 + }
849 +
850 + for (int64_t row = first_row; row < end_row; row += row_stride) {
851 + const scalar_t* x_ptr = input + row * 2 * d;
852 + const scalar_t* y_ptr = x_ptr + d;
853 + scalar_t* out_ptr = out + row * d;
854 + const scalar_t x = VLLM_LDG(&x_ptr[idx]);
855 + const scalar_t y = VLLM_LDG(&y_ptr[idx]);
856 + out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit,
857 + alpha, beta);
858 + }
859 +}
860 +
861 +template <typename scalar_t, bool BATCHED_EXPERTS>
862 +__global__ void masked_swigluoai_and_mul_kernel(
863 + scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
864 + const int* __restrict__ valid_token_counts, const int max_num_tokens,
865 + const int d, const float alpha, const float limit) {
866 + int64_t first_row, end_row, row_stride;
867 + const int idx = get_masked_feature_index();
868 + if (idx >= d ||
869 + !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
870 + first_row, end_row, row_stride)) {
871 + return;
872 + }
873 +
874 + for (int64_t row = first_row; row < end_row; row += row_stride) {
875 + const scalar_t* in_ptr = input + row * 2 * d;
876 + scalar_t* out_ptr = out + row * d;
877 + out_ptr[idx] =
878 + swigluoai_and_mul(in_ptr[2 * idx], in_ptr[2 * idx + 1], alpha, limit);
879 + }
880 +}
881 +
882 +template <typename scalar_t, bool BATCHED_EXPERTS>
883 +__global__ void masked_swiglustep_and_mul_kernel(
884 + scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
885 + const int* __restrict__ valid_token_counts, const int max_num_tokens,
886 + const int d, const float limit) {
887 + int64_t first_row, end_row, row_stride;
888 + const int idx = get_masked_feature_index();
889 + if (idx >= d ||
890 + !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
891 + first_row, end_row, row_stride)) {
892 + return;
893 + }
894 +
895 + for (int64_t row = first_row; row < end_row; row += row_stride) {
896 + const scalar_t* gate_ptr = input + row * 2 * d;
897 + const scalar_t* up_ptr = gate_ptr + d;
898 + scalar_t* out_ptr = out + row * d;
899 + const float gate = (float)VLLM_LDG(&gate_ptr[idx]);
900 + const float up = (float)VLLM_LDG(&up_ptr[idx]);
901 + const float gate_silu = gate / (1.0f + expf(-gate));
902 + const float gate_clamped = fminf(gate_silu, limit);
903 + const float up_clamped = fmaxf(fminf(up, limit), -limit);
904 + out_ptr[idx] = (scalar_t)(gate_clamped * up_clamped);
905 + }
906 +}
907 +
908 +template <typename T>
909 +__device__ __forceinline__ T silu_no_mul_kernel(const T& x) {
910 + return silu_kernel(x, 1.0f);
911 +}
912 +
913 +template <typename T>
914 +__device__ __forceinline__ T gelu_no_mul_kernel(const T& x) {
915 + return gelu_kernel(x, 1.0f);
916 +}
917 +
918 +template <typename T>
919 +__device__ __forceinline__ T gelu_tanh_no_mul_kernel(const T& x) {
920 + return gelu_tanh_kernel(x, 1.0f);
921 +}
922 +
923 +template <typename T>
924 +__device__ __forceinline__ T relu2_no_mul_kernel(const T& x) {
925 + const float value = fmaxf((float)x, 0.0f);
926 + return (T)(value * value);
927 +}
928 +
929 +template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
930 + bool BATCHED_EXPERTS>
931 +__global__ void masked_activation_kernel(
932 + scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
933 + const int* __restrict__ valid_token_counts, const int max_num_tokens,
934 + const int d) {
935 + int64_t first_row, end_row, row_stride;
936 + const int idx = get_masked_feature_index();
937 + if (idx >= d ||
938 + !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
939 + first_row, end_row, row_stride)) {
940 + return;
941 + }
942 +
943 + for (int64_t row = first_row; row < end_row; row += row_stride) {
944 + const int64_t offset = row * d + idx;
945 + out[offset] = ACT_FN(VLLM_LDG(&input[offset]));
946 + }
947 +}
948 +
949 } // namespace vllm
950
951 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
@@ -890,8 +1040,7 @@ void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d]
1040 // through), matching SituAndMul(linear_beta=None) on the Python side.
1041 void situ_and_mul(torch::stable::Tensor& out, // [..., d]
1042 torch::stable::Tensor& input, // [..., 2 * d]
893 - double beta, double linear_beta,
894 - std::optional<torch::stable::Tensor> valid_rows) {
1043 + double beta, double linear_beta) {
1044 int d = input.size(-1) / 2;
1045 int64_t num_tokens = input.numel() / input.size(-1);
1046 if (num_tokens == 0) {
@@ -899,14 +1048,6 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d]
1048 }
1049 dim3 grid(num_tokens);
1050 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 - }
1051 const torch::stable::accelerator::DeviceGuard device_guard(
1052 input.get_device_index());
1053 const cudaStream_t stream = get_current_cuda_stream();
@@ -914,7 +1055,7 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d]
1055 input.scalar_type(), "situ_and_mul_kernel", [&] {
1056 vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
1057 out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
917 - d, (float)beta, (float)linear_beta, valid_rows_ptr);
1058 + d, (float)beta, (float)linear_beta);
1059 });
1060 }
1061
@@ -925,7 +1066,7 @@ void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d]
1066 int num_experts = input.size(0);
1067 int max_num_tokens = input.size(1);
1068 int d = input.size(2) / 2;
928 - if (num_experts == 0 || max_num_tokens == 0) {
1069 + if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
1070 return;
1071 }
1072 constexpr int block_size = 256;
@@ -936,21 +1077,128 @@ void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d]
1077 const cudaStream_t stream = get_current_cuda_stream();
1078 VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1079 input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
939 - vllm::masked_situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
940 - out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
941 - expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
942 - (float)beta, (float)linear_beta);
1080 + vllm::masked_situ_and_mul_kernel<scalar_t, true>
1081 + <<<grid, block, 0, stream>>>(
1082 + out.mutable_data_ptr<scalar_t>(),
1083 + input.const_data_ptr<scalar_t>(),
1084 + expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
1085 + (float)beta, (float)linear_beta);
1086 });
1087 }
1088
1089 +#define LAUNCH_MASKED_ACT_AND_MUL(KERNEL, ACT_FIRST, HAS_CLAMP, LIMIT, ALPHA, \
1090 + BETA) \
1091 + vllm::masked_act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST, \
1092 + HAS_CLAMP, BATCHED_EXPERTS> \
1093 + <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
1094 + input.const_data_ptr<scalar_t>(), \
1095 + valid_token_counts.const_data_ptr<int>(), \
1096 + max_num_tokens, d, LIMIT, ALPHA, BETA)
1097 +
1098 +#define LAUNCH_MASKED_ACTIVATION(KERNEL) \
1099 + vllm::masked_activation_kernel<scalar_t, KERNEL<scalar_t>, BATCHED_EXPERTS> \
1100 + <<<grid, block, 0, stream>>>( \
1101 + out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
1102 + valid_token_counts.const_data_ptr<int>(), max_num_tokens, d)
1103 +
1104 +void masked_moe_activation(
1105 + torch::stable::Tensor& out, // [T, d] or [E, T, d]
1106 + torch::stable::Tensor& input, // [T, d] or [E, T, d], gated has 2 * d
1107 + const torch::stable::Tensor& valid_token_counts,
1108 + const std::string& activation, double clamp_limit, double alpha,
1109 + double beta, double situ_beta, double situ_linear_beta) {
1110 + const bool batched_experts = input.dim() == 3;
1111 + const int num_experts = batched_experts ? input.size(0) : 1;
1112 + const int max_num_tokens = batched_experts ? input.size(1) : input.size(0);
1113 + const int d = out.size(-1);
1114 + if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
1115 + return;
1116 + }
1117 +
1118 + constexpr int block_size = 256;
1119 + const int feature_blocks = (d + block_size - 1) / block_size;
1120 + const int token_blocks =
1121 + std::min(max_num_tokens, vllm::kMaxMaskedTokenBlocks);
1122 + // Batched grid: (feature tile, expert, token lane); flat grid: (feature
1123 + // tile, token lane). Token lanes grid-stride each valid prefix.
1124 + dim3 grid = batched_experts ? dim3(feature_blocks, num_experts, token_blocks)
1125 + : dim3(feature_blocks, token_blocks);
1126 + dim3 block(block_size);
1127 + const torch::stable::accelerator::DeviceGuard device_guard(
1128 + input.get_device_index());
1129 + const cudaStream_t stream = get_current_cuda_stream();
1130 + auto launch = [&]<bool BATCHED_EXPERTS>() {
1131 + VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1132 + input.scalar_type(), "masked_moe_activation", [&] {
1133 + if (activation == "silu") {
1134 + LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, false, 0.0f,
1135 + 1.0f, 0.0f);
1136 + } else if (activation == "silu_with_clamp") {
1137 + LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, true,
1138 + (float)clamp_limit, (float)alpha,
1139 + (float)beta);
1140 + } else if (activation == "gelu") {
1141 + LAUNCH_MASKED_ACT_AND_MUL(vllm::gelu_kernel, true, false, 0.0f,
1142 + 1.0f, 0.0f);
1143 + } else if (activation == "gelu_tanh") {
1144 + LAUNCH_MASKED_ACT_AND_MUL(vllm::gelu_tanh_kernel, true, false, 0.0f,
1145 + 1.0f, 0.0f);
1146 + } else if (activation == "situ") {
1147 + vllm::masked_situ_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
1148 + <<<grid, block, 0, stream>>>(
1149 + out.mutable_data_ptr<scalar_t>(),
1150 + input.const_data_ptr<scalar_t>(),
1151 + valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
1152 + (float)situ_beta, (float)situ_linear_beta);
1153 + } else if (activation == "swigluoai") {
1154 + vllm::masked_swigluoai_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
1155 + <<<grid, block, 0, stream>>>(
1156 + out.mutable_data_ptr<scalar_t>(),
1157 + input.const_data_ptr<scalar_t>(),
1158 + valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
1159 + (float)alpha, (float)clamp_limit);
1160 + } else if (activation == "swigluoai_uninterleave") {
1161 + LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, true,
1162 + (float)clamp_limit, (float)alpha,
1163 + (float)beta);
1164 + } else if (activation == "swiglustep") {
1165 + vllm::masked_swiglustep_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
1166 + <<<grid, block, 0, stream>>>(
1167 + out.mutable_data_ptr<scalar_t>(),
1168 + input.const_data_ptr<scalar_t>(),
1169 + valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
1170 + (float)clamp_limit);
1171 + } else if (activation == "silu_no_mul") {
1172 + LAUNCH_MASKED_ACTIVATION(vllm::silu_no_mul_kernel);
1173 + } else if (activation == "gelu_no_mul") {
1174 + LAUNCH_MASKED_ACTIVATION(vllm::gelu_no_mul_kernel);
1175 + } else if (activation == "gelu_tanh_no_mul") {
1176 + LAUNCH_MASKED_ACTIVATION(vllm::gelu_tanh_no_mul_kernel);
1177 + } else if (activation == "relu2_no_mul") {
1178 + LAUNCH_MASKED_ACTIVATION(vllm::relu2_no_mul_kernel);
1179 + } else {
1180 + STD_TORCH_CHECK(false,
1181 + "Unsupported masked MoE activation: ", activation);
1182 + }
1183 + });
1184 + };
1185 + if (batched_experts) {
1186 + launch.template operator()<true>();
1187 + } else {
1188 + launch.template operator()<false>();
1189 + }
1190 +}
1191 +
1192 +#undef LAUNCH_MASKED_ACT_AND_MUL
1193 +#undef LAUNCH_MASKED_ACTIVATION
1194 // 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.
1195 +// `num_valid_tokens` is the int32 DeepEP token count (psum[-1]); rows = it *
1196 +// `topk`, so padding is excluded on-device and skipped rows receive scale 1.
1197 void situ_and_mul_quant(torch::stable::Tensor& out, // [..., d] (fp8)
1198 torch::stable::Tensor& scale, // [..., 1 or d/128]
1199 torch::stable::Tensor& input, // [..., 2 * d]
1200 double beta, double linear_beta, int64_t group_size,
953 - std::optional<torch::stable::Tensor> valid_rows,
1201 + std::optional<torch::stable::Tensor> num_valid_tokens,
1202 int64_t topk) {
1203 STD_TORCH_CHECK(
1204 out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
@@ -974,16 +1222,17 @@ void situ_and_mul_quant(torch::stable::Tensor& out, // [..., d] (fp8)
1222 }
1223 STD_TORCH_CHECK(out.size(-1) == d && out.numel() == num_tokens * (int64_t)d,
1224 "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");
1225 + const int32_t* num_valid_tokens_ptr = nullptr;
1226 + if (num_valid_tokens.has_value()) {
1227 + STD_TORCH_CHECK(
1228 + num_valid_tokens->is_cuda() &&
1229 + num_valid_tokens->get_device_index() == input.get_device_index(),
1230 + "situ_and_mul_quant: num_valid_tokens must be a CUDA tensor on the "
1231 + "same device as input");
1232 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>();
1233 + num_valid_tokens->scalar_type() == torch::headeronly::ScalarType::Int,
1234 + "situ_and_mul_quant: num_valid_tokens must be int32 (psum count)");
1235 + num_valid_tokens_ptr = num_valid_tokens->const_data_ptr<int32_t>();
1236 }
1237 const torch::stable::accelerator::DeviceGuard device_guard(
1238 input.get_device_index());
@@ -1031,7 +1280,7 @@ void situ_and_mul_quant(torch::stable::Tensor& out, // [..., d] (fp8)
1280 out.mutable_data_ptr<fp8_t>(),
1281 scale.mutable_data_ptr<float>(),
1282 input.const_data_ptr<scalar_t>(), num_tokens,
1034 - valid_rows_ptr, topk);
1283 + num_valid_tokens_ptr, topk);
1284 return;
1285 }
1286 }
@@ -1045,7 +1294,7 @@ void situ_and_mul_quant(torch::stable::Tensor& out, // [..., d] (fp8)
1294 scale.mutable_data_ptr<float>(),
1295 input.const_data_ptr<scalar_t>(), d, num_groups,
1296 (float)beta, (float)linear_beta, num_tokens,
1048 - valid_rows_ptr, topk);
1297 + num_valid_tokens_ptr, topk);
1298 });
1299 });
1300 }
csrc/libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.cu
+8 -7
@@ -146,16 +146,17 @@ __global__ void preprocessTopkIdKernel(int* topk_id_ptr, int size,
146 }
147 __syncthreads();
148
149 - // query global expert id in expert map.
150 - // if global expert id = -1 in exert map, plus n_expert
151 - // else set global expert id = exert map[global expert id]
149 if (offset + tidx < bound) {
150 auto topk_id = topk_id_ptr[offset + tidx];
154 - auto local_expert_idx = smem_expert_map[topk_id];
155 - if (local_expert_idx == -1) {
156 - topk_id += num_experts;
151 + if (topk_id < 0 || topk_id >= num_experts) {
152 + topk_id = num_experts;
153 } else {
158 - topk_id = local_expert_idx;
154 + auto local_expert_idx = smem_expert_map[topk_id];
155 + if (local_expert_idx == -1) {
156 + topk_id += num_experts;
157 + } else {
158 + topk_id = local_expert_idx;
159 + }
160 }
161 __syncwarp();
162 topk_id_ptr[offset + tidx] = topk_id;
csrc/libtorch_stable/ops.h
+10 -5
@@ -530,20 +530,25 @@ void fatrelu_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
530 double threshold);
531 void swigluoai_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
532 double alpha = 1.702, double limit = 7.0);
533 -void situ_and_mul(
534 - torch::stable::Tensor& out, torch::stable::Tensor& input, double beta = 1.0,
535 - double linear_beta = -1.0,
536 - std::optional<torch::stable::Tensor> valid_rows = std::nullopt);
533 +void situ_and_mul(torch::stable::Tensor& out, torch::stable::Tensor& input,
534 + double beta = 1.0, double linear_beta = -1.0);
535 void situ_and_mul_quant(
536 torch::stable::Tensor& out, torch::stable::Tensor& scale,
537 torch::stable::Tensor& input, double beta = 1.0, double linear_beta = -1.0,
538 int64_t group_size = 0,
541 - std::optional<torch::stable::Tensor> valid_rows = std::nullopt,
539 + std::optional<torch::stable::Tensor> num_valid_tokens = std::nullopt,
540 int64_t topk = 1);
541 void masked_situ_and_mul(torch::stable::Tensor& out,
542 torch::stable::Tensor& input,
543 const torch::stable::Tensor& expert_num_tokens,
544 double beta = 1.0, double linear_beta = -1.0);
545 +void masked_moe_activation(torch::stable::Tensor& out,
546 + torch::stable::Tensor& input,
547 + const torch::stable::Tensor& valid_token_counts,
548 + const std::string& activation,
549 + double clamp_limit = 0.0, double alpha = 1.0,
550 + double beta = 0.0, double situ_beta = 1.0,
551 + double situ_linear_beta = -1.0);
552 void gelu_new(torch::stable::Tensor& out, torch::stable::Tensor& input);
553 void gelu_fast(torch::stable::Tensor& out, torch::stable::Tensor& input);
554 void gelu_quick(torch::stable::Tensor& out, torch::stable::Tensor& input);
csrc/libtorch_stable/torch_bindings.cpp
+8 -2
@@ -614,7 +614,7 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
614 // SituGLU implementation used in Kimi models.
615 ops.def(
616 "situ_and_mul(Tensor! out, Tensor input, float beta=1.0, float "
617 - "linear_beta=-1.0, Tensor? valid_rows=None) -> ()");
617 + "linear_beta=-1.0) -> ()");
618 // Fused SituGLU activation + dynamic FP8 quantization for the Humming w2 path
619 // (writes the fp8 down input and its float32 scale). group_size=0 ->
620 // per-token scale [.., 1]; group_size=128 -> k-major block-FP8 scale [..,
@@ -622,10 +622,15 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) {
622 ops.def(
623 "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, "
624 "float beta=1.0, float linear_beta=-1.0, int group_size=0, "
625 - "Tensor? valid_rows=None, int topk=1) -> ()");
625 + "Tensor? num_valid_tokens=None, int topk=1) -> ()");
626 ops.def(
627 "masked_situ_and_mul(Tensor! out, Tensor input, Tensor "
628 "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()");
629 + ops.def(
630 + "masked_moe_activation(Tensor! out, Tensor input, Tensor "
631 + "valid_token_counts, str activation, float clamp_limit=0.0, float "
632 + "alpha=1.0, float beta=0.0, float situ_beta=1.0, float "
633 + "situ_linear_beta=-1.0) -> ()");
634
635 // GELU implementation used in GPT-2.
636 ops.def("gelu_new(Tensor! out, Tensor input) -> ()");
@@ -852,6 +857,7 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) {
857 ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul));
858 ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant));
859 ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul));
860 + ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation));
861 ops.impl("gelu_new", TORCH_BOX(&gelu_new));
862 ops.impl("gelu_fast", TORCH_BOX(&gelu_fast));
863 ops.impl("gelu_quick", TORCH_BOX(&gelu_quick));
tests/kernels/core/test_activation.py
+208
@@ -22,7 +22,13 @@ from vllm.model_executor.layers.activation import (
22 SwigluStepAndMul,
23 swiglustep_and_mul_triton,
24 )
25 +from vllm.model_executor.layers.fused_moe.activation import (
26 + ApplyMoEActivationConfig,
27 + MoEActivation,
28 + apply_moe_activation,
29 +)
30 from vllm.model_executor.layers.fused_moe.utils import swiglu_limit_func
31 +from vllm.platforms import current_platform
32 from vllm.utils.torch_utils import set_random_seed
33
34 DTYPES = [torch.half, torch.bfloat16, torch.float]
@@ -34,6 +40,100 @@ CUDA_DEVICES = [
40 ]
41
42
43 +def test_masked_moe_activation_rejects_unsupported_activation() -> None:
44 + input = torch.empty(1, 1, 2)
45 + output = torch.empty(1, 1, 1)
46 + valid_token_counts = torch.ones(1, dtype=torch.int32)
47 +
48 + with pytest.raises(NotImplementedError, match="relu2"):
49 + apply_moe_activation(
50 + MoEActivation.RELU2,
51 + output,
52 + input,
53 + valid_token_counts=valid_token_counts,
54 + )
55 +
56 +
57 +def test_moe_silu_clamp_uses_native_xpu_fallback(
58 + default_vllm_config, monkeypatch
59 +) -> None:
60 + monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
61 + clamp_limit = 3.0
62 + input = torch.tensor([[12.0, -12.0, 8.0, -8.0], [-2.0, 2.0, -4.0, 4.0]])
63 + output = torch.empty(2, 2)
64 +
65 + apply_moe_activation(
66 + MoEActivation.SILU,
67 + output,
68 + input,
69 + activation_config=ApplyMoEActivationConfig(clamp_limit=clamp_limit),
70 + )
71 +
72 + expected = SiluAndMulWithClamp(clamp_limit, compile_native=False).forward_native(
73 + input
74 + )
75 + torch.testing.assert_close(output, expected)
76 +
77 +
78 +def _assert_masked_moe_activation(
79 + activation: MoEActivation,
80 + activation_config: ApplyMoEActivationConfig,
81 + *,
82 + dtype: torch.dtype,
83 + mask_layout: str,
84 + d: int,
85 + max_num_tokens: int,
86 +) -> None:
87 + device = CUDA_DEVICES[0]
88 + num_experts = 4
89 + input_dim = 2 * d if activation.is_gated else d
90 + if mask_layout == "flat":
91 + input = torch.randn(max_num_tokens, input_dim, dtype=dtype, device=device)
92 + valid_token_counts = torch.tensor(
93 + [max_num_tokens // 2], dtype=torch.int32, device=device
94 + )
95 + output = torch.full((max_num_tokens, d), 42.0, dtype=dtype, device=device)
96 + else:
97 + input = torch.randn(
98 + num_experts, max_num_tokens, input_dim, dtype=dtype, device=device
99 + )
100 + valid_token_counts = torch.tensor(
101 + [0, 1, max_num_tokens // 2, max_num_tokens],
102 + dtype=torch.int32,
103 + device=device,
104 + )
105 + output = torch.full(
106 + (num_experts, max_num_tokens, d), 42.0, dtype=dtype, device=device
107 + )
108 +
109 + apply_moe_activation(
110 + activation,
111 + output,
112 + input,
113 + activation_config=activation_config,
114 + valid_token_counts=valid_token_counts,
115 + )
116 +
117 + batched_input = input.view(-1, max_num_tokens, input_dim)
118 + batched_output = output.view(-1, max_num_tokens, d)
119 + for expert, num_tokens in enumerate(valid_token_counts.cpu().tolist()):
120 + if num_tokens:
121 + expected = torch.empty((num_tokens, d), dtype=dtype, device=device)
122 + apply_moe_activation(
123 + activation,
124 + expected,
125 + batched_input[expert, :num_tokens].clone(),
126 + activation_config=activation_config,
127 + )
128 + torch.testing.assert_close(
129 + batched_output[expert, :num_tokens],
130 + expected,
131 + atol=get_default_atol(output),
132 + rtol=get_default_rtol(output),
133 + )
134 + assert torch.all(batched_output[expert, num_tokens:] == 42.0)
135 +
136 +
137 @pytest.mark.parametrize(
138 "activation",
139 [
@@ -252,6 +352,114 @@ def test_masked_situ_and_mul(
352 )
353
354
355 +@pytest.mark.parametrize(
356 + ("activation", "activation_config"),
357 + [
358 + (MoEActivation.SILU, ApplyMoEActivationConfig()),
359 + (
360 + MoEActivation.SILU,
361 + ApplyMoEActivationConfig(clamp_limit=3.0),
362 + ),
363 + (MoEActivation.GELU, ApplyMoEActivationConfig()),
364 + (MoEActivation.GELU_TANH, ApplyMoEActivationConfig()),
365 + (
366 + MoEActivation.SITU,
367 + ApplyMoEActivationConfig(
368 + activation_situ_beta=1.5,
369 + activation_situ_linear_beta=2.0,
370 + ),
371 + ),
372 + (MoEActivation.SWIGLUOAI, ApplyMoEActivationConfig()),
373 + (
374 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
375 + ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
376 + ),
377 + (MoEActivation.SWIGLUSTEP, ApplyMoEActivationConfig()),
378 + (MoEActivation.SILU_NO_MUL, ApplyMoEActivationConfig()),
379 + (MoEActivation.GELU_NO_MUL, ApplyMoEActivationConfig()),
380 + (MoEActivation.GELU_TANH_NO_MUL, ApplyMoEActivationConfig()),
381 + (MoEActivation.RELU2_NO_MUL, ApplyMoEActivationConfig()),
382 + ],
383 + ids=[
384 + "silu",
385 + "silu_clamp",
386 + "gelu",
387 + "gelu_tanh",
388 + "situ",
389 + "swigluoai",
390 + "swigluoai_uninterleave",
391 + "swiglustep",
392 + "silu_no_mul",
393 + "gelu_no_mul",
394 + "gelu_tanh_no_mul",
395 + "relu2_no_mul",
396 + ],
397 +)
398 +@torch.inference_mode()
399 +def test_masked_moe_activation_dispatch(
400 + default_vllm_config,
401 + activation: MoEActivation,
402 + activation_config: ApplyMoEActivationConfig,
403 +) -> None:
404 + _assert_masked_moe_activation(
405 + activation,
406 + activation_config,
407 + dtype=torch.bfloat16,
408 + mask_layout="batched_experts",
409 + d=513,
410 + max_num_tokens=7,
411 + )
412 +
413 +
414 +@pytest.mark.parametrize(
415 + ("activation", "activation_config", "mask_layout"),
416 + [
417 + pytest.param(
418 + MoEActivation.SILU,
419 + ApplyMoEActivationConfig(),
420 + "flat",
421 + id="flat",
422 + ),
423 + pytest.param(
424 + MoEActivation.SITU,
425 + ApplyMoEActivationConfig(
426 + activation_situ_beta=1.5,
427 + activation_situ_linear_beta=2.0,
428 + ),
429 + "batched_experts",
430 + id="batched-experts",
431 + ),
432 + ],
433 +)
434 +@torch.inference_mode()
435 +def test_masked_moe_activation_grid_stride(
436 + default_vllm_config,
437 + activation: MoEActivation,
438 + activation_config: ApplyMoEActivationConfig,
439 + mask_layout: str,
440 +) -> None:
441 + _assert_masked_moe_activation(
442 + activation,
443 + activation_config,
444 + dtype=torch.half,
445 + mask_layout=mask_layout,
446 + d=513,
447 + max_num_tokens=67,
448 + )
449 +
450 +
451 +@torch.inference_mode()
452 +def test_masked_moe_activation_opcheck(default_vllm_config) -> None:
453 + device = CUDA_DEVICES[0]
454 + input = torch.randn(2, 3, 64, dtype=torch.half, device=device)
455 + output = torch.empty(2, 3, 32, dtype=torch.half, device=device)
456 + valid_token_counts = torch.tensor([1, 3], dtype=torch.int32, device=device)
457 + opcheck(
458 + torch.ops._C.masked_moe_activation,
459 + (output, input, valid_token_counts, "silu", 0.0, 1.0, 0.0, 1.0, -1.0),
460 + )
461 +
462 +
463 @pytest.mark.parametrize(
464 "activation",
465 [
tests/kernels/moe/test_batched_moe.py
+145 -3
@@ -8,6 +8,7 @@ import torch
8
9 from tests.kernels.moe.utils import (
10 batched_moe,
11 + make_dummy_moe_config,
12 make_quantized_test_activations,
13 make_test_weights,
14 naive_batched_moe,
@@ -16,7 +17,13 @@ from tests.kernels.quant_utils import native_batched_masked_quant_matmul
17 from tests.kernels.utils import torch_experts
18 from vllm.config import VllmConfig, set_current_vllm_config
19 from vllm.model_executor.layers.fused_moe import fused_topk
20 +from vllm.model_executor.layers.fused_moe.activation import (
21 + ApplyMoEActivationConfig,
22 + MoEActivation,
23 + apply_moe_activation,
24 +)
25 from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
26 + BatchedTritonExperts,
27 invoke_moe_batched_triton_kernel,
28 )
29 from vllm.platforms import current_platform
@@ -496,11 +503,146 @@ def test_batched_mm_td_zero_expert_tokens():
503
504 # BatchedTritonExperts device enablement (XPU)
505 def test_batched_triton_experts_supports_current_device():
499 - from vllm.model_executor.layers.fused_moe.experts.fused_batched_moe import (
500 - BatchedTritonExperts,
506 + assert BatchedTritonExperts._supports_current_device()
507 +
508 +
509 +@pytest.mark.parametrize(
510 + "activation",
511 + [
512 + MoEActivation.SITU,
513 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
514 + MoEActivation.SWIGLUSTEP,
515 + ],
516 +)
517 +def test_batched_triton_supports_masked_activations(
518 + monkeypatch, activation: MoEActivation
519 +):
520 + monkeypatch.setattr(current_platform, "is_xpu", lambda: False)
521 + assert BatchedTritonExperts._supports_activation(activation)
522 +
523 + monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
524 + assert not BatchedTritonExperts._supports_activation(activation)
525 +
526 +
527 +def _torch_experts_with_shared_activation(
528 + a: torch.Tensor,
529 + w1: torch.Tensor,
530 + w2: torch.Tensor,
531 + topk_weights: torch.Tensor,
532 + topk_ids: torch.Tensor,
533 + activation: MoEActivation,
534 + activation_config: ApplyMoEActivationConfig,
535 +) -> torch.Tensor:
536 + num_tokens, hidden_dim = a.shape
537 + topk = topk_ids.shape[1]
538 + routed_input = a[:, None, :].expand(-1, topk, -1).reshape(-1, hidden_dim)
539 + routed_experts = topk_ids.reshape(-1)
540 + routed_output = torch.zeros_like(routed_input)
541 +
542 + for expert in range(w1.shape[0]):
543 + expert_mask = routed_experts == expert
544 + if not expert_mask.any():
545 + continue
546 + projected = routed_input[expert_mask] @ w1[expert].T
547 + activated = torch.empty(
548 + projected.shape[0], w2.shape[2], dtype=a.dtype, device=a.device
549 + )
550 + apply_moe_activation(
551 + activation,
552 + activated,
553 + projected,
554 + activation_config=activation_config,
555 + )
556 + routed_output[expert_mask] = activated @ w2[expert].T
557 +
558 + return (
559 + (
560 + routed_output.view(num_tokens, topk, hidden_dim).float()
561 + * topk_weights[..., None]
562 + )
563 + .sum(dim=1)
564 + .to(a.dtype)
565 )
566
503 - assert BatchedTritonExperts._supports_current_device()
567 +
568 +_MASKED_BATCHED_ACTIVATION_CASES = [
569 + pytest.param(
570 + MoEActivation.SITU,
571 + ApplyMoEActivationConfig(
572 + activation_situ_beta=1.5,
573 + activation_situ_linear_beta=2.0,
574 + ),
575 + id="situ",
576 + ),
577 + pytest.param(
578 + MoEActivation.SWIGLUOAI_UNINTERLEAVE,
579 + ApplyMoEActivationConfig(clamp_limit=3.0, alpha=1.3, beta=0.5),
580 + id="swigluoai-uninterleave",
581 + ),
582 + pytest.param(
583 + MoEActivation.SWIGLUSTEP,
584 + ApplyMoEActivationConfig(),
585 + id="swiglustep",
586 + ),
587 +]
588 +
589 +
590 +@pytest.mark.parametrize(
591 + ("activation", "activation_config"), _MASKED_BATCHED_ACTIVATION_CASES
592 +)
593 +@torch.inference_mode()
594 +def test_batched_experts_masked_activation_end_to_end(
595 + activation: MoEActivation,
596 + activation_config: ApplyMoEActivationConfig,
597 +):
598 + if current_platform.is_xpu():
599 + pytest.skip("Masked activations are not enabled on XPU")
600 +
601 + from vllm.v1.worker.workspace import init_workspace_manager
602 +
603 + set_random_seed(7)
604 + device = current_platform.device_type
605 + init_workspace_manager(torch.device(f"{device}:0"))
606 + m, n, k, e, topk = 16, 128, 128, 8, 2
607 +
608 + a = torch.randn((m, k), device=device, dtype=torch.bfloat16) / 10
609 + score = torch.randn((m, e), device=device, dtype=torch.bfloat16)
610 + w1 = torch.randn((e, 2 * n, k), device=device, dtype=torch.bfloat16) / 15
611 + w2 = torch.randn((e, k, n), device=device, dtype=torch.bfloat16) / 15
612 + moe_config = make_dummy_moe_config(
613 + num_experts=e,
614 + experts_per_token=topk,
615 + hidden_dim=k,
616 + intermediate_size=n,
617 + in_dtype=a.dtype,
618 + activation=activation,
619 + )
620 + moe_config.swiglu_limit = activation_config.clamp_limit
621 + moe_config.swiglu_alpha = activation_config.alpha
622 + moe_config.swiglu_beta = activation_config.beta
623 + moe_config.activation_situ_beta = activation_config.activation_situ_beta
624 + moe_config.activation_situ_linear_beta = (
625 + activation_config.activation_situ_linear_beta
626 + )
627 +
628 + routing_weights = torch.softmax(score.float(), dim=-1)
629 + topk_weights, topk_ids = torch.topk(routing_weights, topk, dim=-1)
630 + topk_ids = topk_ids.to(torch.int32)
631 + with set_current_vllm_config(vllm_config):
632 + baseline_output = _torch_experts_with_shared_activation(
633 + a,
634 + w1,
635 + w2,
636 + topk_weights,
637 + topk_ids,
638 + activation,
639 + activation_config,
640 + )
641 + triton_output = batched_moe(
642 + a, w1, w2, topk_weights, topk_ids, moe_config=moe_config
643 + )
644 +
645 + torch.testing.assert_close(triton_output, baseline_output, atol=3e-2, rtol=2e-2)
646
647
648 @pytest.mark.parametrize("m,n,k,e,topk", [(32, 512, 512, 8, 2), (45, 1024, 128, 8, 1)])
tests/kernels/moe/test_moe.py
+260 -1
@@ -1377,6 +1377,102 @@ def test_humming_activation_metadata_tracks_shared_apply(activation: MoEActivati
1377 ) == apply_moe_activation_supported(activation)
1378
1379
1380 +@pytest.mark.parametrize(
1381 + ("env_override", "use_ep", "expected"),
1382 + [
1383 + (None, False, "indexed"),
1384 + (None, True, "grouped_contiguous"),
1385 + ("auto", True, "grouped_contiguous"),
1386 + ("indexed", True, "indexed"),
1387 + ("grouped", False, "grouped_contiguous"),
1388 + ],
1389 +)
1390 +def test_humming_selects_gemm_from_parallelism_and_override(
1391 + monkeypatch: pytest.MonkeyPatch,
1392 + env_override: str | None,
1393 + use_ep: bool,
1394 + expected: str,
1395 +):
1396 + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1397 + get_humming_moe_gemm_type,
1398 + )
1399 +
1400 + if env_override is None:
1401 + monkeypatch.delenv("VLLM_HUMMING_MOE_GEMM_TYPE", raising=False)
1402 + else:
1403 + monkeypatch.setenv("VLLM_HUMMING_MOE_GEMM_TYPE", env_override)
1404 + moe_config = make_dummy_moe_config()
1405 + moe_config.moe_parallel_config.use_ep = use_ep
1406 +
1407 + assert get_humming_moe_gemm_type(moe_config) == expected
1408 +
1409 +
1410 +@pytest.mark.parametrize(
1411 + ("dp_token_counts", "topk_shape", "expected"),
1412 + [
1413 + pytest.param(None, (7, 3), 21, id="runtime-layout"),
1414 + pytest.param([3, 5], (48, 1), 48, id="expanded-dp-layout"),
1415 + ],
1416 +)
1417 +def test_humming_global_valid_shape_m(
1418 + monkeypatch: pytest.MonkeyPatch,
1419 + dp_token_counts: list[int] | None,
1420 + topk_shape: tuple[int, int],
1421 + expected: int,
1422 +):
1423 + from types import SimpleNamespace
1424 +
1425 + import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
1426 +
1427 + dp_metadata = (
1428 + None
1429 + if dp_token_counts is None
1430 + else SimpleNamespace(
1431 + num_tokens_across_dp_cpu=torch.tensor(dp_token_counts, dtype=torch.int32)
1432 + )
1433 + )
1434 + monkeypatch.setattr(
1435 + humming,
1436 + "get_forward_context",
1437 + lambda: SimpleNamespace(dp_metadata=dp_metadata),
1438 + )
1439 + experts = SimpleNamespace(moe_config=make_dummy_moe_config(experts_per_token=6))
1440 + topk_ids = torch.empty(topk_shape, dtype=torch.int64)
1441 +
1442 + result = humming.HummingExpertsBase.get_global_valid_shape_m(experts, topk_ids)
1443 +
1444 + assert result == expected
1445 +
1446 +
1447 +def test_humming_permute_scratch_is_keyed_by_runtime_topk(
1448 + monkeypatch: pytest.MonkeyPatch,
1449 +):
1450 + from types import SimpleNamespace
1451 + from unittest.mock import Mock
1452 +
1453 + import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
1454 +
1455 + scratch_topk6 = Mock()
1456 + scratch_topk1 = Mock()
1457 + scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1])
1458 + monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True)
1459 + monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type)
1460 + moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6)
1461 + moe_config.moe_parallel_config.dp_size = 2
1462 + experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config)
1463 +
1464 + assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
1465 + assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1
1466 + assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6
1467 +
1468 + assert scratch_type.call_count == 2
1469 + first_call, second_call = scratch_type.call_args_list
1470 + assert first_call.kwargs["max_num_tokens"] == 1024
1471 + assert first_call.kwargs["topk"] == 6
1472 + assert second_call.kwargs["max_num_tokens"] == 6144
1473 + assert second_call.kwargs["topk"] == 1
1474 +
1475 +
1476 def test_humming_delegates_to_instance_activation():
1477 from types import SimpleNamespace
1478 from unittest.mock import Mock
@@ -1408,10 +1504,173 @@ def test_humming_delegates_to_instance_activation():
1504 activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
1505 input=input,
1506 output=output,
1411 - valid_rows=None,
1507 )
1508
1509
1510 +def test_humming_grouped_apply_forwards_valid_prefix(
1511 + monkeypatch: pytest.MonkeyPatch,
1512 +):
1513 + from types import MethodType, SimpleNamespace
1514 + from unittest.mock import Mock
1515 +
1516 + import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming
1517 + from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
1518 + HummingExpertsBase,
1519 + HummingGroupedExperts,
1520 + )
1521 +
1522 + activation_func = Mock()
1523 + buffers = {
1524 + "gate_up_output": torch.empty(6, 4),
1525 + "activation_output": torch.empty(6, 2),
1526 + "down_output": torch.empty(6, 4),
1527 + }
1528 + experts = SimpleNamespace(
1529 + activation=activation_func,
1530 + num_experts=2,
1531 + estimate_local_valid_shape_m=lambda _: 6,
1532 + prepare_buffers=lambda *_: buffers,
1533 + _get_permute_scratch=lambda _: None,
1534 + quantize_input=lambda _, *, inputs, input_scale=None, quanted_input: (
1535 + inputs,
1536 + input_scale,
1537 + ),
1538 + humming_forward=Mock(),
1539 + compute_config_str="",
1540 + w13_tuning_config_str="",
1541 + w2_tuning_config_str="",
1542 + )
1543 + experts.apply_activation = MethodType(HummingExpertsBase.apply_activation, experts)
1544 +
1545 + hidden_states = torch.empty(3, 4)
1546 + output = torch.empty_like(hidden_states)
1547 + topk_ids = torch.zeros((3, 2), dtype=torch.int64)
1548 + topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
1549 + expected_counts = torch.tensor([4], dtype=torch.int32)
1550 + expert_offsets = torch.tensor([0, 2, 4], dtype=torch.int64)
1551 + monkeypatch.setattr(
1552 + humming,
1553 + "moe_permute",
1554 + lambda **_: (
1555 + torch.empty(6, 4),
1556 + None,
1557 + expert_offsets,
1558 + torch.arange(6),
1559 + None,
1560 + ),
1561 + )
1562 + monkeypatch.setattr(humming, "moe_unpermute", lambda **_: None)
1563 +
1564 + HummingGroupedExperts.apply(
1565 + experts,
1566 + output=output,
1567 + hidden_states=hidden_states,
1568 + w1=torch.empty(0),
1569 + w2=torch.empty(0),
1570 + topk_weights=topk_weights,
1571 + topk_ids=topk_ids,
1572 + activation=MoEActivation.SITU,
1573 + global_num_experts=2,
1574 + expert_map=None,
1575 + a1q_scale=None,
1576 + a2_scale=None,
1577 + workspace13=torch.empty(0),
1578 + workspace2=torch.empty(0),
1579 + expert_tokens_meta=None,
1580 + apply_router_weight_on_input=False,
1581 + )
1582 +
1583 + activation_func.assert_called_once()
1584 + call_kwargs = activation_func.call_args.kwargs
1585 + assert call_kwargs["activation"] == MoEActivation.SITU
1586 + assert call_kwargs["input"].shape == (6, 4)
1587 + assert call_kwargs["output"].shape == (6, 2)
1588 + torch.testing.assert_close(call_kwargs["valid_token_counts"], expected_counts)
1589 +
1590 +
1591 +def test_batched_marlin_activation_uses_expert_token_counts(
1592 + monkeypatch: pytest.MonkeyPatch,
1593 +):
1594 + from types import SimpleNamespace
1595 + from unittest.mock import Mock
1596 +
1597 + import vllm.model_executor.layers.fused_moe.experts.marlin_moe as marlin
1598 + from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
1599 + BatchedMarlinExperts,
1600 + )
1601 + from vllm.model_executor.layers.fused_moe.modular_kernel import (
1602 + ExpertTokensMetadata,
1603 + )
1604 +
1605 + activation_func = Mock()
1606 + experts = SimpleNamespace(
1607 + activation=activation_func,
1608 + w1_bias=None,
1609 + w2_bias=None,
1610 + w1_scale=None,
1611 + w2_scale=None,
1612 + quant_type_id=0,
1613 + a1_gscale=None,
1614 + a2_gscale=None,
1615 + w13_g_idx=None,
1616 + w2_g_idx=None,
1617 + w13_g_idx_sort_indices=None,
1618 + w2_g_idx_sort_indices=None,
1619 + w1_zp=None,
1620 + w2_zp=None,
1621 + input_dtype=None,
1622 + is_k_full=True,
1623 + activation_config=ApplyMoEActivationConfig(
1624 + clamp_limit=3.0, alpha=1.3, beta=0.5
1625 + ),
1626 + )
1627 + hidden_states = torch.empty(2, 3, 4)
1628 + expert_num_tokens = torch.tensor([1, 2], dtype=torch.int32)
1629 + topk_ids = torch.zeros((2, 3), dtype=torch.int64)
1630 + expert_map = torch.arange(2, dtype=torch.int32)
1631 +
1632 + def fake_batched_fused_marlin_moe(**kwargs):
1633 + kwargs["activation_func"](
1634 + kwargs["activation"],
1635 + torch.empty(6, 4),
1636 + torch.empty(6, 8),
1637 + topk_ids=topk_ids,
1638 + expert_map=expert_map,
1639 + )
1640 + return kwargs["output"]
1641 +
1642 + monkeypatch.setattr(
1643 + marlin, "batched_fused_marlin_moe", fake_batched_fused_marlin_moe
1644 + )
1645 + output = torch.empty_like(hidden_states)
1646 +
1647 + BatchedMarlinExperts.apply(
1648 + experts,
1649 + output=output,
1650 + hidden_states=hidden_states,
1651 + w1=torch.empty(0),
1652 + w2=torch.empty(0),
1653 + topk_weights=torch.ones_like(topk_ids, dtype=torch.float32),
1654 + topk_ids=topk_ids,
1655 + activation=MoEActivation.SWIGLUOAI_UNINTERLEAVE,
1656 + global_num_experts=2,
1657 + expert_map=expert_map,
1658 + a1q_scale=None,
1659 + a2_scale=None,
1660 + workspace13=torch.empty(0),
1661 + workspace2=torch.empty(0),
1662 + expert_tokens_meta=ExpertTokensMetadata(expert_num_tokens, None),
1663 + apply_router_weight_on_input=False,
1664 + )
1665 +
1666 + activation_func.assert_called_once()
1667 + call = activation_func.call_args
1668 + assert call.args[0] == MoEActivation.SWIGLUOAI_UNINTERLEAVE
1669 + assert call.args[1].shape == (2, 3, 4)
1670 + assert call.args[2].shape == (2, 3, 8)
1671 + torch.testing.assert_close(call.kwargs["valid_token_counts"], expert_num_tokens)
1672 +
1673 +
1674 @pytest.mark.parametrize(
1675 "activation",
1676 [
tests/kernels/moe/test_moe_permute_unpermute.py
+36
@@ -286,3 +286,39 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):
286 permuted_idx_1.untyped_storage().data_ptr()
287 == scratch.permuted_idx.untyped_storage().data_ptr()
288 )
289 +
290 +
291 +def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:
292 + if not moe_permute_unpermute_supported():
293 + pytest.skip("moe_permute_unpermute is not supported on this platform.")
294 +
295 + hidden_states = torch.arange(5 * 16, dtype=torch.bfloat16, device="cuda").view(
296 + 5, 16
297 + )
298 + topk_ids = torch.tensor([[0], [-1], [1], [4], [2]], device="cuda")
299 + expert_map = torch.tensor([0, 1, -1, -1], dtype=torch.int32, device="cuda")
300 + scratch = MoEPermuteScratch(
301 + max_num_tokens=5,
302 + topk=1,
303 + num_experts=4,
304 + num_local_experts=2,
305 + device=hidden_states.device,
306 + hidden_size=16,
307 + hidden_dtype=hidden_states.dtype,
308 + )
309 +
310 + permuted, _, expert_offsets, _, _ = moe_permute(
311 + hidden_states=hidden_states,
312 + a1q_scale=None,
313 + topk_ids=topk_ids,
314 + n_expert=4,
315 + n_local_expert=2,
316 + expert_map=expert_map,
317 + scratch=scratch,
318 + )
319 +
320 + torch.testing.assert_close(
321 + expert_offsets,
322 + torch.tensor([0, 1, 2], dtype=torch.int64, device="cuda"),
323 + )
324 + torch.testing.assert_close(permuted[:2], hidden_states[[0, 2]])
tests/kernels/moe/test_situ_mul_fp8_quant.py
+7 -7
@@ -5,7 +5,7 @@ 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.
8 +the masked 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
@@ -40,7 +40,7 @@ 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
43 +# (num_tokens, num_valid_tokens). None means all rows are valid. 2048/4096 rows
44 # span multiple grid-stride waves of the persistent grid (GRID_DIM = 132*8).
45 CASES = [
46 (1, None),
@@ -111,7 +111,7 @@ def test_situ_and_mul_quant_pipelined(
111 (num_tokens, d // GROUP_SIZE), -1.0, dtype=torch.float32, device=DEVICE
112 )
113 out_sentinel = out.clone()
114 - valid_rows = (
114 + num_valid_tokens = (
115 None if valid is None else torch.tensor(valid, dtype=torch.int32, device=DEVICE)
116 )
117
@@ -122,7 +122,7 @@ def test_situ_and_mul_quant_pipelined(
122 beta=SITU_BETA,
123 linear_beta=SITU_LINEAR_BETA,
124 group_size=GROUP_SIZE,
125 - valid_rows=valid_rows,
125 + num_valid_tokens=num_valid_tokens,
126 )
127
128 # Padding contract: skipped rows keep their bytes; their scale becomes 1.0.
@@ -157,7 +157,7 @@ def test_situ_and_mul_quant_pipelined(
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."""
160 + """The kernel expands num_valid_tokens by topk to bound rows."""
161 set_random_seed(7)
162 d = SITU_D
163 num_tokens = tokens * topk + 37 # padding tail past the valid rows
@@ -168,7 +168,7 @@ def test_situ_and_mul_quant_topk_row_bound(tokens: int, topk: int) -> None:
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)
171 + num_valid_tokens = torch.tensor(tokens, dtype=torch.int32, device=DEVICE)
172
173 situ_and_mul_quant(
174 out,
@@ -177,7 +177,7 @@ def test_situ_and_mul_quant_topk_row_bound(tokens: int, topk: int) -> None:
177 beta=SITU_BETA,
178 linear_beta=SITU_LINEAR_BETA,
179 group_size=GROUP_SIZE,
180 - valid_rows=valid_rows,
180 + num_valid_tokens=num_valid_tokens,
181 topk=topk,
182 )
183
tests/kernels/moe/utils.py
+3 -1
@@ -128,6 +128,7 @@ def batched_moe(
128 quant_dtype: torch.dtype | None = None,
129 per_act_token_quant: bool = False,
130 block_shape: list[int] | None = None,
131 + moe_config: FusedMoEConfig | None = None,
132 ) -> torch.Tensor:
133 max_num_tokens = round_up(a.shape[0], 64)
134
@@ -141,7 +142,8 @@ def batched_moe(
142 a2_scale=a2_scale,
143 )
144
144 - moe_config = make_dummy_moe_config()
145 + if moe_config is None:
146 + moe_config = make_dummy_moe_config()
147
148 fused_experts = FusedMoEKernel(
149 BatchedPrepareAndFinalize(
vllm/model_executor/layers/fused_moe/activation.py
+139 -22
@@ -9,6 +9,8 @@ from typing import TYPE_CHECKING
9 import torch
10 import torch.nn.functional as F
11
12 +from vllm.platforms import current_platform
13 +
14 if TYPE_CHECKING:
15 from vllm.model_executor.layers.fused_moe.config import (
16 FusedMoEConfig,
@@ -134,11 +136,31 @@ _APPLY_MOE_ACTIVATIONS = frozenset(
136 )
137
138
139 +_MASKED_MOE_ACTIVATION_NAMES: dict[MoEActivation, str] = {
140 + MoEActivation.SILU: "silu",
141 + MoEActivation.GELU: "gelu",
142 + MoEActivation.GELU_TANH: "gelu_tanh",
143 + MoEActivation.SITU: "situ",
144 + MoEActivation.SWIGLUOAI: "swigluoai",
145 + MoEActivation.SWIGLUOAI_UNINTERLEAVE: "swigluoai_uninterleave",
146 + MoEActivation.SWIGLUSTEP: "swiglustep",
147 + MoEActivation.SILU_NO_MUL: "silu_no_mul",
148 + MoEActivation.GELU_NO_MUL: "gelu_no_mul",
149 + MoEActivation.GELU_TANH_NO_MUL: "gelu_tanh_no_mul",
150 + MoEActivation.RELU2_NO_MUL: "relu2_no_mul",
151 +}
152 +
153 +
154 def apply_moe_activation_supported(activation: MoEActivation) -> bool:
155 """Whether ``apply_moe_activation`` supports an activation."""
156 return activation in _APPLY_MOE_ACTIVATIONS
157
158
159 +def apply_moe_activation_masked_supported(activation: MoEActivation) -> bool:
160 + """Whether the masked ``apply_moe_activation`` path supports an activation."""
161 + return activation in _MASKED_MOE_ACTIVATION_NAMES
162 +
163 +
164 @dataclass(frozen=True)
165 class ApplyMoEActivationConfig:
166 """Configuration forwarded to ``apply_moe_activation``."""
@@ -177,6 +199,96 @@ class ApplyMoEActivationConfig:
199 _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG = ApplyMoEActivationConfig()
200
201
202 +def _validate_moe_activation_shapes(
203 + activation: MoEActivation,
204 + output: torch.Tensor,
205 + input: torch.Tensor,
206 + expected_dim: int,
207 +) -> None:
208 + assert input.dim() == expected_dim, f"Input must be {expected_dim}D"
209 + assert output.dim() == expected_dim, f"Output must be {expected_dim}D"
210 + assert input.shape[:-1] == output.shape[:-1], (
211 + f"Input/output leading shapes must match: {input.shape} vs {output.shape}"
212 + )
213 + if activation.is_gated:
214 + assert output.size(-1) * 2 == input.size(-1), (
215 + f"{activation.value} expects 2x ratio: "
216 + f"{output.size(-1) * 2} vs {input.size(-1)}"
217 + )
218 + else:
219 + assert output.size(-1) == input.size(-1), (
220 + f"{activation.value} expects equal sizes: "
221 + f"{output.size(-1)} vs {input.size(-1)}"
222 + )
223 +
224 +
225 +def _apply_moe_activation_masked(
226 + activation: MoEActivation,
227 + output: torch.Tensor,
228 + input: torch.Tensor,
229 + valid_token_counts: torch.Tensor,
230 + config: ApplyMoEActivationConfig,
231 +) -> torch.Tensor:
232 + masked_activation = _MASKED_MOE_ACTIVATION_NAMES.get(activation)
233 + if masked_activation is None:
234 + raise NotImplementedError(
235 + f"Masked MoE activation is not implemented for {activation.value}"
236 + )
237 +
238 + assert input.dim() in (2, 3), "Masked input must be 2D or 3D"
239 + _validate_moe_activation_shapes(activation, output, input, expected_dim=input.dim())
240 + assert input.dtype == output.dtype, "Input and output dtypes must match"
241 + assert input.device == output.device, "Input and output devices must match"
242 + assert input.is_contiguous(), "Input must be contiguous"
243 + assert output.is_contiguous(), "Output must be contiguous"
244 + assert valid_token_counts.dtype == torch.int32, (
245 + "valid_token_counts must use torch.int32"
246 + )
247 + assert valid_token_counts.dim() == 1, "valid_token_counts must be 1D"
248 + assert valid_token_counts.device == input.device, (
249 + "valid_token_counts must be on the input device"
250 + )
251 + assert valid_token_counts.is_contiguous(), "valid_token_counts must be contiguous"
252 + expected_counts = input.size(0) if input.dim() == 3 else 1
253 + assert valid_token_counts.size(0) == expected_counts, (
254 + f"valid_token_counts must have {expected_counts} element(s) for "
255 + f"{input.dim()}D input"
256 + )
257 +
258 + masked_clamp_limit = 0.0 if config.clamp_limit is None else config.clamp_limit
259 + masked_alpha = config.alpha
260 + if activation == MoEActivation.SILU and config.clamp_limit is not None:
261 + masked_activation = "silu_with_clamp"
262 + elif activation == MoEActivation.SITU:
263 + assert config.activation_situ_beta is not None, (
264 + "SITU requires activation_situ_beta from FusedMoEConfig"
265 + )
266 + elif activation == MoEActivation.SWIGLUOAI:
267 + masked_clamp_limit = 7.0
268 + masked_alpha = 1.702
269 + elif activation == MoEActivation.SWIGLUOAI_UNINTERLEAVE:
270 + assert config.clamp_limit is not None, (
271 + "SWIGLUOAI_UNINTERLEAVE requires clamp_limit"
272 + )
273 + elif activation == MoEActivation.SWIGLUSTEP:
274 + masked_clamp_limit = 7.0
275 +
276 + torch.ops._C.masked_moe_activation(
277 + output,
278 + input,
279 + valid_token_counts,
280 + masked_activation,
281 + masked_clamp_limit,
282 + masked_alpha,
283 + config.beta,
284 + 1.0 if config.activation_situ_beta is None else config.activation_situ_beta,
285 + -1.0
286 + if config.activation_situ_linear_beta is None
287 + else config.activation_situ_linear_beta,
288 + )
289 + return output
290 +
291 +
292 def silu_and_mul_with_clamp(
293 output: torch.Tensor,
294 input: torch.Tensor,
@@ -184,7 +296,12 @@ def silu_and_mul_with_clamp(
296 topk_ids: torch.Tensor | None = None,
297 expert_map: torch.Tensor | None = None,
298 ) -> None:
187 - if topk_ids is not None and expert_map is not None:
299 + if current_platform.is_xpu():
300 + d = input.shape[-1] // 2
301 + gate = torch.clamp(input[..., :d], max=clamp_limit)
302 + up = torch.clamp(input[..., d:], min=-clamp_limit, max=clamp_limit)
303 + output.copy_(gate * torch.sigmoid(gate) * up)
304 + elif topk_ids is not None and expert_map is not None:
305 from vllm.model_executor.layers.fused_moe.utils import swiglu_limit_func
306
307 swiglu_limit_func(output, input, clamp_limit, topk_ids, expert_map)
@@ -201,12 +318,15 @@ def apply_moe_activation(
318 activation_config: ApplyMoEActivationConfig | None = None,
319 topk_ids: torch.Tensor | None = None,
320 expert_map: torch.Tensor | None = None,
204 - valid_rows: torch.Tensor | None = None,
321 + valid_token_counts: torch.Tensor | None = None,
322 ) -> torch.Tensor:
323 """Apply MoE activation function.
324
325 The configuration drives specialized activation behavior. Routing tensors
209 - remain per-call inputs because they depend on the current token assignment.
326 + and valid token counts remain per-call inputs because they depend on the
327 + current token assignment. A single token count masks a flat ``[T, D]``
328 + buffer; one count per expert masks each prefix in a padded ``[E, T, D]``
329 + buffer.
330 """
331 config = (
332 _DEFAULT_APPLY_MOE_ACTIVATION_CONFIG
@@ -214,19 +334,17 @@ def apply_moe_activation(
334 else activation_config
335 )
336
217 - assert input.dim() == 2, "Input must be 2D"
218 - assert output.dim() == 2, "Output must be 2D"
219 - if activation.is_gated:
220 - assert output.size(-1) * 2 == input.size(-1), (
221 - f"{activation.value} expects 2x ratio: "
222 - f"{output.size(-1) * 2} vs {input.size(-1)}"
223 - )
224 - else:
225 - assert output.size(-1) == input.size(-1), (
226 - f"{activation.value} expects equal sizes: "
227 - f"{output.size(-1)} vs {input.size(-1)}"
337 + if valid_token_counts is not None:
338 + return _apply_moe_activation_masked(
339 + activation,
340 + output,
341 + input,
342 + valid_token_counts,
343 + config,
344 )
345
346 + _validate_moe_activation_shapes(activation, output, input, expected_dim=2)
347 +
348 # Activations with gated multiplication (gate × activation(up))
349 if activation == MoEActivation.SILU:
350 if config.clamp_limit is not None:
@@ -257,7 +375,6 @@ def apply_moe_activation(
375 -1.0
376 if config.activation_situ_linear_beta is None
377 else config.activation_situ_linear_beta,
260 - valid_rows,
378 )
379 elif activation == MoEActivation.SWIGLUOAI:
380 torch.ops._C.swigluoai_and_mul(output, input)
@@ -298,7 +415,7 @@ def situ_and_mul_quant(
415 beta: float,
416 linear_beta: float | None,
417 group_size: int = 0,
301 - valid_rows: torch.Tensor | None = None,
418 + num_valid_tokens: torch.Tensor | None = None,
419 topk: int = 1,
420 ) -> None:
421 """Fused Kimi SITU activation + dynamic FP8 quantization.
@@ -314,11 +431,11 @@ def situ_and_mul_quant(
431 [M, d // 128], matching humming ``quant_input(group_size=128, float32)``.
432
433 ``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.
434 + ``SituAndMul(linear_beta=None)``. ``num_valid_tokens`` (int32 scalar tensor)
435 + is the DeepEP v2 contiguous-layout token count (e.g. ``psum[-1:]``); the
436 + kernel multiplies it by ``topk`` to bound rows. Padding rows are skipped and
437 + receive a benign scale of 1.0. Kept in this module so that every
438 + ``torch.ops._C.situ*`` call lives in one file.
439 """
440 torch.ops._C.situ_and_mul_quant(
441 output,
@@ -327,6 +444,6 @@ def situ_and_mul_quant(
444 beta,
445 -1.0 if linear_beta is None else linear_beta,
446 group_size,
330 - valid_rows,
447 + num_valid_tokens,
448 topk,
449 )
vllm/model_executor/layers/fused_moe/experts/fused_batched_moe.py
+29 -31
@@ -5,7 +5,10 @@
5 import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation
8 +from vllm.model_executor.layers.fused_moe.activation import (
9 + MoEActivation,
10 + apply_moe_activation_masked_supported,
11 +)
12 from vllm.model_executor.layers.fused_moe.config import (
13 FusedMoEConfig,
14 FusedMoEParallelConfig,
@@ -19,7 +22,6 @@ from vllm.model_executor.layers.fused_moe.utils import (
22 _resize_cache,
23 moe_kernel_quantize_input,
24 normalize_batched_scales_shape,
22 - swiglu_limit_func,
25 )
26 from vllm.model_executor.layers.quantization.utils.quant_utils import (
27 QuantKey,
@@ -839,16 +841,18 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
841
842 @staticmethod
843 def _supports_activation(activation: MoEActivation) -> bool:
842 - return activation in [
843 - MoEActivation.SILU,
844 - MoEActivation.GELU,
845 - MoEActivation.GELU_TANH,
846 - MoEActivation.SWIGLUOAI,
847 - MoEActivation.SILU_NO_MUL,
848 - MoEActivation.GELU_NO_MUL,
849 - MoEActivation.GELU_TANH_NO_MUL,
850 - MoEActivation.RELU2_NO_MUL,
851 - ]
844 + if current_platform.is_xpu():
845 + return activation in [
846 + MoEActivation.SILU,
847 + MoEActivation.GELU,
848 + MoEActivation.GELU_TANH,
849 + MoEActivation.SWIGLUOAI,
850 + MoEActivation.SILU_NO_MUL,
851 + MoEActivation.GELU_NO_MUL,
852 + MoEActivation.GELU_TANH_NO_MUL,
853 + MoEActivation.RELU2_NO_MUL,
854 + ]
855 + return apply_moe_activation_masked_supported(activation)
856
857 @staticmethod
858 def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
@@ -858,20 +862,6 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
862 # Let PrepareAndFinalize::finalize() decide the impl.
863 return TopKWeightAndReduceDelegate()
864
861 - def activation(
862 - self,
863 - activation: MoEActivation,
864 - output: torch.Tensor,
865 - input: torch.Tensor,
866 - **kwargs,
867 - ) -> None:
868 - gemm1_clamp_limit = self.quant_config.gemm1_clamp_limit
869 - if activation == MoEActivation.SILU and gemm1_clamp_limit is not None:
870 - swiglu_limit_func(output, input, float(gemm1_clamp_limit))
871 - return
872 -
873 - super().activation(activation, output, input)
874 -
865 def workspace_shapes(
866 self,
867 M: int,
@@ -998,11 +988,19 @@ class BatchedTritonExperts(mk.FusedMoEExpertsModular):
988 intermediate_cache2.fill_(0)
989
990 # TODO (bnell): use triton utility from batched deep gemm.
1001 - self.activation(
1002 - activation,
1003 - intermediate_cache2.view(-1, activation_out_dim),
1004 - intermediate_cache1.view(-1, N),
1005 - )
991 + if current_platform.is_xpu():
992 + self.activation(
993 + activation,
994 + intermediate_cache2.view(-1, activation_out_dim),
995 + intermediate_cache1.view(-1, N),
996 + )
997 + else:
998 + self.activation(
999 + activation,
1000 + intermediate_cache2,
1001 + intermediate_cache1,
1002 + valid_token_counts=expert_num_tokens,
1003 + )
1004
1005 qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(
1006 intermediate_cache2,
vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py
+52 -51
@@ -35,10 +35,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
35 TopKWeightAndReduceDelegate,
36 TopKWeightAndReduceNoOP,
37 )
38 -from vllm.model_executor.layers.fused_moe.utils import (
39 - _resize_cache,
40 - swiglu_limit_func,
41 -)
38 +from vllm.model_executor.layers.fused_moe.utils import _resize_cache
39 from vllm.model_executor.layers.quantization.utils.quant_utils import (
40 INT4_DTYPE,
41 INT8_DTYPE,
@@ -103,17 +100,18 @@ def _is_supported_wna16_weight_key(weight_key: QuantKey | None) -> bool:
100 )
101
102
106 -def get_humming_moe_gemm_type() -> str:
103 +def get_humming_moe_gemm_type(moe_config: FusedMoEConfig | None = None) -> str:
104 env_gemm_type: str | None = envs.VLLM_HUMMING_MOE_GEMM_TYPE
108 - gemm_type = "indexed"
109 - if env_gemm_type is not None:
105 + if env_gemm_type is not None and env_gemm_type.lower() != "auto":
106 env_gemm_type = env_gemm_type.lower()
111 - if env_gemm_type == "indexed":
112 - gemm_type = env_gemm_type
113 - elif env_gemm_type in ["grouped_contiguous", "grouped"]:
107 + if env_gemm_type == "grouped":
108 gemm_type = "grouped_contiguous"
109 else:
116 - gemm_type = "indexed"
110 + gemm_type = env_gemm_type
111 + elif moe_config is not None and moe_config.moe_parallel_config.use_ep:
112 + gemm_type = "grouped_contiguous"
113 + else:
114 + gemm_type = "indexed"
115
116 logger.info_once(f"Using {gemm_type} gemm for humming moe") # noqa
117 return gemm_type
@@ -178,7 +176,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
176 max_num_tokens=max_num_tokens,
177 num_dispatchers=num_dispatchers,
178 )
181 - self._permute_scratch: MoEPermuteScratch | None = None
179 + self._permute_scratch: dict[int, MoEPermuteScratch] = {}
180
181 def init_humming_moe(self):
182 from vllm.utils.humming import get_heuristics_config
@@ -257,26 +255,36 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
255 **kwargs,
256 )
257
260 - def _get_permute_scratch(self) -> MoEPermuteScratch | None:
261 - if self._permute_scratch is None and moe_permute_unpermute_supported():
262 - self._permute_scratch = MoEPermuteScratch(
263 - max_num_tokens=self.moe_config.max_num_tokens,
264 - topk=self.moe_config.experts_per_token,
258 + def _get_permute_scratch(self, topk: int) -> MoEPermuteScratch | None:
259 + if not moe_permute_unpermute_supported():
260 + return None
261 +
262 + scratch = self._permute_scratch.get(topk)
263 + if scratch is None:
264 + max_expanded_rows = (
265 + self.moe_config.max_num_tokens
266 + * self.moe_config.dp_size
267 + * self.moe_config.experts_per_token
268 + )
269 + scratch = MoEPermuteScratch(
270 + max_num_tokens=math.ceil(max_expanded_rows / topk),
271 + topk=topk,
272 num_experts=self.moe_config.num_experts,
273 num_local_experts=self.moe_config.num_local_experts,
274 device=torch.device(self.moe_config.device),
275 hidden_size=self.moe_config.hidden_dim,
276 hidden_dtype=self.moe_config.in_dtype,
277 )
271 - return self._permute_scratch
278 + self._permute_scratch[topk] = scratch
279 + return scratch
280
281 def get_global_valid_shape_m(self, topk_ids: torch.Tensor):
274 - num_tokens = topk_ids.size(0)
282 ctx = get_forward_context()
283 if ctx.dp_metadata is not None:
284 num_tokens = ctx.dp_metadata.num_tokens_across_dp_cpu.sum().item()
285 + return num_tokens * self.moe_config.experts_per_token
286
279 - return num_tokens * topk_ids.size(1)
287 + return topk_ids.size(0) * topk_ids.size(1)
288
289 def estimate_local_valid_shape_m(self, topk_ids: torch.Tensor):
290 # estimate shape_m for kernel tuning
@@ -645,7 +653,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
653 if supported:
654 assert hasattr(cls, "humming_gemm_type")
655 gemm_type = cls.humming_gemm_type().value.lower()
648 - preferred_gemm_type = get_humming_moe_gemm_type()
656 + preferred_gemm_type = get_humming_moe_gemm_type(moe_config)
657 supported = preferred_gemm_type.lower() == gemm_type
658 if not supported:
659 reason = (
@@ -660,25 +668,16 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
668 activation: MoEActivation,
669 output: torch.Tensor,
670 input: torch.Tensor,
663 - valid_rows: torch.Tensor | None = None,
671 + valid_token_counts: torch.Tensor | None = None,
672 ) -> None:
665 - activation_config = self.activation_config
666 - if (
667 - activation == MoEActivation.SILU
668 - and activation_config.clamp_limit is not None
669 - ):
670 - swiglu_limit_func(
671 - output=output,
672 - input=input,
673 - swiglu_limit=activation_config.clamp_limit,
674 - )
675 - else:
676 - self.activation(
677 - activation=activation,
678 - input=input,
679 - output=output,
680 - valid_rows=valid_rows,
681 - )
673 + activation_kwargs: dict[str, Any] = dict(
674 + activation=activation,
675 + input=input,
676 + output=output,
677 + )
678 + if valid_token_counts is not None:
679 + activation_kwargs["valid_token_counts"] = valid_token_counts
680 + self.activation(**activation_kwargs)
681
682 def fused_situ_quant_enabled(self, activation: MoEActivation) -> bool:
683 """Whether the SITU activation + w2 quant can be fused into one kernel.
@@ -718,7 +717,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
717 self,
718 gate_up_output: torch.Tensor,
719 quanted_down_input: torch.Tensor,
721 - valid_rows: torch.Tensor | None,
720 + num_valid_tokens: torch.Tensor | None,
721 topk: int = 1,
722 ) -> tuple[torch.Tensor, torch.Tensor]:
723 """Run the fused SITU activation + FP8 quant for the w2 input.
@@ -750,7 +749,7 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):
749 beta=cfg.activation_situ_beta,
750 linear_beta=cfg.activation_situ_linear_beta,
751 group_size=group_size,
753 - valid_rows=valid_rows,
752 + num_valid_tokens=num_valid_tokens,
753 topk=topk,
754 )
755 return quanted_down_input, input_scale
@@ -878,8 +877,8 @@ class HummingIndexedExperts(HummingExpertsBase):
877 )
878
879 # psum[-1:] is the DeepEP valid *token* count as a zero-cost int32 view.
881 - # situ (rows = tokens*topk) multiplies by topk on-device; mul_sum bounds
882 - # on tokens, so both read this pointer directly -- no host cast/multiply.
880 + # The fused path and mul_sum consume tokens directly. The unfused
881 + # activation consumes the expanded row count (tokens * topk).
882 valid_tokens = None
883 topk = topk_ids.size(1)
884 if (
@@ -894,21 +893,18 @@ class HummingIndexedExperts(HummingExpertsBase):
893 inputs, input_scale = self.fused_situ_quant(
894 gate_up_output=buffers["gate_up_output"],
895 quanted_down_input=buffers["quanted_down_input"],
897 - valid_rows=valid_tokens,
896 + num_valid_tokens=valid_tokens,
897 topk=topk,
898 )
899 else:
901 - # Fallback situ_and_mul takes int64 row counts (tokens*topk).
902 - valid_rows = (
903 - valid_tokens.to(torch.int64) * topk
904 - if valid_tokens is not None
905 - else None
900 + valid_token_counts = (
901 + valid_tokens * topk if valid_tokens is not None else None
902 )
903 self.apply_activation(
904 activation=activation,
905 input=buffers["gate_up_output"],
906 output=buffers["activation_output"],
911 - valid_rows=valid_rows,
907 + valid_token_counts=valid_token_counts,
908 )
909
910 inputs, input_scale = self.quantize_input(
@@ -1002,7 +998,7 @@ class HummingGroupedExperts(HummingExpertsBase):
998 n_expert=global_num_experts,
999 n_local_expert=self.num_experts,
1000 expert_map=expert_map,
1005 - scratch=self._get_permute_scratch(),
1001 + scratch=self._get_permute_scratch(topk_ids.size(1)),
1002 )
1003
1004 inputs, input_scale = self.quantize_input(
@@ -1028,6 +1024,11 @@ class HummingGroupedExperts(HummingExpertsBase):
1024 activation=activation,
1025 input=buffers["gate_up_output"],
1026 output=buffers["activation_output"],
1027 + valid_token_counts=(
1028 + expert_first_token_offset[-1:].to(torch.int32)
1029 + if expert_tokens_meta is None
1030 + else None
1031 + ),
1032 )
1033
1034 inputs, input_scale = self.quantize_input(
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+13 -16
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.activation import (
13 ApplyMoEActivationConfig,
14 MoEActivation,
15 apply_moe_activation,
16 + apply_moe_activation_masked_supported,
17 apply_moe_activation_supported,
18 )
19 from vllm.model_executor.layers.fused_moe.config import (
@@ -955,6 +956,10 @@ class BatchedMarlinExperts(MarlinExpertsBase):
956 def activation_format() -> mk.FusedMoEActivationFormat:
957 return mk.FusedMoEActivationFormat.BatchedExperts
958
959 + @staticmethod
960 + def _supports_activation(activation: MoEActivation) -> bool:
961 + return apply_moe_activation_masked_supported(activation)
962 +
963 def workspace_shapes(
964 self,
965 M: int,
@@ -1000,26 +1005,18 @@ class BatchedMarlinExperts(MarlinExpertsBase):
1005 act: MoEActivation,
1006 act_output: torch.Tensor,
1007 act_input: torch.Tensor,
1003 - **kwargs,
1008 + *,
1009 + topk_ids: torch.Tensor | None = None,
1010 + expert_map: torch.Tensor | None = None,
1011 ) -> None:
1005 - if act != MoEActivation.SITU:
1006 - self.activation(
1007 - act,
1008 - act_output,
1009 - act_input,
1010 - **kwargs,
1011 - )
1012 - return
1013 -
1012 num_experts, max_num_tokens = hidden_states.shape[:2]
1015 - beta = self.activation_config.activation_situ_beta
1016 - linear_beta = self.activation_config.activation_situ_linear_beta
1017 - torch.ops._C.masked_situ_and_mul(
1013 + self.activation(
1014 + act,
1015 act_output.view(num_experts, max_num_tokens, -1),
1016 act_input.view(num_experts, max_num_tokens, -1),
1020 - expert_tokens_meta.expert_num_tokens,
1021 - 1.0 if beta is None else beta,
1022 - -1.0 if linear_beta is None else linear_beta,
1017 + topk_ids=topk_ids,
1018 + expert_map=expert_map,
1019 + valid_token_counts=expert_tokens_meta.expert_num_tokens,
1020 )
1021
1022 return batched_fused_marlin_moe(
vllm/model_executor/layers/fused_moe/modular_kernel.py
+2 -2
@@ -902,7 +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,
905 + valid_token_counts: torch.Tensor | None = None,
906 ) -> None:
907 apply_moe_activation(
908 activation,
@@ -911,7 +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,
914 + valid_token_counts=valid_token_counts,
915 )
916
917 @abstractmethod