diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index cab2ce96c4..c42d02cb22 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -265,6 +265,23 @@ AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata] PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict +def count_nans_per_row(logits: torch.Tensor) -> torch.Tensor: + """Per-row NaN counts, left on device.""" + return logits.isnan().sum(dim=-1, dtype=torch.int32) + + +def nans_to_dict(counts: list[int], req_id_to_index: dict[str, int]) -> dict[str, int]: + """Map per-row NaN counts onto request ids. + + Rows and requests need not correspond one-to-one, so requests without a + row are reported as 0. + """ + return { + req_id: int(counts[i]) if i < len(counts) else 0 + for req_id, i in req_id_to_index.items() + } + + # Wrapper for ModelRunnerOutput to support overlapped execution. class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): def __init__( @@ -277,6 +294,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): vocab_size: int, routed_experts: RoutedExpertsTensors | None = None, check_ep_fault: bool = False, + num_nans: torch.Tensor | None = None, ): self._model_runner_output = model_runner_output self._invalid_req_indices = invalid_req_indices @@ -291,6 +309,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): self.vocab_size = vocab_size self._logprobs_tensors = logprobs_tensors self._routed_experts = routed_experts + self._num_nans = num_nans self._has_fault: torch.Tensor | None = None # Initiate the copy on a separate stream, but do not synchronize it. @@ -310,6 +329,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): if self._routed_experts is not None else None ) + self._num_nans_cpu = ( + self._num_nans.to("cpu", non_blocking=True) + if self._num_nans is not None + else None + ) if check_ep_fault: has_fault = get_ep_all2all_manager().query_fault() self._has_fault = has_fault.to("cpu", non_blocking=True) @@ -349,6 +373,14 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput): output.routed_experts = self._routed_experts_cpu.tolists() del self._routed_experts + if self._num_nans_cpu is not None: + output.num_nans_in_logits = nans_to_dict( + self._num_nans_cpu.tolist(), output.req_id_to_index + ) + if envs.VLLM_RAISE_ON_LOGIT_NANS: + raise_if_nan_logits(output.num_nans_in_logits) + del self._num_nans + if self._has_fault is not None and self._has_fault.item(): mask = get_ep_all2all_manager().query_active_mask() raise RuntimeError( @@ -3737,6 +3769,7 @@ class GPUModelRunner( num_scheduled_tokens: int, ) -> tuple[ dict[str, int], + torch.Tensor | None, LogprobsLists | None, list[list[int]], dict[str, LogprobsTensors | None], @@ -3744,9 +3777,15 @@ class GPUModelRunner( dict[str, int], list[int], ]: - num_nans_in_logits = {} + num_nans: torch.Tensor | None = None + num_nans_in_logits: dict[str, int] = {} if envs.VLLM_COMPUTE_NANS_IN_LOGITS: - num_nans_in_logits = self._get_nans_in_logits(logits) + if self.use_async_scheduling: + # Keep the counts on device; they ride the async output copy + # stream rather than blocking here. + num_nans = None if logits is None else count_nans_per_row(logits) + else: + num_nans_in_logits = self._get_nans_in_logits(logits) num_reqs = self.input_batch.num_reqs discard_sampled_tokens_req_indices = np.nonzero( @@ -3861,6 +3900,7 @@ class GPUModelRunner( return ( num_nans_in_logits, + num_nans, logprobs_lists, valid_sampled_token_ids, prompt_logprobs_dict, @@ -4713,6 +4753,7 @@ class GPUModelRunner( with record_function_or_nullcontext("gpu_model_runner: bookkeep"): ( num_nans_in_logits, + num_nans_device, logprobs_lists, valid_sampled_token_ids, prompt_logprobs_dict, @@ -4815,6 +4856,7 @@ class GPUModelRunner( vocab_size=self.input_batch.vocab_size, routed_experts=routed_experts_snapshot, check_ep_fault=self.check_ep_fault, + num_nans=num_nans_device, ) with record_function_or_nullcontext( "gpu_model_runner: set_async_sampled_token_ids" @@ -5718,23 +5760,15 @@ class GPUModelRunner( return prompt_logprobs_dict - def _get_nans_in_logits( - self, - logits: torch.Tensor | None, - ) -> dict[str, int]: - try: - if logits is None: - return {req_id: 0 for req_id in self.input_batch.req_ids} + def _get_nans_in_logits(self, logits: torch.Tensor | None) -> dict[str, int]: + """Count NaNs per request, reading the result back to the host. - num_nans_in_logits = {} - num_nans_for_index = logits.isnan().sum(dim=-1).cpu().numpy() - for req_id in self.input_batch.req_ids: - req_index = self.input_batch.req_id_to_index[req_id] - num_nans_in_logits[req_id] = ( - int(num_nans_for_index[req_index]) - if num_nans_for_index is not None and req_index < logits.shape[0] - else 0 - ) + Only used under sync scheduling, The async path keeps the counts + on device instead; see`AsyncGPUModelRunnerOutput`. + """ + try: + counts = [] if logits is None else count_nans_per_row(logits).tolist() + num_nans_in_logits = nans_to_dict(counts, self.input_batch.req_id_to_index) if envs.VLLM_RAISE_ON_LOGIT_NANS: raise_if_nan_logits(num_nans_in_logits) return num_nans_in_logits