From 70e24da543dc387a814150705b30679fe60c857c Mon Sep 17 00:00:00 2001
From: indianspeedster <38083203+indianspeedster@users.noreply.github.com>
Date: Thu, 8 Oct 2026 18:54:54 +0000
Subject: [PATCH] Use int64 addressing in _triton_permute_bwd_kernel

The kernel computed `row * cols` in int32 for both its gradient load and its
scatter store, so once rows * cols >= 2**31 the addresses overflowed. On
MI355X and B200 the access then hits an illegal address. Widen the loaded
destination rows and the gradient row offsets to int64.

_triton_permute_bwd is the backward of permute_mxfp8_fwd_hp_bwd, so this
fixes that path too.

Test: test_triton_permute_bwd_addresses_past_int32 covers the store (a small
gradient scattered past 2**31 output elements) and the load (a gradient with
more than 2**31 elements). It runs the kernel in a child process, so an
illegal address or hang becomes a normal test failure.
---
 .../prototype/moe_training/ep/test_kernels.py | 73 +++++++++++++++++++
 torchao/prototype/moe_training/ep/permute.py  |  5 +-
 2 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/test/prototype/moe_training/ep/test_kernels.py b/test/prototype/moe_training/ep/test_kernels.py
index 97ad73dda89..3f140c4b149 100644
--- a/test/prototype/moe_training/ep/test_kernels.py
+++ b/test/prototype/moe_training/ep/test_kernels.py
@@ -4,9 +4,15 @@
 # This source code is licensed under the BSD 3-Clause license found in the
 # LICENSE file in the root directory of this source tree.
 
+import os
+import subprocess
+import sys
+import textwrap
+
 import pytest
 import torch
 
+import torchao
 from torchao.utils import is_cuda_version_at_least, is_sm_at_least_100
 
 if not (
@@ -102,3 +108,70 @@ def test_triton_permute_bwd(
         atol=0,
         msg="Triton permute backward kernel output does not match PyTorch reference",
     )
+
+
+def _run_in_child(code: str, timeout: int = 120) -> subprocess.CompletedProcess:
+    """Run `code` in a fresh Python process that imports this same torchao.
+
+    Kernels with bad addressing can fault or hang instead of returning wrong
+    values, and either would take down the pytest process.
+    """
+    torchao_root = os.path.dirname(os.path.dirname(torchao.__file__))
+    env = dict(os.environ)
+    env["PYTHONPATH"] = os.pathsep.join(
+        p for p in (torchao_root, env.get("PYTHONPATH")) if p
+    )
+    try:
+        return subprocess.run(
+            [sys.executable, "-c", code],
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+            env=env,
+        )
+    except subprocess.TimeoutExpired:
+        pytest.fail(f"child process hung for more than {timeout} s")
+
+
+_ADDRESS_PAST_INT32_CHILD = textwrap.dedent(
+    """
+    import torch
+    from torchao.prototype.moe_training.ep.permute import _triton_permute_bwd
+
+    cols = 1024
+    big = 2**31 // cols + 1024  # big * cols > 2**31
+    n = 4096
+
+    # Store side: scatter a small gradient to the last n rows of a big output.
+    # These rows straddle the 2**31-element boundary; the last 1024 start past it.
+    idx = torch.arange(big - n, big, device="cuda", dtype=torch.int32)
+    grad = torch.randn(n, cols, device="cuda", dtype=torch.bfloat16)
+    out = _triton_permute_bwd(grad, idx, big, cols)
+    torch.cuda.synchronize()
+    assert torch.equal(out[big - n :], grad), "store past 2**31 elements is wrong"
+    assert not out[: big - n].any(), "store wrote into rows it should not touch"
+    del idx, grad, out
+    torch.cuda.empty_cache()
+
+    # Load side: a big gradient whose last n rows scatter into a small output.
+    idx = torch.full((big,), -1, device="cuda", dtype=torch.int32)
+    idx[big - n :] = torch.arange(n, device="cuda", dtype=torch.int32)
+    grad = torch.zeros(big, cols, device="cuda", dtype=torch.bfloat16)
+    grad[big - n :] = torch.randn(n, cols, device="cuda", dtype=torch.bfloat16)
+    out = _triton_permute_bwd(grad, idx, n, cols)
+    torch.cuda.synchronize()
+    assert torch.equal(out, grad[big - n :]), "load past 2**31 elements is wrong"
+    """
+)
+
+
+def test_triton_permute_bwd_addresses_past_int32():
+    """Kernel addresses must not overflow int32 once rows * cols >= 2**31.
+
+    Covers both the gradient load and the scatter store.
+    """
+    free, _ = torch.cuda.mem_get_info()
+    if free < 6 * 2**30:
+        pytest.skip("needs ~6 GiB of free GPU memory")
+    proc = _run_in_child(_ADDRESS_PAST_INT32_CHILD)
+    assert proc.returncode == 0, proc.stderr[-2000:]
diff --git a/torchao/prototype/moe_training/ep/permute.py b/torchao/prototype/moe_training/ep/permute.py
index b4c93044696..1fb51e6ef58 100644
--- a/torchao/prototype/moe_training/ep/permute.py
+++ b/torchao/prototype/moe_training/ep/permute.py
@@ -303,15 +303,16 @@ def _triton_permute_bwd_kernel(
     row_offsets = row_pid * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
     col_offsets = col_pid * BLOCK_COLS + tl.arange(0, BLOCK_COLS)
 
+    # int64 so `row * cols` does not overflow once rows * cols >= 2**31.
     dest_rows = tl.load(
         permuted_indices_ptr + row_offsets,
         mask=row_offsets < grad_rows,
         other=PADDING_VALUE,
-    )
+    ).to(tl.int64)
 
     read_mask = (row_offsets[:, None] < grad_rows) & (col_offsets[None, :] < grad_cols)
     grad_values = tl.load(
-        grad_ptr + row_offsets[:, None] * grad_cols + col_offsets[None, :],
+        grad_ptr + row_offsets[:, None].to(tl.int64) * grad_cols + col_offsets[None, :],
         mask=read_mask,
         other=PADDING_VALUE,
     )