[Core] Preallocate sampler logits workspace during memory profiling

[glaziermag]

Summary

Classification: PARTIAL_MITIGATION / TARGETED sampler workspace accounting improvement.

This PR preallocates a float32 sampler logits workspace during GPU memory profiling and reuses that workspace in Sampler.forward() when it is compatible with the current logits tensor. If the workspace is absent, too small, on the wrong device, or the logits are already float32, the sampler uses the normal safe path instead of blindly slicing the workspace.

The intended benefit is narrow: account for the sampler fp16-to-fp32 logits workspace when vLLM sizes the KV cache, reducing an unprofiled allocation spike during sampler warmup/serving.

Change

• Allocate sampler_workspace as a static float32 tensor sized to maximum live requests and vocabulary size, only on the last pipeline rank of generation runners, during model memory accounting.
• In Sampler.forward(), reuse that workspace only when dtype, device, row capacity, and vocab dimension match the current logits.
• Leave float32 logits in place and fall back to logits.to(torch.float32) when the preallocated workspace cannot safely satisfy the current call.
• Keep the claim limited to this workspace allocation path.

Validation Scope

Independent GCP validation was run on 2026-05-05 against the original PR head before the final workspace-hardening follow-up:

• base: be12afd284f3f09991a7fcf506553375dc58be36
• validated PR head: 4d580373668c3441c16dde174c1af603d0bb3c40
• historical hardening head: ab70bb5249ef4f93f6958b23aba9641591ca43a6
• model: Qwen/Qwen2.5-0.5B
• L4 24 GB: GCP g2-standard-32, NVIDIA L4, driver 580.126.20, CUDA 13.0
• A100: GCP a2-highgpu-1g, CUDA 13.0
• Python 3.10.12, torch 2.11.0+cu130, transformers 4.57.6

Results on L4 24 GB:

• Base at gpu_memory_utilization=0.99: failed during CUDA graph capture with CUDA OOM.
• PR head at gpu_memory_utilization=0.99: still failed before serving with CUDA OOM while trying to allocate 2 MiB.
• PR head at gpu_memory_utilization=0.95: booted and completed repeated greedy generation probes.
• PR head at 0.95 with prompt logprobs plus chunked prefill: booted and probe passed.
• PR head at 0.95 with max_num_seqs=2 and max_num_batched_tokens=128: booted and boundary probe passed.
• PR head at 0.95 with ngram speculative decoding: booted and probe passed.

Results on A100:

• Base/head at gpu_memory_utilization=0.99 both failed early because free startup memory was below the requested 0.99 reservation.
• PR head at 0.95 passed generation, prompt-logprobs/chunked-prefill, max-seqs/batched-token boundary, and ngram speculative probes.

Current rebase and validation (2026-07-12)

Rebased onto upstream main at 27c3e579f0e5f345a86e512e26e1231d3689931f.

The current head corrects the prior over-reservation and reuse race:

• Sampler.forward() receives one next/bonus logits row per live request; speculative target rows are handled separately by RejectionSampler, so the persistent workspace is max_num_reqs × vocab_size, not max_num_reqs × (1 + num_spec_tokens) × vocab_size.
• The workspace is allocated only on the last pipeline rank for generation runners. Other ranks and pooling/draft runners do not reserve it.
• Full processed-logit outputs are cloned if they alias the reusable workspace, preventing the next default-stream sample from racing an asynchronous D2H output copy.
• Incompatible dtype, device, rank, capacity, or vocabulary width still falls back to the normal float32 conversion.

Current focused validation:

• pytest -q tests/v1/sample/test_sampler.py -k "workspace or full_processed" --confcutdir=tests/v1/sample on CPU -> 4 passed, 36 deselected
• pytest -q tests/v1/worker/test_gpu_model_runner.py -k "sampler_workspace_allocation" --confcutdir=tests/v1/worker on CPU -> 4 passed, 34 deselected
• ruff format --check, ruff check, python -m py_compile, and git diff --check on all four changed files -> passed

