[Kimi-K3] FlyDSL SiTUv2 MoE kernels, strided grouped-topk, tuned configs by carlushuang · Pull Request #4397 · ROCm/aiter

carlushuang · 2026-07-27T15:12:47Z

Summary

Brings up Kimi-K3 (KDA-hybrid, MXFP4 MoE) in aiter, covering both gfx950 (MI355) and gfx1250. All changes are additive and gated behind K3-specific paths (SiTUv2 activation, K3 shapes); default behavior for existing models is unchanged.

Changes

FlyDSL SiTUv2 MoE kernels (a16w4)

flydsl_moe_stage1/2 with SiTUv2 activation (situ_beta, situ_linear_beta): mixed 2-stage a16w4 MoE GEMM, fused route/quant/scatter, opt-sort and shuffle utilities for the MXFP4 (a16w4) path.
flydsl 0.2.4 compatibility: drop loc= from _split_mfma_operands calls, _MIN_FLYDSL_VERSION bumped to 0.2.4.
Observable tile-alignment guard: force tile_n/k 256→128 when inter_dim is not 256-aligned (e.g. 384), with a one-time warning.

A4W4 / A8W4 SiTUv2 (gfx950 + gfx1250)

FlyDSL A4W4 SiTUv2 path and AITER_SITUV2_A8W4 env-gated a8w4 mode; a8w4 fmoe tune configs.

Grouped topk: row-strided inputs

biased_grouped_topk and the opt-sort kernel accept row-strided logits (the K3 fused MoE-front router slice of the [gate_up|experts|routed] buffer) — avoids a per-layer contiguous copy. Coverage folded into the existing test_biased_grouped_topk as a K3-shaped case (no new test file).

Tuned configs (Kimi-K3)

kimik3_fp4_tuned_fmoe.csv, kimik3_fp4_untuned_fmoe.csv, kimik3_bf16_tuned_gemm.csv, kimik3_bf16_untuned_gemm.csv, a16w4 fmoe tuned config (inter_dim=384), a8w4 fmoe tune configs, bf16 dense GEMM tuned config, gfx1250 conv2d (ViT) configs, gfx1250 bf16 tuned config.

Vision / misc

Triton conv2d configs for the K3 ViT tower (gfx1250), Gluon reduce guard, MlaVersion enum in module_aiter_core, CK gfx1250 enable patch (temporary).

* add Situv2 activation for a8w4 MoE stage1 * refactor(flydsl): scope situv2 helpers and fix lint Move situ beta compile-time constants into situ_elem/situ_up_elem; apply ruff/black fixes on the a8w4 situv2 test module. * feat(moe): integrate SiTUv2 into fused_moe API, split-K, tuner/prebuild * test(moe): add SiTUv2 host-ref + a4w4/a8w4 stage1 test (adapted to Situv2 naming) Complements MHYang's aiter/ops/flydsl/test_flydsl_moe_a8w4.py by adding a host-only (no-GPU) SiTUv2 reference sweep plus a4w4 stage1 coverage in addition to a8w4. Adapted to this branch's API: ActivationType.Situv2 enum and situ_beta / situ_linear_beta parameters on torch_moe_stage1 and flydsl_moe_stage1. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): add SiTUv2 default cases + --beta/--linear-beta to test_moe_2stage (Situv2 naming) Co-authored-by: Cursor <cursoragent@cursor.com> * test: consolidate a8w4 SiTUv2 vec4 cases into op_tests/flydsl_tests. Move pytest tile/gate_mode sweep from aiter/ops/flydsl/test_flydsl_moe_a8w4.py into test_flydsl_moe_situv2.py and remove the misplaced test file. * style: black-format silu_and_mul_fq.py for CI pre-checks. --------- Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

