[feat] mega stage2 of gfx1250

[yanboshao]

Motivation

Technical Details

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[github-actions]

🏷️ 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:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
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 4785 --add-label <label>

[github-actions]

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

ruff

---

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

aiter/op_tests/multigpu_tests/test_mega_moe.py

Lines 32 to 46 in cb929e1

	import argparse
	import os
	import torch
	import torch.distributed as dist
	import torch.profiler as tprof
	import aiter
	from aiter import dtypes
	from aiter import ActivationType, QuantType, get_gfx
	from aiter.fused_moe import fused_moe
	from aiter.ops.shuffle import shuffle_weight, moe_shuffle_scale
	from aiter.ops.flydsl.moe_common import GateMode
	from aiter.utility import fp4_utils
	from aiter import get_hip_quant, get_torch_quant, pertoken_quant

---

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

aiter/op_tests/multigpu_tests/test_mega_moe.py

Line 50 in cb929e1

except Exception: # pragma: no cover

---

⚠️ [ruff] <RUF059> _{reported by reviewdog 🐶}
Unpacked variable E is never used

aiter/op_tests/multigpu_tests/test_mega_moe.py

Line 160 in cb929e1

E, two_inter = t.shape[:2]

---

⚠️ [ruff] <BLE001> _{reported by reviewdog 🐶}
Do not catch blind exception: Exception

aiter/op_tests/multigpu_tests/test_mega_moe.py

Line 782 in cb929e1

except Exception as _e:

[github-actions]

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

Suggested change

	ep_scatter_params = dict(
	gather_w=_gather_w_buf,
	tis=stage2_scatter.source_token_map,
	ep_rowmap=ep_rowmap,
	cap_rows=_cap_rows,
	topk=int(topk),
	max_tok=int(stage2_scatter.max_tokens_per_rank),
	slot_stride=int(stage2_scatter.max_tokens_per_rank) * int(topk),
	)
	ep_scatter_params = {
	"gather_w": _gather_w_buf,
	"tis": stage2_scatter.source_token_map,
	"ep_rowmap": ep_rowmap,
	"cap_rows": _cap_rows,
	"topk": int(topk),
	"max_tok": int(stage2_scatter.max_tokens_per_rank),
	"slot_stride": int(stage2_scatter.max_tokens_per_rank) * int(topk),
	}

[github-actions]

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

Suggested change

	dict(
	stage2_scatter=stage2_scatter,
	ep_destination_stride=(int(stage2_scatter.max_tokens_per_rank) * int(topk)),
	ep_row_map=ep_rowmap,
	)
	{
	"stage2_scatter": stage2_scatter,
	"ep_destination_stride": (int(stage2_scatter.max_tokens_per_rank) * int(topk)),
	"ep_row_map": ep_rowmap,
	}

[github-actions]

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

Suggested change

	import os
	import flydsl.compiler as flyc
	import flydsl.expr as fx
	import mori.cco.device.flydsl as cco
	from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops
	from aiter.ops.flydsl.kernels import vector
	from aiter.ops.flydsl.kernels.buffer_ops import (
	buffer_load,
	buffer_store,
	create_buffer_resource_from_addr,
	)
	from flydsl.expr import arith, const_expr, range_constexpr
	from flydsl.expr.typing import Int32, Int64, T
	from . import primitives as P
	from .config import (
	_LANE_MASK as LANE_MASK,
	_LOG2_WAVE_SIZE as LOG2_WAVE,
	_WAVE_SIZE as WAVE,
	)
	# NOTE: the cross-device xdb barrier is kept inlined per kernel (dispatch Phase 2,
	import os
	import flydsl.compiler as flyc
	import flydsl.expr as fx
	import mori.cco.device.flydsl as cco
	from flydsl.expr import arith, const_expr, range_constexpr
	from flydsl.expr.typing import Int32, Int64, T
	from aiter.ops.flydsl.kernels import communication_ops_utils as comm_ops
	from aiter.ops.flydsl.kernels import vector
	from aiter.ops.flydsl.kernels.buffer_ops import (
	buffer_load,
	buffer_store,
	create_buffer_resource_from_addr,
	)
	from . import primitives as P
	from .config import (
	_LANE_MASK as LANE_MASK,
	)
	from .config import (
	_LOG2_WAVE_SIZE as LOG2_WAVE,
	)
	from .config import (
	_WAVE_SIZE as WAVE,
	)
	# NOTE: the cross-device xdb barrier is kept inlined per kernel (dispatch Phase 2,

[github-actions]

