Mask a_s/b_s scale loads in blockwise FP8 training GEMM kernels

[amanyagami]

Summary

Both triton_fp8_gemm_1x128_128x128_kernel and triton_fp8_gemm_1x128_128x1_kernel in torchao/prototype/blockwise_fp8_training/kernels.py load their per-tile a_s/b_s scale values without a bounds mask, while every other load and the final store in these kernels is masked. When the autotuner's tile does not evenly divide M or N, the scale loads for the tail tile dereference memory past the end of a_s/b_s.

The GEMM's numeric output is unaffected — the out-of-range lanes are dropped by c_mask at the store — so this is a memory-safety bug, not a numerical one. The issue reports it via compute-sanitizer as an Invalid __global__ read, which intermittently surfaces as CUDA error: unspecified launch failure depending on allocator slack.

Fix

mask=offs_m < M for the a_s load, mask=offs_n < N for the b_s load — matching the bounds every sibling load (a_mask, b_mask) and the final store (c_mask) already use in the same kernels.

Test

These kernels require CUDA to run (real @triton.jit GEMMs targeting H100/B200), so I can't exercise them in this environment. Instead, added test/prototype/blockwise_fp8_training/test_gemm_scale_bounds.py, which isolates the exact offs_m/offs_n → scale-index arithmetic in plain Python for tile geometries where M/N are not a multiple of the block size (the exact case the issue reports), and checks:

• unmasked indices do go out of [0, M)/[0, N) for these shapes (reproducing the bug's precondition)
• masked indices (mask=offs_m < M / mask=offs_n < N) never do, and still cover every valid row/col at least once

This proves the fix's indices are correct independent of actually running the Triton kernel. End-to-end verification (e.g. under compute-sanitizer, per the issue's own repro) still needs GPU CI — flagging this explicitly.

Fixes #4634

🤖 Generated with Claude Code

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4793

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.