* add Situv2 activation for a8w4 MoE stage1 * refactor(flydsl): scope situv2 helpers and fix lint Move situ beta compile-time constants into situ_elem/situ_up_elem; apply ruff/black fixes on the a8w4 situv2 test module. * feat(moe): integrate SiTUv2 into fused_moe API, split-K, tuner/prebuild * test(moe): add SiTUv2 host-ref + a4w4/a8w4 stage1 test (adapted to Situv2 naming) Complements MHYang's aiter/ops/flydsl/test_flydsl_moe_a8w4.py by adding a host-only (no-GPU) SiTUv2 reference sweep plus a4w4 stage1 coverage in addition to a8w4. Adapted to this branch's API: ActivationType.Situv2 enum and situ_beta / situ_linear_beta parameters on torch_moe_stage1 and flydsl_moe_stage1. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): add SiTUv2 default cases + --beta/--linear-beta to test_moe_2stage (Situv2 naming) Co-authored-by: Cursor <cursoragent@cursor.com> * test: consolidate a8w4 SiTUv2 vec4 cases into op_tests/flydsl_tests. Move pytest tile/gate_mode sweep from aiter/ops/flydsl/test_flydsl_moe_a8w4.py into test_flydsl_moe_situv2.py and remove the misplaced test file. * style: black-format silu_and_mul_fq.py for CI pre-checks. * Fix GUI shuffle_scale k_pad for non-256-aligned MoE K and auto stage2 tile_k. Pad w2/w1 GUI e8m0 scales to k_groups multiple of 8 (DSV4 inter=640), unify stage2 tile_k selection in flydsl_moe_stage2 and fused_moe, and add staged a8w4 regression tests across inter/model K sweeps. * test: integrate a8w4 K-dimension regressions into existing test suites. Move shuffle pad coverage into test_quant_mxfp4 and FlyDSL GUI stage2/e2e into test_flydsl_moe_a8w4; drop standalone op_tests files for repo convention. * test: move flydsl a8w4 GUI regressions to op_tests/flydsl_tests. Relocate test_flydsl_moe_a8w4 under op_tests/flydsl_tests and align with pytest conventions used by other FlyDSL op tests. --------- Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>

* add Situv2 activation for a8w4 MoE stage1 * refactor(flydsl): scope situv2 helpers and fix lint Move situ beta compile-time constants into situ_elem/situ_up_elem; apply ruff/black fixes on the a8w4 situv2 test module. * feat(moe): integrate SiTUv2 into fused_moe API, split-K, tuner/prebuild * test(moe): add SiTUv2 host-ref + a4w4/a8w4 stage1 test (adapted to Situv2 naming) Complements MHYang's aiter/ops/flydsl/test_flydsl_moe_a8w4.py by adding a host-only (no-GPU) SiTUv2 reference sweep plus a4w4 stage1 coverage in addition to a8w4. Adapted to this branch's API: ActivationType.Situv2 enum and situ_beta / situ_linear_beta parameters on torch_moe_stage1 and flydsl_moe_stage1. Co-authored-by: Cursor <cursoragent@cursor.com> * test(moe): add SiTUv2 default cases + --beta/--linear-beta to test_moe_2stage (Situv2 naming) Co-authored-by: Cursor <cursoragent@cursor.com> * test: consolidate a8w4 SiTUv2 vec4 cases into op_tests/flydsl_tests. Move pytest tile/gate_mode sweep from aiter/ops/flydsl/test_flydsl_moe_a8w4.py into test_flydsl_moe_situv2.py and remove the misplaced test file. * style: black-format silu_and_mul_fq.py for CI pre-checks. * Fix GUI shuffle_scale k_pad for non-256-aligned MoE K and auto stage2 tile_k. Pad w2/w1 GUI e8m0 scales to k_groups multiple of 8 (DSV4 inter=640), unify stage2 tile_k selection in flydsl_moe_stage2 and fused_moe, and add staged a8w4 regression tests across inter/model K sweeps. * test: integrate a8w4 K-dimension regressions into existing test suites. Move shuffle pad coverage into test_quant_mxfp4 and FlyDSL GUI stage2/e2e into test_flydsl_moe_a8w4; drop standalone op_tests files for repo convention. * test: move flydsl a8w4 GUI regressions to op_tests/flydsl_tests. Relocate test_flydsl_moe_a8w4 under op_tests/flydsl_tests and align with pytest conventions used by other FlyDSL op tests. * Add FlyDSL MoE a16wfp4 (bf16 x mxfp4) kernels on gfx950. Port a16w4 stage1/stage2 into mixed_moe_gemm_2stage as dedicated _a16w4 builders while keeping HEAD fp8/fp4 compile paths unchanged. Wire dispatch in moe_kernels and add op_tests/flydsl_tests regression for stage1, stage2, and e2e against torch_moe references. * Fix ruff/black lint issues in a16wfp4 PR files. Remove unused _fly imports, move lds_space/types imports to module top, add noqa for test sys.path bootstrap, and apply black formatting. * Remove dead duplicate waves_per_eu None checks in a16w4 stage2. The parameter defaults to int and callers always pass an int after dispatch. * Fix a16w4 stage2: skip second compute when total_tiles==1. The CK-style ping-pong pipeline assumed total_tiles>=2. When inter_dim==tile_k (e.g. inter_dim=256, tile_k=256) total_tiles=1, the HEAD pre-fetched an out-of-bounds A tile into lds_ping, and the TAIL then computed with that garbage data, producing cos~0.71 vs reference. Fix: const_expr(total_tiles==1) guard in the TAIL block--skip the second compute_tile and the ping-buffer load entirely. Also tighten test _check_result: add cosine+rel_L2 as primary gate (cos>0.999) so future stage2 bugs cannot be masked by atol=1.0 on small-magnitude outputs. * Add pipeline guards for a16w4 stage1 total_tiles==1 and odd-tile validation. The stage1 ping-pong pipeline has the same total_tiles>=2 assumption as stage2. Add: - total_tiles==1 guard in stage1 TAIL (same pattern as stage2 fix) - Validation: reject odd total_tiles and K < tile_k at compile time for both stage1 and stage2 a16w4 paths These configs cannot occur with real LLM shapes (model_dim and inter_dim are always multiples of 256), but the guards prevent silent corruption on contrived inputs. * Add SiTUv2 activation to a16w4 (bf16 x mxfp4) stage1. Port the SiTUv2 activation from pr/a8w4-situv2 into the a16w4 stage1 kernel (both direct and cshuffle epilogues): situ_g = beta * tanh(gate/beta) * sigmoid(gate) up_scaled = linear_beta * tanh(up/linear_beta) with gate/up clamped to <=7 / [-7,7] before activation, matching the a8w4 situv2 numerics exactly. Thread situ_beta/situ_linear_beta through compile_flydsl_moe_stage1 and flydsl_moe_stage1; the per-beta cache tag keeps distinct binaries. Verified vs an inline torch SiTUv2 reference (mxfp4 dequant GEMM + clamp + situv2): cos>=0.99996 across beta in {(1,1),(0.5,2),(1.5,0.8)}, shapes (512/256, 3072/256), tile_n {128,256}. silu/swiglu unchanged. * Strip dead generic pipeline from a16w4 stage1 kernel (~2260 lines). compile_mixed_moe_gemm1_a16w4 was created by copying the generic fp8/fp4 builder and gating with is_a16w4_stage1. The live a16w4 path is fully contained in `if const_expr(is_a16w4_stage1): ... return`; everything after that return was an unreachable copy of the generic f8f6f4 kernel body (referenced generic-only vars like _lds_tid_offset_pong / _pipe_* that are never defined on the a16w4 path -- proven dead since a16w4 tests pass without them). Remove the dead generic kernel body, the generic-only setup blocks (postlude pipeline schedule, not-is_a16w4 guards), and the now-unused setup vars. Purely dead-code deletion: a16w4 stage1/stage2/e2e/situv2 + a4w4 all still pass with identical cosine. * Strip dead generic pipeline from a16w4 stage2 kernel (~1140 lines). Same shape as the stage1 cleanup: compile_mixed_moe_gemm2_a16w4's live a16w4 path is fully contained in `if const_expr(is_a16w4): ... return` inside _moe_gemm2_then_body; everything after that return was an unreachable copy of the generic f8f6f4 stage2 body (f8f6f4 MFMA, generic-only scale layouts). Proven dead: a16w4 stage2/e2e/situv2 + a4w4 all pass unchanged after removal. Remove the dead body plus the now-unused generic setup vars (pack_N/pack_K, cbsz/blgp, generic scale layouts, sx_rsrc sentinel, etc.). Purely dead-code deletion. * feat(moe): route a16w4 (bf16 x mxfp4) SiTUv2 through fused_moe Wire the mixed_moe a16w4 kernel into the fused_moe 2-stage path for SiTUv2: get_2stage_cfgs now matches bf16/fp16 x fp4 when activation is SiTUv2 (which uniquely identifies the a16w4 kernel, so GPT-OSS / legacy bf16-Swiglu keep their CK-Tile routing), maps _a_type to bf16/fp16, and infers q_dtype_a=bf16 for SiTUv2+separated so the activation stays bf16 (no fp4 quant). The stage1/ stage2 activation-quant gates also accept SiTUv2. Tests: a16w4 SiTUv2 via fused_moe (cos=1.0 vs torch ref) and an a8w4 SiTUv2 inter=640 vec4 case exercising the fix-k non-256 K-tiling. * feat(moe): support non-256 inter_dim for a16w4 (a8w4 parity) Bring the FlyDSL a16w4 (bf16 x mxfp4) MoE kernels to a8w4 parity so inter_dim need not be a multiple of 256 (e.g. DSV4 inter=640, and arbitrary values like 384 via inter_dim_pad), across stage1/stage2/E2E. Kernel (mixed_moe_gemm_2stage.py): - stage2: allow inter_dim % tile_k == 0 (drop even-only rule), add odd_k_tiles 1-tile tail path, 256-pad scale reads (#3476). - stage1: same odd_k_tiles K-loop + 256-padded w1 scale layout; fix gx launch grid (tile2_pad + full-N when inter_dim_pad > 0). Caller (moe_kernels.py): - pick/resolve_flydsl_stage1_tile_n: use tile_n=128 when inter%256!=0 (fixes stage1 error on cols 0..inter%256 under tile_n=256). - zero-init stage1 out when inter_dim_pad > 0. Tests (test_flydsl_moe_a16wfp4.py): - non-256 regressions (256/384/640) for stage1/stage2/E2E + pick_* tests. - _generate_a16wfp4_data gains activation/situ_beta/situ_linear_beta. - --perf sweep mode (correctness + latency/TFLOPs). * test(moe): fold situv2 tests into dtype files, drop standalone situv2 file Consolidate test_flydsl_moe_situv2.py into the per-dtype test files (matching how test_flydsl_moe_a16wfp4.py already embeds its situv2 tests): - test_flydsl_moe_a8w4.py: add the a8w4 SiTUv2 vec4 stage1 sweep + the host-only situv2 reference test; reuse the file's existing _check_close helper (fp32 cast for the bf16-ref vs f16-out case). Add shuffle_weight / e8m0_shuffle imports. - Delete test_flydsl_moe_situv2.py. Its broken-under-pytest a4w4/a8w4 main()-runner is dropped; a8w4 situv2 is now real parametrized coverage. a16wfp4 situv2 tests already lived in their dtype file and are unchanged. * fix conflict * add timing args --------- Co-authored-by: MHYang <mengyang@amd.com> Co-authored-by: Clement Lin <Clement.Lin@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mh <mh@local>

…n tests Squashed fix + tests (PR#3 commits 50729303, 2f33bd40, 1f0c6485, 46e3e056). flydsl_moe_stage1 only downgraded tile_n (256->128) for non-256-aligned inter_dim on a16w4 (bf16 x mxfp4); a8w4 (fp8 x mxfp4) kept tile_n=256, which in separated gate_mode over-runs the gate/up (N) axis for non-256 inter_dim: ~30% wrong E2E output or GPU memfault at inter=384/640. Extend the resolve to a8w4 (b_dtype fp4/mxfp4, a_dtype in {bf16,fp8}); a4w4 untouched. Callers keep passing tile_n=256; the kernel resolves internally. Tests (numeric vs torch reference, no-pad path, 128-multiple inter_dim): - test_flydsl_e2e_a8w4_situv2 (separated, the production/customer path) - test_flydsl_e2e_a16wfp4_situv2 (separated + interleave) - test_flydsl_e2e_a8w4_gui extended to inter=256/384/640 (interleave, swiglu) Verified on gfx950 (full rebuild): a8w4 20 / a16wfp4 27 / a4w4 4 passed; ruff + black clean.

…tuv2 E2E Squashed (PR#3 c827ddbb + d381aae7). - resolve_flydsl_stage1_tile_n / resolve_flydsl_stage2_tile_k: keep the silent auto-downgrade of a non-dividing tile (256->128) for non-256 inter_dim, but make it observable -- full docstrings noting tile=256 is NOT tunable for such shapes, plus a one-time (deduped) logger.warning on override. - test_flydsl_e2e_a8w4_situv2: enable interleave (a16w4-style shuffle_weight_a16w4 recipe) alongside separated, over 128-multiple inter_dim. Verified all activations (silu/swiglu/situv2) x both gate_modes x non-256 give E2E 0.0000.

Needed by atom-k3 (rocm/atom HEAD imports them from aiter.ops.shuffle); ported verbatim from aiter main. Keeps the single aiter-k3 branch usable by vllm-k3, sglang-k3 and atom-k3.

Bring k3-for-amd up to rocm/aiter main (6c0c6f6, +25 commits) on top of the K3 kernel work (fda97e9 shuffle ops). One conflict in aiter/ops/shuffle.py (GUI-interleave scale padding): rocm's version kept — its if k_%8!=0 block already pads with proper fill, making the K3 block redundant (identical padded shape; scale fill value is irrelevant since the padded weights are zero). interleave_gate_up_rows + moe_shuffle_weight preserved.

…API removed in 0.2.4)

[kimi-K3] extend conv2d support to gfx1250

flydsl 0.2.4 removed the loc= parameter from rocdl._split_mfma_operands. The A16W4 BF16 K32 MFMA helper in mixed_moe_gemm_2stage.py called it as _split_mfma(operands, loc=loc), which raises TypeError under 0.2.4. Drop loc= from both (identical) helper sites; _split_mfma_operands only unwraps operand Values and reads int flags, so it never needed loc. The loc/ip are still threaded to the real op builder (_mfma_k32_raw), so location tracking is preserved. Bump _MIN_FLYDSL_VERSION to 0.2.4 and drop the now-stale 0.2.2 pin note. Verified on gfx950 with flydsl 0.2.4 via test_flydsl_moe_a16wfp4.py (stage1, stage2, e2e all pass, cos>=0.99999). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

[Kimi-K3] fix MoE A16W4 for flydsl 0.2.4 and bump _MIN_FLYDSL_VERSION

[Kimi-K3] add fmoe tuned config

[Kimi-K3] extend a16w4 fmoe tuned config inter_dim=384

Keep Kimi-K3 SiTUv2 on the A4W4 FlyDSL path, add gfx1250 grouped-MoE activation and stability support, preserve ATOM's legacy A4W4 call, and cover the gfx1250 path with focused tests. Co-authored-by: Cursor <cursoragent@cursor.com>

Signed-off-by: Hongxia Yang <hongxia.yang@amd.com>

Enable a8w4 SiTUv2 MoE via AITER_SITUV2_A8W4

patches/ck_gfx1250_enable.patch only served to build native CK kernels (quant/cache/rmsnorm/moe/sample/custom_all_reduce) on gfx1250. Verified that with ENABLE_CK=0 (no patch, triton/hip fallbacks) Kimi-K3 serves correctly on gfx1250 x4 (tp4): full gsm8k 1319 = 0.956 flexible/strict, vs 0.9613 with the patched native-CK path — within stderr (+/-0.0056). Dropping the patch keeps cleanup/k3-minimal minimal for merge into k3-for-amd; gfx1250 runs ENABLE_CK=0 (cost: ~13% decode vs native CK, no accuracy loss).

Kimi-K3 runs full-attn via MLAAttention (MLA-latent) and KDA via fla, so it never dispatches to unified_attention (verified 0 calls across a full gsm8k run). The non-power-of-two head_size / Gluon-reduce guards were added for K3's old MHA-via-unified_attention path, which ATOM has now removed. Reverting unified_attention.py to k3-for-amd keeps cleanup/k3-minimal focused on what K3 actually needs. Full gsm8k 1319 with this + the dropped CK patch = 0.955.

[Kimi-K3] add a8w4 fmoe tune config

Resolve aiter/fused_moe.py SiTUv2 conflict by taking k3-for-amd's a16w4/a8w4 logic: on gfx1250 q_dtype_a is unconditionally overridden below (fp4x2, or a8w4 via AITER_FORCE_A8W4), so Kimi-K3 is unaffected by this branch; the upstream version is more correct for non-gfx1250. PR#14's remaining slimming (drop CK gfx1250 patch, drop unified_attention Gluon guard) is preserved.

# Conflicts: # csrc/kernels/topk_softmax_kernels_group.cu

add gfx1250 bf16 tuned config

github-actions · 2026-07-27T15:13:27Z

🏷️ 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 4397 --add-label <label>

github-actions

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

ruff

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

aiter/aiter/ops/flydsl/moe_kernels.py

Line 33 in bc3c444

except Exception: # pragma: no cover - logging must never break the kernel

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

aiter/aiter/ops/flydsl/moe_kernels.py

Line 141 in bc3c444

def get_flydsl_kernel_params(name: str) -> Optional[Dict]:

⚠️ [ruff] <UP006> _{reported by reviewdog 🐶}
Use dict instead of Dict for type annotation

aiter/aiter/ops/flydsl/moe_kernels.py

Line 141 in bc3c444

def get_flydsl_kernel_params(name: str) -> Optional[Dict]:

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

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1279 in bc3c444

w1_scale: Optional[torch.Tensor] = None,

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

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1280 in bc3c444

a1_scale: Optional[torch.Tensor] = None,

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

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1281 in bc3c444

sorted_weights: Optional[torch.Tensor] = None,

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

aiter/aiter/ops/flydsl/moe_kernels.py

Line 1285 in bc3c444

k_batch_intra_block: Optional[int] = None,

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

aiter/op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py

Lines 31 to 52 in bc3c444

    
           import pytest  # noqa: E402 
        
           import torch  # noqa: E402 
        
           import aiter  # noqa: E402 
        
           from aiter import dtypes, QuantType, ActivationType  # noqa: E402 
        
           from aiter.fused_moe import (  # noqa: E402 
        
               fused_topk, 
        
               moe_sorting, 
        
               torch_moe_stage1, 
        
               torch_moe_stage2, 
        
           ) 
        
           from aiter.jit.utils.chip_info import get_gfx  # noqa: E402 
        
           from aiter.ops.flydsl.utils import is_flydsl_available  # noqa: E402 
        
           from aiter.ops.flydsl.moe_kernels import (  # noqa: E402 
        
               pick_flydsl_stage1_tile_n, 
        
               pick_flydsl_stage2_tile_k, 
        
           ) 
        
           from aiter.ops.shuffle import shuffle_scale_a16w4, shuffle_weight_a16w4  # noqa: E402 
        
           from aiter.ops.quant import (  # noqa: E402 
        
               mxfp4_moe_sort_fwd, 
        
               per_1x32_f8_scale_f8_quant, 
        
           ) 
        
           from aiter.utility.fp4_utils import e8m0_shuffle  # noqa: E402