A high-throughput and memory-efficient inference and serving engine for LLMs
[Security] Fix remote DoS via invalid recovered token reinjection (#44744)
Signed-off-by: jperezde <jperezde@redhat.com>
Juan Pérez de Algaba committed
Jun 10, 2026 at 11:31 UTC
8a5cf1ccd65e8ac7635c402c1ec0b08988bc26ca
2 files changed
+99
-1
tests/v1/sample/test_rejection_sampler.py
+93
@@ -998,6 +998,99 @@ def test_sample_recovered_tokens_uses_fp64_exponential_race_when_requested():
998
assert torch.equal(actual, expected)
999
1000
1001
+@pytest.mark.parametrize("no_draft_probs", [True, False])
1002
+@pytest.mark.parametrize(
1003
+ "vocab_size",
1004
+ [
1005
+ 100, # below BLOCK_SIZE: single partial tile with many padding entries
1006
+ 8193, # BLOCK_SIZE + 1: only 1 valid entry in the last tile
1007
+ 10000, # non-aligned, moderate tail
1008
+ 151936, # real-world Qwen3 vocab size from the CVE report
1009
+ ],
1010
+)
1011
+def test_sample_recovered_tokens_vocab_boundary(vocab_size: int, no_draft_probs: bool):
1012
+ """Regression test for GHSA-8wr5-jm2h-8r4f.
1013
+
1014
+ When vocab_size is not a multiple of BLOCK_SIZE (8192), the last Triton
1015
+ tile extends beyond the vocabulary. If all valid entries in that tail tile
1016
+ have zero target probability, the out-of-range masked positions (score 0)
1017
+ could win the tl.max tie-break, producing recovered_id >= vocab_size.
1018
+ This test forces that scenario and asserts every recovered token is valid.
1019
+ """
1020
+ BLOCK_SIZE = 8192
1021
+ batch_size = 2
1022
+ max_spec_len = 3
1023
+ num_tokens = batch_size * max_spec_len
1024
+
1025
+ last_tile_start = (vocab_size // BLOCK_SIZE) * BLOCK_SIZE
1026
+
1027
+ target_probs = torch.rand(
1028
+ num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE
1029
+ )
1030
+ if last_tile_start > 0:
1031
+ # Zero out valid entries in the last partial tile so the only
1032
+ # non-zero scores come from earlier, fully-covered tiles.
1033
+ target_probs[:, last_tile_start:] = 0.0
1034
+ else:
1035
+ # vocab_size < BLOCK_SIZE: single tile. Concentrate all mass on
1036
+ # entry 0 so the NO_DRAFT_PROBS path (which zeroes the draft
1037
+ # token entry) can drive all valid scores to zero.
1038
+ target_probs = torch.zeros_like(target_probs)
1039
+ target_probs[:, 0] = 1.0
1040
+ # Re-normalize so it's a valid distribution.
1041
+ target_probs = target_probs / target_probs.sum(dim=-1, keepdim=True)
1042
+
1043
+ draft_probs = torch.rand(
1044
+ num_tokens, vocab_size, dtype=torch.float32, device=DEVICE_TYPE
1045
+ )
1046
+ draft_probs = torch.nn.functional.softmax(draft_probs, dim=-1)
1047
+
1048
+ if last_tile_start == 0:
1049
+ # Force draft token to 0 so the NO_DRAFT_PROBS path zeroes the
1050
+ # only non-zero entry, leaving all valid scores at zero.
1051
+ draft_token_ids = torch.zeros(
1052
+ num_tokens, 1, dtype=torch.int32, device=DEVICE_TYPE
1053
+ )
1054
+ else:
1055
+ draft_token_ids = torch.randint(
1056
+ 0, vocab_size, (num_tokens, 1), dtype=torch.int32, device=DEVICE_TYPE
1057
+ )
1058
+
1059
+ temperature = torch.ones(batch_size, dtype=torch.float32, device=DEVICE_TYPE)
1060
+ generators = {
1061
+ i: torch.Generator(device=DEVICE_TYPE).manual_seed(42 + i)
1062
+ for i in range(batch_size)
1063
+ }
1064
+ sampling_metadata = create_sampling_metadata(
1065
+ all_greedy=False, temperature=temperature, generators=generators
1066
+ )
1067
+
1068
+ spec_decode_metadata = create_spec_decode_metadata(
1069
+ draft_token_ids.reshape(batch_size, max_spec_len).tolist(),
1070
+ torch.rand(num_tokens, vocab_size, device=DEVICE_TYPE),
1071
+ )
1072
+
1073
+ recovered = sample_recovered_tokens(
1074
+ max_spec_len,
1075
+ spec_decode_metadata.num_draft_tokens,
1076
+ spec_decode_metadata.cu_num_draft_tokens,
1077
+ draft_token_ids.squeeze(-1),
1078
+ None if no_draft_probs else draft_probs,
1079
+ target_probs,
1080
+ sampling_metadata,
1081
+ device=DEVICE_TYPE,
1082
+ )
1083
+
1084
+ assert (recovered >= 0).all(), (
1085
+ f"Recovered token IDs contain negative values: "
1086
+ f"{recovered[recovered < 0].tolist()}"
1087
+ )
1088
+ assert (recovered < vocab_size).all(), (
1089
+ f"Recovered token IDs >= vocab_size ({vocab_size}): "
1090
+ f"{recovered[recovered >= vocab_size].tolist()}"
1091
+ )
1092
+
1093
+
1094
########################### Tests for Synthetic Rejection Sampling #########
1095
1096
vllm/v1/sample/rejection_sampler.py
+6
-1
@@ -921,12 +921,17 @@ def sample_recovered_tokens_kernel(
921
other=0.0,
922
)
923
924
- # Local tile reduction
924
+ # Local tile reduction.
925
+ # Mask out-of-vocabulary entries to -inf so they can never win
926
+ # the argmax — prevents producing recovered_id >= vocab_size
927
+ # when all valid entries in the last tile have zero probability.
928
score = prob * inv_q
929
+ score = tl.where(vocab_mask, score, float("-inf"))
930
local_max, local_id = tl.max(score, axis=0, return_indices=True)
931
932
if local_max > max_val:
933
max_val = local_max
934
recovered_id = v + local_id
935
936
+ recovered_id = tl.minimum(recovered_id, vocab_size - 1)
937
tl.store(output_token_ids_ptr + token_idx, recovered_id)