[BugFix] Fix topk_softmax assert (#19764)

Signed-off-by: Luka Govedic <lgovedic@redhat.com>

[BugFix] Fix topk_softmax assert (#19764)
Signed-off-by: Luka Govedic <lgovedic@redhat.com>
4f35be10 · Luka Govedič · GitHub · 2b61d2e2 · 4f35be10
Unverified Commit 4f35be10 authored Aug 27, 2025 by Luka Govedič Committed by GitHub Aug 27, 2025
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

csrc/moe/topk_softmax_kernels.cu csrc/moe/topk_softmax_kernels.cu +1 -1

No files found.
--- a/csrc/moe/topk_softmax_kernels.cu
+++ b/csrc/moe/topk_softmax_kernels.cu
@@ -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>(),