Aoti Debug

AOTI Debugging Guide

This skill helps diagnose and fix common AOTInductor issues.

Error Pattern Routing

Check the error message and route to the appropriate sub-guide:

Triton Index Out of Bounds

If the error matches this pattern:

Assertion `index out of bounds: 0 <= tmpN < ksM` failed

→ Follow the guide in `triton-index-out-of-bounds.md`

All Other Errors

Continue with the sections below.

---

First Step: Always Check Device and Shape Matching

For ANY AOTI error (segfault, exception, crash, wrong output), ALWAYS check these first:

1. Compile device == Load device: The model must be loaded on the same device type it was compiled on 2. Input devices match: Runtime inputs must be on the same device as the compiled model 3. Input shapes match: Runtime input shapes must match the shapes used during compilation (or satisfy dynamic shape constraints)

# During compilation - note the device and shapes
model = MyModel().eval()           # What device? CPU or .cuda()?
inp = torch.randn(2, 10)           # What device? What shape?
compiled_so = torch._inductor.aot_compile(model, (inp,))

# During loading - device type MUST match compilation
loaded = torch._export.aot_load(compiled_so, "???")  # Must match model/input device above

# During inference - device and shapes MUST match
out = loaded(inp.to("???"))  # Must match compile device, shape must match

If any of these don't match, you will get errors ranging from segfaults to exceptions to wrong outputs.

Key Constraint: Device Type Matching

AOTI requires compile and load to use the same device type.

• If you compile on CUDA, you must load on CUDA (device index can differ)
• If you compile on CPU, you must load on CPU
• Cross-device loading (e.g., compile on GPU, load on CPU) is NOT supported

Common Error Patterns

1. Device Mismatch Segfault

Symptom: Segfault, exception, or crash during aot_load() or model execution.

Example error messages:

• The specified pointer resides on host memory and is not registered with any CUDA device
• Crash during constant loading in AOTInductorModelBase
• Expected out tensor to have device cuda:0, but got cpu instead

Cause: Compile and load device types don't match (see "First Step" above).

Solution: Ensure compile and load use the same device type. If compiled on CPU, load on CPU. If compiled on CUDA, load on CUDA.

2. Input Device Mismatch at Runtime

Symptom: RuntimeError during model execution.

Cause: Input device doesn't match compile device (see "First Step" above).

Better Debugging: Run with AOTI_RUNTIME_CHECK_INPUTS=1 for clearer errors. This flag validates all input properties including device type, dtype, sizes, and strides:

AOTI_RUNTIME_CHECK_INPUTS=1 python your_script.py

This produces actionable error messages like:

Error: input_handles[0]: unmatched device type, expected: 0(cpu), but got: 1(cuda)

Debugging CUDA Illegal Memory Access (IMA) Errors

If you encounter CUDA illegal memory access errors, follow this systematic approach:

Step 1: Sanity Checks

Before diving deep, try these debugging flags:

AOTI_RUNTIME_CHECK_INPUTS=1
TORCHINDUCTOR_NAN_ASSERTS=1

These flags take effect at compilation time (at codegen time):

• AOTI_RUNTIME_CHECK_INPUTS=1 checks if inputs satisfy the same guards used during compilation
• TORCHINDUCTOR_NAN_ASSERTS=1 adds codegen before and after each kernel to check for NaN

Step 2: Pinpoint the CUDA IMA

CUDA IMA errors can be non-deterministic. Use these flags to trigger the error deterministically:

PYTORCH_NO_CUDA_MEMORY_CACHING=1
CUDA_LAUNCH_BLOCKING=1

These flags take effect at runtime:

• PYTORCH_NO_CUDA_MEMORY_CACHING=1 disables PyTorch's Caching Allocator, which allocates bigger buffers than needed immediately. This is usually why CUDA IMA errors are non-deterministic.
• CUDA_LAUNCH_BLOCKING=1 forces kernels to launch one at a time. Without this, you get "CUDA kernel errors might be asynchronously reported" warnings since kernels launch asynchronously.