⚠️ [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

[github-actions]

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

Suggested change

	if const_expr(_XDB_FLAG_SLOTS > block_num):
	if bid == 0:
	_tail = _XDB_FLAG_SLOTS - block_num
	_nthr = warp_num_per_block * WAVE
	for r in range_constexpr((_tail + _nthr - 1) // _nthr):
	idx = tid + r * _nthr
	if idx < _tail:
	buffer_store(
	phase + arith.constant(1, type=T.i64),
	rsrc_xdb_flag,
	block_num + idx,
	)
	if const_expr(_XDB_FLAG_SLOTS > block_num) and bid == 0:
	_tail = _XDB_FLAG_SLOTS - block_num
	_nthr = warp_num_per_block * WAVE
	for r in range_constexpr((_tail + _nthr - 1) // _nthr):
	idx = tid + r * _nthr
	if idx < _tail:
	buffer_store(
	phase + arith.constant(1, type=T.i64),
	rsrc_xdb_flag,
	block_num + idx,
	)

[github-actions]

⚠️ [ruff] <UP037> _{reported by reviewdog 🐶}
Remove quotes from type annotation

Suggested change

	workgroup_mask: Union[int, "ir.Value"] = 0,
	workgroup_mask: Union[int, ir.Value] = 0,

[github-actions]

⚠️ [ruff] <F811> _{reported by reviewdog 🐶}
Redefinition of unused buffer_ops from line 20: buffer_ops redefined here

Suggested change

	from aiter.ops.flydsl.kernels import buffer_ops
	from aiter.ops.flydsl.kernels.tensor_shim import (
	from aiter.ops.flydsl.kernels.tensor_shim import (

[github-actions]

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

Suggested change

	neg1 = arith.constant(-1, type=i32)
	arith.constant(-1, type=i32)

[github-actions]

⚠️ [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

[github-actions]

⚠️ [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

[github-actions]

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

aiter/aiter/ops/flydsl/kernels/moe_contiguous_psum.py

Lines 15 to 33 in cb929e1

	import flydsl.compiler as flyc
	import flydsl.expr as fx
	from flydsl._mlir import ir
	from flydsl._mlir.dialects import llvm, scf
	from flydsl.expr import arith, const_expr, gpu, ptrtoint, range_constexpr
	from aiter.ops.flydsl.kernels import buffer_ops
	from flydsl.expr.typing import Int32, T
	from flydsl.expr.arith import ArithValue, CmpIPredicate, CmpFPredicate, _to_raw as _raw
	from flydsl.compiler.kernel_function import CompilationContext
	from flydsl.runtime.device import get_rocm_arch
	from flydsl.utils.smem_allocator import SmemAllocator, SmemPtr
	from aiter.ops.flydsl.kernels import buffer_ops
	from aiter.ops.flydsl.kernels.tensor_shim import (
	STensor,
	AITER_FLYDSL_KERNARG_PRELOAD,
	AITER_FLYDSL_KERNARG_PRELOAD_COUNT,
	ptr_rsrc,
	)

[github-actions]

⚠️ [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 44 in cb929e1

	import math
	from collections import namedtuple
	import flydsl.compiler as flyc
	import flydsl.expr as fx
	from flydsl.expr import (
	arith,
	const_expr,
	range_constexpr,
	rocdl,
	tdm_ops,
	)
	from aiter.ops.flydsl.kernels import vector
	from flydsl.expr.typing import Constexpr, T
	from flydsl.expr.typing import Vector as Vec
	from aiter.utility.mx_types import MxDtypeInt as MxDtype
	from .gemm_common_gfx1250 import (
	batched_silu_swiglu,
	batched_situv2,
	fused_silu_swiglu_elem,
	fused_situv2_elem,
	make_lds_copy_ops,
	pipeline_fence,
	situv2_consts,
	workgroup_barrier,
	)
	from .quant_utils import (
	emit_amax_e8m0_native_scale,
	emit_cvt_scalef32_pk8_fp8_f32,
	)
	# tdm_scatter is vendored under MegaMoE (self-contained on the stock FlyDSL
	# wheel's low-level TDM intrinsics) so this kernel needs no FlyDSL-side patch.
	# Once the gather/scatter wrappers land upstream, import them from
	# flydsl.expr.rocdl.tdm_ops instead and delete the local shim.
	from .mega_moe_gfx1250.tdm_gather_shim import make_tensor_gather_descriptor, tensor_store_gather
	from .tensor_shim import AITER_FLYDSL_MOE_EXPERT_SCHEDULING_MODE