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 4, 10:00 UTC
[Perf] Warmup forward_native sampler kernel (#41375)
Signed-off-by: Artem Perevedentsev <aperevedents@nvidia.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Artem Perevedentsev committed
May 1, 2026 at 17:04 UTC
c3e64696cdea5df92eb15e20b18ba979a536c1e3
1 file changed
+20
vllm/v1/worker/gpu_model_runner.py
+20
@@ -5701,6 +5701,26 @@ class GPUModelRunner(
5701
sampler_output = self.sampler(
5702
logits=logits, sampling_metadata=dummy_metadata
5703
)
5704
+ # Also warm forward_native (taken when generators dict is non-empty),
5705
+ # but skip the extra call in 'processed_logits' / 'processed_logprobs'
5706
+ # modes — there TopKTopPSampler binds forward = forward_native at
5707
+ # init time, so the warmup call is redundant and only inflates peak
5708
+ # memory during profile_run.
5709
+ # No .clone() of logits: warmup output is discarded, so any in-place
5710
+ # mutation by forward_native does not affect correctness.
5711
+ if self.sampler.logprobs_mode not in (
5712
+ "processed_logits",
5713
+ "processed_logprobs",
5714
+ ):
5715
+ self.sampler(
5716
+ logits=logits,
5717
+ sampling_metadata=replace(
5718
+ dummy_metadata,
5719
+ generators={
5720
+ 0: torch.Generator(device=self.device).manual_seed(0)
5721
+ },
5722
+ ),
5723
+ )
5724
except RuntimeError as e:
5725
if "out of memory" in str(e):
5726
raise RuntimeError(