[Bugfix][ROCm][MoE] Fix mxfp4 oracle regressions from #37128

[AndreasKaratzas]

Fixes several issues introduced by #37128 that broke gpt-oss on ROCm.

• Restore gfx950 gate for CK mxfp4 backend selection. The old code only picked CK on gfx950 via on_gfx950(), the refactor dropped this and let CK get selected on gfx942 where it crashes.
• Restore CK_MXFP4_MOE_DIM_ALIGNMENT (256) check. Models with intermediate_size not aligned to 256 (like gpt-oss-20b at 2880) hit a reshape error in aiter shuffle_scale_a16w4. Added is_supported_config to AiterExperts so the backend selector falls through to Triton.
• Restore hidden_pad/intermediate_pad for CK path. These were passed to rocm_aiter_ops.fused_moe() in the old code but got lost in the refactor. Added fields to FusedMoEQuantConfig and wired them through.

Tested on MI325X (gfx942):

• test_gpt_oss_speculative_reasoning_leakage passes
• GPQA eval via gpt_oss.evals: 56.76% (1584 questions, effort=low)
• Backend correctly falls back to Triton on non-gfx950

Related:

• [Bigfix]fix lora test by pass padded size back to the layer #37811
• [XPU][MoE Refactor] Refactor xpu mxfp4 support into oracle #37784
• Revert "[MoE Refactor] Mxfp4 oracle rebased" (#37128) #37786
• [MoE Refactor] Mxfp4 oracle rebased #37128

cc @kenroche

[gemini-code-assist]

Code Review

This pull request addresses several regressions related to MoE with mxfp4 quantization, primarily affecting ROCm platforms, which were introduced in a recent refactoring. The fixes include restoring platform-specific checks for the CK backend, correctly handling dimension padding, and resolving an issue with LoRA on NVIDIA. My review identifies a critical contradiction in one of the changes: while the PR description claims to enable mxfp4 LoRA on ROCm, the code continues to raise a NotImplementedError, albeit with a more descriptive message. This discrepancy needs to be resolved.

[mergify]

Hi @AndreasKaratzas, the pre-commit checks have failed. Please run:

uv pip install pre-commit>=4.5.1
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[AndreasKaratzas]

Putting this PR in draft mode, as CI regression seems not to be addressed by the PR. The most straight-forward solution probably is to revert the problematic PR.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @AndreasKaratzas.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[BowenBao]

Please share if there's a discussion with upstream folks. #37128 broke CI, and rather than reverting it to figure out a proper solution, there's pressure to quickly land this PR. The size -> shape change and padding logic feel like quick(hacky) band-aids that will require the ROCm team to do more follow-up work down the line.

I don't feel this aligns with OSS best practices. However, if the maintainers are aware of this and have agreed it's the best way forward, I have nothing to add. cc @tjtanaa, @gshtras, @dllehr-amd, @ChuanLi1101

[AndreasKaratzas]

Please share if there's a discussion with upstream folks. #37128 broke CI, and rather than reverting it to figure out a proper solution, there's pressure to quickly land this PR. The size -> shape change and padding logic feel like quick(hacky) band-aids that will require the ROCm team to do more follow-up work down the line.

I don't feel this aligns with OSS best practices. However, if the maintainers are aware of this and have agreed it's the best way forward, I have nothing to add. cc @tjtanaa, @gshtras, @dllehr-amd, @ChuanLi1101

The .size() in to .shape change isn't a band-aid, it's the correct fix. triton_kernels.tensor.Tensor (used for MXFP4 swizzled weights since #37128) exposes .shape but not .size() or .dim(). .shape is the common interface that works with both tensor types and is the more Pythonic/numpy-standard accessor anyway. There's no follow-up work needed here.

On the padding logic, I hear the concern about it being incremental. The CK MXFP4 kernels on gfx950 require 256-byte aligned dimensions, and the padding was already happening in create_weights but wasn't being properly communicated to the layer/config.

[gshtras]

This appears to correctly fix the fallback to Triton caused by #35893 on gpt-oss-120b
lm_eval scores are back to normal when using CK MoE

[gshtras]

Nit. not is_cuda doesn't automatically mean rocm. Other platforms may get confused by this message.

[AndreasKaratzas]

Done :)