The earlier L4/A100 runtime results remain historical validation of the targeted accounting approach. This follow-up was not rerun on GPU and does not strengthen the original OOM claim.

What This Does Not Prove

This PR should not claim to make gpu_memory_utilization=0.99 safe on 24 GB GPUs. The L4 head run still OOMed at 0.99.

This PR also does not prove a fix for the closest sampler OOM issue I found, #33920. That issue uses GLM-4.7-Flash, tensor parallel size 2, 2x RTX 4090, max_model_len=90000, and an OOM in the top-k/top-p sort path. That original issue was not reproduced by this validation.

Safe Claim

Preallocates/account-profiles the sampler fp32 logits workspace to avoid one unprofiled allocation spike; full 24 GB/0.99 OOM behavior remains unresolved. The current head safely falls back when the workspace is not compatible. On L4, this validated boot/generation at 0.95 and prompt-logprobs, chunked-prefill, max-seqs/batched-token, and ngram speculative probes. It does not make 0.99 utilization safe on 24 GB GPUs and does not prove the GLM-4.7-Flash/top-k-sort issue is fixed.

Review Recommendation

Review this as a targeted sampler workspace accounting improvement. If the goal is to fix a 24 GB 0.99 boot OOM or #33920 specifically, that needs a separate original-condition reproduction and broader memory-sizing fix.

[gemini-code-assist]

Code Review

The pull request correctly identifies and addresses a critical Out-of-Memory (OOM) issue in the sampler by introducing a pre-allocated workspace and using copy_() for in-place casting. This is a good architectural fix for the identified problem. However, a critical sizing mismatch exists in the allocation of the sampler_workspace, which could lead to an index out-of-bounds error during runtime. The workspace is currently sized based on the maximum number of sequences (max_num_seqs), but it needs to accommodate the maximum number of batched tokens (max_num_batched_tokens) to prevent errors during prefill or large decode steps.

[gemini-code-assist]

The logits.size(0) represents the total number of tokens in the current batch (total_num_scheduled_tokens), which can be as large as self.scheduler_config.max_num_batched_tokens. However, self.sampler_workspace is allocated in gpu_model_runner.py with a size based on max_num_seqs (maximum number of sequences).

This mismatch means that if logits.size(0) exceeds max_num_seqs (which is common during prefill or large decode steps), accessing self.sampler_workspace[: logits.size(0)] will result in an IndexError due to out-of-bounds access. The sampler_workspace must be large enough to hold the maximum possible number of tokens in a batch.

[gemini-code-assist]

The sampler_workspace is allocated using max_batch_size = self.scheduler_config.max_num_seqs. As noted in the sampler.py review, the sampler's forward method will attempt to slice this workspace using logits.size(0), which can be up to self.scheduler_config.max_num_batched_tokens.

This allocation should use self.scheduler_config.max_num_batched_tokens (or self.max_num_tokens) instead of max_num_seqs to ensure the workspace is sufficiently large for all possible batch sizes, preventing an IndexError during runtime.

