Devlog: Resolving GPU Memory Corruption & CUDA Assertions on Wide Matvecs

1. The Symptoms

When running inference/benchmarks on unquantized models (like Qwen3.5-9B) using the Triton backend on laptop GPUs (e.g., RTX 5050), the runtime encountered two asynchronous CUDA faults:

• cudaErrorIllegalAddress (Illegal memory access) during cache eviction.
• device-side assert triggered ( indexSelectSmallIndex assertion failed) during embedding vocabulary lookups.

2. The Root Cause (Resource Starvation & Memory Corruption)

• Wide Layer Geometries: In Qwen3.5, the projection matrices (Query, Key, Value) have a columns dimension of cols = 3584 .
• Triton Padded Tiling: Triton pads non-power-of-two shapes to the next power of two ( BLOCK_N = 4096 ).
• Static Tile Size: The default Triton config mapping uses BLOCK_M = 64 for wide matrices. This instructs Triton to compile a kernel loading a massive static tile size of 64 × 4096 = 262,144 elements.
• Register & Shared Memory Overflow: In FP32 accumulation, this tile requires 1 MB of memory per thread block. However, laptop GPU streaming multiprocessors (SMs) have a physical hardware limit of 96KB–100KB
  of shared memory/register files.
• The Failure Cascade: The compiler failed to allocate resources cleanly, resulting in severe register spills and memory read violations. Triton silently returned corrupted/NaN outputs. The subsequent
  argmax logit lookup generated a garbage token index (out of bounds), which triggered the PyTorch device-side assertion on the next step.

3. The Fix (Dynamic Chunked Looping)

To prevent resource starvation, we modified Triton’s launching configuration in triton_kernels.py:

• Wide Column Routing: Added a check to detect if cols > 2048 .
• Looped Execution: Instead of loading the entire 4096 columns block statically, we force the kernel to execute _matvec_looped_kernel using a chunk size of BLOCK_N = 512 .
• Memory Optimization: The kernel loops through the columns sequentially in chunks of 512. This drops the tile size memory requirement per block to under 64 KB, which safely fits inside any laptop GPU’s
  physical SM limits.

4. Results

I am currently running the final profiles ( max-performance and max-max-perf ). Let me check on the status!