[Bugfix][V1] Chunk prompt-logprobs logits to bound the activation peak

tests/v1/worker/test_gpu_prompt_logprobs_chunking.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Chunked prompt-logprobs must match the unchunked result exactly.
+
+`_get_prompt_logprobs_dict` materializes full-vocabulary logits in row
+chunks so the float32 score tensor cannot dominate the activation peak
+(issue #5907). Chunking rewrites the destination row mapping, so the
+regression risk is misaligned rows rather than wrong values -- these
+tests pin the mapping by comparing against a single-chunk run.
+"""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+from vllm import envs
+from vllm.utils import torch_utils
+from vllm.v1.outputs import LogprobsTensors
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+VOCAB_SIZE = 64
+HIDDEN_SIZE = 8
+
+
+@pytest.fixture(autouse=True)
+def _no_pin_memory(monkeypatch):
+    """Run the host-to-device copy unpinned.
+
+    `_get_prompt_logprobs_dict` calls `async_tensor_h2d`, which pins the
+    source buffer when `PIN_MEMORY` is set. Pinning requires a CUDA
+    allocator, so leaving it enabled makes this test unrunnable on
+    CPU-only machines. The chunking logic under test is unaffected.
+    """
+    monkeypatch.setattr(torch_utils, "PIN_MEMORY", False)
+
+
+def _make_runner(logprobs_mode: str, prompt_len: int, num_scheduled: int):
+    """Build the minimal runner state `_get_prompt_logprobs_dict` reads."""
+    torch.manual_seed(0)
+    # A deterministic projection so compute_logits is reproducible across
+    # chunk sizes: identical rows in must give identical logits out.
+    weight = torch.randn(HIDDEN_SIZE, VOCAB_SIZE)
+
+    request = SimpleNamespace(
+        prompt_token_ids=list(range(prompt_len)),
+        num_computed_tokens=0,
+        in_progress_prompt_logprobs_cpu=None,
+    )
+
+    from vllm.v1.sample.sampler import Sampler
+
+    runner = object.__new__(GPUModelRunner)
+    runner.num_prompt_logprobs = {"req0": 3}
+    runner.requests = {"req0": request}
+    runner.device = torch.device("cpu")
+    runner.model = SimpleNamespace(compute_logits=lambda h: h @ weight)
+    runner.sampler = SimpleNamespace(
+        compute_logprobs=Sampler.compute_logprobs,
+        gather_logprobs=Sampler.gather_logprobs,
+    )
+    runner.model_config = SimpleNamespace(logprobs_mode=logprobs_mode)
+    runner.input_batch = SimpleNamespace(req_id_to_index={"req0": 0})
+    runner.query_start_loc = SimpleNamespace(np=np.array([0], dtype=np.int32))
+    runner._sync_device = lambda: None
+
+    hidden_states = torch.randn(num_scheduled, HIDDEN_SIZE)
+    return runner, hidden_states, request
+
+
+def _run(monkeypatch, logprobs_mode, prompt_len, num_scheduled, chunk_size):
+    monkeypatch.setattr(envs, "VLLM_PROMPT_LOGPROBS_CHUNK_SIZE", chunk_size)
+    runner, hidden_states, _request = _make_runner(
+        logprobs_mode, prompt_len, num_scheduled
+    )
+    out = runner._get_prompt_logprobs_dict(hidden_states, {"req0": num_scheduled})
+    tensors = out["req0"]
+    assert isinstance(tensors, LogprobsTensors)
+    return tensors
+
+
+@pytest.mark.parametrize("logprobs_mode", ["raw_logprobs", "raw_logits"])
+@pytest.mark.parametrize(
+    "chunk_size",
+    [
+        1,  # degenerate: one row per chunk
+        3,  # does not divide evenly -> short final chunk
+        4,  # divides evenly
+        1024,  # larger than the batch -> single chunk
+    ],
+)
+def test_chunked_matches_unchunked(monkeypatch, logprobs_mode, chunk_size):
+    # Schedule the whole prompt in one step so prefill completes and the
+    # request's tensors are returned. num_logits is then prompt_len - 1.
+    prompt_len, num_scheduled = 13, 13
+
+    reference = _run(
+        monkeypatch, logprobs_mode, prompt_len, num_scheduled, chunk_size=10**9
+    )
+    chunked = _run(
+        monkeypatch, logprobs_mode, prompt_len, num_scheduled, chunk_size=chunk_size
+    )
+
+    torch.testing.assert_close(chunked.logprobs, reference.logprobs)
+    assert torch.equal(chunked.logprob_token_ids, reference.logprob_token_ids)
+    assert torch.equal(
+        chunked.selected_token_ranks, reference.selected_token_ranks
+    )
+
+
+def test_chunking_bounds_the_materialized_logits(monkeypatch):
+    """No single compute_logits call may exceed the configured chunk size."""
+    chunk_size = 4
+    prompt_len, num_scheduled = 13, 13
+    num_logits = prompt_len - 1
+    monkeypatch.setattr(envs, "VLLM_PROMPT_LOGPROBS_CHUNK_SIZE", chunk_size)
+
+    runner, hidden_states, _ = _make_runner(
+        "raw_logprobs", prompt_len, num_scheduled
+    )
+    seen_rows = []
+    inner = runner.model.compute_logits
+
+    def recording_compute_logits(h):
+        seen_rows.append(h.shape[0])
+        return inner(h)
+
+    runner.model = SimpleNamespace(compute_logits=recording_compute_logits)
+    runner._get_prompt_logprobs_dict(hidden_states, {"req0": num_scheduled})
+
+    assert seen_rows, "compute_logits was never called"
+    assert max(seen_rows) <= chunk_size
+    # Every row is still covered exactly once.
+    assert sum(seen_rows) == num_logits

vllm/envs.py

@@ -166,6 +166,7 @@
     VLLM_RUST_FRONTEND_PATH: str | None = "auto"
     VLLM_SERVER_DEV_MODE: bool = False
     VLLM_V1_OUTPUT_PROC_CHUNK_SIZE: int = 128
+    VLLM_PROMPT_LOGPROBS_CHUNK_SIZE: int = 1024
     VLLM_MLA_DISABLE: bool = False
     VLLM_RAY_PER_WORKER_GPUS: float = 1.0
     VLLM_RAY_BUNDLE_INDICES: str = ""
@@ -1430,6 +1431,15 @@ def _resolve_rust_cli_path() -> str | None:
     "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE": lambda: int(
         os.getenv("VLLM_V1_OUTPUT_PROC_CHUNK_SIZE", "128")
     ),
+    # Maximum number of prompt-token rows whose full-vocabulary logits are
+    # materialized at once while computing prompt logprobs. The prompt
+    # logprobs path upcasts logits to float32 via log_softmax, so an
+    # unchunked [num_scheduled_tokens, vocab_size] tensor can dominate the
+    # activation peak and OOM (see issue #5907). Lower this to trade
+    # throughput for a smaller peak.
+    "VLLM_PROMPT_LOGPROBS_CHUNK_SIZE": lambda: int(
+        os.getenv("VLLM_PROMPT_LOGPROBS_CHUNK_SIZE", "1024")
+    ),
     # If set, vLLM will disable the MLA attention optimizations.
     "VLLM_MLA_DISABLE": lambda: bool(int(os.getenv("VLLM_MLA_DISABLE", "0"))),
     # If set, vLLM will pick up the provided Flash Attention MLA
@@ -2346,6 +2356,7 @@ def compile_factors() -> dict[str, object]:
         "VLLM_WORKER_MULTIPROC_METHOD",
         "VLLM_ENABLE_V1_MULTIPROCESSING",
         "VLLM_V1_OUTPUT_PROC_CHUNK_SIZE",
+        "VLLM_PROMPT_LOGPROBS_CHUNK_SIZE",
         "VLLM_CPU_KVCACHE_SPACE",
         "VLLM_CPU_MOE_PREPACK",
         "VLLM_ZENTORCH_WEIGHT_PREPACK",

vllm/v1/worker/gpu_model_runner.py

@@ -5758,33 +5758,52 @@ def _get_prompt_logprobs_dict(
             req_idx = self.input_batch.req_id_to_index[req_id]
             offset = self.query_start_loc.np[req_idx].item()
             prompt_hidden_states = hidden_states[offset : offset + num_logits]
-            logits = self.model.compute_logits(prompt_hidden_states)
 
             # Get the "target" tokens for each index. For prompt at index i,
             # the token at prompt index i+1 is the "sampled" token we want
             # to gather the logprob for.
             tgt_token_ids = prompt_token_ids[start_tok : start_tok + num_logits]
 
-            # Compute prompt scores respecting logprobs_mode.
-            # NOTE: prompt tokens skip sampling processors, so
-            # processed_* and raw_* yield the same scores here.
-            if self.model_config.logprobs_mode in ("raw_logits", "processed_logits"):
-                scores = logits.to(torch.float32)
-            else:
-                scores = self.sampler.compute_logprobs(logits)
-            token_ids, logprobs, ranks, *_ = self.sampler.gather_logprobs(
-                scores, num_prompt_logprobs, tgt_token_ids
-            )
+            # Materialize the full-vocabulary logits in row chunks. Prompt
+            # logprobs upcast to float32, so an unchunked
+            # [num_logits, vocab_size] tensor plus its float32 copy can
+            # dominate the activation peak: at 8192 rows and a 152k
+            # vocabulary that is ~5 GiB for the float32 scores alone. That
+            # peak is not exercised by profile_run (_dummy_sampler_run runs
+            # with max_num_logprobs=None), so it is not reflected in the KV
+            # cache sizing and shows up as an OOM at serving time (#5907).
+            chunk_size = envs.VLLM_PROMPT_LOGPROBS_CHUNK_SIZE
+            for pos in range(0, num_logits, chunk_size):
+                end = min(pos + chunk_size, num_logits)
+                logits = self.model.compute_logits(prompt_hidden_states[pos:end])
+
+                # Compute prompt scores respecting logprobs_mode.
+                # NOTE: prompt tokens skip sampling processors, so
+                # processed_* and raw_* yield the same scores here.
+                if self.model_config.logprobs_mode in (
+                    "raw_logits",
+                    "processed_logits",
+                ):
+                    scores = logits.to(torch.float32)
+                else:
+                    scores = self.sampler.compute_logprobs(logits)
+                del logits
+                token_ids, logprobs, ranks, _ = self.sampler.gather_logprobs(
+                    scores, num_prompt_logprobs, tgt_token_ids[pos:end]
+                )
+                del scores
 
-            # Transfer GPU->CPU async.
-            chunk_slice = slice(start_idx, start_idx + num_logits)
-            logprobs_tensors.logprob_token_ids[chunk_slice].copy_(
-                token_ids, non_blocking=True
-            )
-            logprobs_tensors.logprobs[chunk_slice].copy_(logprobs, non_blocking=True)
-            logprobs_tensors.selected_token_ranks[chunk_slice].copy_(
-                ranks, non_blocking=True
-            )
+                # Transfer GPU->CPU async.
+                chunk_slice = slice(start_idx + pos, start_idx + end)
+                logprobs_tensors.logprob_token_ids[chunk_slice].copy_(
+                    token_ids, non_blocking=True
+                )
+                logprobs_tensors.logprobs[chunk_slice].copy_(
+                    logprobs, non_blocking=True
+                )
+                logprobs_tensors.selected_token_ranks[chunk_slice].copy_(
+                    ranks, non_blocking=True
+                )
 
         # Remove requests that have completed prefill from the batch
         # num_prompt_logprobs_dict.