Suggested change

	max_batch_size = self.scheduler_config.max_num_seqs
	self.sampler.sampler_workspace = torch.empty(
	max_batch_size = self.scheduler_config.max_num_batched_tokens

[glaziermag]

@gemini-code-assist Thanks for the review. We actually don't need to expand this workspace to max_num_batched_tokens—the original bounds based on max_num_seqs (accounting for speculative draft tokens) are sufficient.

The assertion that we need max_num_batched_tokens to 'prevent errors during prefill' is incorrect due to how vLLM handles prompt logprobs. During chunked prefill, the prompt logprobs generation bypasses sampler.forward(). Instead, gpu_model_runner.py extracts the logits and manually passes them to the static method self.sampler.compute_logprobs(logits) which triggers a native PyTorch log_softmax without using the sampler_workspace.

The only logits that reach sampler.forward() (and therefore trigger logits_fp32.copy_(logits)) are the generation target tokens filtered by logits_indices. Because there is exactly 1 target token per sequence (plus speculative draft tokens), logits.size(0) will not exceed max_num_seqs * (1 + num_spec_tokens).

I have pushed a revision that reverts the memory bounds back to track the actual required target sequence limits. This keeps the copy_() safe from out-of-bounds exceptions while saving VRAM.

[gemini-code-assist]

Thank you for the detailed explanation and for clarifying the memory allocation requirements for sampler_workspace. Your explanation regarding how prompt logprobs generation bypasses sampler.forward() during chunked prefill, and that only generation target tokens hit sampler.forward(), makes perfect sense. The logic that logits.size(0) will not exceed max_num_seqs * (1 + num_spec_tokens) is well-reasoned. I appreciate you pushing the revision to strictly track the actual required target sequence limits, which indeed saves a significant block of VRAM. This is a valuable optimization.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @glaziermag.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[glaziermag]

Independent validation addendum (2026-05-05)

I reran the memory behavior validation on GCP against the exact base and PR head commits:

• base: be12afd284f3f09991a7fcf506553375dc58be36
• head: 4d580373668c3441c16dde174c1af603d0bb3c40
• model: Qwen/Qwen2.5-0.5B
• L4 run: GCP g2-standard-32, 1x NVIDIA L4 24 GB, driver 580.126.20, CUDA 13.0
• A100 run: GCP a2-highgpu-1g, 1x, CUDA 13.0

L4 results:

• Base at gpu_memory_utilization=0.99: failed during CUDA graph capture with CUDA OOM.
• PR head at gpu_memory_utilization=0.99: still failed before serving with CUDA OOM while trying to allocate 2 MiB.
• PR head at gpu_memory_utilization=0.95: booted and completed repeated greedy generation probes.
• PR head at 0.95 with prompt logprobs plus chunked prefill: booted and probe passed.
• PR head at 0.95 with max_num_seqs=2 and max_num_batched_tokens=128: booted and boundary probe passed.
• PR head at 0.95 with ngram speculative decoding: booted and probe passed.

A100 results:

• Base/head at gpu_memory_utilization=0.99 both failed early because free startup memory was below the requested 0.99 reservation.
• PR head at 0.95 passed generation, prompt-logprobs/chunked-prefill, max-seqs/batched-token boundary, and ngram speculative probes.

Classification: UNSAFE_CLAIM as originally worded, with underlying evidence TARGETED.

I amended the PR title/body to narrow the claim. I did not find a complete fix for a 24 GB 0.99 boot OOM. This PR validates a sampler workspace allocation accounting improvement at 0.95, but head still OOMs at the original L4 0.99 condition. If this PR is intended to fix #33920, that still needs the original GLM-4.7-Flash, TP=2, 2x RTX 4090, top-k/top-p sort OOM reproduction.

[glaziermag]

Follow-up pushed: ab70bb5249ef4f93f6958b23aba9641591ca43a6.

What changed:

• Sampler.forward() now reuses the profiled float32 workspace only when dtype/device/shape are compatible;
• float32 logits are left in place;
• missing, too-small, wrong-device, or wrong-width workspaces fall back to the normal logits.to(torch.float32) path instead of slicing out of bounds or dereferencing None;
• workspace sizing continues to account for speculative token rows via max_num_reqs * (1 + num_spec_tokens).

Validation of the follow-up:

• python -m pytest tests/v1/sample/test_sampler.py -q -k "sampler_workspace_reused or sampler_workspace_falls_back or sampler_workspace_leaves" -> 3 passed
• python -m py_compile vllm/v1/sample/sampler.py vllm/v1/worker/gpu_model_runner.py tests/v1/sample/test_sampler.py -> passed

Classification remains TARGETED, not ACTUAL. The original L4 24 GB gpu_memory_utilization=0.99 boot condition still OOMed on the validated PR head, so this PR should be reviewed as a sampler workspace accounting/hardening change, not as a complete 24 GB 0.99 boot-OOM fix.