Skip to content

moe rebase and refactor - #4394

Merged
coderfeli merged 27 commits into
mainfrom
dev/450ep_felix_tdm_port_rebase
Jul 31, 2026
Merged

moe rebase and refactor#4394
coderfeli merged 27 commits into
mainfrom
dev/450ep_felix_tdm_port_rebase

Conversation

@yadaish

@yadaish yadaish commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

yadaish and others added 5 commits July 27, 2026 14:30
@yadaish @Zzz9990
@lalala-sh @XingerZhu @coderfeli @yanboshao
rebase
9a9c3bf 
Co-authored-by: zanzhang <zanzhang@amd.com>
Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: XingerZhu <233663585+XingerZhu@users.noreply.github.com>
Co-authored-by: Felix Li <felix.li@amd.com>
Co-authored-by: yanboshao <yashao@amd.com>
@HaonanWang98 @Zzz9990
@lalala-sh @XingerZhu @coderfeli @yanboshao @yadaish
fix
b370e62 
Co-authored-by: zanzhang <zanzhang@amd.com>
Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: XingerZhu <233663585+XingerZhu@users.noreply.github.com>
Co-authored-by: Felix Li <felix.li@amd.com>
Co-authored-by: yanboshao <yashao@amd.com>
@yadaish @Zzz9990
@lalala-sh @XingerZhu @coderfeli @yanboshao
support test csv
7c3e3bb 
Co-authored-by: zanzhang <zanzhang@amd.com>
Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: XingerZhu <233663585+XingerZhu@users.noreply.github.com>
Co-authored-by: Felix Li <felix.li@amd.com>
Co-authored-by: yanboshao <yashao@amd.com>
@HaonanWang98 @Zzz9990
@lalala-sh @XingerZhu @coderfeli @yanboshao @yadaish
fix
105f856 
Co-authored-by: zanzhang <zanzhang@amd.com>
Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: XingerZhu <233663585+XingerZhu@users.noreply.github.com>
Co-authored-by: Felix Li <felix.li@amd.com>
Co-authored-by: yanboshao <yashao@amd.com>
@yadaish @Zzz9990
@lalala-sh @XingerZhu @coderfeli @yanboshao
fix
cafce66 
Co-authored-by: zanzhang <zanzhang@amd.com>
Co-authored-by: lalala-sh <Jiaxing.Wen@amd.com>
Co-authored-by: XingerZhu <233663585+XingerZhu@users.noreply.github.com>
Co-authored-by: Felix Li <felix.li@amd.com>
Co-authored-by: yanboshao <yashao@amd.com>
@yadaish
yadaish requested a review from a team July 27, 2026 07:25
@github-actions

github-actions Bot commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

🏷️ 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:triton-300x Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
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 4394 --add-label <label>

github-actions[bot]
github-actions Bot reviewed Jul 27, 2026
View reviewed changes

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit

ruff

⚠️ [ruff] <B008> reported by reviewdog 🐶
Do not perform function call fx.Stream in argument defaults; instead, perform the call within the function, or read the default from a module-level singleton variable

aiter/aiter/ops/flydsl/kernels/moe_route_maps.py

Line 518 in cafce66

stream: fx.Stream = fx.Stream(None),

⚠️ [ruff] <B008> reported by reviewdog 🐶
Do not perform function call fx.Stream in argument defaults; instead, perform the call within the function, or read the default from a module-level singleton variable

aiter/aiter/ops/flydsl/kernels/moe_route_maps.py

Line 759 in cafce66

stream: fx.Stream = fx.Stream(None),

⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

aiter/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

Lines 6 to 29 in cafce66

import math
from collections import namedtuple
import flydsl.compiler as flyc
import flydsl.expr as fx
from flydsl.expr import const_expr, range_constexpr, rocdl, tdm_ops, vector
from flydsl.expr.typing import Constexpr, T
from flydsl.expr.typing import Vector as Vec
from .gemm_common_gfx1250 import (
lds_load_b32_raw,
lds_load_b128_raw,
lds_store_b32_raw,
lds_store_b128_raw,
lds_store_b64_raw,
pipeline_fence,
workgroup_barrier,
fused_silu_swiglu_elem,
batched_silu_swiglu,
)
from .quant_utils import (
emit_amax_e8m0_native_scale,
emit_cvt_scalef32_pk8_fp8_f32,
)
from aiter.utility.mx_types import MxDtypeInt as MxDtype

⚠️ [ruff] <F841> reported by reviewdog 🐶
Local variable SA_SUPERS is assigned to but never used

aiter/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

Line 95 in cafce66

SA_SUPERS, SB_SUPERS = tile_m // 32, tile_n // 32

