Use int64 addressing in _triton_permute_bwd_kernel

[indianspeedster]

Summary

_triton_permute_bwd_kernel computes row * cols in int32 for both its gradient
load and its scatter store. Once rows * cols >= 2**31 the address overflows and
the access hits an illegal address. This widens the loaded destination rows and
the gradient row offsets to int64.

The size is reachable in practice: DeepSeek-V3-16B with EP=8 and imbalanced
routing sent 1,141,504 rows x 2048 cols = 2,337,800,192 elements through this
kernel. _triton_permute_bwd is also the backward of permute_mxfp8_fwd_hp_bwd,
so this fixes that path too.

Split out of #4623 as requested in review. #4623 is stacked on this PR, because
it routes the bf16 permute_and_pad backward through this kernel.

Test plan

New test test_triton_permute_bwd_addresses_past_int32 in
test/prototype/moe_training/ep/test_kernels.py covers both sides of the fix:

• Store: a 4096-row gradient scattered into the last rows of an output with more
  than 231 elements (the last 1024 of those rows start past 231).
• Load: a gradient with more than 2**31 elements whose last rows scatter into a
  small output.

It runs the kernel in a child process with a 120 s timeout, so an illegal address
or hang fails the test instead of crashing the pytest process. The child is
pointed at the same torchao the test imported. Each case needs about 4 GiB of GPU
memory; the test skips with less than 6 GiB free.

	MI355X (gfx950)	B200 (SM100)
test_permute.py + ep/test_kernels.py with this change	11 passed	11 passed
New test against main's kernel	fails: illegal memory access	fails: illegal memory access
New test with only the load widening reverted	fails	not run
New test with only the store widening reverted	fails	not run

Test command on both machines: pytest test/prototype/moe_training/ep/test_permute.py test/prototype/moe_training/ep/test_kernels.py. MI355X: torch 2.15.0.dev+rocm7.14, triton 3.8.0; because main gates these files to SM100, the MI355X runs override that gate locally. B200: torch 2.14.1+cu130, triton 3.8.0.

No upstream CI runner has SM100, and these test files are gated to SM100 on main, so today these tests run in no CI job. After #4987 (which enables them on gfx950) they run in the ROCm regression job (regression_test_rocm.yml, MI350 runner, triggered by ciflow/rocm or a push to main).

This PR and #4989 both append tests to the end of ep/test_kernels.py, so
whichever lands second needs a trivial rebase.

---

Authored with assistance from an AI coding agent.

[pytorch-bot]

🔗 Helpful Links

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

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