karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 2 Packages Activity

skip cudagraph/DP padding in topk (#48979)

Browse Source 
Signed-off-by: gnovack <novackgm@gmail.com>
This commit is contained in:
gnovack
2026-07-21 15:20:30 -07:00
committed by GitHub
parent 08e5067561
commit 85f638a2b8
10 changed files with 302 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/moe/moe_ops.h
+6 -3
View File
@@ -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<torch::stable::Tensor> bias);
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> 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<torch::stable::Tensor> bias,
double routed_scaling_factor);
double routed_scaling_factor,
std::optional<torch::stable::Tensor> 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<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid);
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& is_padding);
void moe_sum(torch::stable::Tensor& input, torch::stable::Tensor& output,
std::optional<torch::stable::Tensor> topk_ids,
csrc/libtorch_stable/moe/topk_softmax_kernels.cu
+45 -24
View File
@@ -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<int, float>;
@@ -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<IndType>(-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 <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG, int WA
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGating(const InputType* input, const bool* finished, float* output, const int num_rows, IndType* indices,
int* source_rows, const int k, const int start_expert, const int end_expert, const bool renormalize,
const float* bias, const double routed_scaling_factor)
const float* bias, const double routed_scaling_factor, const bool* is_padding)
{
static_assert(std::is_same_v<InputType, float> || std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -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<IndType>(-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 <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM, int MAX_BYTES_PER_LDG, typename IndType, typename InputType, ScoringFunc SF>
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<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM, InputType>;
@@ -616,7 +621,7 @@ void topkGatingLauncherHelper(const InputType* input, const bool* finished, floa
dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);
topkGating<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG, WARP_SIZE_PARAM, IndType, InputType, SF><<<num_blocks, block_dim, 0, stream>>>(
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<NUM_EXPERTS, WARPS_PER_TB, 32, MAX_BYTES, \
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 { \
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<TPB><<<num_tokens, TPB, 0, stream>>>(
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<torch::stable::Tensor> bias,
double routed_scaling_factor,
cudaStream_t stream)
cudaStream_t stream,
std::optional<torch::stable::Tensor> 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<float>();
}
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<bool>();
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
vllm::moe::topkGatingKernelLauncher<int, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -777,7 +796,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
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<uint32_t, ComputeType, SF>(
reinterpret_cast<const ComputeType*>(gating_output.const_data_ptr()),
@@ -786,7 +805,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
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<int64_t, ComputeType, SF>(
@@ -796,7 +815,7 @@ void dispatch_topk_launch(
token_expert_indices.mutable_data_ptr<int>(),
softmax_workspace.mutable_data_ptr<float>(),
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<torch::stable::Tensor> bias)
std::optional<torch::stable::Tensor> bias,
std::optional<torch::stable::Tensor> 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<float, 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::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<torch::stable::Tensor> bias,
double routed_scaling_factor)
double routed_scaling_factor,
std::optional<torch::stable::Tensor> 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<float, 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::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());
}
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+100 -49
View File
@@ -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<int>(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<OutIndType>(expert);
indices[offset] = !is_pad_row ? static_cast<OutIndType>(expert)
: static_cast<OutIndType>(-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<OutIndType, HashIndType>;
cudaLaunchConfig_t config = {};
@@ -134,7 +141,7 @@ void launchDsv4HashTopk(const float* input, float* output, OutIndType* indices,
config.numAttrs = 1;
const float scale = static_cast<float>(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<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -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<InputType, float>) {
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<InputType, float>) {
using VecType = AlignedArray<float, ELTS_PER_LDG>;
VecType* row_chunk_vec_ptr = reinterpret_cast<VecType*>(&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<int64_t>(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<IndType>(expert);
indices[idx] = !is_pad_row ? static_cast<IndType>(expert)
: static_cast<IndType>(-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<IndType>(-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<<<num_blocks, block_dim, 0, stream>>>(
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<NUM_EXPERTS, WARPS_PER_TB, 32, \
MAX_BYTES>( \
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<InputType, float>) {
if (use_hash && topk == 6 && renormalize &&
(num_experts == 256 || num_experts == 384)) {
launchDsv4HashTopk<IndType, HashIndType>(
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<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream) {
const std::optional<torch::stable::Tensor>& tid2eid, cudaStream_t stream,
const std::optional<torch::stable::Tensor>& is_padding) {
const float* bias_ptr = nullptr;
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
@@ -737,6 +770,22 @@ void dispatch_topk_softplus_sqrt_launch(
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::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<bool>();
}
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<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int64_t>(),
tid2eid.value().const_data_ptr<int64_t>(), stream);
tid2eid.value().const_data_ptr<int64_t>(), 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<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, true, input_ids.value().const_data_ptr<int>(),
tid2eid.value().const_data_ptr<int>(), stream);
tid2eid.value().const_data_ptr<int>(), stream, is_padding_ptr);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
@@ -769,7 +818,7 @@ void dispatch_topk_softplus_sqrt_launch(
topk_indices_ptr, token_expert_indices.mutable_data_ptr<int>(),
num_tokens, num_experts, topk, renormalize, routed_scaling_factor,
bias_ptr, false, static_cast<const OutIndType*>(nullptr),
static_cast<const OutIndType*>(nullptr), stream);
static_cast<const OutIndType*>(nullptr), stream, is_padding_ptr);
}
};
@@ -793,7 +842,8 @@ void topk_softplus_sqrt(
bool renormalize, double routed_scaling_factor,
const std::optional<torch::stable::Tensor>& correction_bias,
const std::optional<torch::stable::Tensor>& input_ids,
const std::optional<torch::stable::Tensor>& tid2eid) {
const std::optional<torch::stable::Tensor>& tid2eid,
const std::optional<torch::stable::Tensor>& 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<float>(
gating_output.const_data_ptr<float>(), 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<const __half*>(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<const __nv_bfloat16*>(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());
csrc/libtorch_stable/moe/torch_bindings.cpp
+3 -3
View File
@@ -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
tests/kernels/moe/test_topk_softplus_sqrt.py
+117
View File
@@ -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)
vllm/_custom_ops.py
+6
View File
@@ -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,
)
vllm/envs.py
+3 -4
View File
@@ -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(
vllm/model_executor/layers/fused_moe/modular_kernel.py
-16
View File
@@ -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
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+11
View File
@@ -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
vllm/model_executor/layers/fused_moe/router/fused_topk_router.py
+11
View File
@@ -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