Repository navigation
[FlyDSL] fix(flydsl): support kmk3 ep moe & opt for ep moe gemm2 - #5804
Merged
Merged
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags & labels: |
github-actions
Bot
changed the title
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Add regression coverage for dynamic-token EP GEMM2 atomic writes.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 1
Open (1)
What changed in this PR
Fixes FlyDSL EP GEMM2 atomic writes by bounding them to valid tokens, with dispatch/AOT propagation and Kimi K3 configurations.
Changes:
- Adds runtime valid-token write limits.
- Propagates and caches kernel variants through dispatch and AOT paths.
- Adds Kimi K3 tuning configurations.
| File | Summary |
|---|---|
aiter/ops/flydsl/kernels/mxmoe_gemm_v2.py |
Bounds atomic stores by valid tokens. |
aiter/ops/flydsl/kernels/mxmoe_dispatcher.py |
Adds runtime dispatch and kernel variants. |
aiter/fused_moe.py |
Propagates valid-token metadata through EP stage 2. |
aiter/configs/model_configs/kimik3_a8w4_tuned_fmoe.csv |
Adds Kimi K3 EP configurations. |
aiter/aot/flydsl/mxfp4_moe.py |
Precompiles token-bound variants. |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Bernard-Liu
changed the title
- fix(flydsl): bound EP GEMM2 atomics by valid tokens - test(flydsl): clarify EP benchmark throughput metrics - add kmk3 ep config
Bernard-Liu
force-pushed
the
opt/support_and_opt_for_ep_kk3
branch
from
October 8, 2026 05:19
c710d91 to
bff0e9c
Compare
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
Add a regression test covering padded inputs and valid-token-bounded atomic GEMM2 writes.
1 open finding
🧠 Review effort: Lite
Give feedback about Copilot approvals in this survey to enter a drawing for a $150 gift card.
Qubitium
added a commit
to Qubitium/aiter
that referenced
this pull request
Oct 12, 2026
Single conflict in aiter/aot/flydsl/mxfp4_moe.py: this branch's include_bias_variants gate and main's use_valid_token_count variants (ROCm#5804) both rewrote the same stage-2 option list in parse_csv(). Resolution keeps both behaviours: include_bias_variants still selects enable_bias_options, and count_options still adds the atomic-epilog use_valid_token_count variants inside that loop, so the parser emits the union of both job sets and the tuner path keeps its no-bias fast scan.
zhiding512
added a commit
that referenced
this pull request
Oct 13, 2026
mxfp4_moe.py stays deleted. #5804 taught it the use_valid_token_count GEMM2 variant, which the runtime selects when the caller passes num_local_tokens; the next commit has the replay enumerate that caller choice instead. Signed-off-by: zhimding <zhimding@amd.com>
zhiding512
added a commit
that referenced
this pull request
Oct 13, 2026
Callers behind an EP dispatch pass the received token count as num_local_tokens. Since #5804 the runtime then bounds the layout-v2 GEMM2 atomics by it, a compile-time choice the old mxfp4_moe.py had to be taught by hand; on inline-sort rows the runtime instead drops the tuned config for heuristic FlyDSL kernels that no AOT compiled. Each row is now replayed with num_local_tokens too. Of the old script's use_valid_token_count kernels, the replay has every one the runtime reaches; the 227 it lacks come from bias variants of Silu/SiTUv2 rows (the runtime enables bias for Swiglu only since #6376), from inline-sort rows the runtime abandons under num_local_tokens, and from a token=3 row that get_padded_M never selects. Signed-off-by: zhimding <zhimding@amd.com>
2 tasks done
Draft
4 tasks done
6 of 11 tasks
zhiding512
added a commit
that referenced
this pull request
Nov 4, 2026
…hem in test_moe_ep A caller behind an expert-parallel dispatch passes fused_moe an expert mask, and the launches it masks compile for the mask's length: the gfx950 masked stage-2 reduction and route-guarded GEMM2, the gfx1250 g2l route and bf16 gather-reduce. No tuned row records that length, so the AOT build never compiled them and EP serving JIT-compiled them on the request path. aiter/aot/flydsl/moe_ep_profiles.csv lists EP callers as what one rank passes fused_moe: local experts, routing columns, expert-mask length, largest M. moe.py replays each profile with the mask at both ends of every get_padded_M bucket up to that M (bisecting between them as for rows), with and without num_local_tokens, against the tuned tables serving reads, so a profile needs no tuned row. --ep-profile selects profile tables; without --csv and --ep-profile both defaults are replayed, and run_aot replays both. The default table covers gfx950 Kimi-K3 EP8, DeepSeek-V4-Pro EP8 (EPLB r0, r32, r64) and EP16, GLM-5.3-Flash EP8/EP4/EP2, and gfx1250 Kimi-K3 EP8/EP16 and DeepSeek-V4-Pro EP4, from the EP rows tuned in #5804, #6928, #6435, op_tests/test_moe_ep.py -t aot_ep_profiles, first in the defaults, runs every profile of the GPU's arch through the EP fused_moe call in FlyDSL run-only mode and requires logits_diff < 0.01. The MXFP4 cases now pin quant_dtype_a instead of setting env switches, accept SiTUv2, and compare against a stage-wise reference that quantizes a4w4 activations as the kernels do. Co-authored-by: Cursor <cursoragent@cursor.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist