Unify FlyDSL W4A4/G1U0 updates and tuning fixes by rujiacai · Pull Request #2565 · ROCm/aiter

rujiacai · 2026-04-01T03:27:02Z

add G1U0 support for stage1 GEMM path
enable GELU path for W4A4 kernels
fix stage1 selection and GELU behavior
refresh tuned FP4 config and remove obsolete code

Motivation

flydsl w4a4 moe support g1u0

Technical Details

Test Plan

Test Result

Submission Checklist

Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

github-actions · 2026-04-01T03:27:51Z

🏷️ CI Guide

Runs automatically on every PR:

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

Extended tests (opt-in via labels):

Label	Tests
`ci:triton-355`	Run Triton tests on MI355 in addition to MI325
`ci:sglang`	SGLang integration tests
`ci:atom`	ATOM benchmark (DeepSeek-R1 + GPT-OSS)
`ci:vllm`	vLLM benchmark
`ci:all`	All of the above

Add labels via the sidebar or gh pr edit 2565 --add-label <label>

Copilot

Pull request overview

This PR updates the FlyDSL MoE 2-stage (stage1/stage2) FP4 (W4A4) paths to support G1U0 in stage1, enable/propagate GELU activation, and refresh tuning/shuffle behavior for FP4 configs.

Changes:

Add g1u0 plumbing across FlyDSL stage1 compilation/launch paths and tuning/task generation.
Add GELU activation support for FlyDSL stage1 (including test coverage and activation mapping).
Update FP4 (a16w4) stage2 preshuffle/scale shuffling and relax tuner dtype-casting failures.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py	Extends FlyDSL tuning tasks for g1u0, adds activation mapping, updates FP4 stage2 shuffles/scales.
aiter/utility/base_tuner.py	Makes dtype normalization in `tune_summary` more tolerant by casting per-column.
aiter/ops/flydsl/test_flydsl_moe_a4w4.py	Adds CLI/config for `--g1u0` and `--act`, updates reference generation accordingly.
aiter/ops/flydsl/moe_kernels.py	Adds `g1u0` to stage1 compile/cache/launch path and adjusts inter-dim interpretation.
aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage.py	Implements g1u0 layout handling and GELU epilogue logic for stage1 FP4 kernels.
aiter/fused_moe.py	Wires activation mapping to include GELU and forwards `g1u0` into FlyDSL stage1 wrapper.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-01T03:34:26Z

+        tunedf_subset = tunedf[self.untunedf.columns]
+        for col in tunedf_subset.columns:
+            try:
+                tunedf_subset[col] = tunedf_subset[col].astype(self.untunedf[col].dtype)


tunedf_subset = tunedf[self.untunedf.columns] followed by per-column assignment can trigger pandas chained-assignment behavior and may not reliably apply the dtype conversions used for the tuple-based membership test. To make the dtype normalization deterministic, take an explicit copy (e.g., via .loc[:, cols].copy()) and cast using .loc[:, col] = ... (or build a dtype map and call astype(..., errors="ignore")).

Suggested change

tunedf_subset = tunedf[self.untunedf.columns]

for col in tunedf_subset.columns:

try:

tunedf_subset[col] = tunedf_subset[col].astype(self.untunedf[col].dtype)

tunedf_subset = tunedf.loc[:, self.untunedf.columns].copy()

for col in tunedf_subset.columns:

try:

tunedf_subset.loc[:, col] = tunedf_subset[col].astype(self.untunedf[col].dtype)

Copilot · 2026-04-01T03:34:26Z

    token_num = a.shape[0]
    E = w1.shape[0]
