[BugFix] Fix topk_softmax assert (#19764)

ProExpertProg · web-flow · commit 4f35be10a96f · 2025-08-27T09:47:28.000-07:00
Signed-off-by: Luka Govedic &lt;lgovedic@redhat.com&gt;
diff --git 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>(),

Original file line number	Diff line number	Diff line change
`@@ -573,7 +573,7 @@ void topk_softmax(`
`573`	`573`	`stream);`
`574`	`574`	`}`
`575`	`575`	`else {`
`576`		`- assert(topk_indices.scalar_type() == at::ScalarType::Int64);`
	`576`	`+ TORCH_CHECK(topk_indices.scalar_type() == at::ScalarType::Long);`
`577`	`577`	`vllm::moe::topkGatingSoftmaxKernelLauncher(`
`578`	`578`	`gating_output.data_ptr<float>(),`
`579`	`579`	`topk_weights.data_ptr<float>(),`