[Model Runner V2] Support num NaNs in logits

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Author: WoosukKwon

vllm/v1/worker/gpu/async_utils.py

@@ -8,8 +8,8 @@
     AsyncModelRunnerOutput,
     LogprobsTensors,
     ModelRunnerOutput,
-    SamplerOutput,
 )
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
 
 
 class AsyncOutput(AsyncModelRunnerOutput):
@@ -54,6 +54,10 @@ def __init__(
                 )
             else:
                 self.logprobs_tensors = None
+            if sampler_output.num_nans is not None:
+                self.num_nans = sampler_output.num_nans.to("cpu", non_blocking=True)
+            else:
+                self.num_nans = None
             self.num_sampled_tokens_cpu = num_sampled_tokens.to(
                 "cpu", non_blocking=True
             )
@@ -80,6 +84,13 @@ def get_output(self) -> ModelRunnerOutput:
             del sampled_token_ids[i][num_sampled_tokens_np[i] :]
         self.model_runner_output.sampled_token_ids = sampled_token_ids
 
+        if self.num_nans is not None:
+            num_nans_np = self.num_nans.numpy()
+            self.model_runner_output.num_nans_in_logits = {
+                req_id: int(num_nans_np[i])
+                for i, req_id in enumerate(self.model_runner_output.req_ids)
+            }
+
         if self.logprobs_tensors is not None:
             self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
         self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict

vllm/v1/worker/gpu/metrics/__init__.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+from torch._inductor.runtime.triton_helpers import libdevice
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _num_nans_kernel(
+    logits_ptr,
+    logits_stride,
+    num_nans_ptr,
+    vocab_size,
+    BLOCK_SIZE: tl.constexpr,
+):
+    req_idx = tl.program_id(0)
+    num_nans = 0
+    for i in range(0, vocab_size, BLOCK_SIZE):
+        block = i + tl.arange(0, BLOCK_SIZE)
+        mask = block < vocab_size
+        logits = tl.load(
+            logits_ptr + req_idx * logits_stride + block, mask=mask, other=0
+        )
+        logits = logits.to(tl.float32)
+        is_nan = libdevice.isnan(logits).to(tl.int1)
+        num_nans += tl.sum(is_nan).to(tl.int32)
+    tl.store(num_nans_ptr + req_idx, num_nans)
+
+
+def get_num_nans(logits: torch.Tensor) -> torch.Tensor:
+    num_reqs, vocab_size = logits.shape
+    BLOCK_SIZE = 8192
+    num_nans = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
+    _num_nans_kernel[(num_reqs,)](
+        logits,
+        logits.stride(0),
+        num_nans,
+        vocab_size,
+        BLOCK_SIZE=BLOCK_SIZE,
+    )
+    return num_nans

vllm/v1/worker/gpu/metrics/logits.py

@@ -25,7 +25,6 @@
     LogprobsTensors,
     ModelRunnerOutput,
 )
-from vllm.v1.sample.sampler import SamplerOutput
 from vllm.v1.worker.gpu.async_utils import AsyncOutput, async_barrier
 from vllm.v1.worker.gpu.attn_utils import (
     build_attn_metadata,
@@ -53,6 +52,7 @@
     SamplingMetadata,
     expand_sampling_metadata,
 )
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 from vllm.v1.worker.gpu.spec_decode import init_speculator
 from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample

vllm/v1/worker/gpu/model_runner.py

@@ -39,9 +39,7 @@ def _min_p_kernel(
         tl.store(logits_ptr + req_idx * logits_stride + block, logits, mask=mask)
 
 
-def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor | None) -> None:
-    if min_p is None:
-        return
+def apply_min_p(logits: torch.Tensor, min_p: torch.Tensor) -> None:
     num_reqs, vocab_size = logits.shape
     BLOCK_SIZE = 1024
     _min_p_kernel[(num_reqs,)](

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

@@ -0,0 +1,14 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import torch
+
+from vllm.v1.outputs import LogprobsTensors
+
+
+@dataclass
+class SamplerOutput:
+    sampled_token_ids: torch.Tensor
+    logprobs_tensors: LogprobsTensors | None
+    num_nans: torch.Tensor | None

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

@@ -3,13 +3,15 @@
 
 import torch
 
+import vllm.envs as envs
 from vllm.config.model import LogprobsMode
-from vllm.v1.outputs import SamplerOutput
 from vllm.v1.sample.ops.topk_topp_sampler import apply_top_k_top_p
+from vllm.v1.worker.gpu.metrics.logits import get_num_nans
 from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample
 from vllm.v1.worker.gpu.sample.logprob import compute_topk_logprobs
 from vllm.v1.worker.gpu.sample.metadata import SamplingMetadata
 from vllm.v1.worker.gpu.sample.min_p import apply_min_p
+from vllm.v1.worker.gpu.sample.output import SamplerOutput
 from vllm.v1.worker.gpu.sample.penalties import apply_penalties_and_temperature
 
 
@@ -21,6 +23,7 @@ def __init__(
         if logprobs_mode not in ["processed_logprobs", "raw_logprobs"]:
             raise NotImplementedError(f"Unsupported logprobs_mode: {logprobs_mode}")
         self.logprobs_mode = logprobs_mode
+        self.compute_nans = envs.VLLM_COMPUTE_NANS_IN_LOGITS  # False by default.
 
     def __call__(
         self,
@@ -42,13 +45,15 @@ def __call__(
         else:
             logprobs_tensors = None
 
+        num_nans = get_num_nans(processed_logits) if self.compute_nans else None
         # These are GPU tensors.
         sampler_output = SamplerOutput(
             # The sampled tokens are expanded to 2D tensor with shape
             # [num_requests, 1], where each row represents one generated
             # token per request.
             sampled_token_ids=sampled.view(-1, 1),
             logprobs_tensors=logprobs_tensors,
+            num_nans=num_nans,
         )
         return sampler_output
 
@@ -63,7 +68,8 @@ def sample(
         # Apply penalties and temperature in place.
         apply_penalties_and_temperature(logits, sampling_metadata)
         # Apply min_p in place.
-        apply_min_p(logits, sampling_metadata.min_p)
+        if sampling_metadata.min_p is not None:
+            apply_min_p(logits, sampling_metadata.min_p)
         # Apply top_k and/or top_p. This might return a new tensor.
         logits = apply_top_k_top_p(
             logits, sampling_metadata.top_k, sampling_metadata.top_p