fix bug in tl.store mask for kernel _to_fp8_row_major_t_and_non_t

[danielvegamyhre]

Discovered this bug while prototyping torchtitan integration with float8nocompile in pytorch/torchtitan#778, the symptom was the grads went to NaN during training.

The issue was a mismatch between how output offsets for writing the transposed tensor are calculated, versus the offsets used in creating the output mask. Output offsets for writing were correct but the offsets used in the output mask were backwards accidentally.

[danielvegamyhre]

Stack from ghstack (oldest at bottom):

• float8nocompile: add e2e fsdp test #1523
• support for activation checkpointing in float8nocompile #1517
• -> fix bug in tl.store mask for kernel _to_fp8_row_major_t_and_non_t #1516

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 96ee5ee with merge base f86fda9 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test (CUDA 2.4, linux.g5.12xlarge.nvidia.gpu, torch==2.4.0, cuda, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t fb4d6860a3727f171085677a64a804406265c4269e6eefe27625371655fc4133 /exec failed with exit code 1

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

[vkuzo]

looks like there is a gap in test coverage, is it doable to add a test to this PR which fails before the fix and passes after the fix?

[danielvegamyhre]

looks like there is a gap in test coverage, is it doable to add a test to this PR which fails before the fix and passes after the fix?

ah thanks for the reminder, just updated the e2e training test with a new test case with large inputs, which I've confirmed fails without this change and passes with this change.