[Core] Take the top-k/top-p Triton scratch from WorkspaceManager

[aoshen02]

Overview

Take the per-program float32 scratch of apply_top_k_top_p_triton from WorkspaceManager, reserved at its maximum size by the GPU model runners before capture and the workspace lock. Part of #59055; with #59156 the memory is then released by /sleep.

Claims

• The sampler scratch (min(num_sm, batch) × vocab × 4 B; 77 MiB per GPU on DeepSeek-V4-Flash, 92 MiB on GLM-5.3) no longer stays resident while the engine sleeps once [Feature] Release WorkspaceManager scratch on sleep #59156 lands, and no longer grows lazily under traffic after warm-up.
• Kernel behavior is unchanged: the same (NUM_PROGRAMS, vocab) float32 buffer is handed to _topk_topp.
• Callers without a reservation (CPU, tests, diffusion_gemma, the watermark sampler, batches larger than the reservation) keep the module-level cached buffer, so nothing can hit the workspace lock.

Validation

• Unit tests (H200): tests/v1/sample/test_topk_topp_sampler.py, test_sampler.py, test_rejection_sampler.py: 236 passed, 1 skipped, 7 failed; the 7 failures (TestFlashInferDistributionMatch, No module named 'scipy') are identical on unmodified main in that environment.
• pre-commit (ruff, ruff format, mypy, etc.) passes on the commit.
• End-to-end, DeepSeek-V4-Flash TP4 with MTP, CUDA graphs and --enable-sleep-mode on 4× GB300, same image and node, on top of [Feature] Release WorkspaceManager scratch on sleep #59156 + [Feature] Offload KV-init runtime state on sleep #59158 + [Core] Release the FlashInfer allreduce workspace (MNNVL or TRT-LLM) with --enable-nccl-comm-suspend #59360:
  • Memory per GPU after /sleep (nvidia-smi): 6313/6439 MiB without this PR (an earlier 100-cycle run of the same stack on the same node, not a back-to-back A/B), 6023 MiB on all 4 GPUs with it. Only about 75 MiB of that is attributable to this change: torch-allocated memory outside the KV cache drops by 75 MiB (the lazily grown sampler buffer); the rest is caching-allocator reserved space (-162 MiB) and the +126 MiB the baseline shows on two GPUs from run to run.
  • Post-wake parity 0/0/0/0 MiB; 3 sleep/wake cycles (full, selective, full) with 195/195 outputs identical to the pre-sleep reference; no Workspace is locked errors.

Details

• reserve_topk_topp_workspace(vllm_config, device) reserves min(num_sm, max_num_seqs × (num_speculative_tokens + 1)) × vocab_size float32 elements, the bound the kernel's own warm-up already uses (_max_sampler_batch_size), and records it per device. Both GPU model runners call it right after building their sampler, i.e. before kernel warm-up, capture and lock_workspace().
• apply_top_k_top_p_triton requests ((NUM_PROGRAMS, vocab), float32) from the manager when that fits the recorded reservation, else falls back to the existing cache. The reservation is a floor on the shared scratch slot, which the MoE workspaces already exceed on the models above.
• The sampler runs on the main thread in lane 0 after the forward; by the WorkspaceManager contract no view from the forward may be live across the next get_simultaneous, and none of the inputs (logits, k, p) is a workspace view.
• The split-row top-p cache (_TRITON_SPLIT_CACHE, ~1 MiB) is left alone.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

AI assistance: written with Claude Code (Claude Fable 5.1), reviewed and tested by the human submitter.

Re-check on current main (bdd31c31d), GB300 (supersedes the earlier end-to-end figure)

Re-checked on main bdd31c3 (applies cleanly), 4x GB300, DeepSeek-V4-Flash TP4 + MTP,
--enable-sleep-mode, default V2 model runner; each launch sends a 160-request seeded
top-k/top-p burst before the first sleep (seeded requests take the Triton path):

• Unit tests: 243 passed, 1 skipped, 3 runs (same as main).
• Memory per GPU: -290 MiB both awake and asleep vs the main launch in the same
  allocator mode (11031 vs 11321 and 10579 vs 10869 MiB asleep); KV cache unchanged
  (201.51 GiB); no "Workspace is locked" errors.
• 3 sleep/wake cycles: 28/28 greedy and 28/28 seeded top-k/top-p outputs identical to
  the awake reference after every wake; wake 0.97-1.37 s (main 1.06-1.31 s).

🤖 Generated with Claude Code