diff --git a/csrc/libtorch_stable/moe/moe_ops.h b/csrc/libtorch_stable/moe/moe_ops.h index 9ea1dae7cae..03525a05f9d 100644 --- a/csrc/libtorch_stable/moe/moe_ops.h +++ b/csrc/libtorch_stable/moe/moe_ops.h @@ -9,14 +9,16 @@ void topk_softmax(torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, torch::stable::Tensor& token_expert_indices, torch::stable::Tensor& gating_output, bool renormalize, - std::optional bias); + std::optional bias, + std::optional is_padding); void topk_sigmoid(torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, torch::stable::Tensor& token_expert_indices, torch::stable::Tensor& gating_output, bool renormalize, std::optional bias, - double routed_scaling_factor); + double routed_scaling_factor, + std::optional is_padding); void topk_softplus_sqrt( torch::stable::Tensor& topk_weights, torch::stable::Tensor& topk_indices, @@ -25,7 +27,8 @@ void topk_softplus_sqrt( double routed_scaling_factor, const std::optional& correction_bias, const std::optional& input_ids, - const std::optional& tid2eid); + const std::optional& tid2eid, + const std::optional& is_padding); void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output, std::optional topk_ids, diff --git a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu index b4bcd9479e9..098644a1b90 100644 --- a/csrc/libtorch_stable/moe/topk_softmax_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softmax_kernels.cu @@ -174,7 +174,8 @@ __launch_bounds__(TPB) __global__ void moeTopK( const int end_expert, const bool renormalize, const float* bias, - const double routed_scaling_factor) + const double routed_scaling_factor, + const bool* is_padding) { using cub_kvp = cub::KeyValuePair; @@ -228,12 +229,14 @@ __launch_bounds__(TPB) __global__ void moeTopK( const int expert = result_kvp.key; const bool node_uses_expert = expert >= start_expert && expert < end_expert; const bool should_process_row = row_is_active && node_uses_expert; + const bool is_pad_row = is_padding != nullptr && is_padding[block_row]; const int idx = k * block_row + k_idx; // Return the unbiased scores for output weights output[idx] = inputs_after_softmax[thread_read_offset + expert]; - indices[idx] = should_process_row ? (expert - start_expert) : num_experts; - assert(indices[idx] >= 0); + indices[idx] = is_pad_row ? static_cast(-1) + : (should_process_row ? (expert - start_expert) : num_experts); + assert(is_pad_row || indices[idx] >= 0); source_rows[idx] = k_idx * num_rows + block_row; if (renormalize) { selected_sum += inputs_after_softmax[thread_read_offset + expert]; @@ -277,7 +280,7 @@ template || std::is_same_v || std::is_same_v, @@ -545,12 +548,14 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // Add a guard to ignore experts not included by this node const bool node_uses_expert = expert >= start_expert && expert < end_expert; const bool should_process_row = row_is_active && node_uses_expert; + const bool is_pad_row = is_padding != nullptr && is_padding[thread_row]; // The lead thread from each sub-group will write out the final results to global memory. (This will be a // single) thread per row of the input/output matrices. const int idx = k * thread_row + k_idx; output[idx] = max_val; - indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS; + indices[idx] = is_pad_row ? static_cast(-1) + : (should_process_row ? (expert - start_expert) : NUM_EXPERTS); source_rows[idx] = k_idx * num_rows + thread_row; if (renormalize) { selected_sum += max_val; @@ -605,7 +610,7 @@ struct TopkConstants template void topkGatingLauncherHelper(const InputType* input, const bool* finished, float* output, IndType* indices, int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, - const float* bias, const double routed_scaling_factor, cudaStream_t stream) + const float* bias, const double routed_scaling_factor, cudaStream_t stream, const bool* is_padding) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = detail::TopkConstants; @@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB); topkGating<<>>( - input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor); + input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, bias, routed_scaling_factor, is_padding); } #ifndef USE_ROCM @@ -627,7 +632,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); #else #define LAUNCH_TOPK(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -635,13 +640,13 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa IndType, InputType, SF>( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, routed_scaling_factor, stream); \ + bias, routed_scaling_factor, stream, is_padding); \ } else if (WARP_SIZE == 32) { \ topkGatingLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ - bias, routed_scaling_factor, stream); \ + bias, routed_scaling_factor, stream, is_padding); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -661,7 +666,8 @@ void topkGatingKernelLauncher( const bool renormalize, const float* bias, const double routed_scaling_factor, - cudaStream_t stream) { + cudaStream_t stream, + const bool* is_padding) { static constexpr int WARPS_PER_TB = 4; static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16; #ifndef USE_ROCM @@ -736,7 +742,7 @@ void topkGatingKernelLauncher( } moeTopK<<>>( workspace, nullptr, topk_weights, topk_indices, token_expert_indices, - num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor); + num_experts, topk, 0, num_experts, renormalize, bias, routed_scaling_factor, is_padding); } } } @@ -755,7 +761,8 @@ void dispatch_topk_launch( int num_tokens, int num_experts, int topk, bool renormalize, std::optional bias, double routed_scaling_factor, - cudaStream_t stream) + cudaStream_t stream, + std::optional is_padding) { const float* bias_ptr = nullptr; if (bias.has_value()) { @@ -769,6 +776,18 @@ void dispatch_topk_launch( bias_ptr = bias_tensor.const_data_ptr(); } + const bool* is_padding_ptr = nullptr; + if (is_padding.has_value()) { + const torch::stable::Tensor& is_padding_tensor = is_padding.value(); + STD_TORCH_CHECK(is_padding_tensor.scalar_type() == torch::headeronly::ScalarType::Bool, + "is_padding tensor must be bool"); + STD_TORCH_CHECK(is_padding_tensor.dim() == 1, "is_padding tensor must be 1D"); + STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens, + "is_padding size mismatch, expected: ", num_tokens); + STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), "is_padding tensor must be contiguous"); + is_padding_ptr = is_padding_tensor.const_data_ptr(); + } + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { vllm::moe::topkGatingKernelLauncher( reinterpret_cast(gating_output.const_data_ptr()), @@ -777,7 +796,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, routed_scaling_factor, stream); + bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } else if (topk_indices.scalar_type() == torch::headeronly::ScalarType::UInt32) { vllm::moe::topkGatingKernelLauncher( reinterpret_cast(gating_output.const_data_ptr()), @@ -786,7 +805,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, routed_scaling_factor, stream); + bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } else { STD_TORCH_CHECK(topk_indices.scalar_type() == torch::headeronly::ScalarType::Long); vllm::moe::topkGatingKernelLauncher( @@ -796,7 +815,7 @@ void dispatch_topk_launch( token_expert_indices.mutable_data_ptr(), softmax_workspace.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, - bias_ptr, routed_scaling_factor, stream); + bias_ptr, routed_scaling_factor, stream, is_padding_ptr); } } @@ -806,7 +825,8 @@ void topk_softmax( torch::stable::Tensor& token_expert_indices, // [num_tokens, topk] torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, - std::optional bias) + std::optional bias, + std::optional is_padding) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -825,15 +845,15 @@ void topk_softmax( if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, 1.0, stream); + bias, 1.0, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, 1.0, stream); + bias, 1.0, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SOFTMAX>(gating_output, topk_weights, topk_indices, token_expert_indices, softmax_workspace, num_tokens, num_experts, topk, renormalize, - bias, 1.0, stream); + bias, 1.0, stream, is_padding); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } @@ -846,7 +866,8 @@ void topk_sigmoid( torch::stable::Tensor& gating_output, // [num_tokens, num_experts] bool renormalize, std::optional bias, - double routed_scaling_factor) + double routed_scaling_factor, + std::optional is_padding) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; @@ -865,15 +886,15 @@ void topk_sigmoid( if (gating_output.scalar_type() == torch::headeronly::ScalarType::Float) { dispatch_topk_launch(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_launch<__half, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_launch<__nv_bfloat16, vllm::moe::SCORING_SIGMOID>(gating_output, topk_weights, topk_indices, token_expert_indices, workspace, num_tokens, num_experts, topk, renormalize, - bias, routed_scaling_factor, stream); + bias, routed_scaling_factor, stream, is_padding); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index b6878eb2d2f..976eebac3dd 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -80,22 +80,27 @@ __launch_bounds__(128) __global__ OutIndType* indices, int num_rows, int num_experts, float routed_scaling_factor, const HashIndType* input_ids, - const HashIndType* tid2eid) { + const HashIndType* tid2eid, + const bool* is_padding) { const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32; const int lane = threadIdx.x % 32; if (warp >= num_rows) return; const int64_t token_id = load_index_as_int64(input_ids, warp); + const bool is_pad_row = is_padding != nullptr && is_padding[warp]; #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) cudaGridDependencySynchronize(); #endif int expert = 0; float weight = 0.f; - if (lane < 6) { + if (lane < 6 && !is_pad_row) { // only load and calculate for 6 experts expert = static_cast(tid2eid[token_id * 6 + lane]); const float x = input[warp * num_experts + expert]; weight = sqrtf(fmaxf(x, 0.f) + __logf(1.f + __expf(-fabsf(x)))); + if (isnan(weight)) { + weight = 0.f; + } } float weight_sum = weight; #pragma unroll @@ -111,7 +116,8 @@ __launch_bounds__(128) __global__ const int offset = warp * 6 + lane; output[offset] = weight * routed_scaling_factor / (weight_sum > 0.f ? weight_sum : 1.f); - indices[offset] = static_cast(expert); + indices[offset] = !is_pad_row ? static_cast(expert) + : static_cast(-1); } } @@ -120,7 +126,8 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, int num_rows, int num_experts, double routed_scaling_factor, const HashIndType* input_ids, - const HashIndType* tid2eid, cudaStream_t stream) { + const HashIndType* tid2eid, cudaStream_t stream, + const bool* is_padding) { if (num_rows == 0) return; auto* kernel = &dsv4HashTopkSoftplusSqrt; cudaLaunchConfig_t config = {}; @@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices, config.numAttrs = 1; const float scale = static_cast(routed_scaling_factor); cudaLaunchKernelEx(&config, kernel, input, output, indices, num_rows, - num_experts, scale, input_ids, tid2eid); + num_experts, scale, input_ids, tid2eid, is_padding); } #endif @@ -166,7 +173,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ const int num_rows, IndType* indices, int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, - const HashIndType* input_ids, const HashIndType* tid2eid) { + const HashIndType* input_ids, const HashIndType* tid2eid, + const bool* is_padding) { static_assert(std::is_same_v || std::is_same_v || std::is_same_v, @@ -231,6 +239,7 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ return; } const bool row_is_active = finished ? !finished[thread_row] : true; + const bool is_pad_row = is_padding != nullptr && is_padding[thread_row]; // We finally start setting up the read pointers for each thread. First, each // thread jumps to the start of the row it will read. @@ -249,9 +258,12 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ cudaGridDependencySynchronize(); #endif - // NOTE(zhuhaoran): dispatch different input types loading, BF16/FP16 convert - // to float - if constexpr (std::is_same_v) { + if (is_pad_row) { +#pragma unroll + for (int ii = 0; ii < VPT; ++ii) { + row_chunk[ii] = 0.f; + } + } else if constexpr (std::is_same_v) { using VecType = AlignedArray; VecType* row_chunk_vec_ptr = reinterpret_cast(&row_chunk); const VecType* vec_thread_read_ptr = @@ -315,12 +327,22 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ if constexpr (USE_HASH) { const int64_t token_id = load_index_as_int64(input_ids, thread_row); const int64_t token_expert_offset = token_id * static_cast(k); + if (!is_pad_row) { #pragma unroll - for (int ii = 0; ii < VPT; ++ii) { - float val = row_chunk[ii]; - float val_b = val * beta; - val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; - row_chunk[ii] = sqrtf(val); + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + float val_b = val * beta; + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; + val = sqrtf(val); + + // Dummy/padding tokens can result in NaN values, so + // clamp them to 0.0. Note: this clamp could likely be removed if + // 'is_padding' is made mandatory + if (isnan(val)) { + val = 0.f; + } + row_chunk[ii] = val; + } } float selected_sum = 0.f; #pragma unroll @@ -335,7 +357,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ group_id * THREADS_PER_ROW * ELTS_PER_LDG + local_id; if (expert == expert_idx) { - indices[idx] = static_cast(expert); + indices[idx] = !is_pad_row ? static_cast(expert) + : static_cast(-1); selected_sum += row_chunk[ii]; break; } @@ -379,23 +402,31 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ #endif return; } else { + if (!is_pad_row) { #pragma unroll - for (int ii = 0; ii < VPT; ++ii) { - float val = row_chunk[ii]; - float val_b = val * beta; - // Compute softplus: log(1 + exp(val)) with numerical stability - // When val > threshold, softplus(x) ≈ x to avoid exp overflow - val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; - val = sqrtf(val); - if (correction_bias) { - const int group_id = ii / ELTS_PER_LDG; - const int local_id = ii % ELTS_PER_LDG; - const int expert_idx = first_elt_read_by_thread + - group_id * THREADS_PER_ROW * ELTS_PER_LDG + - local_id; - val = val + correction_bias[expert_idx]; + for (int ii = 0; ii < VPT; ++ii) { + float val = row_chunk[ii]; + float val_b = val * beta; + // Compute softplus: log(1 + exp(val)) with numerical stability + // When val > threshold, softplus(x) ≈ x to avoid exp overflow + val = (val_b > threshold) ? val : (__logf(1.0f + __expf(val_b))) / beta; + val = sqrtf(val); + // Dummy/padding tokens can result in NaN values, so + // clamp them to 0.0. Note: this clamp could likely be removed if + // 'is_padding' is made mandatory + if (isnan(val)) { + val = 0.f; + } + if (correction_bias) { + const int group_id = ii / ELTS_PER_LDG; + const int local_id = ii % ELTS_PER_LDG; + const int expert_idx = first_elt_read_by_thread + + group_id * THREADS_PER_ROW * ELTS_PER_LDG + + local_id; + val = val + correction_bias[expert_idx]; + } + row_chunk[ii] = val; } - row_chunk[ii] = val; } // Original TopK path: find top-k experts by score @@ -450,18 +481,19 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__ // Add a guard to ignore experts not included by this node const bool node_uses_expert = expert >= start_expert && expert < end_expert; - const bool should_process_row = row_is_active && node_uses_expert; + const bool should_process_row = + row_is_active && node_uses_expert && !is_pad_row; // The lead thread from each sub-group will write out the final results // to global memory. (This will be a single) thread per row of the // input/output matrices. const int idx = k * thread_row + k_idx; - if (correction_bias != nullptr) { + if (correction_bias != nullptr && should_process_row) { max_val -= correction_bias[expert]; } output[idx] = max_val; indices[idx] = - should_process_row ? (expert - start_expert) : NUM_EXPERTS; + !is_pad_row ? expert - start_expert : static_cast(-1); source_rows[idx] = k_idx * num_rows + thread_row; if (renormalize) { selected_sum += max_val; @@ -544,7 +576,7 @@ void topkGatingSoftplusSqrtLauncherHelper( const int start_expert, const int end_expert, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, - const HashIndType* tid2eid, cudaStream_t stream) { + const HashIndType* tid2eid, cudaStream_t stream, const bool* is_padding) { static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS); using Constants = @@ -573,12 +605,12 @@ void topkGatingSoftplusSqrtLauncherHelper( cudaLaunchKernelEx(&config, kernel, input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid); + input_ids, tid2eid, is_padding); #else kernel<<>>( input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid); + input_ids, tid2eid, is_padding); #endif }) } @@ -592,7 +624,7 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, tid2eid, \ - stream); + stream, is_padding); #else #define LAUNCH_SOFTPLUS_SQRT(NUM_EXPERTS, WARPS_PER_TB, MAX_BYTES) \ if (WARP_SIZE == 64) { \ @@ -601,14 +633,14 @@ void topkGatingSoftplusSqrtLauncherHelper( gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream); \ + tid2eid, stream, is_padding); \ } else if (WARP_SIZE == 32) { \ topkGatingSoftplusSqrtLauncherHelper( \ gating_output, nullptr, topk_weights, topk_indices, \ token_expert_indices, num_tokens, topk, 0, num_experts, renormalize, \ routed_scaling_factor, correction_bias, use_hash, input_ids, \ - tid2eid, stream); \ + tid2eid, stream, is_padding); \ } else { \ assert(false && \ "Unsupported warp size. Only 32 and 64 are supported for ROCm"); \ @@ -622,14 +654,14 @@ void topkGatingSoftplusSqrtKernelLauncher( const int topk, const bool renormalize, double routed_scaling_factor, const float* correction_bias, const bool use_hash, const HashIndType* input_ids, const HashIndType* tid2eid, - cudaStream_t stream) { + cudaStream_t stream, const bool* is_padding) { #ifndef USE_ROCM if constexpr (std::is_same_v) { if (use_hash && topk == 6 && renormalize && (num_experts == 256 || num_experts == 384)) { launchDsv4HashTopk( gating_output, topk_weights, topk_indices, num_tokens, num_experts, - routed_scaling_factor, input_ids, tid2eid, stream); + routed_scaling_factor, input_ids, tid2eid, stream, is_padding); return; } } @@ -728,7 +760,8 @@ void dispatch_topk_softplus_sqrt_launch( int num_experts, int topk, bool renormalize, double routed_scaling_factor, const std::optional& correction_bias, const std::optional& input_ids, - const std::optional& tid2eid, cudaStream_t stream) { + const std::optional& tid2eid, cudaStream_t stream, + const std::optional& is_padding) { const float* bias_ptr = nullptr; if (correction_bias.has_value()) { bias_ptr = correction_bias.value().const_data_ptr(); @@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch( auto launch = [&](auto* topk_indices_ptr) { using OutIndType = typename std::remove_pointer::type; + + const bool* is_padding_ptr = nullptr; + if (is_padding.has_value()) { + const torch::stable::Tensor& is_padding_tensor = is_padding.value(); + STD_TORCH_CHECK(is_padding_tensor.scalar_type() == + torch::headeronly::ScalarType::Bool, + "is_padding tensor must be bool"); + STD_TORCH_CHECK(is_padding_tensor.dim() == 1, + "is_padding tensor must be 1D"); + STD_TORCH_CHECK(is_padding_tensor.size(0) == num_tokens, + "is_padding size mismatch, expected: ", num_tokens); + STD_TORCH_CHECK(is_padding_tensor.is_contiguous(), + "is_padding tensor must be contiguous"); + is_padding_ptr = is_padding_tensor.const_data_ptr(); + } + if (tid2eid.has_value()) { STD_TORCH_CHECK(input_ids.has_value(), "input_ids is required for hash MoE"); @@ -751,7 +800,7 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr); } else { STD_TORCH_CHECK(tid2eid.value().scalar_type() == torch::headeronly::ScalarType::Int); @@ -761,7 +810,7 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, true, input_ids.value().const_data_ptr(), - tid2eid.value().const_data_ptr(), stream); + tid2eid.value().const_data_ptr(), stream, is_padding_ptr); } } else { vllm::moe::topkGatingSoftplusSqrtKernelLauncher( @@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch( topk_indices_ptr, token_expert_indices.mutable_data_ptr(), num_tokens, num_experts, topk, renormalize, routed_scaling_factor, bias_ptr, false, static_cast(nullptr), - static_cast(nullptr), stream); + static_cast(nullptr), stream, is_padding_ptr); } }; @@ -793,7 +842,8 @@ void topk_softplus_sqrt( bool renormalize, double routed_scaling_factor, const std::optional& correction_bias, const std::optional& input_ids, - const std::optional& tid2eid) { + const std::optional& tid2eid, + const std::optional& is_padding) { const int num_experts = gating_output.size(-1); const auto num_tokens = gating_output.numel() / num_experts; const int topk = topk_weights.size(-1); @@ -806,21 +856,22 @@ void topk_softplus_sqrt( dispatch_topk_softplus_sqrt_launch( gating_output.const_data_ptr(), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, - routed_scaling_factor, correction_bias, input_ids, tid2eid, stream); + routed_scaling_factor, correction_bias, input_ids, tid2eid, stream, + is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::Half) { dispatch_topk_softplus_sqrt_launch<__half>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream); + input_ids, tid2eid, stream, is_padding); } else if (gating_output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { dispatch_topk_softplus_sqrt_launch<__nv_bfloat16>( reinterpret_cast(gating_output.const_data_ptr()), topk_weights, topk_indices, token_expert_indices, num_tokens, num_experts, topk, renormalize, routed_scaling_factor, correction_bias, - input_ids, tid2eid, stream); + input_ids, tid2eid, stream, is_padding); } else { STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); diff --git a/csrc/libtorch_stable/moe/torch_bindings.cpp b/csrc/libtorch_stable/moe/torch_bindings.cpp index a0296adca78..c6467f1c1ff 100644 --- a/csrc/libtorch_stable/moe/torch_bindings.cpp +++ b/csrc/libtorch_stable/moe/torch_bindings.cpp @@ -8,19 +8,19 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_moe_C, m) { m.def( "topk_softmax(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, Tensor? " - "bias) -> ()"); + "bias, Tensor? is_padding) -> ()"); // Apply topk sigmoid to the gating outputs. m.def( "topk_sigmoid(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, " - "Tensor? bias, float routed_scaling_factor) -> ()"); + "Tensor? bias, float routed_scaling_factor, Tensor? is_padding) -> ()"); m.def( "topk_softplus_sqrt(Tensor! topk_weights, Tensor! topk_indices, Tensor! " "token_expert_indices, Tensor gating_output, bool renormalize, float " "routed_scaling_factor, Tensor? " - "bias, Tensor? input_ids, Tensor? tid2eid) -> ()"); + "bias, Tensor? input_ids, Tensor? tid2eid, Tensor? is_padding) -> ()"); // Calculate the result of moe by summing up the partial results // from all selected experts. topk_ids/expert_map are optional and, when diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 5c6691bbc32..c5f00253e78 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -6,6 +6,7 @@ import pytest import torch import torch.nn.functional as F +import vllm._custom_ops as ops from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, @@ -231,3 +232,119 @@ def test_dsv4_fast_topk( atol=2e-5, rtol=2e-5, ) + + +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="This test is skipped on non-CUDA platform.", +) +@pytest.mark.parametrize("use_hash", [False, True]) +@pytest.mark.parametrize("use_bias", [False, True]) +@pytest.mark.parametrize("use_padding_mask", [False, True]) +@pytest.mark.parametrize("pad_with_nan", [False, True]) +@pytest.mark.parametrize("num_experts", [128, 256, 384]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half, torch.float32]) +def test_fused_topk_softplus_sqrt_padding( + use_hash: bool, + use_bias: bool, + use_padding_mask: bool, + pad_with_nan: bool, + num_experts: int, + dtype: torch.dtype, +): + """Verify explicit padding and NaN-padded rows do not affect real rows.""" + torch.manual_seed(0) + num_tokens = 8 + topk = 6 + indices_dtype = torch.int32 + + gating_output = torch.randn((num_tokens, num_experts), dtype=dtype, device="cuda") + + padding_rows = torch.zeros(num_tokens, dtype=torch.bool, device="cuda") + padding_rows[1::2] = True + if pad_with_nan: + gating_output[padding_rows] = float("nan") + is_padding = padding_rows if use_padding_mask else None + + # A negative correction bias makes explicit pad rows look selectable unless + # the kernel uses the is_padding guard. + e_score_correction_bias = None + if use_bias: + e_score_correction_bias = ( + -torch.rand((num_experts,), dtype=torch.float32, device="cuda") - 1.0 + ) + + input_ids = None + hash_indices_table = None + if use_hash: + vocab_size = 64 + hash_indices_table = torch.stack( + [torch.randperm(num_experts)[:topk] for _ in range(vocab_size)] + ).to(device="cuda", dtype=indices_dtype) + input_ids = torch.randint( + 0, vocab_size, (num_tokens,), dtype=indices_dtype, device="cuda" + ) + + topk_weights = torch.empty(num_tokens, topk, dtype=torch.float32, device="cuda") + topk_ids = torch.empty(num_tokens, topk, dtype=indices_dtype, device="cuda") + token_expert_indices = torch.empty( + num_tokens, topk, dtype=torch.int32, device="cuda" + ) + + ops.topk_hash_softplus_sqrt( + topk_weights, + topk_ids, + token_expert_indices, + gating_output, + renormalize=True, + routed_scaling_factor=1.0, + e_score_correction_bias=e_score_correction_bias, + input_tokens=input_ids, + hash_indices_table=hash_indices_table, + is_padding=is_padding, + ) + + if use_padding_mask: + pad_ids = topk_ids[padding_rows] + pad_weights = topk_weights[padding_rows] + assert torch.equal(pad_ids, torch.full_like(pad_ids, -1)), ( + f"Explicit pad rows should contain only -1 ids, got {pad_ids.tolist()}" + ) + assert (pad_weights == 0).all(), ( + "Explicit pad rows should have all-zero weights, " + f"got {pad_weights.tolist()}" + ) + + if pad_with_nan: + nan_pad_weights = topk_weights[padding_rows] + assert torch.isfinite(nan_pad_weights).all(), ( + f"NaN-padded rows have non-finite weights, got {nan_pad_weights.tolist()}" + ) + assert (nan_pad_weights == 0).all(), ( + "NaN-padded rows should have all-zero weights, " + f"got {nan_pad_weights.tolist()}" + ) + + topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt( + gating_output=gating_output, + topk=topk, + renormalize=True, + routed_scaling_factor=1.0, + e_score_correction_bias=e_score_correction_bias, + input_ids=input_ids, + hash_indices_table=hash_indices_table, + ) + + rows_to_compare = torch.ones(num_tokens, dtype=torch.bool, device="cuda") + if use_padding_mask or pad_with_nan: + rows_to_compare = ~padding_rows + + sorted_ref_ids, idx_ref = topk_ids_ref[rows_to_compare].sort(dim=-1) + sorted_ids, idx_ops = topk_ids[rows_to_compare].sort(dim=-1) + torch.testing.assert_close( + sorted_ref_ids, sorted_ids.to(sorted_ref_ids.dtype), atol=0, rtol=0 + ) + + sorted_w_ref = topk_weights_ref[rows_to_compare].gather(1, idx_ref) + sorted_w = topk_weights[rows_to_compare].gather(1, idx_ops) + torch.testing.assert_close(sorted_w_ref, sorted_w, atol=2e-2, rtol=1e-2) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index df57bc0f77d..f635077a1d5 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2389,6 +2389,7 @@ def topk_softmax( gating_output: torch.Tensor, renormalize: bool = False, e_score_correction_bias: torch.Tensor | None = None, + is_padding: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.topk_softmax( topk_weights, @@ -2397,6 +2398,7 @@ def topk_softmax( gating_output, renormalize, e_score_correction_bias, + is_padding, ) @@ -2408,6 +2410,7 @@ def topk_sigmoid( renormalize: bool = False, e_score_correction_bias: torch.Tensor | None = None, routed_scaling_factor: float = 1.0, + is_padding: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.topk_sigmoid( topk_weights, @@ -2417,6 +2420,7 @@ def topk_sigmoid( renormalize, e_score_correction_bias, routed_scaling_factor, + is_padding, ) @@ -2430,6 +2434,7 @@ def topk_hash_softplus_sqrt( e_score_correction_bias: torch.Tensor | None = None, input_tokens: torch.Tensor | None = None, hash_indices_table: torch.Tensor | None = None, + is_padding: torch.Tensor | None = None, ) -> None: torch.ops._moe_C.topk_softplus_sqrt( topk_weights, @@ -2441,6 +2446,7 @@ def topk_hash_softplus_sqrt( e_score_correction_bias, input_tokens, hash_indices_table, + is_padding, ) diff --git a/vllm/envs.py b/vllm/envs.py index ff6b78b8a41..9bc2ad1769a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -193,7 +193,7 @@ if TYPE_CHECKING: "relax", ] = "relax" VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True - VLLM_MOE_SKIP_PADDING: bool = False + VLLM_MOE_SKIP_PADDING: bool = True VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1519,9 +1519,8 @@ environment_variables: dict[str, Callable[[], Any]] = { ), # Skip cudagraph/DP padding tokens in the MoE path by forcing their expert # ids to -1 so the dispatch and experts drop them. Requires a MoE kernel that - # treats topk_id == -1 as a skip sentinel; off by default because not all - # kernels support it yet. - "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "0"))), + # treats topk_id == -1 as a skip sentinel + "VLLM_MOE_SKIP_PADDING": lambda: bool(int(os.getenv("VLLM_MOE_SKIP_PADDING", "1"))), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 0eb411dabdf..2c8ba874f0c 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -10,7 +10,6 @@ from typing import final import torch import vllm.envs as envs -from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import ( MoEActivation, @@ -1194,21 +1193,6 @@ class FusedMoEKernelModularImpl: The _prepare method is a wrapper around self.prepare_finalize.prepare that handles DBO and async. """ - # Skip cudagraph/DP padding tokens uniformly across all a2a backends: - # forcing padded rows' expert ids to -1 makes every prepare_finalize drop - # them (not dispatched / not computed by the experts). The V2 model runner - # marks them in forward_context.is_padding; it is None for runners that do - # not populate it, leaving topk_ids unchanged. - # Gated by VLLM_MOE_SKIP_PADDING (off by default) because this requires the - # experts kernel to treat topk_id == -1 as a skip sentinel, which not all - # MoE backends support yet. - is_padding = None - if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): - is_padding = get_forward_context().is_padding - if is_padding is not None: - n = topk_ids.shape[0] - # TODO: Properly support DBO (padding lives at the batch tail). - topk_ids = torch.where(is_padding[:n].unsqueeze(1), -1, topk_ids) if not self.prepare_finalize.supports_async(): # We shouldn't be running an a2a kernel that doesn't diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py index f89eb4910b3..bc019a58e2e 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py @@ -9,6 +9,7 @@ import vllm._custom_ops as ops import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, @@ -20,6 +21,13 @@ from vllm.model_executor.layers.fused_moe.router.dsv4_topk import ( ) +def _get_padding_mask(num_tokens: int) -> torch.Tensor | None: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + return is_padding[:num_tokens] if is_padding is not None else None + return None + + def vllm_topk_softmax( topk_weights: torch.Tensor, topk_indices: torch.Tensor, @@ -35,6 +43,7 @@ def vllm_topk_softmax( gating_output, renormalize, e_score_correction_bias, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices @@ -57,6 +66,7 @@ def vllm_topk_sigmoid( renormalize, e_score_correction_bias, routed_scaling_factor, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices @@ -144,6 +154,7 @@ def vllm_topk_softplus_sqrt( e_score_correction_bias, input_tokens, hash_indices_table, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices diff --git a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py index 855fa606565..28f67b5c693 100644 --- a/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py +++ b/vllm/model_executor/layers/fused_moe/router/fused_topk_router.py @@ -5,8 +5,10 @@ from collections.abc import Callable import torch import vllm._custom_ops as ops +import vllm.envs as envs from vllm._aiter_ops import rocm_aiter_ops from vllm.distributed.eplb.eplb_state import EplbLayerState +from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.model_executor.layers.fused_moe.config import ( RoutingMethodType, get_routing_method_type, @@ -14,6 +16,13 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter +def _get_padding_mask(num_tokens: int) -> torch.Tensor | None: + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): + is_padding = get_forward_context().is_padding + return is_padding[:num_tokens] if is_padding is not None else None + return None + + def vllm_topk_softmax( topk_weights: torch.Tensor, topk_indices: torch.Tensor, @@ -27,6 +36,7 @@ def vllm_topk_softmax( token_expert_indices, gating_output, renormalize, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices @@ -45,6 +55,7 @@ def vllm_topk_sigmoid( token_expert_indices, gating_output, renormalize, + is_padding=_get_padding_mask(topk_indices.shape[0]), ) return topk_weights, topk_indices