Repository navigation
[Security] Fix remote DoS via invalid recovered token reinjection - #44744
MatthewBonanni merged 3 commits into
Conversation
Mask out-of-vocabulary entries to -inf in sample_recovered_tokens_kernel so they can never win the tiled argmax when all valid entries in the last tile have zero probability. Add a belt-and-suspenders clamp on recovered_id before store. Add regression test that forces the zero-tail-tile scenario for non-aligned vocab sizes. Signed-off-by: Juan Pérez de Algaba <jperezde@redhat.com> Signed-off-by: jperezde <jperezde@redhat.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
MatthewBonanni
left a comment
MatthewBonanni
left a comment
There was a problem hiding this comment.
Thanks for the fix, this seems reasonable to me. Can you confirm that the test fails on the unpatched kernel though? Shouldn't the failure happen when vocab_size < BLOCK_SIZE?
Add vocab_size=100 (below BLOCK_SIZE) to the parametrized test to cover the single partial tile scenario where padding entries outnumber valid entries. Force draft_token_id=0 in this case so the NO_DRAFT_PROBS path zeroes the only non-zero entry, reproducing the tie-break condition. Signed-off-by: jperezde <jperezde@redhat.com>
|
@MatthewBonanni Yes, the test fails on the unpatched kernel -- without the -inf masking, the zero-score out-of-range entries in the last partial tile can win the tl.max tie-break, producing recovered_id >= vocab_size. The bug applies when vocab_size < BLOCK_SIZE too (single partial tile, same mechanism). The fix handles both cases since vocab_mask = vocab_offset < vocab_size works regardless of whether vocab_size is above or below BLOCK_SIZE. I have added also a test case for that specific use case |
MatthewBonanni
added
the
ready
…lm-project#44744) Signed-off-by: jperezde <jperezde@redhat.com>
Cherry-pick 62 bugfix/security PRs from upstream vllm-project/vllm main (2026-05-03 to 2026-06-17), covering scheduler, engine core, model runner, worker, attention, KV cache, compilation, and structured output fixes. Security (4): vllm-project#43286 vllm-project#44744 vllm-project#45118 vllm-project#45252 Bugfix (56): vllm-project#35536 vllm-project#36616 vllm-project#38895 vllm-project#39155 vllm-project#39324 vllm-project#39562 vllm-project#39805 vllm-project#40398 vllm-project#40726 Runner fix (2): vllm-project#44568 vllm-project#44603 Skipped: vllm-project#43781 (ROCm-specific, not applicable to Ascend NPU) Conflict resolutions: - Manual merge: vllm-project#43286 vllm-project#45118 vllm-project#42112 vllm-project#43160 vllm-project#43719 vllm-project#44560 - Upstream-preferred (-X theirs): vllm-project#43808 vllm-project#43988 vllm-project#42967 vllm-project#35536 vllm-project#45195 - Test files (--theirs): vllm-project#44744 vllm-project#41895 vllm-project#42040 vllm-project#41233 vllm-project#45345 vllm-project#43982 Co-authored-by: GitHub Copilot Signed-off-by: MingqiWang-coder <mingqiwang@hust.edu.cn>
ekagra-ranjan
left a comment
ekagra-ranjan
left a comment
There was a problem hiding this comment.
Thanks for the PR @jperezdealgaba !
I have a followup on a pending issue which is similar to what was observed a year ago but hasn't been fixed on vLLM and is related to the fix in this PR so curious to know your thoughts on below @MatthewBonanni @jperezdealgaba
| max_val = local_max | ||
| recovered_id = v + local_id | ||
|
|
||
| recovered_id = tl.minimum(recovered_id, vocab_size - 1) |
There was a problem hiding this comment.
wondering what's the original motivation for this when the code earlier maks the OOV tokens with -inf so they are not sampled?
On the contrary, I think at this line we should make any OOV token as 0 as per this year old PR. That PR found that nans in the logits will make the target sample tokenid 0 whereas triton will sample OOV since argmax of [nan, nan, -inf, -inf] is 2. The current vLLM code can still lead to OOV from draft until we add this fix. The year old PR doesn't fix the vLLM's behavior of outputting tokenid 0 when nans are in logits but rather makes draft behave like target. As of now, nans in logits gets unnoticed in non SD inference and without this fix it gives an impression that SD is breaking something.
As to why nans can be present in model, I can think of 2 scenarios: some quant method OR some early stage RL checkpoints.
…lm-project#44744) Signed-off-by: jperezde <jperezde@redhat.com>
…ckport vllm#44744) Backport of vllm-project/vllm#44744 (GHSA-8wr5-jm2h-8r4f): sample_recovered_tokens_kernel masks the lanes past the vocabulary to -inf before the tiled argmax and clamps recovered_id to vocab_size - 1, together with its regression test test_sample_recovered_tokens_vocab_boundary. On 1Cat main the kernel returns vocab_size for an all-NaN target row on the path without draft probabilities (DSpark, MTP, EAGLE): tl.max skips NaN and picks a padded lane of the last tile. The next step feeds that id to the embedding lookup, a device-side assert (the chain 1CatAI#658 describes). The added test covers that case with vocab sizes 1000, 129280 and 151936. Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Peuqui <peuqui@github.com>
…ckport vllm#44744) Backport of vllm-project/vllm#44744 (GHSA-8wr5-jm2h-8r4f): sample_recovered_tokens_kernel masks the lanes past the vocabulary to -inf before the tiled argmax and clamps recovered_id to vocab_size - 1, together with its regression test test_sample_recovered_tokens_vocab_boundary. On 1Cat main the kernel returns vocab_size for an all-NaN target row on the path without draft probabilities (DSpark, MTP, EAGLE): tl.max skips NaN and picks a padded lane of the last tile. The next step feeds that id to the embedding lookup, a device-side assert (the chain 1CatAI#658 describes). The added test covers that case with vocab sizes 1000, 129280 and 151936. Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Peuqui <peuqui@github.com>
Summary
Fixes GHSA-8wr5-jm2h-8r4f
When
vocab_sizeis not a multiple ofBLOCK_SIZE(8192), the last Triton tile extends beyond the vocabulary. If all valid entries in that tail tile have zero target probability, the out-of-range masked positions (score 0) could win thetl.maxtie-break, producingrecovered_id >= vocab_size. This results in a remote denial of service via invalid token reinjection.Fix:
-infinsample_recovered_tokens_kernelso they can never win the tiled argmaxrecovered_idbefore storeTest plan
test_sample_recovered_tokens_vocab_boundaryintests/v1/sample/test_rejection_sampler.pythat forces the zero-tail-tile scenario for non-aligned vocab sizes (8193, 10000, 151936)[0, vocab_size)Notes
Made with Cursor