Problem Description
Summary
We are seeing performance issues with the latest AITER FlyDSL MXFP4 MoE pipeline:
- Large-shape MXFP4 MoE has low MFU compared with the expected kernel efficiency.
- Small-token cases show long latency, which makes the pipeline less competitive for decode / small-batch workloads.
- The current implementation is functional, but needs tuning across token regimes and representative MoE shapes.
Context
Pipeline: latest MXFP4 MoE path in AITER using FlyDSL mixed MoE GEMM 2-stage kernels.
Target behavior:
- Support MXFP4 MoE efficiently for both GEMM1 and GEMM2.
- Keep the current scale shuffle / weight shuffle layout compatible with the AITER MXFP4 path.
- Improve both large-token throughput and small-token latency.
Observed Problem
For large shapes, MFU is lower than expected. This likely indicates that the current tile / pipeline configuration is not fully utilizing the hardware.
For small token counts, latency is high. This may be caused by launch overhead sensitivity, suboptimal tile selection, insufficient specialization for skinny M, or pipeline parameters tuned more toward large-M throughput.
Expected Behavior
The FlyDSL MXFP4 MoE path should have:
- Better MFU on large prefill-style shapes.
- Lower latency on small-token / decode-style shapes.
- Tuned configs for representative MXFP4 MoE GEMM1 and GEMM2 shapes.
- Stable performance across the token sweep used by
test_moe_2stage.py.
Requested Work
Please help tune the latest AITER FlyDSL MXFP4 MoE pipeline, especially:
- Tile configs for large-shape MFU.
- Small-token latency-sensitive configs.
- GEMM1 and GEMM2 separately if needed.
- Stage / preload / async-copy / persistent scheduling parameters.
- Any required shape-specific dispatch or heuristic changes.
Target Shapes
Validation
A good validation target would be:
test_moe_2stage.py passes correctness.- Performance is checked across the full token sweep on gfx950.
- MFU improves for large shapes.
- Latency improves for small token counts.
Notes
This is not a correctness issue. The main problem is performance tuning for the latest MXFP4 MoE pipeline in AITER.
GPU
gfx950
ROCm Version
7.2
Tuning Suggestions
Support Slice K which make different waves to calc different K and reduce in LDS.
Add Weight loading through LDS and enable 2stages, 3stages pipeling.
Make sure no lds bank conflicts for all shapes.
Tune pipelines.
Problem Description
Summary
We are seeing performance issues with the latest AITER FlyDSL MXFP4 MoE pipeline:
Context
Pipeline: latest MXFP4 MoE path in AITER using FlyDSL mixed MoE GEMM 2-stage kernels.
Target behavior:
Observed Problem
For large shapes, MFU is lower than expected. This likely indicates that the current tile / pipeline configuration is not fully utilizing the hardware.
For small token counts, latency is high. This may be caused by launch overhead sensitivity, suboptimal tile selection, insufficient specialization for skinny M, or pipeline parameters tuned more toward large-M throughput.
Expected Behavior
The FlyDSL MXFP4 MoE path should have:
test_moe_2stage.py.Requested Work
Please help tune the latest AITER FlyDSL MXFP4 MoE pipeline, especially:
Target Shapes
Validation
A good validation target would be:
test_moe_2stage.pypasses correctness.Notes
This is not a correctness issue. The main problem is performance tuning for the latest MXFP4 MoE pipeline in AITER.
GPU
gfx950
ROCm Version
7.2
Tuning Suggestions
Support Slice K which make different waves to calc different K and reduce in LDS.
Add Weight loading through LDS and enable 2stages, 3stages pipeling.
Make sure no lds bank conflicts for all shapes.
Tune pipelines.