Unverified Commit 4f35be10 authored by Luka Govedič's avatar Luka Govedič Committed by GitHub
Browse files

[BugFix] Fix topk_softmax assert (#19764)


Signed-off-by: default avatarLuka Govedic <lgovedic@redhat.com>
parent 2b61d2e2
...@@ -573,7 +573,7 @@ void topk_softmax( ...@@ -573,7 +573,7 @@ void topk_softmax(
stream); stream);
} }
else { else {
assert(topk_indices.scalar_type() == at::ScalarType::Int64); TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
vllm::moe::topkGatingSoftmaxKernelLauncher( vllm::moe::topkGatingSoftmaxKernelLauncher(
gating_output.data_ptr<float>(), gating_output.data_ptr<float>(),
topk_weights.data_ptr<float>(), topk_weights.data_ptr<float>(),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment