PyTorch Asynchronous Assert

PyTorch's torch._assert_async API enables asynchronous assertion on GPU device streams, avoiding graph breaks in torch.compile. The API accepts a boolean tensor and reports failures only when the GPU stream synchronizes with the CPU thread, as demonstrated in a model that validates softmax output sums.

PyTorch Asynchronous Assert Introduction When using PyTorch, sometimes we would like to check if certain variable conditions are met during the forward pass. If a PyTorch program runs on GPU, we would like to have this check to be performed asynchronously on the GPU device stream, without blocking the CPU thread. Otherwise, it is a graph break operation which can disturb the optimization of the computation graph from a neural network compiler, such as torch.compile . In PyTorch, such asynchronous assertion can be performed using the torch. assert async API. In this blog post, I would like to quickly discuss how to use the torch. assert async API and how it is implemented in PyTorch. PyTorch Asynchronous Assert The torch. assert async is not well documented. What’s different from the torch. assert https://github.com/pytorch/pytorch/blob/cf30153c4c131c8164ee7798e5022d810682e2cb/torch/ init .py L2270 API is that torch. assert async accepts a boolean tensor whereas torch. assert accepts a Python boolean value. When the boolean tensor is on GPU, the assertion will be performed asynchronously on the GPU device stream. Since the assertion is performed asynchronously, if the assertion fails, the error will be reported at a later time only when the GPU stream is synchronized with the CPU thread.In the following example, we inserted a torch. assert async assertion in a PyTorch model. We will test what happens when the assertion fails and how it is reported asynchronously when torch.compile is used or not used. 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215 | python import argparseimport loggingfrom pathlib import Pathfrom typing import Optionalimport torchimport torch.nn as nnlogging.basicConfig level=logging.INFO, format=" % levelname s % message s" class NormalizedLinear nn.Module : """ Applies a linear transformation followed by Softmax, then validates that output probabilities sum to 1.0 asynchronously on the device stream. """ def init self, in features: int, out features: int, trigger fail: bool = False - None: super . init self.linear = nn.Linear in features, out features self.softmax = nn.Softmax dim=-1 self.trigger fail = trigger fail def forward self, x: torch.Tensor - torch.Tensor: probs = self.softmax self.linear x if self.trigger fail: Force condition to False to demonstrate deferred stream failure cond = torch.tensor False, device=probs.device else: prob sum = probs.sum dim=-1 Pure GPU Tensor Ops abs - le - all : Returns a 0-D boolean CUDA Tensor WITHOUT calling .item or causing CPU-GPU host synchronization. cond = torch.abs prob sum - 1.0 .le 1e-5 .all Non-blocking async check enqueued into the GPU execution stream torch. assert async cond, "Assertion Failed: Probabilities must sum to 1.0 " return probs 10.0def create profiler device: torch.device, trace path: Path, stacks path: Optional Path = None, warmup steps: int = 3, active steps: int = 5, - torch.profiler.profile: """Configures and returns a PyTorch profiler using torch.profiler.schedule for warmup and active steps.""" Ensure parent directories exist trace path.parent.mkdir parents=True, exist ok=True if stacks path: stacks path.parent.mkdir parents=True, exist ok=True activities = torch.profiler.ProfilerActivity.CPU if device.type == "cuda": activities.append torch.profiler.ProfilerActivity.CUDA def on trace ready prof: torch.profiler.profile - None: prof.export chrome trace str trace path logging.info f"Saved Chrome trace to: {trace path.resolve }" if stacks path: metric = "self cuda time total" if device.type == "cuda" else "self cpu time total" prof.export stacks str stacks path , metric=metric logging.info f"Saved Flamegraph stacks to: {stacks path.resolve }" return torch.profiler.profile activities=activities, schedule=torch.profiler.schedule wait=0, warmup=warmup steps, active=active steps, repeat=1 , record shapes=True, profile memory=True, with stack=stacks path is not None, on trace ready=on trace ready, def run benchmark args: argparse.Namespace - None: device = torch.device args.device logging.info "=" 60 logging.info f"Device : {device.type.upper }" logging.info f"Compile : {args.compile}" logging.info f"Profile : {args.profile}" logging.info f"Trigger Fail : {args.trigger fail}" if args.profile: logging.info f"Profiler Config: {args.warmup} Warmup Steps | {args.iters} Active Steps" logging.info "=" 60 Initialize model model = NormalizedLinear args.in features, args.out features, trigger fail=args.trigger fail .to device x = torch.randn args.batch size, args.in features, device=device if args.compile: logging.info "Compiling model via torch.compile ..." model = torch.compile model try: if args.profile: total steps = args.warmup + args.iters logging.info f"Running Profiler {args.warmup} warmup + {args.iters} active = {total steps} total steps ..." with create profiler device=device, trace path=args.trace path, stacks path=args.stacks path, warmup steps=args.warmup, active steps=args.iters as prof: for step in range total steps : = model x prof.step Advances schedule: warmup - active - trace ready if device.type == "cuda": torch.cuda.synchronize else: Standalone Warmup logging.info f"Warming up {args.warmup} iterations ..." for in range args.warmup : = model x if device.type == "cuda": torch.cuda.synchronize Standalone Benchmark Loop logging.info f"Executing {args.iters} benchmark iterations..." for in range args.iters : = model x if device.type == "cuda": torch.cuda.synchronize logging.info "Run finished successfully " except RuntimeError as err: logging.error "Caught Exception from GPU Stream Sync Expected if --trigger-fail was set :" logging.error f"-- {err}" def parse args - argparse.Namespace: parser = argparse.ArgumentParser description= "PyTorch Async Assertion & Compiler Benchmark with Profiler Schedule", formatter class=argparse.ArgumentDefaultsHelpFormatter, Feature Flags parser.add argument "--compile", action="store true", help="Compile model graph with torch.compile" parser.add argument "--profile", action="store true", help="Enable torch.profiler tracing" parser.add argument "--trigger-fail", action="store true", help="Force condition failure to test deferred stream errors" File Paths parser.add argument "--trace-path", type=Path, default=Path "my trace.json" , help="Output JSON path for Chrome/Perfetto trace" parser.add argument "--stacks-path", type=Path, default=None, help="Optional output text path for Flamegraph stack trace" Model Hyperparameters & Profiler Steps parser.add argument "--device", type=str, default="cuda" if torch.cuda.is available else "cpu", help="Target device" parser.add argument "--batch-size", type=int, default=128, help="Batch size" parser.add argument "--in-features", type=int, default=1024, help="Input feature size" parser.add argument "--out-features", type=int, default=2048, help="Output feature size" parser.add argument "--warmup", type=int, default=3, help="Warmup iterations before active profiling" parser.add argument "--iters", type=int, default=5, help="Number of active benchmark iterations to capture in profile" return parser.parse args if name == " main ": run benchmark parse args | To profile the PyTorch program with and without using torch.compile , we could run the following commands, which will generate Perfetto profiling traces for both cases. 12 | bash $ python assert async.py --profile --trace-path assert async trace.json$ python assert async.py --compile --profile --trace-path assert async compiled trace.json | In the profiling trace https://ui.perfetto.dev/ /?url=https://raw.githubusercontent.com/leimao/PyTorch-Asynchronous-Assert/refs/heads/main/assert async trace.json of the run that does not use torch.compile , we could see that the assertion is performed asynchronously on the GPU stream. In the profiling trace https://ui.perfetto.dev/ /?url=https://raw.githubusercontent.com/leimao/PyTorch-Asynchronous-Assert/refs/heads/main/assert async compiled trace.json of the run that uses torch.compile , we could see that the assertion is fused into the compiled graph with other operations, which is friendly to the optimization of the computation graph. If an assertion fails, the error will be reported asynchronously when the GPU stream is synchronized with the CPU thread. For example, if we run the following command to trigger an assertion failure. 123456789101112131415 | bash $ python assert async.py --trigger-fail INFO ============================================================ INFO Device : CUDA INFO Compile : False INFO Profile : False INFO Trigger Fail : True INFO ============================================================ INFO Warming up 3 iterations .../opt/pytorch/pytorch/aten/src/ATen/native/cuda/TensorCompare.cu:109: assert async cuda kernel: block: 0,0,0 , thread: 0,0,0 Assertion Assertion Failed: Probabilities must sum to 1.0 failed. ERROR Caught Exception from GPU Stream Sync Expected if --trigger-fail was set : ERROR -- CUDA error: device-side assert triggeredSearch for cudaErrorAssert' in https://docs.nvidia.com/cuda/cuda-runtime-api/group CUDART TYPES.html for more information.CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.For debugging consider passing CUDA LAUNCH BLOCKING=1Compile with TORCH USE CUDA DSA to enable device-side assertions. | On caveat is that an assertion failure will cause CUDA context to be poisoned, which will cause all subsequent CUDA calls to fail. Without restarting the CUDA context, the PyTorch program will not be able to continue on GPU. That is why it is an assertion rather than an exception which can be caught and handled. PyTorch Asynchronous Assert Implementation Because the purpose of assertion failure is to terminate the program, the implementation of torch. assert async is designed to poison the CUDA context. In the CUDA operation, the trap instruction, which translates to asm volatile "trap;" I believe, is used to terminate the program when an assertion fails. In the Triton compilation, the tl.device assert https://triton-lang.org/main/python-api/generated/triton.language.device assert.html instruction is used to terminate the program when an assertion fails. Conclusions The torch. assert async calls are not completely free and it can poison the CUDA context when an assertion fails. Therefore, ideally it should be used in development, and should not be used in production because assertion should be expected to always pass. In C++, similarly, the assert will be optimized away by the compiler in release or production builds when the macro NDEBUG is defined. PyTorch Asynchronous Assert