[ModelRunner V2] Reduce peak memory of chunked prompt logprobs

tests/v1/sample/test_logprobs.py

@@ -1630,6 +1630,38 @@ def test_token_logprobs_large_batch_int64_row_offset():
     )
 
 
+@pytest.mark.parametrize("num_prompt_logprobs", [-1, 3])
+def test_chunked_prompt_logprobs_match_unchunked(num_prompt_logprobs: int):
+    """Chunks written into preallocated outputs must match one unchunked pass."""
+    if not current_platform.is_cuda_alike():
+        pytest.skip("the logprob kernels are Triton")
+    from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores
+    from vllm.v1.worker.gpu.sample.prompt_logprob import (
+        CHUNK_SIZE,
+        compute_prompt_logprobs_with_chunking,
+    )
+
+    device = torch.device(current_platform.device_type)
+    vocab_size = 512
+    num_tokens = 2 * CHUNK_SIZE + 3
+    logits = torch.randn(num_tokens, vocab_size, device=device)
+    prompt_token_ids = torch.randint(vocab_size, (num_tokens,), device=device)
+
+    token_ids, logprobs, ranks = compute_prompt_logprobs_with_chunking(
+        prompt_token_ids,
+        logits,
+        lambda hidden_states: hidden_states,
+        num_prompt_logprobs,
+    )
+
+    num_logprobs = vocab_size if num_prompt_logprobs == -1 else num_prompt_logprobs
+    expected = compute_topk_scores(logits, num_logprobs, prompt_token_ids)
+    assert token_ids.dtype == torch.int32
+    torch.testing.assert_close(token_ids, expected.logprob_token_ids.to(torch.int32))
+    torch.testing.assert_close(logprobs, expected.logprobs)
+    torch.testing.assert_close(ranks, expected.selected_token_ranks)
+
+
 def test_trace_replay_with_spec_decode(monkeypatch):
     monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
 

vllm/v1/worker/gpu/sample/prompt_logprob.py

@@ -293,12 +293,13 @@ def compute_prompt_logprobs_with_chunking(
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     # Since materializing the full prompt logits can take too much memory,
     # we compute it in chunks.
-    token_ids = []
-    scores = []
-    ranks = []
+    num_tokens = prompt_token_ids.shape[0]
+    token_ids: torch.Tensor | None = None
+    scores: torch.Tensor | None = None
+    ranks: torch.Tensor | None = None
     logits_mode = logprobs_mode in ("raw_logits", "processed_logits")
     prompt_token_ids = prompt_token_ids.to(torch.int64)
-    for start_idx in range(0, prompt_token_ids.shape[0], CHUNK_SIZE):
+    for start_idx in range(0, num_tokens, CHUNK_SIZE):
         end_idx = start_idx + CHUNK_SIZE
         # NOTE(woosuk): logits_fn can be slow because it involves all-gather.
         prompt_logits = logits_fn(prompt_hidden_states[start_idx:end_idx])
@@ -313,11 +314,22 @@ def compute_prompt_logprobs_with_chunking(
             prompt_token_ids[start_idx:end_idx],
             logits_mode=logits_mode,
         )
-        token_ids.append(result.logprob_token_ids)
-        scores.append(result.logprobs)
-        ranks.append(result.selected_token_ranks)
-
-    token_ids = torch.cat(token_ids, dim=0) if len(token_ids) > 1 else token_ids[0]
-    scores = torch.cat(scores, dim=0) if len(scores) > 1 else scores[0]
-    ranks = torch.cat(ranks, dim=0) if len(ranks) > 1 else ranks[0]
+        # Return token_ids as int32 to reduce peak memory usage.
+        if num_tokens <= CHUNK_SIZE:
+            logprob_token_ids = result.logprob_token_ids.to(torch.int32)
+            return logprob_token_ids, result.logprobs, result.selected_token_ranks
+
+        if token_ids is None:
+            width = result.logprobs.shape[1]
+            token_ids = prompt_token_ids.new_empty(
+                (num_tokens, width), dtype=torch.int32
+            )
+            scores = result.logprobs.new_empty((num_tokens, width))
+            ranks = result.selected_token_ranks.new_empty(num_tokens)
+        assert scores is not None and ranks is not None
+        token_ids[start_idx:end_idx] = result.logprob_token_ids
+        scores[start_idx:end_idx] = result.logprobs
+        ranks[start_idx:end_idx] = result.selected_token_ranks
+
+    assert token_ids is not None and scores is not None and ranks is not None
     return token_ids, scores, ranks