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
@yadaish yadaish self-assigned this Aug 4, 2026
@bingxche bingxche mentioned this pull request Aug 7, 2026
Closed
1 of 4 tasks
@amd-bot amd-bot mentioned this pull request Aug 9, 2026
Open
@amd-bot amd-bot mentioned this pull request Aug 16, 2026
Open
@jhinpan jhinpan mentioned this pull request Aug 20, 2026
Closed
@amd-bot amd-bot mentioned this pull request Aug 28, 2026
Open
@jin-amd jin-amd mentioned this pull request Sep 30, 2026
Merged
3 of 4 tasks
jin-amd added a commit to jin-amd/aiter that referenced this pull request Sep 30, 2026
@jin-amd
[Bugfix][CK] Keep two K tiles per split in blockscale MoE stage-1 spl…
75c2199 
…it-K

The default fp8-blockscale 2-stage fused MoE returns wrong output whenever
get_ksplit() asks for model_dim / 256 splits, one 256-wide K tile each. This
happens at small decode batches on gfx942 (cu_num 304), for example at one
token for (model_dim, inter_dim) = (4096, 256) with 288 experts and top-8,
(7168, 128) with 257 experts and top-9, and (2048, 384) with 128 experts and
top-8. Every output element differs from the reference, and the values
change from run to run.

ck_moe_stage1_gemm turns the requested number of splits into CK's KBatch,
the number of KPerBlock tiles each split covers: K / (splitk * KPerBlock).
With one tile per split that is 1, which DeviceMoeGemmBlockScale reads as
"no split": it runs a single z-block over the whole K and skips zeroing the
output. The split-K instance still accumulates into that output atomically,
and fused_moe allocates it with torch.empty. StrideE also keys the split-K
layout (2 * N per row) on KBatch > 1, so the rows overlap as well.

ROCm#3997 fixed this in ck_moe_stage1() by turning split-K off for such
requests, and the MoE refactor in ROCm#4394 dropped that guard. This fix sits at
the KBatch conversion instead, so no caller reaches KBatch == 1 through
split-K, whether the split comes from get_ksplit(), a tuned CSV or
AITER_KSPLIT. It also keeps the split rather than turning it off: when
splitk would leave fewer than two tiles per split, stage 1 uses fewer splits
of at least two tiles, KBatch = 2 or the smallest divisor of the tile count
that is at least 2. At the (4096, 256, E=288, top-8) one-token shape the
operator then takes 32.1 us, against 41.1 us without split-K. A K of a
single tile cannot be split and now fails with an error. Requests that
already give two or more tiles per split are unchanged.

Signed-off-by: Jin Tao <jintao12@amd.com>
@jin-amd jin-amd mentioned this pull request Sep 30, 2026
Merged
4 of 5 tasks
yifehuan pushed a commit that referenced this pull request Oct 2, 2026
@jin-amd
[Bugfix][CK] Keep two K tiles per split in blockscale MoE stage-1 spl…
635c10f 
…it-K (#5979)

The default fp8-blockscale 2-stage fused MoE returns wrong output whenever
get_ksplit() asks for model_dim / 256 splits, one 256-wide K tile each. This
happens at small decode batches on gfx942 (cu_num 304), for example at one
token for (model_dim, inter_dim) = (4096, 256) with 288 experts and top-8,
(7168, 128) with 257 experts and top-9, and (2048, 384) with 128 experts and
top-8. Every output element differs from the reference, and the values
change from run to run.

ck_moe_stage1_gemm turns the requested number of splits into CK's KBatch,
the number of KPerBlock tiles each split covers: K / (splitk * KPerBlock).
With one tile per split that is 1, which DeviceMoeGemmBlockScale reads as
"no split": it runs a single z-block over the whole K and skips zeroing the
output. The split-K instance still accumulates into that output atomically,
and fused_moe allocates it with torch.empty. StrideE also keys the split-K
layout (2 * N per row) on KBatch > 1, so the rows overlap as well.

#3997 fixed this in ck_moe_stage1() by turning split-K off for such
requests, and the MoE refactor in #4394 dropped that guard. This fix sits at
the KBatch conversion instead, so no caller reaches KBatch == 1 through
split-K, whether the split comes from get_ksplit(), a tuned CSV or
AITER_KSPLIT. It also keeps the split rather than turning it off: when
splitk would leave fewer than two tiles per split, stage 1 uses fewer splits
of at least two tiles, KBatch = 2 or the smallest divisor of the tile count
that is at least 2. At the (4096, 256, E=288, top-8) one-token shape the
operator then takes 32.1 us, against 41.1 us without split-K. A K of a
single tile cannot be split and now fails with an error. Requests that
already give two or more tiles per split are unchanged.

Signed-off-by: Jin Tao <jintao12@amd.com>
@ntgiang71096 ntgiang71096 mentioned this pull request Oct 9, 2026
Open
4 of 5 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