Repository navigation
[Triton/Gluon] Add a16w4 MoE tuned dispatch and gfx942 DSv4.1 tables - #5667
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
There was a problem hiding this comment.
🟡 Changes recommended
Add gfx942 test coverage and validate or constrain the MI300X rollout.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds gfx942-specific, per-shape Triton dispatch tuning for A16W4 MoE GEMMs while preserving fallback behavior.
Changes:
- Adds architecture- and shape-specific dispatch lookup.
- Adds eight gfx942 tuning configurations.
- Preserves defaults for unmatched shapes and architectures.
File summaries
| File | Summary |
|---|---|
aiter/ops/triton/moe/moe_op_gemm_a16w4.py |
Applies tuned dispatch entries with fallback defaults. Moderate finding (3 votes): gfx942 coverage is skipped by existing CI. Moderate finding (1 vote): MI300X applicability is unvalidated. |
aiter/ops/triton/configs/gfx942/triton/moe/a16w4/DEFAULT.json |
Defines eight gfx942 Triton tuning entries. |
Review details
Suppressed comments (1)
aiter/ops/triton/moe/moe_op_gemm_a16w4.py:98
- [inferred] The lookup keys only on
get_arch(), so this table applies to every gfx942 device, while the supplied measurements cover only MI325X. Because MI300X shares gfx942 but has different memory bandwidth, these tiles may become a throughput regression there even though the path is enabled unconditionally. Author must validate the table on MI300X or provide a justified device-specific rollout before applying it to all gfx942 devices.
tuned = get_moe_dispatch("A16W4", get_arch(), "triton").get(
f"bm{block_m}_n{n}_k{k}"
)
- Files reviewed: 2/2 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
c411de2 to
147a54c
Compare
|
Pushed a fix after the last CI run came back with 11 failures, all of them in the new resolver test. Tracing them led to a gap in shared moe_config_utils: get_moe_dispatch() takes an arch argument but never forwards it to resolve_config_dir(), so arch only keyed the lru_cache while the table kept resolving against whatever arch was actually running. The test patches get_arch to gfx942 so it can run anywhere, but that patch was effectively ignored so on my MI325X the running arch is gfx942 so it passed locally, while on CI's non-gfx942 runner no a16w4 table was found, the lookup returned {}, and all eight measured tiles fell back to stock defaults. Forwarding the arch fixes it |
|
Hi @Boss2002n Could you please take a look when you get a chance? This is ready for review and I’d like to get it merged soon as currently this moe backend is performing worse than trition unfused on gfx942. I have tried to follow the existing A8W4 dispatch pattern and this gives a sizeable+14.9% E2E output throughput as measured on MI325X. Thanks! I saw that it was added to some board but is it so that the items are reviewed in some order? I am relatively newbie to AITER PRs so dont have full context on the board so decided to send a ping 😅 |
|
@vgokhale @brunomazzottiamd can you review this? vLLM's |
Review: PR #5667 — a16w4 MoE GEMM gfx942 tuned configThe config table and wrapper wiring are good: My concern is What actually needs testing hereThere is one real gap: the existing numerics test ( What the test actually doesOf the six test functions, only two are about this config:
The other four are testing the shared dispatch framework, not the a16w4 table:
Roughly half the file ( No other kernel does this
So this introduces a bespoke per-kernel config-test pattern with no precedent, and inconsistently: if a16w4 needs it, a4w4/a8w4 need it too. Either the framework is trusted (this is unnecessary) or it needs testing once, centrally (this is the wrong place). Suggested direction
Net: the table is fine; the test is testing shared plumbing through one kernel's keyhole. Trim it to the a16w4-specific check and move the dispatch-contract coverage to the framework, so we don't set a heavy per-kernel precedent for something the numerics tests already cover everywhere else. |
|
Instead of standalone UT we can just extend the current a16w4 UT to cover gfx942. If you comment out the is_fp4_available skip and run the a16w4 UTs on gfx942 does everything pass? |
|
Thanks a lot for the comments and feedbacks, now pushed
Since there was a lot of feedback 😅 , showing a quick action points from me on each of these
Happy to expand on any of them more or further refine this PR! @lburzawa one small heads up though! The existing shapes never hit the tuned N,K or |
|
@akii96 Thanks for the changes. The tuned table can become large with time so running every case might result in long CI times in the future. Instead please add the shapes you care about to existing a16w4 UTs and remove the skip for gfx942. To calculate M to use in test cases based on block_m use this formula: block_m = ceil(M * num_experts_active / num_experts_total). |
|
@lburzawa Moved the 8 shapes into the existing test as normal cases and turned the gfx942 skip off, so the number of cases stays fixed as the table grows. Thanks again for guiding with the UTs! I will track the CI and see if anything pops up and address it immediately! Edit: One small note rather than removing the arch check I widened it to is_fp4_avail() or gfx942, since dropping it entirely would also run archs the kernel doesn't support. Happy to remove it outright if you'd prefer. |
|
The red shards are unrelated to this PR. Every failure is in test_pa_mqa_logits_mxfp4.py and test_rope_mxfp4.py, which both go through paged_mxfp4_mqa_logits. That kernel asserts get_arch() == "gfx950", so it fails on the MI300X runners, and those shards are red on main too. Our test is in shard 6, which passed on both MI300X and MI35X. |
e047bc5 to
84b55ce
Compare
|
@Boss2002n FYI, the remaining MI300X failures come from main and not this PR. Once the fixes in #5958 and #5984 land, CI should clear up. Thanks for tracking this PR 🙏 |
1f541cf to
4ca3a30
Compare
get_kernel_config_triton hardcodes block_k=256/block_n=512/num_stages=1
for every arch. That fits CDNA4's 160 KiB of LDS, not CDNA3's 64 KiB, so
the gfx942 prefill tile runs at very low occupancy.
Consult get_moe_dispatch("A16W4", ...) before those defaults, as
moe_op_gemm_a8w4 does, and ship a gfx942 table for DeepSeek-V4.1-Flash at
TP4 and TP8. Arches with no table get {} and are unaffected; on gfx942
only the eight listed shapes change. Decode is left alone: one
(block_m, N, K) key would cover 1-1024 tokens, where the stock path is
m-adaptive and the measured optimum shifts across the range.
test_moe_gemm_a16w4.py skips unless arch_info.is_fp4_avail(), which is
gfx950/gfx1250 only, so it cannot cover this table. Add a resolver test
that pins the measured tiles independently of the shipped JSON and runs
on any arch.
MI325X: 3.67-4.22x on the MoE GEMMs at block_m=128, +14.9% output
throughput and -19.6% TPOT end to end (65536 in / 512 out, conc 8, TP4).
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
get_moe_dispatch takes an arch argument but drops it: the directory is
built by resolve_config_dir() from the *running* arch, so arch only keys
the lru_cache. The parameter therefore promises a control it does not
deliver, and a caller that resolves a table for an arch other than the
one it is running on silently gets that other arch's table, or {}.
resolve_config_dir() already accepts an arch override and
mhc_config_utils passes it; pass it here too.
This is a no-op in production: all five call sites (A16W4, A8W4 triton
and gluon, A4W4) already pass the running arch, so the resolved
directory is unchanged. Only callers that deliberately override see a
difference.
It does fix CI for test_moe_gemm_a16w4_config.py. Those tests patch
get_arch to gfx942 so the gfx942 table resolves on any host; with arch
dropped the patch moved the cache key only, the lookup fell back to the
runner's arch, found no a16w4 table and returned {}, and all eight
measured tiles read back as the stock defaults. It passed on an MI325X,
where the running arch happens to be the one being asserted.
Also clear the lru_cache around the fixture and the untuned-arch probe
in _stock_config: get_moe_dispatch is cached on
(config_name, arch, backend), so without it the suite passes in
declaration order but not under reordering.
Verified on MI325X (gfx942): 25 passed on the real host, and 25 passed
with arch_info.get_arch faked to a non-gfx942 value, which reproduced
the 11 CI failures test-for-test beforehand.
The standalone resolver test asserted shared dispatch behaviour through one kernel's keyhole, which no other tuned kernel does. Replace it with a single numerics test in the existing file, and revert the shared-code change it needed. test_op skips unless is_fp4_avail() (gfx950/gfx1250), so nothing exercised the gfx942 table. The kernel carries no fp4 gate and runs on gfx942, so the entries can be covered there directly: one case per shipped key, parametrized off the JSON so the test cannot drift from the table. Each case asserts selection as well as numerics. A key no real shape produces would fall back to the stock tile and still compute the right answer, so numerics alone cannot catch a dead entry. The table load is required rather than optional for the same reason: a missing file would collect zero cases and report green. get_moe_dispatch no longer forwards arch to resolve_config_dir. Every caller passes the running arch, so it was a no-op at runtime that existed only for the deleted test's monkeypatch. Verified on MI325X (gfx942, ROCm 6.4.2): 8 passed; worst max_err 0.0070 against maxtol 0.4, worst rms_err 0.00002 against rmstol 0.04. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
block_n=256 with block_k=256 needs 67584 B of LDS at block_m=64 (32 KiB A + 32 KiB packed B + 2 KiB mx scales), over gfx942's 65536 B limit, so bm64_n1024_k5120 and bm64_n1536_k5120 failed to launch with OutOfResources. They were measured under triton 3.7.1, which accepted them; 3.8.0 is what CI installs and it rejects the launch. Re-tuned both under 3.8.0 on the same MI325X, sweeping only tiles that fit: bm64_n1024_k5120 bn=256 bk=256 -> bn=128 bk=256 nkd=16 wpe=2 (50176 B) bm64_n1536_k5120 bn=256 bk=256 -> bn=128 bk=128 nkd=16 wpe=2 (25088 B) Both are the fastest fitting tile for their shape. The other six entries are unchanged. Verified on MI325X (gfx942) under triton 3.8.0: 8 passed. Worst max_err 0.0070 against maxtol 0.4, worst rms_err 0.00002 against rmstol 0.04. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
The tuned entries were covered by a dedicated test parametrized off the shipped JSON, so the case count grew with the table. Move the eight shapes into test_op as plain Case rows and widen the arch gate so gfx942 runs them; the count is now fixed. gfx1200/gfx1201 stay skipped, the kernel has no support there. Keep one selection check: a stock tile computes the same answer as a tuned one, so numerics alone cannot tell whether an entry was reached. A deliberately wrong tile passes all 48 numerics variants of a case. Verified on gfx942 under triton 3.8.0: 1153 passed, 576 skipped. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
4ca3a30 to
44d76da
Compare
|
Since this PR is time-sensitive for us, I took a look at the intermittent MI300X Shard 7 failure (test_batched_gemm_a8w8 memory access fault, also seen on main) and put up a 2 line small fix in #6082. Happy to adjust it if needed. Once it lands, I'll update this branch from main. cc @Boss2002n |
@akii96 Thank you 😃 |
Cover block_m=128 and block_m=64 with the n1536/k5120 shape instead of all eight table entries, cutting the added test_op cases from 384 to 96. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
|
@Boss2002n Trimmed it to 2 cases, one per tuned block_m (128 and 64), in b59cdb0. The new commit dismissed the earlier approvals however. Once the CI passes could you help get this merged ? Thanks so much! |
Summary
get_kernel_config_tritonhardcodesblock_k=256,block_n=512andnum_stages=1for every architecture. Those suit CDNA4's 160 KiB of LDS; CDNA3 has 64 KiB, so the gfx942 prefill tile runs at very low occupancy.This routes the kernel through the tuned dispatch that
a8w4already uses and ships a gfx942 table for DeepSeek-V4.1-Flash at TP4 and TP8.get_moe_dispatch("A16W4", get_arch(), "triton"), keyed bybm<block_m>_n<N>_k<K>configs/gfx942/triton/moe/a16w4/DEFAULT.json, 8 entriesget_moe_dispatchreturns{}when no file is shipped, so untuned architectures are unaffected. On gfx942 only the 8 listed prefill shapes change; decode (block_m16/32) and unlisted shapes keep the existing defaults.Results
MI325X (gfx942), ROCm 7.2.3, DeepSeek-V4.1-Flash, E=384, topk=6.
bm128_n1536_k5120bm128_n5120_k768bm128_n1024_k5120bm128_n5120_k512bm64_n1536_k5120bm64_n5120_k768bm64_n1024_k5120bm64_n5120_k512End to end via vLLM (
--moe-backend aiter_triton_mxfp4_bf16, random 65536 in / 512 out, concurrency 8, TP=4, 3 repeats, CV 0.14%):Notes on Tuning
No a16w4 tile tuner exists in-tree, so the entries come from a sweep over all six knobs the kernel forwards to the launch.
Worth knowing when adding shapes:
block_k=64compiles only withmatrix_instr_nonkdim=32, and it wins atblock_m=128on most shapes. Sweeping withnonkdimfixed hides it entirely.block_n=256won on every shape; the stock 512 never did.Tests
block_kchanging the K-split.block_n,num_stagesandnum_warpsare bit-identical.Limits