From f9ff209223d3f3d34c6ef7881793309ede2aaf93 Mon Sep 17 00:00:00 2001 From: indianspeedster <38083203+indianspeedster@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:55:36 +0000 Subject: [PATCH] Make _triton_permute_bwd handle non-contiguous inputs and bad indices - The kernel addresses rows as `row * cols`, so a non-contiguous grad_output or permuted_indices gave wrong results with no error. Make both contiguous inside the op (a no-op when they already are). This covers the permute_mxfp8_fwd_hp_bwd backward, which passes grad_output through as-is. - The scatter only skipped the -1 padding value, so an index >= rows or any other negative index wrote outside output_buffer. Skip any index outside [0, original_rows). generate_permute_indices never emits one; this only matters if a caller passes inconsistent counts. Tests: test_triton_permute_bwd_non_contiguous_inputs covers a transposed gradient and a strided index tensor. test_triton_permute_bwd_kernel_out_of_range_indices launches the kernel into the middle of a larger guard buffer and checks the guard rows are untouched for indices past the end and below zero. --- .../prototype/moe_training/ep/test_kernels.py | 66 ++++++++++++++++++- torchao/prototype/moe_training/ep/permute.py | 11 +++- 2 files changed, 74 insertions(+), 3 deletions(-) diff --git a/test/prototype/moe_training/ep/test_kernels.py b/test/prototype/moe_training/ep/test_kernels.py index 97ad73dda89..c11937c9897 100644 --- a/test/prototype/moe_training/ep/test_kernels.py +++ b/test/prototype/moe_training/ep/test_kernels.py @@ -6,6 +6,7 @@ import pytest import torch +import triton from torchao.utils import is_cuda_version_at_least, is_sm_at_least_100 @@ -17,7 +18,10 @@ pytest.skip("Test requires CUDA 12.8+ with SM >= 100", allow_module_level=True) from torchao.prototype.moe_training.ep.kernels import generate_permute_indices -from torchao.prototype.moe_training.ep.permute import _triton_permute_bwd +from torchao.prototype.moe_training.ep.permute import ( + _triton_permute_bwd, + _triton_permute_bwd_kernel, +) @pytest.mark.parametrize( @@ -102,3 +106,63 @@ def test_triton_permute_bwd( atol=0, msg="Triton permute backward kernel output does not match PyTorch reference", ) + + +def test_triton_permute_bwd_non_contiguous_inputs(): + """The backward op must handle a non-contiguous gradient and index tensor.""" + device = "cuda" + rows, cols = 300, 384 + idx = torch.randperm(rows, device=device, dtype=torch.int32) + idx = torch.cat([idx, torch.full((20,), -1, device=device, dtype=torch.int32)]) + grad = torch.randn(cols, idx.numel(), device=device, dtype=torch.bfloat16).t() + # Every other element of a twice-as-long tensor: same values, stride 2. + idx_strided = torch.stack([idx, torch.zeros_like(idx)], dim=1)[:, 0] + assert not grad.is_contiguous() and not idx_strided.is_contiguous() + + ref = _triton_permute_bwd(grad.contiguous(), idx, rows, cols) + assert torch.equal(_triton_permute_bwd(grad, idx, rows, cols), ref) + assert torch.equal( + _triton_permute_bwd(grad.contiguous(), idx_strided, rows, cols), ref + ) + + +def test_triton_permute_bwd_kernel_out_of_range_indices(): + """Indices outside [0, rows) are skipped like the -1 padding value. + + generate_permute_indices never emits them, but a bad index must not write + outside the output. The kernel writes into rows [guard, guard + rows) of a + larger buffer, so a stray write lands in the guard rows and is detected. + """ + device = "cuda" + rows, cols, guard = 64, 256, 8 + idx = [0, 5, rows, rows + 3, -1, -2, -7, 63] + idx = torch.tensor(idx, device=device, dtype=torch.int32) + grad = torch.randn(idx.numel(), cols, device=device, dtype=torch.bfloat16) + + sentinel = 7.0 + buf = torch.full( + (guard + rows + guard, cols), sentinel, device=device, dtype=torch.bfloat16 + ) + out = buf[guard : guard + rows] + out.zero_() + block_rows, block_cols = 256, 256 + grid = (triton.cdiv(idx.numel(), block_rows), triton.cdiv(cols, block_cols)) + _triton_permute_bwd_kernel[grid]( + grad, + idx, + out, + idx.numel(), + cols, + rows, + cols, + BLOCK_ROWS=block_rows, + BLOCK_COLS=block_cols, + PADDING_VALUE=-1, + ) + + assert (buf[:guard] == sentinel).all(), "wrote before the output" + assert (buf[guard + rows :] == sentinel).all(), "wrote past the output" + valid = (idx >= 0) & (idx < rows) + expected = torch.zeros_like(out) + expected[idx[valid].long()] = grad[valid] + assert torch.equal(out, expected) diff --git a/torchao/prototype/moe_training/ep/permute.py b/torchao/prototype/moe_training/ep/permute.py index b4c93044696..e9b49761abf 100644 --- a/torchao/prototype/moe_training/ep/permute.py +++ b/torchao/prototype/moe_training/ep/permute.py @@ -264,6 +264,9 @@ def _triton_permute_bwd( Returns: grad_input: bf16 gradient tensor (unpermuted) """ + # The kernel addresses rows as `row * cols`, so it needs row-major inputs. + grad_output = grad_output.contiguous() + permuted_indices = permuted_indices.contiguous() grad_rows, grad_cols = grad_output.shape output_buffer = grad_output.new_zeros((original_rows, original_cols)) grid = lambda meta: ( @@ -316,8 +319,12 @@ def _triton_permute_bwd_kernel( other=PADDING_VALUE, ) - write_mask = (dest_rows[:, None] != PADDING_VALUE) & ( - col_offsets[None, :] < original_cols + # Skip padding (-1) and any index outside [0, original_rows), so a bad index + # never writes outside output_buffer. + write_mask = ( + (dest_rows[:, None] >= 0) + & (dest_rows[:, None] < original_rows) + & (col_offsets[None, :] < original_cols) ) tl.store( output_buffer_ptr + dest_rows[:, None] * original_cols + col_offsets[None, :],