SIGN IN SIGN UP

[Bugfix] Normalize disabled top_k rows to vocab_size in mixed sampling batches (#646) (#648)

This PR is:

- To normalize disabled top_k rows to the vocab_size sentinel in mixed 
- To match vLLM's per-row top_k convention in GPUInputBatch and the v1
sampler
- To keep disabled rows sampling from the full vocabulary while enabled
rows keep their
- To stop a mixed top_k=0 / top_k=N batch from crashing EngineCore.

Before: _make_top_k() passed raw per-request top_k values to vLLM's
sampler, so a disabled row carried 0. The PyTorch top-k path computes a
gather index of vocab_size - top_k, indexing out of bounds and
terminating EngineCore (issue #646).

After: _make_top_k() sends top_k when 0 < top_k < vocab_size and the
vocab_size sentinel otherwise. A regression test covers the mixed batch
through the real Sampler.forward() path.

Signed-off-by: Mihal Dimo <mihal@kakao.com>
M
Mihal Dimo committed
813e738d95840bd66b60248ad4a557485320d896
Parent: db4d7c5
Committed by GitHub <noreply@github.com> on 8/27/2026, 10:47:19 AM