Repository navigation
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
6cb9ebc to
d3ae81d
Compare
github-actions
Bot
changed the title
aiter PR #5059 — [JIT] Add gfx1201 BF16 G1U1 selected large-M Triton MoE pathRoutes ten measured Qwen3.6 35B-A3B token counts on gfx1201 to the existing Triton MoE kernels, reaching parity with the vLLM baseline that the current AITER path loses to by 2–2.7x. Review (advisory): Worth stating plainly, because it changes how the rest should be read: this PR adds no kernel. Every kernel it calls —
More importantly, that workflow's triggers are
Under real serving M moves continuously with batch composition, so the fast path will fire far less often than a ten-row table suggests, and the 2–2.7x gap remains for everything in between. 📝 Two silences stack. 📝 📝 |
github-actions
Bot
changed the title
Thanks for the review, this is helpful. Here's the updated fix point-by-point:
Agreed. This was a real gap.
So this test is now both:
Agreed. This was also a real issue.
This behavior is intentional for this first step. The current design here is an exact-match allow-list of validated gfx1201 rows, not a generic bucket/range policy. That is why the lookup key includes the exact token count and returns We agree this is not a general solution for all nearby
We do plan to revisit bucket/range policy in follow-up work, but only after we have broader measurements and test coverage to justify it.
Agreed for We tightened this:
We kept
We reviewed this path. In the Triton helper, stage1 = torch.empty((M * topk, inter_dim), ...)Here The padding only exists in the routing metadata: sorted_token_ids
sorted_expert_ids
So the distinction is:
In other words, there are no extra padded payload rows allocated in |
67f03eb to
bc189b1
Compare
|
No activity for 15 days, so this is now labelled |
|
[verified] This needs to be reworked using the current AITER Triton MoE architecture. The PR currently creates a parallel config loader and manually assembles routing, GEMM1, GEMM2, and reduction inside Please implement the BF16 G1U1 backend under The tests must follow the existing AITER test structure as well. Do not place a Triton-specific test at the top level of |
|
No activity for 15 days, so this is now labelled |
github-actions
Bot
changed the title
github-actions
Bot
changed the title
github-actions
Bot
changed the title
github-actions
Bot
changed the title
keneoneth
changed the title
github-actions
Bot
changed the title
# Conflicts: # csrc/py_itfs_cu/moe_sorting_opus_kernels.cu
github-actions
Bot
changed the title
Motivation
Enable Qwen3.6 35B-A3B BF16 AITER MoE serving on gfx1201/RDNA4 systems such as Navi48/R9700S.
The existing AITER MoE stack could route Qwen3.6 BF16 G1U1 shapes into paths that were unsupported or not comparable on gfx1201. Small decode shapes are handled by the related small-M PR; this PR adds the selected large-M Triton path and makes the existing CK2Stage fallback baseline runnable for comparison.
Related work:
fmoe_g1u1finds no suitable kernel #4492This PR
1024 <= M <= 4096.topk_ids/topk_weight.get_moe_dispatch()style, with exact keys for the launched GEMM shapes.cu_num=64.token=4096,inter_dim=128CK2Stage row needed by the fallback path when larger M values pad to4096.The original gfx1201 BF16 CK2Stage rows used
cu_num=32, which appears to come from PyTorch reporting the Navi48/RDNA4 WGP count. AITER runtime config lookup usesget_cu_num(), which reports64on the tested R9700S/Navi48 system, matching therocminfoCompute Unit count. With rows keyed as32, runtime misses the tuned rows unlessCU_NUM=32is forced manually. A follow-up may be needed in the tuning/config-generation path so future RDNA configs write AITER's runtime CU convention instead of PyTorch's WGP count.The new Triton path is selected only when all of these are true:
gfx120164quant_type == QuantType.Noactivation == ActivationType.Silumodel_dim == 2048inter_dim == 128num_experts == 256topk == 81024 <= M <= 4096w1andw2Unsupported inputs fall back to the existing AITER path.
Test Plan
Focused correctness and dispatch coverage:
Benchmark smoke:
Full correctness validation compared public AITER
fused_moeagainsttorch_moefor every integerM=1024..4096.Full performance validation compared public AITER
fused_moeagainst vLLMfused_expertsfor every integerM=1024..4096withwarmup=10,iters=30.A forced-policy diagnostic sweep compared the Triton large-M implementation against the corrected CK2Stage fallback for every integer
M=17..4096. That diagnostic was used to verify the large-M implementation's range and to show why lower M values are not included in the production policy.Test Results
Latest focused gfx1201 validation at reviewed branch head
84b421d7e72c:Benchmark smoke:
Production interval vs vLLM,
M=1024..4096:Production interval vs corrected CK2Stage fallback,
M=1024..4096:Forced diagnostic sweep vs CK2Stage,
M=17..4096:The CK2Stage-winning rows are below the production gate or in low/transition regions:
17..202, 204..219, 221..235, 237..239, 241..244, 247, 250, 252..255, 270..271, 282, 484. This is why the production policy starts atM=1024.Reference correctness for the new Triton path over the production interval:
Representative per-M rows from the full validation table. The sampled rows include every 32nd production M plus the tested policy boundary/transition rows, including the 4x / 8x / 16x rows
1024,2048, and4096.The full per-M table for
M=17..4096remains in localpr.md; it was not pasted here because the complete table is several thousand rows and too large for a practical GitHub PR body.PR Self-Review
Reviewed revision
84b421d7e72ccf6abf2fad0e58a7019a6e40c4aeBehavior and intent
This PR changes gfx1201 64-CU BF16 G1U1 MoE dispatch for Qwen3.6 35B-A3B large-M rows in the validated
1024 <= M <= 4096range. The old #5059 implementation was no longer compatible with current AITER because it used removed legacy Triton MoE wrappers. This revision provides a current-architecture Triton MoE route for the validated production range while keeping unsupported inputs on the existing AITER path.Validation and unverified scope
Validated on a Radeon AI Pro R9700S /
gfx1201system using a ROCm vLLM container withflydsl==0.2.4installed inside the disposable container. Public-path correctness, dispatch-gate coverage, route-alignment wrapper coverage, CK2Stage fallback config coverage, and benchmark smoke passed.The full correctness sweep over every integer
M=1024..4096found zero runtime errors and zero AITER MoE criterion failures. The full performance sweep over every integerM=1024..4096found zero mean-latency losers against the vLLM rerun baseline. The forced diagnostic sweep overM=17..4096showed CK2Stage can win in low/transition rows below the production gate, while the new Triton path wins or ties every row in1024..4096.Unverified scope: non-gfx1201 devices, non-64-CU gfx1201 devices, quantized MoE, expert-parallel scatter/combine, shuffled weights, non-contiguous weights, alternate gate modes, padded hidden/intermediate dimensions,
inter_dim=256,M < 1024, andM > 4096. These are intentionally outside the gate.PR-specific question
The changed large-M kernel path is selected only when runtime architecture is
gfx1201,get_cu_num() == 64, inputs/weights are BF16, quantization is disabled, activation is Silu, layout is G1U1 with separated gate mode, shape ismodel_dim=2048,inter_dim=128,num_experts=256,topk=8, weights are contiguous and non-shuffled, EP/scatter metadata is absent, padding/bias/scales/local-token overrides are absent, and1024 <= M <= 4096. Unsupported inputs are prevented from reaching this backend by_can_use_bf16_g1u1_gfx1201_triton_large_m(...); otherwise dispatch falls back to the existing AITER path.Changed op tests
op_tests/triton_tests/moe/test_moe_gemm_a16w16.py: covers the publicaiter.fused_moe(...)path for the gfx1201 BF16 G1U1 current-architecture large-M backend, compares against independenttorch_moe, checks caller-ownedoutput=, and validates the1024..4096policy gate plus neighboring rejected rows. This is not redundant with #4740's small-M direct HIP test because this PR uses a different Triton backend, different route-alignment path, and selected large-M dispatch gate.op_tests/triton_tests/moe/moe_routing/test_topk_align.py: covers the route-alignment wrapper used by the new backend. It validates histograms, raw offsets, padded/tile offsets, block maps, inverse indices, sorted weights, empty experts, uneven expert loads, and per-stream scratch isolation. This is not redundant with public MoE correctness because public output checks can pass on a single stream while shared routing metadata remains unsafe for concurrent streams or graph replay.op_tests/test_moe_ck2stage_bf16.py: covers the corrected gfx1201 BF16 CK2Stage fallback config convention and 4096 padded-M coverage, verifies that the Opus sorting source has the GFX10+ safe fallback shape needed by current main, and runs a representative gfx1201 CK2Stage 4096-bucket path. This is not redundant because the large-M PR uses the fallback as a comparison baseline and the config convention bug would otherwise make that baseline unavailable.op_tests/op_benchmarks/triton/bench_moe_gemm_a16w16.py: provides a standalone public-wrapper benchmark for the new Triton MoE backend. This is not a correctness test; it is included so the launchable Triton kernel family has a local benchmark entry point.Submission Checklist