[BugFix] Fix topk_softmax assert (#19764)

Signed-off-by: Luka Govedic <lgovedic@redhat.com>
This commit is contained in:
Luka Govedič
2025-08-27 12:47:28 -04:00
committed by GitHub
parent 2b61d2e22f
commit 4f35be10a9

View File

@ -573,7 +573,7 @@ void topk_softmax(
stream);
}
else {
assert(topk_indices.scalar_type() == at::ScalarType::Int64);
TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);
vllm::moe::topkGatingSoftmaxKernelLauncher(
gating_output.data_ptr<float>(),
topk_weights.data_ptr<float>(),