From 96fc92d5a7f575df2045e8c71decb7c01b2ec124 Mon Sep 17 00:00:00 2001 From: indianspeedster <38083203+indianspeedster@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:56:03 +0000 Subject: [PATCH] Run EP permute tests on ROCm gfx950 test/prototype/moe_training/ep/test_permute.py and ep/test_kernels.py were gated to CUDA 12.8+ with SM >= 100, so they never ran on ROCm. The kernels they test (generate_permute_indices, _triton_permute_bwd, and the MXFP8 permute) are Triton and also work on gfx950. Allow MI350-series (gfx950) as well, using the existing torchao.utils.is_MI350(). --- test/prototype/moe_training/ep/test_kernels.py | 10 ++++++---- test/prototype/moe_training/ep/test_permute.py | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/test/prototype/moe_training/ep/test_kernels.py b/test/prototype/moe_training/ep/test_kernels.py index 97ad73dda89..ca2db377b84 100644 --- a/test/prototype/moe_training/ep/test_kernels.py +++ b/test/prototype/moe_training/ep/test_kernels.py @@ -7,14 +7,16 @@ import pytest import torch -from torchao.utils import is_cuda_version_at_least, is_sm_at_least_100 +from torchao.utils import is_cuda_version_at_least, is_MI350, is_sm_at_least_100 if not ( torch.cuda.is_available() - and is_sm_at_least_100() - and is_cuda_version_at_least(12, 8) + and ((is_sm_at_least_100() and is_cuda_version_at_least(12, 8)) or is_MI350()) ): - pytest.skip("Test requires CUDA 12.8+ with SM >= 100", allow_module_level=True) + pytest.skip( + "Test requires CUDA 12.8+ with SM >= 100, or ROCm gfx950", + 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 diff --git a/test/prototype/moe_training/ep/test_permute.py b/test/prototype/moe_training/ep/test_permute.py index 601270658bb..ccc6ad44f09 100644 --- a/test/prototype/moe_training/ep/test_permute.py +++ b/test/prototype/moe_training/ep/test_permute.py @@ -1,14 +1,16 @@ import pytest import torch -from torchao.utils import is_cuda_version_at_least, is_sm_at_least_100 +from torchao.utils import is_cuda_version_at_least, is_MI350, is_sm_at_least_100 if not ( torch.cuda.is_available() - and is_sm_at_least_100() - and is_cuda_version_at_least(12, 8) + and ((is_sm_at_least_100() and is_cuda_version_at_least(12, 8)) or is_MI350()) ): - pytest.skip("Test requires CUDA 12.8+ with SM >= 100", allow_module_level=True) + pytest.skip( + "Test requires CUDA 12.8+ with SM >= 100, or ROCm gfx950", + allow_module_level=True, + ) from torchao.prototype.moe_training.ep import permute_mxfp8_fwd_hp_bwd from torchao.prototype.moe_training.ep.permute import permute_and_pad