⚠️ [ruff] <C408> reported by reviewdog 🐶
Unnecessary dict() call (rewrite as a literal)

aiter/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

Line 243 in cafce66

pad_kw = dict(pad_interval=pad[0], pad_amount=pad[1]) if pad else {}

⚠️ [ruff] <F841> reported by reviewdog 🐶
Local variable N_ELEM is assigned to but never used

aiter/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

Line 420 in cafce66

N_ELEM = N_MX_BLKS * WN_PER_MX_BLOCK * 4

⚠️ [ruff] <SIM102> reported by reviewdog 🐶
Use a single if statement instead of nested if statements

aiter/aiter/ops/flydsl/kernels/mxfp4_preshuffle_gfx1250_tdm.py

Lines 467 to 468 in cafce66

if row_rel < mn_oob:
if is_kgrp0:

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1760 in cafce66

g2l_lut: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1761 in cafce66

expert_mask: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1762 in cafce66

gather_w: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1763 in cafce66

weight_in: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1764 in cafce66

counter: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1765 in cafce66

num_local_tokens: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1766 in cafce66

num_valid_routes: Optional[torch.Tensor] = None,

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Lines 1896 to 1898 in cafce66

grouped_a1_scale: Optional[
torch.Tensor
] = None, # (out_E, out_max_m//wmma_rep, (model_dim//32)*wmma_rep) uint8 out

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1899 in cafce66

g2l_lut: Optional[torch.Tensor] = None, # (E_global,) int32 global->local

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1900 in cafce66

gather_w: Optional[torch.Tensor] = None, # (token_num, topk) out; kept->cast,drop->0

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1901 in cafce66

weight_in: Optional[torch.Tensor] = None, # (token_num, topk) f32 route weights in

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1902 in cafce66

counter: Optional[torch.Tensor] = None, # (E,) int32 pre-zeroed slot counter

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 2281 in cafce66

out_payload: Optional[torch.Tensor] = None, # (E, max_m, Pb) uint8

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 2282 in cafce66

out_scale: Optional[torch.Tensor] = None, # (E, max_m//wmma_rep, Ws*wmma_rep)

⚠️ [ruff] <UP045> reported by reviewdog 🐶
Use X | None for type annotations

aiter/aiter/ops/flydsl/moe_kernels.py

Line 2283 in cafce66

num_valid_routes: Optional[torch.Tensor] = None, # (1,) int32; route-branch only: skip routes >= this (EP dead-tail)

yadaish and others added 5 commits July 27, 2026 16:11
@yadaish
fix ruff and black
228dfe9
@yadaish
fix aot
77b5d9f
@yadaish
fix aot
60c8704
@yadaish
fix ruff
ce61a78
@azaidy @cursoragent
refactor(flydsl): rewrite moe_g2l_lut with high-level flydsl API
f4a1e17 
Replace the low-level MLIR/scf/ArithValue construction with flydsl expr
if/select, SharedAllocator struct for the double-buffered scan, and
_lds_load/_lds_store helpers. Behaviour unchanged.

Co-authored-by: Cursor <cursoragent@cursor.com>
@zufayu
zufayu requested a review from coderfeli July 28, 2026 01:44
@yadaish yadaish changed the title Dev/450ep felix tdm port rebase moe rebase and refactor Jul 28, 2026
yadaish and others added 13 commits July 28, 2026 14:14
@yadaish
fix black
9b08354
@yadaish
fix csv mode
b51f161
@yadaish
tiny fix
75a2ae7
@yadaish
update
1246545
@azaidy
remove useless codes
83fcc9c
@yadaish
Merge branch 'main' of github.com:ROCm/aiter into dev/450ep_felix_tdm…
c92b004 
…_port_rebase

