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 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.

assert_async.py

import argparse
import logging
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn

logging.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.0


def 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.

$ 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 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 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.

$ 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 triggered
Search 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=1
Compile 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 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.