b1e12d142d · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -265,6 +265,23 @@ AttnMetadataDict: TypeAlias = dict[str, AttentionMetadata]

265 PerLayerAttnMetadata: TypeAlias = list[AttnMetadataDict] | AttnMetadataDict

266

267

268 +def count_nans_per_row(logits: torch.Tensor) -> torch.Tensor:

269 + """Per-row NaN counts, left on device."""

270 + return logits.isnan().sum(dim=-1, dtype=torch.int32)

271 +

272 +

273 +def nans_to_dict(counts: list[int], req_id_to_index: dict[str, int]) -> dict[str, int]:

274 + """Map per-row NaN counts onto request ids.

275 +

276 + Rows and requests need not correspond one-to-one, so requests without a

277 + row are reported as 0.

278 + """

279 + return {

280 + req_id: int(counts[i]) if i < len(counts) else 0

281 + for req_id, i in req_id_to_index.items()

282 + }

283 +

284 +

285 # Wrapper for ModelRunnerOutput to support overlapped execution.

286 class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

287 def __init__(

@@ -277,6 +294,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

294 vocab_size: int,

295 routed_experts: RoutedExpertsTensors | None = None,

296 check_ep_fault: bool = False,

297 + num_nans: torch.Tensor | None = None,

298 ):

299 self._model_runner_output = model_runner_output

300 self._invalid_req_indices = invalid_req_indices

@@ -291,6 +309,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

309 self.vocab_size = vocab_size

310 self._logprobs_tensors = logprobs_tensors

311 self._routed_experts = routed_experts

312 + self._num_nans = num_nans

313 self._has_fault: torch.Tensor | None = None

314

315 # Initiate the copy on a separate stream, but do not synchronize it.

@@ -310,6 +329,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

329 if self._routed_experts is not None

330 else None

331 )

332 + self._num_nans_cpu = (

333 + self._num_nans.to("cpu", non_blocking=True)

334 + if self._num_nans is not None

335 + else None

336 + )

337 if check_ep_fault:

338 has_fault = get_ep_all2all_manager().query_fault()

339 self._has_fault = has_fault.to("cpu", non_blocking=True)

@@ -349,6 +373,14 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

373 output.routed_experts = self._routed_experts_cpu.tolists()

374 del self._routed_experts

375

376 + if self._num_nans_cpu is not None:

377 + output.num_nans_in_logits = nans_to_dict(

378 + self._num_nans_cpu.tolist(), output.req_id_to_index

379 + )

380 + if envs.VLLM_RAISE_ON_LOGIT_NANS:

381 + raise_if_nan_logits(output.num_nans_in_logits)

382 + del self._num_nans

383 +

384 if self._has_fault is not None and self._has_fault.item():

385 mask = get_ep_all2all_manager().query_active_mask()

386 raise RuntimeError(

@@ -3737,6 +3769,7 @@ class GPUModelRunner(

3769 num_scheduled_tokens: int,

3770 ) -> tuple[

3771 dict[str, int],

3772 + torch.Tensor | None,

3773 LogprobsLists | None,

3774 list[list[int]],

3775 dict[str, LogprobsTensors | None],

@@ -3744,9 +3777,15 @@ class GPUModelRunner(

3777 dict[str, int],

3778 list[int],

3779 ]:

3747 - num_nans_in_logits = {}

3780 + num_nans: torch.Tensor | None = None

3781 + num_nans_in_logits: dict[str, int] = {}

3782 if envs.VLLM_COMPUTE_NANS_IN_LOGITS:

3749 - num_nans_in_logits = self._get_nans_in_logits(logits)

3783 + if self.use_async_scheduling:

3784 + # Keep the counts on device; they ride the async output copy

3785 + # stream rather than blocking here.

3786 + num_nans = None if logits is None else count_nans_per_row(logits)

3787 + else:

3788 + num_nans_in_logits = self._get_nans_in_logits(logits)

3789

3790 num_reqs = self.input_batch.num_reqs

3791 discard_sampled_tokens_req_indices = np.nonzero(

@@ -3861,6 +3900,7 @@ class GPUModelRunner(

3900

3901 return (

3902 num_nans_in_logits,

3903 + num_nans,

3904 logprobs_lists,

3905 valid_sampled_token_ids,

3906 prompt_logprobs_dict,

@@ -4713,6 +4753,7 @@ class GPUModelRunner(

4753 with record_function_or_nullcontext("gpu_model_runner: bookkeep"):

4754 (

4755 num_nans_in_logits,

4756 + num_nans_device,

4757 logprobs_lists,

4758 valid_sampled_token_ids,

4759 prompt_logprobs_dict,

@@ -4815,6 +4856,7 @@ class GPUModelRunner(

4856 vocab_size=self.input_batch.vocab_size,

4857 routed_experts=routed_experts_snapshot,

4858 check_ep_fault=self.check_ep_fault,

4859 + num_nans=num_nans_device,

4860 )

4861 with record_function_or_nullcontext(

4862 "gpu_model_runner: set_async_sampled_token_ids"

@@ -5718,23 +5760,15 @@ class GPUModelRunner(

5760

5761 return prompt_logprobs_dict

5762

5721 - def _get_nans_in_logits(

5722 - self,

5723 - logits: torch.Tensor | None,

5724 - ) -> dict[str, int]:

5725 - try:

5726 - if logits is None:

5727 - return {req_id: 0 for req_id in self.input_batch.req_ids}

5763 + def _get_nans_in_logits(self, logits: torch.Tensor | None) -> dict[str, int]:

5764 + """Count NaNs per request, reading the result back to the host.

5765

5729 - num_nans_in_logits = {}

5730 - num_nans_for_index = logits.isnan().sum(dim=-1).cpu().numpy()

5731 - for req_id in self.input_batch.req_ids:

5732 - req_index = self.input_batch.req_id_to_index[req_id]

5733 - num_nans_in_logits[req_id] = (

5734 - int(num_nans_for_index[req_index])

5735 - if num_nans_for_index is not None and req_index < logits.shape[0]

5736 - else 0

5737 - )

5766 + Only used under sync scheduling, The async path keeps the counts

5767 + on device instead; see`AsyncGPUModelRunnerOutput`.

5768 + """

5769 + try:

5770 + counts = [] if logits is None else count_nans_per_row(logits).tolist()

5771 + num_nans_in_logits = nans_to_dict(counts, self.input_batch.req_id_to_index)

5772 if envs.VLLM_RAISE_ON_LOGIT_NANS:

5773 raise_if_nan_logits(num_nans_in_logits)

5774 return num_nans_in_logits

[V1] Copy NaN-in-logits counts to host asynchronously (#51304)