# Conflicts:
#	aiter/aot/flydsl/grouped_moe.py
#	aiter/ops/flydsl/grouped_moe_gfx1250.py
#	aiter/ops/flydsl/kernels/gemm_common_gfx1250.py
#	aiter/ops/flydsl/kernels/gemm_mxscale_gfx1250.py
#	aiter/ops/flydsl/kernels/kernels_common.py
#	aiter/ops/flydsl/kernels/moe_contiguous_psum.py
#	aiter/ops/flydsl/kernels/moe_fused_route_quant_scatter.py
#	aiter/ops/flydsl/kernels/moe_gather_reduce.py
#	aiter/ops/flydsl/kernels/moe_grouped_gemm_mxscale_gfx1250.py
#	aiter/ops/flydsl/kernels/moe_route_maps.py
#	aiter/ops/flydsl/moe_kernels.py
#	op_tests/test_flydsl_grouped_gemm_gfx1250.py
#	op_tests/test_moe_ep.py
@yadaish
Merge branch 'dev/450ep_felix_tdm_port_rebase' of github.com:ROCm/ait…
3e38725 
…er into dev/450ep_felix_tdm_port_rebase
@yadaish
use new lds interface
66fd25b
@yadaish
Merge branch 'dev/450ep_felix_tdm_port_rebase' of github.com:ROCm/ait…
babe528 
…er into dev/450ep_felix_tdm_port_rebase
@yadaish
Merge remote-tracking branch 'origin/main' into dev/450ep_felix_tdm_p…
7662bf5 
…ort_rebase
@XiaobingSuper
Merge branch 'main' into dev/450ep_felix_tdm_port_rebase
73f2f29
@yadaish
use AITER_FLYDSL_MOE_EXPERT_SCHEDULING_MODE to control moe expert mode
0c18686
@yadaish
refactor (#4459)
dc465f2
@coderfeli
coderfeli merged commit b0b6945 into main Jul 31, 2026
42 checks passed
@coderfeli
coderfeli deleted the dev/450ep_felix_tdm_port_rebase branch July 31, 2026 05:56
@XiaobingSuper XiaobingSuper mentioned this pull request Jul 31, 2026
Merged
This was referenced Aug 1, 2026
Open
Open
Open
yanboshao added a commit that referenced this pull request Aug 3, 2026
@yanboshao @cursoragent
ep scatter: port fused gemm2 P2P scatter-combine onto main's TDM MoE …
7723bc0 
…dispatch

Wire the branch's ep_scatter (combine_mode="scatter_fused") feature onto main's
post-#4394 dispatch, whose grouped_moe/fused_moe were reorganized (entry
grouped_gemm_gfx1250_a8w4 -> _grouped_a8w4_tdm_moe) relative to the pre-rebase
branch. The kernel + launcher layer already carries the feature after the rebase
(mxfp4_preshuffle_gfx1250_tdm ep_p2p_write epilogue + batched_gemm_mxfp4
flydsl_grouped_gemm_a8w4_masked ep_* / arg_ep_rowmap forwarding); this commit
adds only the dispatch-layer wiring that the "keep main's grouped_moe/fused_moe"
rebase strategy deferred:

  grouped_moe_gfx1250.py
    - add _get_compiled_contiguous_psum_remap_ep / _get_compiled_ep_rowmap /
      build_ep_rowmap helpers (compile the ep remap + ep_rowmap kernels that our
      kept moe_contiguous_psum.py already provides)
    - contiguous_psum_remap: add ep= branch (fused ep_rowmap build in the remap)
    - _grouped_a8w4_tdm_moe: ep_* params; build ep_rowmap + _ep_remap; pass
      ep=_ep_remap to contiguous_psum_remap; build _ep_gemm2_kwargs; pass
      **_ep_gemm2_kwargs to the gemm2 flydsl_grouped_gemm_a8w4_masked; ep_scatter
      early-return (combine reads comb_inp, no gather-reduce)
    - grouped_gemm_gfx1250_a8w4: ep_* params forwarded to _grouped_a8w4_tdm_moe
  fused_moe.py
    - fused_moe / fused_moe_ / _fused_moe_impl: ep_* params threaded down to the
      grouped_gemm_gfx1250_a8w4 call

Co-authored-by: yanboshao <yashao@amd.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
yanboshao added a commit that referenced this pull request Aug 3, 2026
@yanboshao
csv: trim tuned_grouped_fmoe to main + 28 genuinely-new tuning points
e234a95 
The rebase-onto-main merge naively unioned old-schema (38-col) rows with new
main's 36-col schema, inflating the file to 257 rows: it resurrected 83 rows
inherited from the pre-#4394 main that new main had dropped, and (because the
38<->36 column shift misaligned the shape key) failed to dedup, keeping 49
re-benchmark duplicates (config identical to a main row, only us/tflops/bw
differ) plus 26 same-shape re-tunes.

Reduce to new main's 71 rows verbatim + only the 28 rows whose full config key
(gfx..max_m) is genuinely absent from main (new token/act/dtype coverage of the
same DeepSeek-family shapes). Drops 158 rows; keeps every shape validated on
gfx1250 (deferring to main's tuning on all overlapping shapes).
@yadaish yadaish self-assigned this Aug 4, 2026
@bingxche bingxche mentioned this pull request Aug 7, 2026
Open
4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@github-actions github-actions[bot] github-actions[bot] left review comments
@coderfeli coderfeli Awaiting requested review from coderfeli

Assignees

@yadaish yadaish

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

5 participants

@yadaish @coderfeli @HaonanWang98 @azaidy @XiaobingSuper