Fix axis-swapped bounds mask in _to_fp8_row_major_t transposed store

[amanyagami]

Summary

The transposed store address in _to_fp8_row_major_t (torchao/prototype/float8nocompile/kernels/fp8_dynamic_tensorwise.py) is:

out_offs = (
    block_col_offs[:, None] * output_stride_row
    + block_row_offs[None, :] * output_stride_col
)

i.e. the tile's first axis addresses the output row via block_col_offs, and the second axis addresses the output col via block_row_offs (this kernel writes the transpose of its input tile). The bounds mask instead compared block_row_offs against output_num_rows and block_col_offs against output_num_cols — axes swapped relative to the address. For rectangular / non-tile-aligned inputs this masks off valid output elements (leaving torch.empty garbage in the result) and/or lets writes through for positions that are actually out of bounds.

Fix

Swap the mask to check the same axes, in the same order, as out_offs.

Test

This kernel's own test (fp8_dynamic_tensorwise_test.py) hard-requires CUDA (assert torch.cuda.is_available()), so I can't exercise it in this environment (no GPU access). Instead, added test/prototype/float8nocompile/test_to_fp8_row_major_t_bounds_mask.py, which isolates the exact block_row_offs/block_col_offs → out_offs/out_mask index arithmetic in plain Python (mirroring precisely what the kernel computes per program instance, for any tile size / non-tile-aligned shape) and checks it against several rectangular shapes:

• the fixed mask covers exactly the valid transposed output region for every shape tested
• the buggy (axis-swapped) mask both misses valid output positions and, for some shapes, writes to positions entirely outside the valid output region (e.g. for a (17,9) input with block_size=8: 72 valid positions missed, 63 out-of-bounds positions written)

This proves the fix is logically correct independent of actually running the Triton kernel. End-to-end verification (e.g. against the existing test_fp8_hp_to_fp8_row_major_t in fp8_dynamic_tensorwise_test.py) still needs GPU CI — flagging this explicitly rather than claiming more than I could verify.

Fixes #4584

🤖 Generated with Claude Code

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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.