diff --git a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu index 095a7667831..785bbf2f6e0 100644 --- a/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu +++ b/csrc/libtorch_stable/moe/topk_softplus_sqrt_kernels.cu @@ -44,6 +44,12 @@ typedef __hip_bfloat162 __nv_bfloat162; namespace vllm { namespace moe { +template +__device__ __forceinline__ int64_t load_index_as_int64(const HashIndType* ptr, + int64_t offset) { + return static_cast(ptr[offset]); +} + /// Aligned array type template + 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 || std::is_same_v || std::is_same_v, @@ -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(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( + 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(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( + 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 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; + 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 +template 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(); } - 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(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), 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::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( + gating_output, topk_weights.mutable_data_ptr(), + 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); + } else { + STD_TORCH_CHECK(tid2eid.value().scalar_type() == + torch::headeronly::ScalarType::Int); + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + 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); + } + } else { + vllm::moe::topkGatingSoftplusSqrtKernelLauncher( + gating_output, topk_weights.mutable_data_ptr(), + 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); + } + }; + + if (topk_indices.scalar_type() == torch::headeronly::ScalarType::Int) { + launch(topk_indices.mutable_data_ptr()); } 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(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), 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()); } 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(); - tid2eid_ptr = tid2eid.value().const_data_ptr(); - } - - vllm::moe::topkGatingSoftplusSqrtKernelLauncher( - gating_output, topk_weights.mutable_data_ptr(), - topk_indices.mutable_data_ptr(), - token_expert_indices.mutable_data_ptr(), 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()); } } @@ -738,4 +747,4 @@ void topk_softplus_sqrt( STD_TORCH_CHECK(false, "Unsupported gating_output data type: ", gating_output.scalar_type()); } -} \ No newline at end of file +} diff --git a/tests/kernels/moe/test_topk_softplus_sqrt.py b/tests/kernels/moe/test_topk_softplus_sqrt.py index 1b68213fafe..46ca934c146 100644 --- a/tests/kernels/moe/test_topk_softplus_sqrt.py +++ b/tests/kernels/moe/test_topk_softplus_sqrt.py @@ -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( 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 058a96ed6b5..1ddcaa50e83 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 @@ -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