Trident: Unifying Guarded Dispatch and Host Execution for PyTorch Triton Workloads

2026-09-29

AI Summary

This study addresses the significant end-to-end latency incurred by host-side scheduling overhead when executing Triton kernels within PyTorch. To mitigate this issue, we propose Trident, a compiler backend built upon Torch-MLIR that unifies guard evaluation, argument preparation, and host execution into a single natively compiled code module. This approach eliminates redundant runtime overhead, falling back to the Python layer only when a new specialization is triggered. When integrated with ATen operator optimizations, Trident achieves 1.47× and 1.68× end-to-end speedups over eager mode and torch.compile, respectively, on large language model inference tasks.

Abstract

User-written Triton kernels enable high-performance GPU computation within PyTorch, but their end-to-end latency can remain dominated by host-side orchestration, especially when device execution is short. Although torch.compile can generate native host wrappers for captured graphs, each invocation still passes through runtime-managed specialization lookup, guard evaluation, and preparation before reaching the wrapper. We present Trident, a compiler backend that removes this recurring overhead from the specialization cache-hit path. Trident introduces the Specialization Cache Module (SCM), which compiles guarded specialization selection, argument and execution-environment preparation, and host execution for multiple specializations into a single executable module. An invocation enters the SCM once, remains in compiled code when a specialization matches, and returns to Python only when a new specialization must be compiled. Built on Torch-MLIR, Trident lowers guards and host-side orchestration to native code while retaining calls to optimized runtime implementations of supported ATen operators. Our evalu- ation on two LLMs shows that Trident achieves up to a 1.47x speedup in model-level end-to-end latency over eager execution and up to 1.68x over torch.compile.

Problem

• Triton kernels
• host-side orchestration overhead
• end-to-end latency
• guard evaluation
• specialization lookup

Innovation

• Trident
• Specialization Cache Module
• Torch-MLIR
• Triton kernels
• host-side orchestration