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

[Perf] Optimize fused_topk_bias for DSv4, 1.5~2x kernel performance improvement (#47463)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-07-15 23:40:55 +00:00
committed by GitHub
parent 015b0320de
commit 2dab187f75
3 changed files with 78 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu
+69 -60
View File
@@ -44,6 +44,12 @@ typedef __hip_bfloat162 __nv_bfloat162;
namespace vllm {
namespace moe {
template <typename HashIndType>
__device__ __forceinline__ int64_t load_index_as_int64(const HashIndType* ptr,
int64_t offset) {
return static_cast<int64_t>(ptr[offset]);
}
/// Aligned array type
template <typename T,
/// Number of elements in the array
@@ -86,14 +92,14 @@ __device__ __forceinline__ float toFloat(T value) {
template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG,
int WARP_SIZE_PARAM, bool USE_HASH, typename IndType,
typename InputType = float>
typename HashIndType, typename InputType = float>
__launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
void topkGatingSoftplusSqrt(
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,
double routed_scaling_factor, const float* correction_bias,
const IndType* input_ids, const IndType* tid2eid) {
const HashIndType* input_ids, const HashIndType* tid2eid) {
static_assert(std::is_same_v<InputType, float> ||
std::is_same_v<InputType, __nv_bfloat16> ||
std::is_same_v<InputType, __half>,
@@ -240,8 +246,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
// Hash MoE path: indices are predetermined from lookup table
if constexpr (USE_HASH) {
const IndType token_id = input_ids[thread_row];
const IndType* expert_indices_for_token = tid2eid + token_id * k;
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);
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
float val = row_chunk[ii];
@@ -252,7 +258,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
float selected_sum = 0.f;
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int expert = expert_indices_for_token[k_idx];
const int expert = static_cast<int>(
load_index_as_int64(tid2eid, token_expert_offset + k_idx));
const int idx = k * thread_row + k_idx;
for (int ii = 0; ii < VPT; ++ii) {
const int group_id = ii / ELTS_PER_LDG;
@@ -261,7 +268,7 @@ __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] = expert;
indices[idx] = static_cast<IndType>(expert);
selected_sum += row_chunk[ii];
break;
}
@@ -285,7 +292,8 @@ __launch_bounds__(WARPS_PER_CTA* WARP_SIZE_PARAM) __global__
#pragma unroll
for (int k_idx = 0; k_idx < k; ++k_idx) {
const int expert = expert_indices_for_token[k_idx];
const int expert = static_cast<int>(
load_index_as_int64(tid2eid, token_expert_offset + k_idx));
const int idx = k * thread_row + k_idx;
for (int ii = 0; ii < VPT; ++ii) {
const int group_id = ii / ELTS_PER_LDG;
@@ -461,14 +469,15 @@ struct TopkConstants {
}
template <int EXPERTS, int WARPS_PER_TB, int WARP_SIZE_PARAM,
int MAX_BYTES_PER_LDG, typename IndType, typename InputType>
int MAX_BYTES_PER_LDG, typename IndType, typename HashIndType,
typename InputType>
void topkGatingSoftplusSqrtLauncherHelper(
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,
double routed_scaling_factor, const float* correction_bias,
const bool use_hash, const IndType* input_ids, const IndType* tid2eid,
cudaStream_t stream) {
const bool use_hash, const HashIndType* input_ids,
const HashIndType* tid2eid, cudaStream_t stream) {
static constexpr int BYTES_PER_LDG =
MIN(MAX_BYTES_PER_LDG, sizeof(InputType) * EXPERTS);
using Constants =
@@ -481,7 +490,8 @@ void topkGatingSoftplusSqrtLauncherHelper(
DISPATCH_HASH(use_hash, USE_HASH, {
auto* kernel =
&topkGatingSoftplusSqrt<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG,
WARP_SIZE_PARAM, USE_HASH, IndType, InputType>;
WARP_SIZE_PARAM, USE_HASH, IndType, HashIndType,
InputType>;
#ifndef USE_ROCM
cudaLaunchConfig_t config = {};
config.gridDim = num_blocks;
@@ -538,13 +548,14 @@ void topkGatingSoftplusSqrtLauncherHelper(
}
#endif
template <typename IndType, typename InputType>
template <typename IndType, typename InputType, typename HashIndType = IndType>
void topkGatingSoftplusSqrtKernelLauncher(
const InputType* gating_output, float* topk_weights, IndType* topk_indices,
int* token_expert_indices, const int num_tokens, const int num_experts,
const int topk, const bool renormalize, double routed_scaling_factor,
const float* correction_bias, const bool use_hash, const IndType* input_ids,
const IndType* tid2eid, cudaStream_t stream) {
const float* correction_bias, const bool use_hash,
const HashIndType* input_ids, const HashIndType* tid2eid,
cudaStream_t stream) {
static constexpr int WARPS_PER_TB = 4;
static constexpr int BYTES_PER_LDG_POWER_OF_2 = 16;
// for bfloat16 dtype, we need 4 bytes loading to make sure num_experts
@@ -644,57 +655,55 @@ void dispatch_topk_softplus_sqrt_launch(
if (correction_bias.has_value()) {
bias_ptr = correction_bias.value().const_data_ptr<float>();
}
bool use_hash = false;
if (tid2eid.has_value()) {
STD_TORCH_CHECK(input_ids.has_value(),
"input_ids is required for hash MoE");
use_hash = true;
}
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
const int* input_ids_ptr = nullptr;
const int* tid2eid_ptr = nullptr;
if (tid2eid.has_value()) {
input_ids_ptr = input_ids.value().const_data_ptr<int>();
tid2eid_ptr = tid2eid.value().const_data_ptr<int>();
}
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int, ComputeType>(
gating_output, topk_weights.mutable_data_ptr<float>(),
topk_indices.mutable_data_ptr<int>(),
token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
input_ids_ptr, tid2eid_ptr, stream);
auto launch = [&](auto* topk_indices_ptr) {
using OutIndType =
typename std::remove_pointer<decltype(topk_indices_ptr)>::type;
if (tid2eid.has_value()) {
STD_TORCH_CHECK(input_ids.has_value(),
"input_ids is required for hash MoE");
STD_TORCH_CHECK(
input_ids.value().scalar_type() == tid2eid.value().scalar_type(),
"input_ids and tid2eid must have the same dtype");
if (tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Long) {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
int64_t>(
gating_output, topk_weights.mutable_data_ptr<float>(),
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);
} else {
STD_TORCH_CHECK(tid2eid.value().scalar_type() ==
torch::headeronly::ScalarType::Int);
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType,
int>(
gating_output, topk_weights.mutable_data_ptr<float>(),
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);
}
} else {
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<OutIndType, ComputeType>(
gating_output, topk_weights.mutable_data_ptr<float>(),
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);
}
};
if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) {
launch(topk_indices.mutable_data_ptr<int>());
} else if (topk_indices.scalar_type() ==
torch::headeronly::ScalarType::UInt32) {
const uint32_t* input_ids_ptr = nullptr;
const uint32_t* tid2eid_ptr = nullptr;
if (tid2eid.has_value()) {
input_ids_ptr = input_ids.value().const_data_ptr<uint32_t>();
tid2eid_ptr = tid2eid.value().const_data_ptr<uint32_t>();
}
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<uint32_t, ComputeType>(
gating_output, topk_weights.mutable_data_ptr<float>(),
topk_indices.mutable_data_ptr<uint32_t>(),
token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
input_ids_ptr, tid2eid_ptr, stream);
launch(topk_indices.mutable_data_ptr<uint32_t>());
} else {
STD_TORCH_CHECK(topk_indices.scalar_type() ==
torch::headeronly::ScalarType::Long);
const int64_t* input_ids_ptr = nullptr;
const int64_t* tid2eid_ptr = nullptr;
if (tid2eid.has_value()) {
input_ids_ptr = input_ids.value().const_data_ptr<int64_t>();
tid2eid_ptr = tid2eid.value().const_data_ptr<int64_t>();
}
vllm::moe::topkGatingSoftplusSqrtKernelLauncher<int64_t, ComputeType>(
gating_output, topk_weights.mutable_data_ptr<float>(),
topk_indices.mutable_data_ptr<int64_t>(),
token_expert_indices.mutable_data_ptr<int>(), num_tokens, num_experts,
topk, renormalize, routed_scaling_factor, bias_ptr, use_hash,
input_ids_ptr, tid2eid_ptr, stream);
launch(topk_indices.mutable_data_ptr<int64_t>());
}
}
@@ -738,4 +747,4 @@ void topk_softplus_sqrt(
STD_TORCH_CHECK(false, "Unsupported gating_output data type: ",
gating_output.scalar_type());
}
}
}
tests/kernels/moe/test_topk_softplus_sqrt.py
+2 -2
View File
@@ -153,9 +153,9 @@ def test_fused_topk_softplus_sqrt_hash(
# experts.
hash_indices_table = torch.stack(
[torch.randperm(num_experts)[:topk] for _ in range(vocab_size)]
).to(device="cuda", dtype=torch.int32)
).to(device="cuda", dtype=torch.long)
input_ids = torch.randint(
0, vocab_size, (num_tokens,), dtype=torch.int32, device="cuda"
0, vocab_size, (num_tokens,), dtype=torch.long, device="cuda"
)
topk_weights_ref, topk_ids_ref = _torch_topk_softplus_sqrt(
vllm/model_executor/layers/fused_moe/router/fused_topk_bias_router.py
+7 -7
View File
@@ -170,13 +170,12 @@ def fused_topk_bias(
hash_indices_table: torch.Tensor | None = None,
routed_scaling_factor: float = 1.0,
):
# The topk kernel dispatches dtype based on topk_ids (set by
# indices_type) and assumes input_tokens/hash_indices_table match.
if indices_type is not None:
if input_tokens is not None and input_tokens.dtype != indices_type:
input_tokens = input_tokens.to(dtype=indices_type)
if hash_indices_table is not None and hash_indices_table.dtype != indices_type:
hash_indices_table = hash_indices_table.to(dtype=indices_type)
if (
input_tokens is not None
and hash_indices_table is not None
and input_tokens.dtype != hash_indices_table.dtype
):
input_tokens = input_tokens.to(dtype=hash_indices_table.dtype)
if not rocm_aiter_ops.is_fused_moe_enabled():
assert hidden_states.size(0) == gating_output.size(0), (
@@ -304,6 +303,7 @@ def fused_topk_bias(
scores_for_choice = scores.view(-1, n_routed_experts)
# For batch invariance, use sorted=True to ensure deterministic expert selection
if hash_indices_table is not None:
assert input_tokens is not None
topk_indices = hash_indices_table[input_tokens]
else:
use_sorted = envs.VLLM_BATCH_INVARIANT