-    inter_dim = w1.shape[1] // 2
+    inter_dim = w1.shape[1] if g1u0 else (w1.shape[1] // 2)
    model_dim = a.shape[1]


When g1u0=True, inter_dim is inferred from w1.shape[1], but there’s no validation that the provided w1 matches the selected mode (G1U0 expects w1.shape[1] == inter_dim, G1U1 expects w1.shape[1] == 2*inter_dim) or that act is compatible (e.g., swiglu requires G1U1). Adding explicit shape/parameter checks here would prevent silent misconfiguration producing incorrect outputs.

Copilot · 2026-04-01T03:34:26Z

+                                if g1u0:
+                                    if enable_bias:
+                                        gate_bias_list = epilogue_pf
+                                        vg = vg + gate_bias_list[ni]
+                                    if act == "gelu":
+                                        y = gelu(vg)
+                                    else:
+                                        y = silu(vg)


In the g1u0 epilogue path, any act value other than "gelu" falls back to SiLU, which can silently produce wrong results if callers accidentally pass "swiglu" (or any typo) with g1u0=True. Consider validating act early (compile-time) based on g1u0 (e.g., allow only {"silu","gelu"} for g1u0, and {"silu","gelu","swiglu"} for g1u1) and raising a clear error on unsupported values.

Copilot · 2026-04-01T03:34:27Z

            # per_1x32 fp4 sorting requires block size to be a multiple of 32, so
            # the tile_m=16 FlyDSL candidate is invalid for this tuning path.
-            if blockM not in [32, 64, 128] or not use_g1u1:
+            if blockM not in [32, 64, 128, 256]: # or not use_g1u1:
                continue


The condition includes a commented-out clause (# or not use_g1u1:) that no longer applies and makes it unclear which configurations are intended to be tuned. Please remove the dead comment or replace it with an explanatory comment about why these specific blockM values are required for this tuning path.

Copilot · 2026-04-01T03:34:27Z

    )


 def _check_result(ref_out, test_out, label, atol=1.0, rtol=0.05, pass_pct=95.0):
    """Compare outputs and print result. Returns (passed, max_delta, pct_close)."""
+    print(f"ref_out: {ref_out.shape}, test_out: {test_out.shape}")


This unconditional print(f"ref_out: {ref_out.shape}, test_out: {test_out.shape}") adds noisy output to every comparison and can make test logs hard to read. Consider removing it or guarding it behind an explicit verbose/debug flag.

coderfeli · 2026-04-01T05:58:40Z

+                    * x
+                    * (1.4426950408889634)
+                )
+                emu = llvm.call_intrinsic(f32, "llvm.amdgcn.exp2.f32", [t], [], [])


use rocdl.exp2?

Thanks for the suggestion. At the moment, this path consistently uses llvm.call_intrinsic(...) for these math ops. We can follow up with a cleanup pass to standardize this style and replace it with rocdl.exp2 (and related rocdl wrappers) where equivalent.

…r stage1 GEMM path - enable GELU path for W4A4 kernels - fix stage1 selection and GELU behavior - refresh tuned FP4 config and remove obsolete code Made-with: Cursor

zufayu · 2026-06-29T05:21:36Z

@coderfeli This PR (FlyDSL W4A4/G1U0 updates) looks like it falls in your area. A few things to address before it can merge:

CI is stale (ran April 3, logs expired) — please rebase onto main and re-run CI.
PR description is incomplete — technical details, test plan and test results sections are all empty.

Passing this to you to drive.

zufayu · 2026-07-13T05:29:04Z

Hi @rujiacai — this PR has been idle since early April; the branch is now behind main by a large margin and shows CONFLICTING. The only CI record is from the 2026-04-03 run, where "Standard Tests (1 GPU)" on MI35X shards 2 and 3 failed — but those logs have since expired (GitHub retention), so the specific failing tests are no longer retrievable, and a 3-month-old result against today's main isn't meaningful anyway.

Could you rebase onto the latest origin/main and re-trigger CI? Once it's green (especially the MI35X FP4/W4A4 shards, since that's this PR's target arch), I'll review the fresh results. If MI35X still fails after rebase, please paste the failing test names so we can pin it down.

A couple of code-level points to look at while you're rebasing (independent of the CI):

In gemm_moe_tune.py, the fp4x2 branch now shuffles w2 via shuffle_weight_a16w4 while w1 still uses shuffle_weight((16,16)) + e8m0_shuffle. If that asymmetry is intentional (stage1 vs stage2 kernels expecting different layouts) that's fine — just want to confirm w2's a16w4 layout matches what the stage2 flydsl kernel expects.
In the g1u0 path, the epilogue only recognizes act=="gelu" and otherwise falls back to silu silently; a caller with use_g1u1=False + Swiglu would get silu with no error. Consider asserting the activation is one of {silu, gelu} on the g1u0 path.

