Make _triton_permute_bwd handle non-contiguous grads and bad indices

[indianspeedster]

Summary

Two robustness fixes to _triton_permute_bwd:

• Non-contiguous inputs. The kernel addresses rows as row * cols, so a
  non-contiguous grad_output or permuted_indices produced wrong results with
  no error. The op now makes both contiguous itself (a no-op when they already
  are). This covers the permute_mxfp8_fwd_hp_bwd backward, which passes
  grad_output through as-is.
• Out-of-range indices. The scatter only skipped the -1 padding value, so
  an index >= rows, or any other negative index, wrote outside output_buffer.
  It now skips any index outside [0, original_rows).
  generate_permute_indices never emits one, so this only matters if a caller
  passes inconsistent counts.

Split out of #4623 as requested in review.

Test plan

New tests in test/prototype/moe_training/ep/test_kernels.py:

• test_triton_permute_bwd_non_contiguous_inputs: a transposed gradient and a
  strided index tensor must give the same result as their contiguous copies.
• test_triton_permute_bwd_kernel_out_of_range_indices: launches the kernel into
  rows [guard, guard + rows) of a larger buffer filled with a sentinel, with
  indices rows, rows + 3, -2 and -7. The guard rows must stay
  untouched and the valid rows must be scattered correctly.

	MI355X (gfx950)	B200 (SM100)
test_permute.py + ep/test_kernels.py with this change	12 passed	12 passed
New tests against main's kernel	both fail	both fail
Without the upper bound (< original_rows)	out-of-range test fails	not run
Without the lower bound (>= 0)	out-of-range test passes	out-of-range test fails
Without either .contiguous()	non-contiguous test fails	not run

On MI355X the stray writes for negative indices don't land in the guard rows even
without the lower bound, so that case is only observable on B200.

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 #4988 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/4989

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