mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-26 16:35:05 +00:00
27 lines
953 B
Python
27 lines
953 B
Python
"""CUDA outlier detection and PyTorch reference."""
|
|
|
|
import torch
|
|
|
|
from deepinfer.model_executor.layers.tilelang_ops import ops
|
|
|
|
MAX_ABS = 1e11
|
|
|
|
|
|
def detect_outlier_and_abort(data: torch.Tensor, id: int = 0) -> None:
|
|
"""Queue a CUDA assertion for NaN, infinity, or values outside (-MAX_ABS, MAX_ABS).
|
|
|
|
The check runs asynchronously. A failed assertion surfaces on a subsequent
|
|
CUDA operation or synchronization and leaves the CUDA context unusable
|
|
until the process exits. Failures print the id, fp32 value, and element
|
|
offset in the flattened tensor.
|
|
|
|
Args:
|
|
data: Any shape, fp16/bf16/fp32 CUDA tensor. Noncontiguous inputs are
|
|
copied before checking; comparisons use fp32.
|
|
id: Runtime int64 identifier included in failure diagnostics.
|
|
"""
|
|
if data.numel() == 0:
|
|
return
|
|
flat_data = data.contiguous().view(-1)
|
|
ops.detect_outlier_and_abort_tl(data.dtype, MAX_ABS)(flat_data, id)
|