Step 3: Identify Problematic Kernels with Intermediate Value Debugger

Use the AOTI Intermediate Value Debugger to pinpoint the problematic kernel:

AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3

This prints kernels one by one at runtime. Together with previous flags, this shows which kernel was launched right before the error.

To inspect inputs to a specific kernel:

AOT_INDUCTOR_FILTERED_KERNELS_TO_PRINT="triton_poi_fused_add_ge_logical_and_logical_or_lt_231,_add_position_embeddings_kernel_5" AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=2

If inputs to the kernel are unexpected, inspect the kernel that produces the bad input.

Additional Debugging Tools

Logging and Tracing

• tlparse / TORCH_TRACE: Provides complete output codes and records guards used
• TORCH_LOGS: Use TORCH_LOGS="+inductor,output_code" to see more PT2 internal logs
• TORCH_SHOW_CPP_STACKTRACES: Set to 1 to see more stack traces

Common Sources of Issues

• Dynamic shapes: Historically a source of many IMAs. Pay special attention when debugging dynamic shape scenarios.
• Custom ops: Especially when implemented in C++ with dynamic shapes. The meta function may need to be Symint'ified.

API Notes

Deprecated API

torch._export.aot_compile()  # Deprecated
torch._export.aot_load()     # Deprecated

Current API

torch._inductor.aoti_compile_and_package()
torch._inductor.aoti_load_package()

The new API stores device metadata in the package, so aoti_load_package() automatically uses the correct device type. You can only change the device index (e.g., cuda:0 vs cuda:1), not the device type.

Environment Variables Summary

AOTI Triton Index Out of Bounds Debug Guide

This guide helps debug AOTI Triton kernel assertion errors with the index out of bounds pattern.

Error Pattern

This guide applies when you see errors like:

/var/tmp/torchinductor_*/.../*.py:NN: unknown: block: [X,Y,Z], thread: [X,Y,Z]
Assertion `index out of bounds: 0 <= tmpN < ksM` failed.

Key Information from Error

Understanding the Assertion

• tmpN: A computed index value in the Triton kernel
• ksM: A dynamic kernel size parameter (runtime value)
• The assertion fails when tmpN < 0 or tmpN >= ksM

---

Step 1: Collect AOTI Package

You need access to the AOTI package that was compiled. This is typically a .pt2 package or extracted archive containing a wrapper.cpp file.

Key File: *.wrapper.cpp contains:

• All Triton kernel source code (embedded as comments)
• Kernel launch configurations
• Input/output tensor mappings
• Dynamic shape variable definitions

---

Step 2: Locate the Failing Kernel in C++ Wrapper

Search for the Assertion Pattern

Extract the assertion pattern from the error (e.g., tmp18 < ks0) and search:

# Search for the specific assertion
grep -n "tmpN < ksM" /path/to/*.wrapper.cpp

# Get context around the assertion (80 lines before, 20 after)
grep -n -B80 -A20 "tmpN < ksM" /path/to/*.wrapper.cpp

Find the Full Kernel Definition

The kernel is embedded as a Python docstring comment in the C++ wrapper:

    /*
    async_compile.triton('triton_red_fused_...', '''
    import triton
    import triton.language as tl
    ...
    def triton_red_fused_...(in_ptr0, out_ptr1, ks0, xnumel, r0_numel, ...):

---

Step 3: Understand the Kernel Logic

Analyze the code path leading to the assertion. Common patterns that cause index out of bounds:

Pattern: Empty Tensor with ks0 = 0

When a dynamic shape ks0 = 0: 1. tmp13 = (-1) + 0 = -1 2. Index wrapping logic produces -1 3. Assertion 0 <= -1 < 0 fails

Example Kernel Pattern

tmp13 = (-1) + ks0           # ks0 - 1
tmp14 = tl.where(tmp12, tmp10, tmp13)  # if condition: use tmp10, else: ks0-1
tmp15 = ks0
tmp16 = tmp14 + tmp15        # wrap-around for negative indices
tmp17 = tmp14 < 0
tmp18 = tl.where(tmp17, tmp16, tmp14)  # if negative: add ks0

# ASSERTION: 0 <= tmp18 < ks0
tl.device_assert(((0 <= tmp18) & (tmp18 < ks0)), "index out of bounds")

---

Step 4: Identify the Dynamic Shape Variable

Find Where the Kernel is Called

grep -n "call_triton_KERNEL_NAME" /path/to/*.wrapper.cpp

Example Output

call_triton_red_fused_...(arg1415_1, buf696, s607, 1L, s13, ...);

Parameter Mapping

Find the Definition of the Shape Variable

grep -n "int64_t s607 = " /path/to/*.wrapper.cpp

This shows which input tensor dimension defines the shape:

int64_t s607 = arg1416_1_size[0];

---

Step 5: Trace Back to Model Input

Find Input Index

Inputs are numbered sequentially. Find which input the argument corresponds to:

grep -n 'inputs_info_\[INDEX\].name = "argNNN_1"' /path/to/*.wrapper.cpp

Check Input Constraints

grep -n "argNNN_1_size\[0\]" /path/to/*.wrapper.cpp

Look for guards like:

if (arg_size[0] > 230400) {  // Upper bound check only - no lower bound!

Common Issue: Upper bound checks exist but no lower bound checks for >= 1.

---

Step 6: Map to Model Code

Use Source Node Comments

The C++ wrapper includes comments showing which PyTorch operations generated each kernel:

grep -n -B5 "call_triton_KERNEL_NAME" /path/to/*.wrapper.cpp | grep "Source Nodes"

Example Output

// Topologically Sorted Source Nodes: [slice_1, sub_89, cumsum, ge_231, where_2, index_copy]

Map Operations to Python Code

---

Root Cause Analysis

Common Root Causes

1. Empty tensor at runtime: A jagged/variable-length tensor has size 0 at runtime but wasn't tested during compilation 2. Missing lower bound guards: AOTI only generates upper bound checks, not lower bound checks 3. Edge case not in sample inputs: Sample inputs during AOTI export never included the edge case

---

Fix Recommendations

Option 1: Add Guard in Forward Method

def forward(self, lengths: torch.Tensor, ...) -> torch.Tensor:
    if lengths.numel() == 0:
        device = lengths.device
        return torch.empty(0, self.output_dim, device=device)
    # ... rest of method

Option 2: Fix the Specific Operation

Add handling for empty tensors in the problematic operation:

def process_events(self, lengths: torch.Tensor, ...):
    if lengths.numel() == 0:
        return torch.empty(0, self.emb_dim, device=lengths.device)
    # ... rest of method

Option 3: Include Edge Cases in AOTI Export

Ensure sample inputs during AOTI export include:

• Empty tensors (size 0)
• Minimum size tensors (size 1)
• Maximum expected sizes

---

Useful Commands Summary

Searching in AOTI Wrapper

# Find kernel by assertion pattern
grep -n "tmpN < ksM" *.wrapper.cpp

# Get full kernel context
grep -n -B80 -A20 "ASSERTION_PATTERN" *.wrapper.cpp

# Find kernel call site
grep -n "call_KERNEL_NAME" *.wrapper.cpp

# Find dynamic shape definition
grep -n "int64_t SHAPE_VAR = " *.wrapper.cpp

# Find input mapping
grep -n 'inputs_info_\[INDEX\].name' *.wrapper.cpp

# Find size constraints
grep -n "SHAPE_VAR_size\[0\]" *.wrapper.cpp

Environment Variables for Debugging

# Enable debug output during torch.compile
export TORCH_COMPILE_DEBUG=1

# Save generated kernels to persistent location
export TORCHINDUCTOR_CACHE_DIR=/path/to/save/kernels

# Enable CUDA launch blocking for accurate stack traces
export CUDA_LAUNCH_BLOCKING=1