[Triton/Gluon] [OPUS] [FlyDSL] Gfx1250/microbench

[JiaoliangYu]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/TheRock/blob/main/GOVERNANCE.md#pull-requests.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5391 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[valarLip]

Reviewed commit 4a896606efab1ab23c2bdb50ff5baea0cc3b8929. I recommend addressing the following before merging, particularly the measurement accuracy and separation of responsibilities.

1. [P2] SMI can monitor the wrong GPU. In smi_monitor.py:405, the device defaults to AITER_SMI_DEVICE=0, while perftest does not pass the workload device. With all devices visible and torch.cuda.set_device(3), the workload executes on GPU 3 but monitoring selects GPU 0. I reproduced this using a real GPU 3 tensor and a monitor spy. Please pass the workload device explicitly, retaining the environment override if needed. The distributed Mega MoE path already passes current_device() correctly.

2. [P2] run_perftest loses callable identity, breaking SMI filtering and labels. The new hook in test_common.py:218 filters by function name, but run_perftest wraps every callable in worker(*args, **kwargs). With AITER_SMI_SKIP_FUNCTIONS=run_torch, run_perftest(run_torch, ...) still performs the replay and emits benchmark_case/worker#1; signature-based tagging also loses backend selectors in the variadic wrapper. Please preserve the original callable metadata/signature or pass explicit backend identifiers. This was reproduced; direct @perftest usage correctly honors the filter.

3. [P2] GEMM1 bandwidth accounting omits top-k output multiplicity. In test_flydsl_grouped_gemm.py:590, FLOPs include topk, but output bytes use only tokens * stage1_n * 2. Under the documented logical GEMM model, the output has tokens * topk rows. Executing the helper confirms that GEMM1 bytes remain unchanged when top-k increases from 1 to 8. Please correct the routed-row count. If the intention is to model actual fused-kernel traffic instead, activation and quantized output formats need separate accounting; logical bytes should not be interpreted as measured HBM traffic.

4. [P3] The data-init CLI accepts empty inputs and silently discards requested distributions. The shared parser uses nargs="*", allowing a bare --data-init to produce an empty list. test_opus_a16w16_gemm.py:539 then indexes the first element and raises IndexError, while entry points that iterate the list can execute zero cases. Conversely, --data-init uniform norm silently runs only uniform in OPUS. Please reject empty lists and either honor every distribution or explicitly accept a single value.

The code organization also needs attention:

• Keep the Mega MoE implementations in separate modules. test_mega_moe.py:17 combines two complete runners in a 2,024-line file under architecture branches, retaining separate CLIs, reference implementations, communication setup/teardown, and reporting, with formatting disabled. A thin architecture dispatcher would provide a unified entry point while preserving backend boundaries.
• Keep test_common.py focused. Input generation and packed-format encoding, JSON reporting, and SMI metadata/state should have dedicated modules with narrow interfaces, leaving a small integration hook in the timing utility. The separate _make_data implementation and different CLI vocabulary added to the QK norm test should also be reconciled with the shared helpers.

Validation in ljin_dev on MI355X/gfx950, using the reviewed source snapshot: 114 passed, 10 skipped for the new MLA test file, including both FP8 cases exercising strides beyond 2 GiB. Another 27 input-generation determinism checks passed. The SMI findings were reproduced with actual GPU work and a monitor spy, not measured power/utilization data. gfx1250 hardware was unavailable, so these results do not validate gfx1250 FlyDSL/OPUS performance; the full distributed Mega MoE benchmark was not executed.

The shared seeded generators, SWA sentinel checks, and large-stride regression tests are useful additions. They can be retained while correcting the measurement semantics and tightening module boundaries.

[JiaoliangYu]

Thanks for the detailed review. Addressed in 5fff12bd4 and 7422c7770:

1. SMI device: regular UTs and SMI both default to logical HIP device 0. --smi-device explicitly selects the monitored logical device, and the BDF mapping handles visibility remapping. Mega MoE already uses each rank's current device. A mismatch occurs only if the workload device is changed without updating --smi-device, so we kept the existing explicit contract.
2. run_perftest: fixed with functools.wraps, preserving names, signatures, filters, and labels.
3. Grouped GEMM bytes: fixed GEMM1 logical input/output accounting to use tokens * topk routed rows.
4. Data-init CLI: changed multi-value options to require at least one value; OPUS now rejects multiple distributions instead of silently taking the first.
5. Mega MoE organization: restored test_mega_moe_v2.py and moved the end-to-end benchmark intact to bench_mega_moe.py, removing the combined architecture-branch file.
6. test_common.py: moved data initialization and packed-format encoding to aiter/benchmark_data_init.py, JSON output to aiter/benchmark_reporting.py, and SMI metadata/state to aiter/smi_monitor.py. The timing utility now keeps only the small hooks needed to establish benchmark label context and hand off prepared replays. QK norm also uses the shared data helper and --data-init/--seed vocabulary.