PyTorch异步断言
内容提要
PyTorch的`torch._assert_async` API可在GPU上异步执行断言，不阻塞CPU线程，避免图断裂。它接受布尔张量，失败时在流同步后报错。实现通过CUDA的`__trap`指令终止程序并污染上下文。建议仅用于开发，生产环境应移除，类似C++的`assert`在发布版中被优化掉。
延伸解读
异步断言与图编译的兼容性
torch._assert_async 的关键优势在于它不会导致图断裂，从而不会干扰 torch.compile 等编译器对计算图的优化。在未使用 torch.compile 时，断言在 GPU 流上异步执行；而在使用 torch.compile 时，断言会被融合进编译后的图中，这有助于保持图优化的完整性。因此，在需要调试且希望保持编译性能的场景下，该 API 是一个合适的选择。
失败后果：CUDA 上下文污染
需要注意的是，一旦断言失败，CUDA 上下文会被污染，导致后续所有 CUDA 调用失败，且无法通过常规异常处理恢复。这意味着断言失败等同于程序终止，而非可捕获的异常。因此，该 API 更适合用于开发阶段的调试，而不适合在生产环境中作为错误处理机制。
开发与生产环境的定位
文章建议将 torch._assert_async 仅用于开发阶段，因为断言在正常情况下应当始终通过，且其执行并非完全无开销。这与 C++ 中 assert 在发布版（定义 NDEBUG）时被优化掉的做法类似。在生产环境中，应移除或禁用此类断言，以避免不必要的性能损耗和潜在的上下文污染风险。
Q&A
PyTorch中torch._assert_async和torch._assert有什么区别？
torch._assert_async接受布尔张量，而torch._assert接受Python布尔值。当布尔张量在GPU上时，断言会异步执行在GPU设备流上，不会阻塞CPU线程。
torch._assert_async在GPU上执行断言时，如果失败，错误何时报告？
错误会在GPU流与CPU线程同步时报告，因此是异步的，可能延迟到后续某个同步点才报错。
使用torch._assert_async时，断言失败会导致什么后果？
断言失败会污染CUDA上下文，导致后续所有CUDA调用失败，程序无法继续在GPU上运行，必须重启CUDA上下文。
torch._assert_async在torch.compile下如何工作？
在torch.compile下，断言会被融合到编译后的计算图中，与其他操作一起优化，不会导致图断裂。
torch._assert_async的实现原理是什么？
实现通过CUDA的__trap指令（汇编为asm volatile("trap;")）在断言失败时终止程序并污染CUDA上下文。在Triton编译中，使用tl.device_assert指令。
torch._assert_async适合在生产环境中使用吗？
不适合。因为断言失败会污染CUDA上下文，且断言本身有开销，应仅在开发阶段使用，生产环境应移除，类似C++中assert在发布版中被优化掉。