…/inferred finding discipline From a full day of real reviews (aiter#3802/#4098/#2565/#4029/#3593/#4171): - grep .cu+.cuh+.h, not just diff files, before claiming missing sync/branch - classify CI failures as infra-flake vs real regression before blaming the PR - tag findings [verified]/[inferred]; never ship an unconfirmed root cause

@valarLip

…error, D10b FlyDSL (#4171) * review-pr skill v2: dispatch gap rule, arch-string FP fix, P5 timing error, D10b FlyDSL Key changes based on 191-PR test sweep (88 aiter + 93 ATOM): - B4 (new): "New dispatch value not handled by all paths, no warning" Highest-frequency finding (~18% of PRs). Covers: new dtype/arch/flag added to a multi-way dispatch where some branches silently fall through to wrong behavior with no assert/warning. Includes FP self-check to filter upstream guards. Old B4 (over-conservative assert) moved to B7. - C4 (new): "New GPU arch string literal in dispatch condition" FP self-check front-loaded: search unchanged lines first; if same arch string already exists → skip (pre-existing style). Fires only on genuinely new introductions. - D10b (new): "FlyDSL arith.bitcast requires arith.unwrap() on operand" Raw DSL value passed to arith.bitcast causes JIT type error at first invocation of that shape/dtype combo. Silent in Python, crashes at kernel JIT time. Real example: aiter#3944 bf16/f16 output path. - P5 (new): "Benchmark timing excludes one-time setup cost" Distinct from P1 (missing numbers). Covers: shuffle_weight, first-call JIT compile, precompile excluded from timing window → claimed speedup is actually a regression. Real examples: aiter#4166 (1.14x claim → 0.83x actual), aiter#3944 (7 precompile kernels on live stream). - Step 3 classification: wired B4/P5/D10b into trigger lines so new constexpr/routing-flag PRs and FlyDSL PRs auto-prompt the right checks. - Step 8 output format: findings now require three parts (Problem + Impact + Action verb). "Author must" / "Reviewer should ask" required; no verb = do not include in output. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * review-pr skill: B6→parametric family, G1b, D1b, C4 FlyDSL exemption - B6 restructured into a 6-row parametric table (param-discard, param-removed, repr-key, arch-discard, dispatch-silent, rename) — unified mental model for all API propagation incompleteness sub-types - G1b: blocking queue.get() without timeout in production serving code - D1b: Python-side UnboundLocalError from conditional variable assignment (ATOM#860 real example; static analysis blind spot) - C4: add FlyDSL capability guard exemption — get_gfx() inside a kernel-specific wrapper is not centralized dispatch hardcoding (aiter#3870) - Step 3 classification: new trigger entries for G1b and D1b Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * review-pr: add cross-file verification + CI-failure triage + verified/inferred finding discipline From a full day of real reviews (aiter#3802/#4098/#2565/#4029/#3593/#4171): - grep .cu+.cuh+.h, not just diff files, before claiming missing sync/branch - classify CI failures as infra-flake vs real regression before blaming the PR - tag findings [verified]/[inferred]; never ship an unconfirmed root cause * review-pr: add E4 (downstream CI coverage / ci:all) + E5 (stable-API owner sign-off) + Step-3 trigger E4: downstream tests (ATOM/SGLang/vLLM) are skip-by-default; a change to a downstream-consumed op can pass aiter CI and break downstream after merge (aiter#3459). Adds dispatch-reachability + model->label triage. E5: behavior/signature/dispatch changes to stable core APIs (fused_moe/mla/attention/quant/jit-core) need de-facto owner sign-off before merge, not a post-merge revert (aiter#3593). * review-pr: add global 🔴 gate — must name concrete triggering input before firing 🔴 Generalizes the per-rule FP self-check (present on B4/C4/B7) into a mandatory gate over all 🔴 findings, closing the gap on rules like D9 that omit one. A 🔴 with no demonstrable triggering shape/scale/input must be downgraded or dropped. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: fix P5 — #4166 was a false positive, not a real example aiter#4166 preshuffles the static weight once outside the timing loop and honestly reports geomean 0.69x; it never claimed 1.14x. The old example taught the amortization error (charging a one-time cost per call). Reframe P5 to fire only on costs that recur per call / per cold start, keep #4166 as a counter-example, and drop the unverifiable #3944 specifics. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: Step 6 — add structural verification for AI-generated code The description-smell table only pre-filters; AI fails structurally. Add six action-forcing checks that each yield a verified/inferred finding: hallucinated -symbol sweep, twin divergence (copy-paste half-adapted), claim/comment↔code + number provenance, safety theater, test-calibrated-to-pass, magic constant without derivation. Motivated by this skill's own fabricated-1.14x P5 bug. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: close Step-3 routing gaps + add FP self-checks + staleness guards ① Route D6/D7 (fake fns), D8 (contiguous), D9 (int32), D4 (invariant reversal), B7 (over-conservative assert), F1 (double HBM) into Step 3 — they were only reachable by scanning all of Section D. ② Add FP self-checks to B2 (tl.load: is the dim provably aligned?) and C2/C3 (search unchanged lines before firing a "new" hardcode), matching C4/D9. ③ Staleness guards on E4 label→model map and P2 production shapes — a stale snapshot produces confidently-wrong recommendations. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: fix B6 (API-modification rule) — route it + widen public-API scope B6 (API propagation) was the rule most about API changes yet had NO Step-3 entry, reachable only by scanning Section B. Route it via a new "API signature change" row + the refactor/rename row. Also fix param-removed/rename scope: a public-op param removal/rename breaks cross-repo consumers (ATOM/SGLang/vLLM), not just same-file call sites — widened, and linked to E1/E5 for public contracts. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: add FP self-check to B6 rename/param-removed (found by live test on #4227) Live-testing the skill on aiter#4227 showed B6's rename/param-removed rows had no own FP self-check: a rename shielded by a same-named wrapper + fc_name would read as a 🔴 break if followed verbatim (only the global 🔴 gate rescued it). Add a self-check to confirm no compatibility shim (same-named wrapper / alias / fc_name binding pin) before firing, with #4227 as the real non-example. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: E5 — reviewer must proactively @-mention the API owner, not wait Strengthen E5's action from passive "get sign-off" to active: the reviewer must post a PR comment @-mentioning the de-facto owner (@valarLip for fused_moe) with a one-line summary + explicit approval request before merge. Finding is unresolved until the owner has been actively pinged and responded. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * review-pr: reflect that aiter rejects permanent env-var knobs (D2 scope + HK9) aiter maintainers generally reject new permanent env-var flags: AITER_MOE_FORCE_BF16_ACT (#3593) was reverted by #4225. Reconcile with D2 by scoping D2 to a *temporary rollback kill-switch* only, and upgrade HK9 from 📝 "document it" to ⚠️ "prefer auto-derivation; permanent knobs face pushback". ATOM unchanged — it manages env vars via atom/envs.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: root <root@hjbog-srdc-24.amd.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

github-actions · 2026-10-07T04:24:53Z

No activity for 15 days, so this is now labelled stale. A push or a comment clears it; keep-open exempts it.

zufayu · 2026-10-09T11:42:57Z

Closing due to inactivity: no pushes, reviews, or comments for 88 days. This is housekeeping, not a judgement on the change — if you want to continue this work, comment here and we'll reopen, or open a fresh PR rebased on the latest main.

rujiacai requested review from a team and Copilot April 1, 2026 03:27

Copilot started reviewing on behalf of rujiacai April 1, 2026 03:28 View session

Copilot AI reviewed Apr 1, 2026

View reviewed changes

coderfeli reviewed Apr 1, 2026

View reviewed changes

rujiacai force-pushed the moe_flydsl_w4a4_pr branch 2 times, most recently from 73c06e6 to c0fa653 Compare April 3, 2026 05:53

Unify FlyDSL W4A4/G1U0 updates and tuning fixes - add G1U0 support fo…

b27a715

…r stage1 GEMM path - enable GELU path for W4A4 kernels - fix stage1 selection and GELU behavior - refresh tuned FP4 config and remove obsolete code Made-with: Cursor

rujiacai force-pushed the moe_flydsl_w4a4_pr branch from c0fa653 to b27a715 Compare April 3, 2026 06:11

zufayu requested review from yzhou103 and zufayu and removed request for yzhou103 and zufayu June 26, 2026 08:58

zufayu self-requested a review June 30, 2026 02:50

zufayu removed their request for review July 13, 2026 05:03

github-actions Bot added the stale The PR hasn't been updated in 2 weeks label Oct 7, 2026

zufayu closed this Oct 9, 2026

Navigation Menu

Repository navigation

Unify FlyDSL W4A4/G1U0 updates and tuning fixes - #2565

Unify FlyDSL W4A4/G1U0 updates and tuning fixes#2565
rujiacai wants to merge 1 commit into
ROCm:mainfrom
rujiacai:moe_flydsl_w4a4_pr

rujiacai commented Apr 1, 2026

Uh oh!

github-actions Bot commented Apr 1, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Apr 1, 2026

Uh oh!

Copilot AI Apr 1, 2026

Uh oh!

Copilot AI Apr 1, 2026

Uh oh!

Copilot AI Apr 1, 2026

Uh oh!

Copilot AI Apr 1, 2026

Uh oh!

coderfeli Apr 1, 2026

Uh oh!

rujiacai Apr 3, 2026

Uh oh!

zufayu commented Jun 29, 2026

Uh oh!

zufayu commented Jul 13, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

zufayu commented Oct 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

rujiacai commented Apr 1, 2026

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

Uh oh!

github-actions Bot commented Apr 1, 2026

🏷️ CI Guide

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

coderfeli Apr 1, 2026

Choose a reason for hiding this comment

Uh oh!

rujiacai Apr 3, 2026

Choose a reason for hiding this comment

Uh oh!

zufayu commented Jun 29, 2026

Uh oh!

zufayu commented Jul 13, 2026

Uh oh!

github-actions Bot commented Oct 7, 2026

Uh oh!

zufayu commented Oct 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants