A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 03:30 UTC
[CI] Reject out-of-vocabulary before they reach the GPU logprob path (#44042)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Jun 2, 2026 at 22:27 UTC
53b88d1dfc1f43cc9c7ead2cdc1586fb01ce89a9
4 files changed
+82
-9
tests/v1/sample/test_logprobs.py
+31
@@ -4,6 +4,7 @@
4
import itertools
5
import math
6
from collections.abc import Generator
7
+from types import SimpleNamespace
8
from typing import get_args
9
10
import pytest
@@ -20,6 +21,7 @@ from tests.v1.sample.utils import (
21
from vllm import SamplingParams
22
from vllm.config.model import LogprobsMode
23
from vllm.distributed import cleanup_dist_env_and_memory
24
+from vllm.exceptions import VLLMValidationError
25
from vllm.platforms import current_platform
26
27
from ...conftest import HfRunner, VllmRunner
@@ -78,6 +80,14 @@ def hf_model(hf_runner) -> Generator[HfRunner, None, None]:
80
yield hf_model
81
82
83
+def _model_config(vocab_size: int = 10):
84
+ return SimpleNamespace(
85
+ max_logprobs=20,
86
+ logits_processors=None,
87
+ get_vocab_size=lambda: vocab_size,
88
+ )
89
+
90
+
91
def _repeat_logprob_config(
92
test_prompts,
93
logprob_prompt_logprob_list: BatchLogprobsSpecType,
@@ -397,6 +407,27 @@ def test_max_logprobs():
407
runner.generate(["Hello world"], sampling_params=bad_sampling_params)
408
409
410
+@pytest.mark.parametrize("token_ids", [[0], [0, 9]])
411
+def test_logprob_token_ids_validate_vocab_bounds_valid(token_ids: list[int]):
412
+ SamplingParams(logprob_token_ids=token_ids).verify(
413
+ _model_config(),
414
+ speculative_config=None,
415
+ structured_outputs_config=None,
416
+ tokenizer=None,
417
+ )
418
+
419
+
420
+@pytest.mark.parametrize("token_ids", [[-1], [10], [-35, 1873042417]])
421
+def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]):
422
+ with pytest.raises(VLLMValidationError, match="logprob_token_ids"):
423
+ SamplingParams(logprob_token_ids=token_ids).verify(
424
+ _model_config(),
425
+ speculative_config=None,
426
+ structured_outputs_config=None,
427
+ tokenizer=None,
428
+ )
429
+
430
+
431
def test_none_logprobs(vllm_model, example_prompts):
432
"""Engine should return `logprobs` and `prompt_logprobs` as `None`
433
vllm/sampling_params.py
+14
@@ -737,6 +737,20 @@ class SamplingParams(
737
parameter="logprob_token_ids",
738
value=n,
739
)
740
+ vocab_size = model_config.get_vocab_size()
741
+ invalid_token_ids = [
742
+ token_id
743
+ for token_id in self.logprob_token_ids
744
+ if token_id < 0 or token_id >= vocab_size
745
+ ]
746
+ if invalid_token_ids:
747
+ raise VLLMValidationError(
748
+ f"token_id(s) {invalid_token_ids} in logprob_token_ids "
749
+ f"contain out-of-vocab token ids. Vocabulary size: "
750
+ f"{vocab_size}",
751
+ parameter="logprob_token_ids",
752
+ value=invalid_token_ids,
753
+ )
754
if self.logprobs is not None and self.logprobs != n:
755
raise VLLMValidationError(
756
f"When both logprobs and logprob_token_ids are set, "
vllm/v1/sample/thinking_budget_state.py
+27
-9
@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Any
6
7
import torch
8
9
+from vllm.platforms import current_platform
10
from vllm.utils.torch_utils import async_tensor_h2d
11
from vllm.v1.sample.logits_processor.interface import (
12
BatchUpdate,
@@ -511,14 +512,31 @@ class ThinkingBudgetStateHolder:
512
513
if active_indices_cpu:
514
device = logits.device
514
- active_indices = async_tensor_h2d(
515
- active_indices_cpu, dtype=torch.long, device=device
516
- )
517
- force_tokens = async_tensor_h2d(
518
- force_tokens_cpu, dtype=torch.long, device=device
519
- )
520
- # Avoid CPU->GPU sync.
521
- fill = logits.new_full((len(active_indices_cpu),), 1e9)
522
- logits.index_put_((active_indices, force_tokens), fill)
515
+ if current_platform.is_rocm() and logits.is_contiguous():
516
+ # Flattened index_fill avoids ROCm faults seen with 2-D
517
+ # advanced-indexing writes on the thinking-budget path.
518
+ vocab_size = logits.shape[1]
519
+ flat_indices_cpu = [
520
+ row * vocab_size + token
521
+ for row, token in zip(active_indices_cpu, force_tokens_cpu)
522
+ ]
523
+ flat_indices = async_tensor_h2d(
524
+ flat_indices_cpu, dtype=torch.long, device=device
525
+ )
526
+ logits.view(-1).index_fill_(0, flat_indices, 1e9)
527
+ elif current_platform.is_rocm():
528
+ fill = logits.new_tensor(1e9)
529
+ for row, token in zip(active_indices_cpu, force_tokens_cpu):
530
+ logits[row, token] = fill
531
+ else:
532
+ active_indices = async_tensor_h2d(
533
+ active_indices_cpu, dtype=torch.long, device=device
534
+ )
535
+ force_tokens = async_tensor_h2d(
536
+ force_tokens_cpu, dtype=torch.long, device=device
537
+ )
538
+ # Avoid CPU->GPU sync.
539
+ fill = logits.new_full((len(active_indices_cpu),), 1e9)
540
+ logits.index_put_((active_indices, force_tokens), fill)
541
542
return logits
vllm/v1/worker/gpu_model_runner.py
+10
@@ -6278,11 +6278,21 @@ class GPUModelRunner(
6278
6279
# Calls torch.accelerator.synchronize()
6280
self._cleanup_profiling_kv_cache()
6281
+ if current_platform.is_rocm():
6282
+ # Drop captured graphs before distributed teardown. On ROCm, delayed
6283
+ # graph destruction can surface HSA faults in the next engine startup.
6284
+ CUDAGraphWrapper.clear_all_graphs()
6285
+ BreakableCUDAGraphWrapper.clear_all_graphs()
6286
+ self.encoder_cudagraph_manager = None
6287
self.compilation_config.static_forward_context.clear()
6288
self.model = None # type: ignore[assignment]
6289
_ROPE_DICT.clear()
6290
6291
reset_workspace_manager()
6292
+ if current_platform.is_rocm():
6293
+ gc.collect()
6294
+ torch.accelerator.empty_cache()
6295
+ torch.accelerator.synchronize()
6296
6297
def _cleanup_profiling_kv_cache(self) -> None:
6298
torch.accelerator.synchronize()