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 11, 18:45 UTC
[Model Runner V2] Support num NaNs in logits (#30187)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Dec 9, 2025 at 10:00 UTC
d471b2aff09028f9c62e861f760a74fd8f99081d
7 files changed
+89
-26
vllm/v1/worker/gpu/async_utils.py
+21
-20
@@ -2,14 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from contextlib import contextmanager
4
5
+import numpy as np
6
import torch
7
8
from vllm.v1.outputs import (
9
AsyncModelRunnerOutput,
10
LogprobsTensors,
11
ModelRunnerOutput,
11
- SamplerOutput,
12
)
13
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
14
15
16
class AsyncOutput(AsyncModelRunnerOutput):
@@ -34,29 +35,18 @@ class AsyncOutput(AsyncModelRunnerOutput):
35
with torch.cuda.stream(self.copy_stream):
36
self.copy_stream.wait_stream(default_stream)
37
37
- # NOTE(woosuk): We must ensure that CPU tensors are not freed
38
- # before the device-to-host copy is fully completed. For instance,
39
- # operations like
40
- # self.sampled_token_np = ...to("cpu", non_blocking=True).numpy()
41
- # are unsafe because the underlying CPU tensor can be prematurely freed and
42
- # reused by other tensors before the asynchronous copy finishes, potentially
43
- # causing race conditions. To prevent this, we delay freeing by holding
44
- # references until the copy event signals completion.
45
- # Likewise, we also need to keep the reference to the GPU tensors.
46
- # This is done by keeping the reference to sampler_output and
47
- # model_runner_output.
48
- self.sampled_token_ids = sampler_output.sampled_token_ids.to(
49
- "cpu", non_blocking=True
50
- )
38
+ self.sampled_token_ids = async_copy_to_np(sampler_output.sampled_token_ids)
39
if sampler_output.logprobs_tensors is not None:
40
self.logprobs_tensors: LogprobsTensors | None = (
41
sampler_output.logprobs_tensors.to_cpu_nonblocking()
42
)
43
else:
44
self.logprobs_tensors = None
57
- self.num_sampled_tokens_cpu = num_sampled_tokens.to(
58
- "cpu", non_blocking=True
59
- )
45
+ if sampler_output.num_nans is not None:
46
+ self.num_nans = async_copy_to_np(sampler_output.num_nans)
47
+ else:
48
+ self.num_nans = None
49
+ self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
50
self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
51
if self.model_runner_output.prompt_logprobs_dict:
52
for k, v in self.model_runner_output.prompt_logprobs_dict.items():
@@ -68,7 +58,6 @@ class AsyncOutput(AsyncModelRunnerOutput):
58
59
def get_output(self) -> ModelRunnerOutput:
60
self.copy_event.synchronize()
71
- num_sampled_tokens_np = self.num_sampled_tokens_cpu.numpy()
61
62
# NOTE(woosuk): The following code is to ensure compatibility with
63
# the existing model runner.
@@ -76,10 +65,18 @@ class AsyncOutput(AsyncModelRunnerOutput):
65
# rather than Python lists.
66
sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
67
num_reqs = len(sampled_token_ids)
68
+ num_sampled_tokens = self.num_sampled_tokens_np.tolist()
69
for i in range(num_reqs):
80
- del sampled_token_ids[i][num_sampled_tokens_np[i] :]
70
+ del sampled_token_ids[i][num_sampled_tokens[i] :]
71
self.model_runner_output.sampled_token_ids = sampled_token_ids
72
73
+ if self.num_nans is not None:
74
+ num_nans = self.num_nans.tolist()
75
+ self.model_runner_output.num_nans_in_logits = {
76
+ req_id: num_nans[i]
77
+ for i, req_id in enumerate(self.model_runner_output.req_ids)
78
+ }
79
+
80
if self.logprobs_tensors is not None:
81
self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
82
self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
@@ -95,3 +92,7 @@ def async_barrier(event: torch.cuda.Event | None):
92
finally:
93
if event is not None:
94
event.record()
95
+
96
+
97
+def async_copy_to_np(x: torch.Tensor) -> np.ndarray:
98
+ return x.to("cpu", non_blocking=True).numpy()
vllm/v1/worker/gpu/metrics/__init__.py
vllm/v1/worker/gpu/metrics/logits.py
new
+42
@@ -0,0 +1,42 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+from torch._inductor.runtime.triton_helpers import libdevice
5
+
6
+from vllm.triton_utils import tl, triton
7
+
8
+
9
+@triton.jit
10
+def _num_nans_kernel(
11
+ logits_ptr,
12
+ logits_stride,
13
+ num_nans_ptr,
14
+ vocab_size,
15
+ BLOCK_SIZE: tl.constexpr,
16
+):
17
+ req_idx = tl.program_id(0)
18
+ num_nans = 0
19
+ for i in range(0, vocab_size, BLOCK_SIZE):
20
+ block = i + tl.arange(0, BLOCK_SIZE)
21
+ mask = block < vocab_size
22
+ logits = tl.load(
23
+ logits_ptr + req_idx * logits_stride + block, mask=mask, other=0
24
+ )
25
+ logits = logits.to(tl.float32)
26
+ is_nan = libdevice.isnan(logits).to(tl.int1)
27
+ num_nans += tl.sum(is_nan).to(tl.int32)
28
+ tl.store(num_nans_ptr + req_idx, num_nans)
29
+
30
+
31
+def get_num_nans(logits: torch.Tensor) -> torch.Tensor:
32
+ num_reqs, vocab_size = logits.shape
33
+ BLOCK_SIZE = 8192
34
+ num_nans = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
35
+ _num_nans_kernel[(num_reqs,)](
36
+ logits,
37
+ logits.stride(0),
38
+ num_nans,
39
+ vocab_size,
40
+ BLOCK_SIZE=BLOCK_SIZE,
41
+ )
42
+ return num_nans
vllm/v1/worker/gpu/model_runner.py
+1
-1
@@ -25,7 +25,6 @@ from vllm.v1.outputs import (
25
LogprobsTensors,
26
ModelRunnerOutput,
27
)
28
-from vllm.v1.sample.sampler import SamplerOutput
28
from vllm.v1.worker.gpu.async_utils import AsyncOutput, async_barrier
29
from vllm.v1.worker.gpu.attn_utils import (
30
build_attn_metadata,
@@ -53,6 +52,7 @@ from vllm.v1.worker.gpu.sample.metadata import (
52
SamplingMetadata,
53
expand_sampling_metadata,
54
)
55
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
56
from vllm.v1.worker.gpu.sample.sampler import Sampler
57
from vllm.v1.worker.gpu.spec_decode import init_speculator
58
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
vllm/v1/worker/gpu/sample/min_p.py
+1
-3
@@ -39,9 +39,7 @@ def _min_p_kernel(
39
tl.store(logits_ptr + req_idx * logits_stride + block, logits, mask=mask)
40
41
42
-def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor | None) -> None:
43
- if min_p is None:
44
- return
42
+def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor) -> None:
43
num_reqs, vocab_size = logits.shape
44
BLOCK_SIZE = 1024
45
_min_p_kernel[(num_reqs,)](
vllm/v1/worker/gpu/sample/output.py
new
+14
@@ -0,0 +1,14 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from dataclasses import dataclass
4
+
5
+import torch
6
+
7
+from vllm.v1.outputs import LogprobsTensors
8
+
9
+
10
+@dataclass
11
+class SamplerOutput:
12
+ sampled_token_ids: torch.Tensor
13
+ logprobs_tensors: LogprobsTensors | None
14
+ num_nans: torch.Tensor | None
vllm/v1/worker/gpu/sample/sampler.py
+10
-2
@@ -3,13 +3,15 @@
3
4
import torch
5
6
+import vllm.envs as envs
7
from vllm.config.model import LogprobsMode
7
-from vllm.v1.outputs import SamplerOutput
8
from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
9
+from vllm.v1.worker.gpu.metrics.logits import get_num_nans
10
from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
11
from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
12
from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
13
from vllm.v1.worker.gpu.sample.min_p import apply_min_p
14
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
15
from vllm.v1.worker.gpu.sample.penalties import apply_penalties_and_temperature
16
17
@@ -21,12 +23,16 @@ class Sampler:
23
if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
24
raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")
25
self.logprobs_mode = logprobs_mode
26
+ self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS # False by default.
27
28
def __call__(
29
self,
30
logits: torch.Tensor,
31
sampling_metadata: SamplingMetadata,
32
) -> SamplerOutput:
33
+ # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear
34
+ # that num_nans is computed before applying penalties and temperature.
35
+ num_nans = get_num_nans(logits) if self.compute_nans else None
36
sampled, processed_logits = self.sample(logits, sampling_metadata)
37
if sampling_metadata.max_num_logprobs is not None:
38
logits = (
@@ -49,6 +55,7 @@ class Sampler:
55
# token per request.
56
sampled_token_ids=sampled.view(-1, 1),
57
logprobs_tensors=logprobs_tensors,
58
+ num_nans=num_nans,
59
)
60
return sampler_output
61
@@ -63,7 +70,8 @@ class Sampler:
70
# Apply penalties and temperature in place.
71
apply_penalties_and_temperature(logits, sampling_metadata)
72
# Apply min_p in place.
66
- apply_min_p(logits, sampling_metadata.min_p)
73
+ if sampling_metadata.min_p is not None:
74
+ apply_min_p(logits, sampling_metadata.min_p)
75
# Apply top_k and/or top_p. This might return a new tensor.
76
logits = apply_top_k_top_p(
77
logits, sampling_metadata.top_k, sampling_metadata.top_p