Adjusting invalid k values in top-k selection beyond vocabulary limits

[dashanji]

Clamp k to [1, vocab_size-1] when it's out of vocab size.

FIX #14181

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[njhill]

Thanks @dashanji! Would it be possible to move this to where the requests are added to the batch?

[DarkLight1337]

Do you know how can the logits shape exceed the vocab size in the first place?

[dashanji]

Thanks @dashanji! Would it be possible to move this to where the requests are added to the batch?

Hi @njhill. Thanks for the suggestion, I have tried to move the same code to here but it doesn't work. Could you please take a look at this?

[dashanji]

Do you know how can the logits shape exceed the vocab size in the first place?

Sorry, I have no idea at present. Maybe I need to explore further.

[njhill]

Do you know how can the logits shape exceed the vocab size in the first place?

Sorry, I have no idea at present. Maybe I need to explore further.

@DarkLight1337 IIUC this isn't related to the logits shape exceeding the vocab size but rather guarding against arbitrary provided top_k values which may be larger than the vocab size (and thus larger than the logits token dimension).

[DarkLight1337]

Do you know how can the logits shape exceed the vocab size in the first place?

Sorry, I have no idea at present. Maybe I need to explore further.

@DarkLight1337 IIUC this isn't related to the logits shape exceeding the vocab size but rather guarding against arbitrary provided top_k values which may be larger than the vocab size (and thus larger than the logits token dimension).

I see. Still, shouldn't we prohibit (or at least warn about) such values in the first place instead of silently clamping them?

[njhill]

@DarkLight1337 I think it's kind of unlikely/obscure for someone to pass such large values and if they do the clamping behaviour is technically still correct so I don't think there's any need for a warning.

@dashanji actually when I looked closer at this I found that there's a more significant bug. I've opened another PR for that #14301 which should address this too. I'll include you as a coauthor on that.

[njhill]

Closing in favour of #14301. Thanks again @dashanji.

[dashanji]

Thanks to @njhill for the quick fix, it works!