Skip to content

Run EP permute tests on ROCm gfx950 - #4987

Open
indianspeedster wants to merge 1 commit into
pytorch:mainfrom
indianspeedster:ep-permute-tests-gfx950
Open

indianspeedster wants to merge 1 commit into
pytorch:mainfrom
indianspeedster:ep-permute-tests-gfx950

Conversation

@indianspeedster

@indianspeedster indianspeedster commented Oct 8, 2026 •

Copy link
Copy Markdown

Summary

test/prototype/moe_training/ep/test_permute.py and ep/test_kernels.py are
gated to CUDA 12.8+ with SM >= 100, so they never run on ROCm. The code they test
(generate_permute_indices, _triton_permute_bwd, and the MXFP8 permute) is
Triton and also works on gfx950. This allows MI350-series (gfx950) as well,
using the existing torchao.utils.is_MI350(). Test-only change.

Split out of #4623 as requested in review.

Test plan

pytest test/prototype/moe_training/ep/test_permute.py test/prototype/moe_training/ep/test_kernels.py:

main this PR
MI355X (gfx950) 2 modules skipped 10 passed
B200 (SM100) 10 passed 10 passed

The MI355X runs used a local environment where the flydsl package didn't match
the repo, so flydsl was hidden and torchao took its non-flydsl path.

CI

No upstream CI runner has SM100, so on main these files run in no CI job. With
this PR they run in the ROCm regression job (regression_test_rocm.yml, which
runs pytest test on the MI350 runner when triggered by ciflow/rocm or a push
to main).

MI355X: torch 2.15.0.dev+rocm7.14, triton 3.8.0. B200: torch 2.14.1+cu130,
triton 3.8.0.


Authored with assistance from an AI coding agent.

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().
@pytorch-bot

pytorch-bot Bot commented Oct 8, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Oct 8, 2026
@pytorch-bot

pytorch-bot Bot commented Oct 8, 2026

Copy link
Copy Markdown

The following ciflow label(s) have been added but CI has not been triggered yet because the workflows are awaiting approval:

  • ciflow/rocm

Once a maintainer approves the workflows (scroll to the bottom of the PR page), the corresponding CI jobs will be triggered automatically. Please ping one of the reviewers if you do not have access to approve and run workflows.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/rocm CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. module: rocm

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant