AOTInductor Triton SASS Inspection

AOTInductor, PyTorch's ahead-of-time inference backend compiler, emits Triton kernels as cubin files stored in the pt2 zip archive, which developers can unzip and disassemble with nvdisasm or cuobjdump --dump-sass to verify that 128-bit LDG.E.128 and STG.E.128 vectorized memory instructions were generated. The technique matters because PyTorch tensors from the cached allocator are 256-byte aligned by default, and the Triton compiler emits those 128-bit vectorized loads and stores only when stride offsets are also 16-byte aligned, so SASS inspection confirms whether the expected vectorization actually occurred.

AOTInductor Triton SASS Inspection Triton is a language and compiler for creating high-performance GPU kernels. AOTInductor is a PyTorch backend compiler that can generate ahead-of-time optimized deep learning inference engines and is compatible with custom Triton kernels and performance autotuning. In order to confirm if the generated Triton kernels have certain behaviors, sometimes it is necessary to inspect the SASS Streaming Assembly of the generated CUDA kernels. In this blog post, I would like to demonstrate how to inspect SASS of Triton kernels from AOTInductor artifacts. Based on the base address of the input tensor and the static offset for data access, the Triton compiler can deduce the appropriate vectorization instructions to generate efficient memory access patterns. In PyTorch, all the tensors allocated from PyTorch cached allocator are 256-byte aligned by default. If the stride offsets are also aligned to 16 bytes, the Triton compiler can generate 128-bit LDG.E.128 / STG.E.128 vectorized memory instructions. In this example, we will verify if 128-bit vectorized memory instructions are generated in the SASS of CUDA kernel of interest. To verify this, because AOTInductor generates Triton kernels as cubin files, we can inspect the generated SASS to check if the 128-bit vectorized memory instructions are indeed emitted. The cubin files are saved in the pt2 zip archive and SASS can be extracted for inspection using nvdisasm or cuobjdump --dump-sass after unzipping the archive. | 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214 | python import osimport shutilimport subprocessimport zipfileimport torchimport torch.nn as nnfrom torch.library import triton op, wrap tritonimport tritonimport triton.language as tl --------------------------------------------------------------------------- 1. Define Triton Vector Copy Kernel ---------------------------------------------------------------------------@triton.jitdef triton copy kernel in ptr, out ptr, n elements, BLOCK SIZE: "tl.constexpr", : pid = tl.program id axis=0 block start = pid BLOCK SIZE offsets = block start + tl.arange 0, BLOCK SIZE 16-byte alignment hint to trigger 128-bit LDG.E.128 / STG.E.128 vectorization offsets = tl.multiple of offsets, 16 mask = offsets < n elements x = tl.load in ptr + offsets, mask=mask tl.store out ptr + offsets, x, mask=mask --------------------------------------------------------------------------- 2. Transparently Traceable Triton Operator ---------------------------------------------------------------------------@triton op "custom ops::triton copy", mutates args= def triton copy op x: torch.Tensor - torch.Tensor: out = torch.empty like x if x.is cuda: n elements = x.numel def grid meta : return triton.cdiv n elements, meta "BLOCK SIZE" , wrap triton triton copy kernel grid x, out, n elements, BLOCK SIZE=1024 else: out.copy x return out --------------------------------------------------------------------------- 3. PyTorch Model ---------------------------------------------------------------------------class VectorCopyModel nn.Module : def forward self, x: torch.Tensor - torch.Tensor: return triton copy op x --------------------------------------------------------------------------- 4. Main Execution Pipeline ---------------------------------------------------------------------------def main : device = "cuda" if torch.cuda.is available else "cpu" if device = "cuda": raise RuntimeError "CUDA device required for AOTInductor Triton compilation." print "\n" + "=" 75 print " PIPELINE: Export - AOTI Compile - Bitwise Test - SASS Inspection" print "=" 75 model = VectorCopyModel .to device .eval example args = torch.randn 1024 1024, device=device, dtype=torch.float32 , -------------------------------------------------------------------------- Step 1: Export Graph -------------------------------------------------------------------------- print "\n STEP 1 Exporting graph with torch.export.export ..., strict=True " ep = torch.export.export model, args=example args, strict=True decomposed ep = ep.run decompositions print " ✅ ExportedProgram graph successfully captured and decomposed." -------------------------------------------------------------------------- Step 2: Compile & Package via AOTInductor -------------------------------------------------------------------------- print "\n STEP 2 Compiling to AOTI Package .pt2 artifact " output dir = os.path.abspath "./aoti output" os.makedirs output dir, exist ok=True pkg path = os.path.join output dir, "model.pt2" compiled pkg = torch. inductor.aoti compile and package decomposed ep, package path=pkg path print f" ✅ Compiled Package Generated: {compiled pkg}" -------------------------------------------------------------------------- Step 3: Run Executable & Verify Bitwise Identity -------------------------------------------------------------------------- print "\n STEP 3 Executing Compiled Model and Verifying Bitwise Identity" runner = torch. inductor.aoti load package compiled pkg test input = torch.randn 1024 1024, device=device, dtype=torch.float32 output = runner test input assert torch.equal test input, output , "FAIL: Tensors are not equal " input bits = test input.view torch.int32 output bits = output.view torch.int32 mismatches = input bits = output bits .sum .item assert mismatches == 0, f"FAIL: Found {mismatches} bitwise mismatch es " print " ✅ Bitwise Identity Confirmed: Input and Output are 100% bitwise identical " -------------------------------------------------------------------------- Step 4: Extract All Artifacts from .pt2 Package -------------------------------------------------------------------------- print "\n STEP 4 Unpacking .pt2 Archive and Scanning Artifacts" extract dir = os.path.join output dir, "extracted" os.makedirs extract dir, exist ok=True extracted cubins = extracted so = None extracted cpp = None with zipfile.ZipFile compiled pkg, "r" as zip ref: zip ref.extractall extract dir for root, , files in os.walk extract dir : for f in files: full path = os.path.join root, f if f.endswith ".cubin" : extracted cubins.append full path elif f.endswith ".wrapper.so" or f.endswith ".so" and not extracted so : extracted so = full path elif f.endswith ".wrapper.cpp" or f.endswith ".cpp" and not extracted cpp : extracted cpp = full path print f" ✅ Found {len extracted cubins } .cubin file s in package." if extracted cpp: print f" ✅ Found wrapper C++: {extracted cpp}" if extracted so: print f" ✅ Found wrapper .so: {extracted so}" -------------------------------------------------------------------------- Step 5: SASS Disassembly on Extracted .cubin via nvdisasm -------------------------------------------------------------------------- print "\n STEP 5 Running SASS Disassembly nvdisasm on Extracted .cubin" if extracted cubins and shutil.which "nvdisasm" : for cubin path in extracted cubins: cubin name = os.path.basename cubin path res = subprocess.run "nvdisasm", "-g", cubin path , capture output=True, text=True, sass output = res.stdout if "LDG.E.128" in sass output or "STG.E.128" in sass output: print f" ✅ nvdisasm CONFIRMED in {cubin name} Found 128-bit vector instructions:" for line in sass output.splitlines : if "LDG.E.128" in line or "STG.E.128" in line: print f" {line.strip }" else: print f" -- Disassembly completed for {cubin name}." elif not shutil.which "nvdisasm" : print " -- SKIPPED 'nvdisasm' not found in system PATH." -------------------------------------------------------------------------- Step 6: SASS Disassembly on Extracted .cubin via cuobjdump -------------------------------------------------------------------------- print "\n STEP 6 Running SASS Disassembly cuobjdump on Extracted .cubin" if extracted cubins and shutil.which "cuobjdump" : for cubin path in extracted cubins: cubin name = os.path.basename cubin path res = subprocess.run "cuobjdump", "-sass", cubin path , capture output=True, text=True, cubin sass = res.stdout if "LDG.E.128" in cubin sass or "STG.E.128" in cubin sass: print f" ✅ cuobjdump CONFIRMED in {cubin name} Found 128-bit vector instructions:" for line in cubin sass.splitlines : if "LDG.E.128" in line or "STG.E.128" in line: print f" {line.strip }" else: print f" -- Disassembly of {cubin name} completed." elif not shutil.which "cuobjdump" : print " -- SKIPPED 'cuobjdump' not found in system PATH." if name == " main ": main | | 12345678910111213141516171819202122232425262728293031323334353637 | bash $ python triton copy.py=========================================================================== PIPELINE: Export - AOTI Compile - Bitwise Test - SASS Inspection=========================================================================== STEP 1 Exporting graph with torch.export.export ..., strict=True /usr/lib/python3.12/copyreg.py:99: FutureWarning: isinstance treespec, LeafSpec is deprecated, use isinstance treespec, TreeSpec and treespec.is leaf instead. return cls. new cls, args ✅ ExportedProgram graph successfully captured and decomposed. STEP 2 Compiling to AOTI Package .pt2 artifact /usr/lib/python3.12/copyreg.py:99: FutureWarning: isinstance treespec, LeafSpec is deprecated, use isinstance treespec, TreeSpec and treespec.is leaf instead. return cls. new cls, args ✅ Compiled Package Generated: /mnt/aoti output/model.pt2 STEP 3 Executing Compiled Model and Verifying Bitwise Identity ✅ Bitwise Identity Confirmed: Input and Output are 100% bitwise identical STEP 4 Unpacking .pt2 Archive and Scanning Artifacts ✅ Found 1 .cubin file s in package. ✅ Found wrapper C++: /mnt/aoti output/extracted/model/data/aotinductor/model/c4si55h2cct2cbdghpljslclh43nkx45u6pefe4b6nymbkyhgoqi.wrapper.cpp ✅ Found wrapper .so: /mnt/aoti output/extracted/model/data/aotinductor/model/c4si55h2cct2cbdghpljslclh43nkx45u6pefe4b6nymbkyhgoqi.wrapper.so STEP 5 Running SASS Disassembly nvdisasm on Extracted .cubin ✅ nvdisasm CONFIRMED in crx777x7kmdkuv7227akgtzfojh4osuapriqhz3jesxvwqpwmfnk.cubin Found 128-bit vector instructions: / 0110 / @ P0 LDG.E.128 R8, desc UR4 R2.64 ; / 0140 / @ P1 LDG.E.128 R12, desc UR4 R2.64+0x800 ; / 0150 / @ P0 STG.E.128 desc UR4 R4.64 , R8 ; / 0170 / STG.E.128 desc UR4 R4.64+0x800 , R12 ; STEP 6 Running SASS Disassembly cuobjdump on Extracted .cubin ✅ cuobjdump CONFIRMED in crx777x7kmdkuv7227akgtzfojh4osuapriqhz3jesxvwqpwmfnk.cubin Found 128-bit vector instructions: / 0110 / @ P0 LDG.E.128 R8, desc UR4 R2.64 ; / 0x0000000402088981 / / 0140 / @ P1 LDG.E.128 R12, desc UR4 R2.64+0x800 ; / 0x00080004020c9981 / / 0150 / @ P0 STG.E.128 desc UR4 R4.64 , R8 ; / 0x0000000804008986 / / 0170 / STG.E.128 desc UR4 R4.64+0x800 , R12 ; / 0x0008000c04007986 / | AOTInductor Triton SASS Inspection https://leimao.github.io/blog/AOTInductor-Triton-SASS-Inspection/ https://leimao.github.io/blog/AOTInductor-Triton-SASS-Inspection/