Skip to content

Commit 336593c

Browse files
authored
[None][fix] Fix topk outIndices when using vectorized_process (#9404)
Signed-off-by: Yue Weng <25103990+yweng0828@users.noreply.github.com>
1 parent f95edb5 commit 336593c

File tree

1 file changed

+9
-1
lines changed

1 file changed

+9
-1
lines changed

cpp/tensorrt_llm/kernels/indexerTopK.cu

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -570,7 +570,15 @@ static __device__ void topKPerRowJob(int const* indices, float const* logits, in
570570
}
571571
else
572572
{
573-
outIndices[i] = smemOutput[i] - rowStart;
573+
if (stride1 == 1)
574+
{
575+
// stride1 == 1 will use vectorized_process, which indexes already skip the rowStart.
576+
outIndices[i] = smemOutput[i];
577+
}
578+
else
579+
{
580+
outIndices[i] = smemOutput[i] - rowStart;
581+
}
574582
}
575583
}
576584
}

0 commit comments

Comments
 (0)