Description
Triton is an accelerator-independent programming language that has become a central
component of modern AI inference systems. In the PyTorch compilation stack, high-
level tensor programs are lowered into optimized Triton programs that execute directly
on AI accelerators, placing Triton between framework-level model code and low-level
accelerator execution. Although Triton raises the level of abstraction compared to
CUDA and Ascend-C, kernel programmers still manipulate pointers, offsets, and
masks explicitly. Incorrect reasoning about these values can cause exploitable spatial
memory vulnerabilities, enabling information disclosure and remote code execution in
accelerator programs. This thesis provides the first comprehensive security analysis of
memory corruption bugs in Triton.
We present TritonShield, the first runtime protection system for AI inference
workloads that execute Triton programs. TritonShield translates pointer and mask
computations into symbolic expressions at compile time, derives assertions that capture
the required spatial memory-safety conditions, and validates these assertions at launch
time before the Triton program executes. This design blocks unsafe launches before
they reach the accelerator, while parameter caching avoids repeated validation work
for common inference workloads.
Our evaluation shows that TritonShield discovers 11 previously unknown bugs and
identifies two previously unknown vulnerability classes in the TritonBench_G dataset.
Across a corpus of 7410 Triton programs, TritonShield protects 99.7% of all memory
accesses. On real-world AI inference workloads based on Llama, DeepSeek, Gemma,
and Qwen models, TritonShield introduces low overhead, averaging approximately
0.7% in our measurements. These results show that precise runtime protection for
Triton programs is practical and can substantially improve the memory safety of AI
inference systems.
|
