Skip to content

MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4) - #4502

Open
coderfeli wants to merge 42 commits into
mainfrom
flydsl-a16w4-mxfp-moe
Open

MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4)#4502
coderfeli wants to merge 42 commits into
mainfrom
flydsl-a16w4-mxfp-moe

Conversation

@coderfeli

@coderfeli coderfeli commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator

Summary

Replaces aiter's a16w4 (bf16 activation x MXFP4 weight) SiTUv2 fused-MoE
kernel with a numerically-correct, faster ported FlyDSL 2-stage kernel
(aiter/ops/flydsl/kernels/moe_2stage_a16wmix), and removes the old
compile_mixed_moe_gemm{1,2}_a16w4 path entirely.

aiter's previous a16w4 SiTUv2 kernel fails aiter's own strict accuracy gate
(op_tests/test_moe_2stage.py, bf16 SiTUv2 reference): logits_diff ~1.0
(essentially wrong output). The port passes the same gate at logits_diff ~1.5e-5 across the full token sweep on both tuned Kimi shapes (3584x384,
3584x512, E896, topk16).

Reproduction (independent re-measurement, gfx950)

Re-measured the shipped port against its own tuned CSV
(kimik3_fp4_tuned_fmoe.csv, abf16_wfp4) on gfx950 (cu=256) via
test_moe_2stage.py --kernel --no-legacy --csv-filter abf16_wfp4, both Kimi
shapes. The kernel runs correctly and meets or beats its tuned latency at
every token
(0.80–1.06×), monotonic across M. Ratio = measured / tuned-CSV
(lower = faster than the recorded tune).

Correctness: op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py 6/6 pass
(2 shapes × tok 1/16/128) vs the authoritative torch SiTUv2 reference.

3584x512

tok measured total (µs) tuned-CSV (µs) ratio
1 19.9 24.0 0.83×
2 29.7 33.3 0.89×
3 38.9 42.5 0.92×
4 41.4 48.2 0.86×
8 62.2 77.5 0.80×
16 111.8 140.1 0.80×
32 200.0 222.4 0.90×
64 307.0 330.1 0.93×
128 424.8 428.1 0.99×
256 479.2 483.2 0.99×
512 470.4 508.2 0.93×
1024 524.5 541.6 0.97×
2048 843.9 792.8 1.06×
4096 1232.2 1215.6 1.01×
8192 2021.1 2025.7 1.00×
16384 3669.9 3694.2 0.99×

3584x384

tok measured total (µs) tuned-CSV (µs) ratio
1 19.0 21.6 0.88×
2 25.8 28.5 0.91×
3 34.0 38.9 0.87×
4 36.9 42.7 0.86×
8 56.8 59.8 0.95×
16 91.5 107.3 0.85×
32 162.6 177.0 0.92×
64 230.9 260.6 0.89×
128 314.0 336.0 0.93×
256 354.1 374.7 0.95×
512 354.7 392.7 0.90×
1024 391.3 422.5 0.93×
2048 644.4 651.3 0.99×
4096 994.1 1087.3 0.91×
8192 1689.6 1994.4 0.85×
16384 3081.7 3600.3 0.86×

What changed

  • Vendored kernel package aiter/ops/flydsl/kernels/moe_2stage_a16wmix/
    (3 files: gemm1.py, gemm2.py, __init__.py): gemm1 (gate+up GEMM + SiTUv2
    -> sorted bf16 intermediate), gemm2 (down-proj + routing-weighted atomic bf16
    scatter). Non-scaled MFMA(16,16,32,bf16), CDNA (gfx942/gfx950). Consolidated
    layout (shared helpers in gemm1, launch glue + CSV tile resolvers in
    __init__); scale loads use the make_buffer_tensor/BufferCopy32b layout API.
  • fused_moe_ dispatch (aiter/ops/flydsl/moe_2stage_a16w4_dispatch.py):
    serves bf16-a16w4 SiTUv2 (the kernel is gate-mode agnostic). Converts the
    incoming guinterleave (a8w4/mxfp8-shared) weight/scale layout to the port's
    standard shuffle_weight/e8m0_shuffle layout via pure byte permutations
    (verified round-trip). gemm2 tiles come from aiter's tuned CSV
    (resolve_a16w4_gemm2_config, on-par-to-faster); gemm1 keeps its own tuned
    default — a perf-gate showed the CSV gemm1 tiles regress mid/high-M by
    1.1x–2.0x, so they are not taken.
  • Removed the old a16w4 path: compile_mixed_moe_gemm{1,2}_a16w4, the compat
    wrapper compile_a16w4_moe_gemm2, and the a16w4-only _decode_e8m0_byte_to_f32
    helper (~4076 lines) from mixed_moe_gemm_2stage.py; the a16w4 branches of
    compile_flydsl_moe_stage1/2; and the dead a16w4 sub-clause / _a_type
    bf16-fp16 branches in get_2stage_cfgs. Generic compile_mixed_moe_gemm1/2
    and all other dtypes (a8w4 / int4 / mxfp8 / a4w4) are untouched.
  • Fallback handling: features the port does not implement (non-default SiTUv2
    beta/linear_beta, per-expert bias, expert-parallel masking) raise a clear
    NotImplementedError rather than routing to the numerically-broken old kernel.
  • Tests: aiter's original op_tests/test_moe_2stage.py a16w4 (abf16_wfp4)
    SiTUv2 rows now exercise the new kernel through the production dispatch —
    _effective_gate_mode returns SEPARATED for bf16-a16w4 (previously INTERLEAVE,
    which silently downgraded those rows to a8w4/fp8), and the a16w4 SiTUv2 path is
    bias-free (the port has no per-expert bias). No bespoke test file.

@coderfeli
coderfeli requested a review from a team August 1, 2026 08:29
@github-actions

github-actions Bot commented Aug 1, 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 4502 --add-label <label>

@coderfeli coderfeli changed the title MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel MoE: correct + faster a16w4 (bf16 A x MXFP4 W) SiTUv2 kernel (replaces old a16w4) Aug 1, 2026
coderfeli and others added 3 commits August 1, 2026 11:13
Replace the a16w4 SiTUv2 SEPARATED fused-MoE path with a numerically-correct
ported FlyDSL 2-stage kernel (aiter/ops/flydsl/kernels/moe_2stage_a16wmix).
aiter's previous compile_mixed_moe_gemm{1,2}_a16w4 fails aiter's own strict
accuracy gate on this path (logits_diff ~1.0); the port passes at ~1.5e-5.

- Vendored kernel package: gemm1 (gate+up GEMM + SiTUv2 -> sorted bf16
  intermediate) and gemm2 (down-proj + routing-weighted atomic bf16 scatter),
  non-scaled MFMA(16,16,32,bf16), CDNA (gfx942/gfx950).
- fused_moe_ a16w4-SiTUv2-SEPARATED branch: converts the incoming guinterleave
  (a8w4/mxfp8-shared) weight/scale layout to the ported kernel's standard
  shuffle_weight/e8m0 layout (pure byte permutations, verified round-trip) and
  routes to the port. Gated + additive; a8w4/mxfp8/INTERLEAVE paths untouched.
- Focused test: op_tests/flydsl_tests/test_a16w4_flydsl_situv2.py.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Step 1 (perf, CSV-driven, --kernel isolated, balanced routing, median-of-3,
gfx950) confirmed the ported kernel is on-par-or-better than aiter's old a16w4
across the throughput band: OURS wins the entire M>=16 range on both Kimi shapes
(0.83-0.99x), with the old kernel ahead only on the tok<=8 fixed-launch-overhead
points and the single 3584x512 tok16384 cell.

Given that (and that the old kernel is numerically broken, 0/N on the strict
gate), remove aiter's a16w4 path entirely so the port is the sole a16w4 kernel:

- mixed_moe_gemm_2stage.py: drop compile_mixed_moe_gemm1_a16w4,
  compile_mixed_moe_gemm2_a16w4, the compat wrapper compile_a16w4_moe_gemm2, and
  the a16w4-only _decode_e8m0_byte_to_f32 helper (~4076 lines). Generic
  compile_mixed_moe_gemm1/2 and all other dtypes (a8w4/int4/mxfp8/a4w4) are
  untouched; now-unused imports pruned.
- moe_kernels.py: remove the a16w4 branch of compile_flydsl_moe_stage1/2; it now
  raises NotImplementedError (the port is dispatched directly from fused_moe_).
- fused_moe.py: broaden the a16w4 dispatch to serve bf16-a16w4 SiTUv2 in both
  gate modes (the kernel is gate-mode agnostic; only activation-dtype routing
  decides bf16 vs fp8), and remove the dead a16w4 sub-clause / _a_type bf16-fp16
  branches from get_2stage_cfgs' use_mxfp4_flydsl path.
- Fallback handling: features the port does not implement (non-default SiTUv2
  beta/linear_beta, per-expert bias, expert-parallel masking) now raise a clear
  NotImplementedError instead of routing to the numerically-broken old kernel
  (explicit-unsupported over silently-wrong).
- Remove op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py (tested the removed
  low-level a16w4 kernel API); the production path is covered by
  test_a16w4_flydsl_situv2.py.

Correctness gate re-run (aiter test_fmoe, a16w4 SiTUv2 SEPARATED, strict, gfx950,
no-bias supported path): PASS at logits_diff ~1.5e-5 across the token sweep on
both shapes. a8w4 FlyDSL tests still pass (shared mixed_moe path intact).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…riginal test

- Port FlyDSL PR #948 consolidation (commit 3d1e03d6) into the vendored package:
  common.py + host.py removed; shared helpers folded into gemm1.py, launch glue +
  CSV resolvers + a16wi4_recommend_block_m into __init__.py. Carried over the
  buffer_ops modernization (scale-load -> make_buffer_tensor/BufferCopy32b). aiter
  adaptations preserved (kernels.common.* -> aiter.ops.flydsl.kernels.*, the
  CDNA3-encoded s_waitcnt shim). Also dropped the unused vendored
  _gemm_common_utils.py (never imported). Public API (flydsl_a16w4_gemm1/2)
  unchanged.
- Condense moe_2stage_a16w4_dispatch.py to the essential dispatch (guinterleave->
  standard-shuffle inversion, CSV gemm2 tiles); keep load-bearing notes.
- gemm2 now takes aiter's tuned CSV tiles via resolve_a16w4_gemm2_config
  (on-par-to-faster: 0.96-1.02x vs kernel default). gemm1 keeps its own tuned
  default (tile_n=128 + M>=16 TILE_K=128/xcd lever): perf-gate showed the CSV
  gemm1 tiles regress mid/high-M by 1.09x-2.02x, so they are NOT taken.
- Route aiter's ORIGINAL a16w4 test through the new kernel: _effective_gate_mode
  returns SEPARATED for bf16-a16w4 (was INTERLEAVE, which downgraded abf16_wfp4
  rows to a8w4/fp8 and never exercised the a16w4 kernel), and the a16w4 SiTUv2
  path is bias-free (the port has no per-expert bias). Removed the bespoke
  test_a16w4_flydsl_situv2.py; the original test_moe_2stage.py abf16_wfp4 rows now
  cover the new kernel.

Gate: aiter test_fmoe a16w4 SiTUv2 SEPARATED strict (gfx950): PASS at
logits_diff ~1.5e-5 across the token sweep, both shapes. a8w4 FlyDSL tests still
pass (25/25). black-88 + ruff clean.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
- Sync the FlyDSL CSV refactor (a16wmix_tuned.csv + resolve_a16wmix_gemm1/2_config
  + _default_tiles_fallback) into the vendored package. The default path of
  flydsl_a16w4_gemm1/2 now resolves tiles from our tuned CSV (data-vs-code;
  provably perf-neutral: CSV values equal the removed inline heuristics, IR
  byte-identical). CSV auto-located in the package dir via
  os.path.dirname(__file__).
- Reroute op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py to test the NEW a16w-mix
  a16w4 kernel through the production fused_moe dispatch (bf16 A x mxfp4 W, SiTUv2 /
  SEPARATED) with a strict cos/logits_diff gate -- explicit FlyDSL-side coverage
  complementing the routed test_moe_2stage.py rows. Does NOT use the removed
  low-level compile_mixed_moe_gemm1_a16w4 API.

Gate: a16w4 SiTUv2 strict (gfx950): PASS at logits_diff ~1.5e-5 (unchanged --
CSV == old defaults). Rerouted test 6/6 pass. black-88 + ruff clean.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli added a commit that referenced this pull request Aug 1, 2026
… resolvers)

Sync the FlyDSL CSV refactor into the vendored package to keep it identical with
PR #4502: add a16wmix_tuned.csv (our tuned tiles, incl. int4 rows) and the
resolve_a16wmix_gemm1/2_config + _default_tiles_fallback resolvers; the
flydsl_a16w4_gemm1/2 default path is now CSV-driven (data-vs-code, perf-neutral:
CSV values equal the removed inline heuristics). CSV auto-located in the package
dir.

The a16wi4 dispatch keeps reading aiter's tuned int4 CSV (kimik2_i4_tuned_fmoe.csv,
the head-to-head-validated on-par config) via resolve_a16w4_gemm1/2_config; the
package sync does not change its behavior.

Gate: a16wi4 strict (gfx950): PASS at logits_diff ~1.5e-5 across the Kimi int4
shapes. black-88 + ruff clean. Package byte-identical with #4502.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli and others added 20 commits August 1, 2026 15:22
…age1/2, drop bespoke layer

Rewire the a16w4 (bf16 A x MXFP4 W) SiTUv2 MoE path so the ported FlyDSL kernel is a
genuine drop-in, removing the per-forward weight relayout and the redundant tuned CSV.

Native layouts (no relayout). The ported gemm1 gains a w_layout="guinterleave" mode that
consumes aiter's native GUGU stage1 W1+scale (shuffle_weight_a16w4 / shuffle_scale_a16w4)
directly. gemm2's native gate_up=False W2+scale is byte-identical to the standard layout
it already reads (E*model_dim % 256 == 0), so no mode is needed. This deletes the four
per-forward _unshuffle_guinterleave -> standard-shuffle passes (hundreds of MB for Kimi
E896) the old dispatch ran on every call.

Route through compile_flydsl_moe_stage1/2. The a16w4 branches of
compile_flydsl_moe_stage{1,2} now build the ported kernel (was: raise). The a16w4 run
path (run_flydsl_a16w4_moe in moe_kernels.py, replacing the deleted
moe_2stage_a16w4_dispatch.py) sorts, resolves gemm2 tiles from aiter's tuned kimik3 CSV
(resolve_a16w4_gemm2_config), and launches both stages; fused_moe_'s a16w4 case calls it.

Drop the redundant tuned config. Delete a16wmix_tuned.csv (321 rows) and the
_load_ours_tuned_csv / _resolve_ours_tuned stack; resolve_a16wmix_gemm{1,2}_config now
return the documented heuristic directly. gemm1 keeps that heuristic (aiter's kimik3
gemm1 tiles regress the ported body); gemm2 uses the kimik3 CSV.

Verified cold on gfx950 (MI355X, cu=256):
  - op_tests/test_moe_2stage.py --csv-filter abf16_wfp4: 32/32 kimik3 a16w4 rows pass
    strict accuracy (logits_diff <= 1.6e-5); perf 54us@tok1 .. 3779us@tok16384 (inter512),
    strictly faster than the relayout dispatch.
  - Native-vs-standard kernel parity bit-identical (max|d|=0.0).
  - op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.py: 6/6 passed.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Comment-only cleanup of the a16w4 shared-kernel package: collapse ASCII
banner blocks to one-line section headers and drop per-line narration that
restated adjacent code (the k_wave=1 else-branch note, the BufferCopy128b
restatement, verbose route-through and use_csv_config comments).

No code change: the Python AST is byte-identical on all three files, so
codegen/ISA and perf are unchanged. All load-bearing "why" comments are
kept -- the @flyc.jit LOAD-BEARING / ~13x-regression warning above the
atomic epilog, the bf16-W num_records i64-fold ABI note, offset-unit
gotchas, and the preshuffle layout/stride math.

black-88 + ruff clean.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Trim the verbose a16w4 SiTUv2 dispatch/guard comments in _fused_moe_impl.
No code change (token-identical to parent); a16w4 + a8w4 rows still pass
the strict test_moe_2stage accuracy gate.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…16w4_moe

Reuse moe_sorting's already-zeroed moe_buf as the gemm2 atomic-scatter target
(matches the a8w4/mxfp4 sibling path) instead of allocating+zeroing a separate
out_buf; make inter_sorted torch.empty (gemm1 masked-store + gemm2 token_id<M
gate mean padding never reaches the output); and drop the two unconditional
torch.cuda.synchronize() calls (g1->g2 order is guaranteed on one stream).

Removes 2 host syncs + 2 device memsets per call. End-to-end fused_moe latency
(gfx950, kimik3 3584x512 a16w4): tok1 -26%, tok32 -7%, tok4096 neutral.
Accuracy unchanged (logits_diff ~1.5e-5). The syncs also blocked CUDA-graph
capture.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
a16wmix_use_k16 now takes the target arch as a required param (callers pass
get_rocm_arch()) instead of reading it internally, and the FLYDSL_A16WMIX_FORCE_K16
test-hook env is removed. Behavior on gfx950 is unchanged (use_k16=False);
validated logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4). Removes the only os use in
gemm1.py.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Move the built-in tile heuristic (resolve_a16wmix_gemm{1,2}_config,
_default_tiles_fallback, _default_tile_n) and the aiter tuned-CSV resolvers
(resolve_a16w4_gemm{1,2}_config, pick_a16w4_config, kernelName decoders, CSV
loader) out of the package __init__ into csv_dispatch.py. __init__ drops from
810 to ~505 lines and keeps only the compile dispatch + launch glue. Pure code
move (re-imported into __init__); tile selection and the compiled kernels are
unchanged -- validated logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Extract _silu_mul_batch/_sigmoid_f32/_tanh_f32/_situ_mul_batch (and LOG2E) from
gemm1.py into a leaf act.py that uses tensor_shim._to_raw (equivalent to gemm1's
local _raw for these fx inputs, so no circular import). gemm1 imports the two
batch helpers it calls. Pure move; validated logits_diff ~1.5e-5 (a16w4) /
2e-5 (a8w4).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace gemm1's local _raw IR-unwrap shim with tensor_shim._to_raw (aliased as
_raw), so the whole package (gemm1, gemm2 via import, act) shares the one
tensor_shim helper. Verified equivalent at every call site (all pass fx values):
accuracy smoke unchanged at logits_diff ~1.5e-5 (a16w4) / 2e-5 (a8w4), pyflakes
clean. Remaining raw llvm/rocdl (MFMA, BufferCopy, masked buffer_store, ptr/gep
shims, side-effecting cvt_pk) are load-bearing or have no tensor_shim equivalent
and are intentionally kept.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…inside)

compile_gemm{1,2}_a16w4_port now take use_k16 as a required keyword arg instead
of calling a16wmix_use_k16(get_rocm_arch()) internally. The dispatch callers
(compile_flydsl_moe_stage{1,2} and _get_compiled_gemm{1,2}_a16w4) resolve it once
and pass it in; the arch query no longer lives in the kernel-building path.
Drops the now-unused get_rocm_arch import from gemm1/gemm2. Validated
logits_diff ~1.4e-5 (a16w4) / 1.9e-5 (a8w4); pyflakes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Remove the use_csv_config/csv_path opt-in from flydsl_a16w4_gemm{1,2}: no caller
  ever set it True (production uses the heuristic for gemm1 and passes CSV-resolved
  tiles to gemm2 explicitly), so the branch was dead. The heuristic
  resolve_a16wmix_gemm{1,2}_config path is now unconditional.
- Drop resolve_a16w4_gemm1_config and _kw_tile_k_for (only reachable via the dead
  use_csv_config branch). resolve_a16w4_gemm2_config stays (run_flydsl_a16w4_moe
  calls it directly).
- Drop _kwave_from_kbatch + the _kb kernelName tag decode: a16w4 tuned rows never
  carry _kb, so k_wave always came from _kw (or 1).

Dead-code removal; tile selection and compiled kernels unchanged (CPU-verified:
heuristic configs + kernelName decode identical). pyflakes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…sl_a16w4_moe

Route a16w4 (bf16 A x MXFP4 W, SiTUv2) through get_2stage_cfgs -> fused_moe_2stages
like a4w4/a8w4, instead of the bespoke pre-dispatch intercept:

- get_2stage_cfgs: re-add a16w4 to use_mxfp4_flydsl (bf16 A, mxfp4 W, Situv2) with
  _a_type="bf16" (kernelName abf16_wfp4, no _gui) and fuse_quant="" -- reusing the
  existing skip-quant branch (a2_scale=None, bf16 intermediate, no inter-stage quant).
- _flydsl_moe_stage{1,2}_impl: replace the stale a16w4 handling (which built the
  standard _s1/_s2 args -- wrong ABI for the port kernel) with early branches that
  call the ported flydsl_a16w4_gemm{1,2} launchers. stage1 returns the bf16
  [sorted_size, inter_dim] sorted intermediate (threaded to stage2 unchanged);
  stage2 atomic-scatters into the moe_sorting-zeroed moe_buf. gemm1 keeps the tuned
  heuristic tiles (tile_n=None); gemm2 uses the CSV tiles (resolve_a16w4_gemm2_config).
- _fused_moe_impl: drop the run_flydsl_a16w4_moe intercept; keep the NotImplementedError
  guards for features the port can't serve (beta/linear_beta, bias, expert_mask, gfx).
- stage2_uses_route_reduce: treat any a16w4 (bf16 A) stage2 as atomic, so a "reduce"-
  named CSV config still gets accumulate=True sorting (moe_buf zeroed) -- the ported
  gemm2 only supports atomic scatter. Fixes a NaN at small-M reduce-named rows.
- Remove run_flydsl_a16w4_moe (its per-stage logic now lives in the impls).

Validated (gfx950, kimik3): logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest
6/6; e2e perf parity with the bespoke path (tok4/8/256/16384 within ~2%); a4w4/a8w4
metadata unchanged. pyflakes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… guard raises

- _flydsl_moe_stage1_impl a16w4 branch: allocate inter_sorted with torch.zeros when
  inter_dim_pad > 0 (torch.empty otherwise) -- restores the pre-fold pad-zeroing that
  was dropped when the alloc moved into the early branch. No-op for unpadded shapes.
- Trim the verbose fold-branch comments in _flydsl_moe_stage{1,2}_impl.
- Collapse the five repetitive a16w4 NotImplementedError guards in _fused_moe_impl
  into one (cond, why) loop.

Behavior-preserving for the tuned (unpadded) a16w4 path; pyflakes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… kernels dir

#2 (dead CSV code): after the fold, a16w4 gemm2 tiles are already carried by the
get_2stage_cfgs kernelName2 (parsed into tile_n/tile_k/b_nt/waves_per_eu/xcd_swizzle
by _flydsl_stage2_wrapper), so the stage2 impl re-reading the CSV via
resolve_a16w4_gemm2_config was redundant. Switch the a16w4 stage2 branch to the
passed tiles (same as a4w4/a8w4) and delete the a16w4 CSV family
(resolve_a16w4_gemm2_config, pick_a16w4_config, _load_a16w4_csv, _decode_a16w4_kname,
_default_a16w4_csv_path, kernelName regexes, FLYDSL_A16W4_TUNED_CSV env). csv_dispatch.py
is now heuristic-only -> renamed tile_heuristic.py.

#4: move act.py (SiLU/SiTUv2 helpers) out of moe_2stage_a16wmix/ to the shared
aiter/ops/flydsl/kernels/act.py so other FlyDSL MoE kernels can reuse it.

Validated (gfx950 kimik3): logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest
6/6; a16w4 e2e perf parity (tok4/8/256/16384 within noise; get_padded_M is identity
at those tokens so kernelName tiles == the old CSV-resolved tiles). pyflakes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…branches)

The flydsl_a16w4_gemm{1,2} launchers (the only callers of resolve_a16wmix_*) are
mxfp4-only; a16wi4 (int4) dispatches via its own get_2stage_cfgs block + kimik2_i4
tuned CSV, never through this heuristic. So the int4/bf16 branches in
_default_tile_n/_default_tiles_fallback (and the w_dtype param) were dead. Remove
them and the w_dtype plumbing from the resolver calls. mxfp4 heuristic outputs are
byte-identical (CPU-verified: gemm1 tok2/8/256 tiles unchanged); the load-bearing
gemm1 heuristic stays (CSV gemm1 tiles regress). pyflakes clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…pret

Replace the fx.Float32(_raw(x).bitcast(T.f32)) boundary in the int4/mxfp4 scale
loads (gemm1 load_b_scale_int4 + _e8m0_byte_to_f32, gemm2 load_b_scale_int4) with
the native .bitcast(fx.Float32) on the already-typed fx.Int32 expr (same emitted
arith.bitcast; drops the _raw unwrap + raw T.f32 + redundant fx.Float32 re-wrap),
and drop one redundant fx.Float32 re-wrap of `eff` in _int4_nibble_to_bf16x8.

Load-bearing raw ops kept: _udiv/_umod (fx.Int32 is signed -> // / % emit
floordivsi/remsi != divui/remui), the masked buffer_store output scatter (fx.copy
can't express runtime-num_records masked scalar scatter), _buffer_i32_scalar_read
(already on the layout API). Validated: pytest 6/6, logits_diff a16w4 ~1.5e-5 /
a8w4 ~1.9e-5, perf parity.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… a16wi4 helpers

#6: _get_compiled_gemm{1,2}_a16w4 is @functools.cache'd, so the builder runs once and
the returned @flyc.jit launcher is dispatched via _run_compiled. Drop the
compile_flydsl_moe_stage{1,2} round-trip in _get_compiled and call
compile_gemm{1,2}_a16w4_port directly -- the kernel name/cache is set by the builder's
name_suffix, so this yields the identical cached kernel (AOT, which still uses
compile_flydsl_moe_stage{1,2}, matches). One fewer hop.

Also remove a16wi4_recommend_block_m and a16wi4_scale_to_kernel_layout: dead on this
branch (only in __all__; a16wi4 dispatches via its own get_2stage_cfgs block + kimik2_i4
CSV, not these helpers).

Validated: pytest 6/6, logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…istic

The gemm1 tile heuristic (tile_heuristic.py resolve_a16wmix_gemm1_config) existed
only because it bypassed the kernelName registry to use tile_k=128 / xcd=1 tiles
that weren't registered. Register those variants and drive gemm1 from the CSV like
gemm2 / a4w4 / a8w4, then delete the heuristic:

- get_flydsl_stage1_kernels: for a16w4 (bf16 A x fp4 W) only, register tile_k in
  {128,256} and xcd_swizzle in {0,1,4} (other dtypes unchanged).
- kimik3_fp4_tuned_fmoe.csv: rewrite the 32 gfx950 kernelName1 to the tuned tiles
  (tok1-2 t32x64x128_bnt0_kw4; tok3-8 t32x128x256_bnt0; tok16-1024 t32x128x128_xcd1;
  tok2048+ t32x128x128_bnt0_xcd1) -- both shapes.
- _flydsl_moe_stage1_impl a16w4 branch: pass the wrapper-parsed CSV tiles
  (tile_n/tile_k/k_wave/b_nt/xcd_swizzle) to flydsl_a16w4_gemm1 instead of tile_n=None;
  force waves_per_eu=None (kernel default) since a no-_w name parses to wpe=1.
- flydsl_a16w4_gemm1/2: drop the resolve_a16wmix_* / _default_tile_n fallback; use the
  caller tiles. Delete tile_heuristic.py + its imports/__all__.

Perf-neutral (CSV tiles == the old heuristic, wpe=None matched): a16w4 e2e 3584x512
tok8=83 / tok256=458 / tok16384=3758 us (ref 84/451/3749, within ~3%). Correctness
logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5; dedicated pytest 6/6 (incl. token=1 kw4);
pyflakes clean; a8w4/a4w4/other-dtype registration unchanged.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- CI: black-reformat moe_kernels.py; wrap the implicit f-string concat in the
  a16w4 guard (ruff ISC004).
- Align the "fp4"/"mxfp4" dtype naming to "fp4": b_dtype is only ever "fp4" at
  runtime (parsed from the abf16_wfp4 kernelName), so the "mxfp4" alias in the
  b_dtype checks was dead -> drop it; rename the port w_dtype "mxfp4"->"fp4"
  (byte-exact: name_suffix maps both to ""). Covers moe_kernels, fused_moe, aot,
  and the moe_2stage_a16wmix package. No CSV change needed (kernelName is wfp4).
- Trim the verbose a16w4 comment blocks in moe_kernels (compile entries + stage
  impl branches).

Behavior-preserving: dedicated pytest 6/6, logits_diff a16w4 ~1.5e-5 / a8w4 ~2e-5;
black + ruff clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Merge the byte-identical _gep1/_gep3 into one _gep. The GEP op is polymorphic
  in the base ptr's address space (the ptr already carries ptr<1>/ptr<3>), so the
  two names were pure duplication.
- Drop the duplicated fx.PointerType.get(i32, Global) construction: move
  _global_i32_ptr above _global_i32_buffer_view and have the buffer view reuse it.

Pure refactor, byte-behavior unchanged: dedicated pytest 6/6, logits_diff
a16w4 ~1.5e-5 / a8w4 ~2.1e-5; black + ruff clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… gap)

The decode-latency tok1/tok2 cells used t32x64x128_kw4 / t32x128x256, which lost
to the (numerically-broken) old a16w4 kernel by ~10-20% fixed launch overhead.
A tile_n=32 gemm1 tile (t32x32x128_kw4) is 25-30% faster there:

  3584x384 tok1 gemm1  17.4 -> 12.2 us   (total ~24.6 -> ~19.4, now beats old ~22)
  3584x512 tok1 gemm1  17.9 -> 13.5 us   (total ~24.8 -> ~20.1, now beats old ~24)
  3584x384 tok2 gemm1  19.6 -> 17.3 us

tok2/512 and tok3+ keep their current tiles (tile_n=32 regresses there as M grows
and the N tile starves). t32x32x128_kw4 is a registry-valid kernel (enum
4*tn<=tk holds: 128<=128), so it stays AOT-cacheable. Correctness unchanged:
logits_diff ~1.4-1.8e-5 at every changed cell; dedicated pytest 6/6.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli and others added 4 commits August 2, 2026 14:19
The AOT precompile drove the a16w4 (bf16 A x fp4 W) port through the generic MX
gemm launch-arg builder (_s1_args_fp4), but the ported moe_2stage_a16wmix kernels
have a different launch ABI (raw fx.Int64 device pointers) -> "too many positional
arguments" at flyc.compile -> the artifact never persisted -> every a16w4 kernel
reported "produced no kernel" (128 MOE AOT failures, CI break).

Add a dedicated _precompile_a16w4_to_cache that drives the SAME runtime launchers
(flydsl_a16w4_gemm{1,2}) the fused-MoE op uses, with real CPU dummies under
COMPILE_ONLY (outside FakeTensorMode, since the port takes raw .data_ptr()). This
guarantees the AOT cache key is identical to runtime by construction:
waves_per_eu=None (both stages), persist=False, w_layout="guinterleave", and the
runtime g2 tile downgrade are all applied inside the launcher.

Verified: kimik3_fp4 AOT now 128 ok / 0 failed; run-only (FLYDSL_RUNTIME_RUN_ONLY=1)
loads every a16w4 artifact for all 32 rows x both shapes (0 cache misses ->
keys match runtime, no JIT fallback). Non-a16w4 dtypes unchanged (gated on
a_dtype=="bf16" and b_dtype=="fp4").

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The artifact is keyed only on the port kernel's constexpr params (shapes / tiles /
topk / act), never on the launch pointers, grid, or n_tokens — and COMPILE_ONLY
never dereferences them. So the per-buffer dummy tensors (correctly-shaped a / w1 /
w1_scale / inter / w2 / ... plus the routing-size math) are unnecessary: a single
1-elem placeholder covers every pointer arg, and n_tokens/M_logical/max_sorted can
be 1. Shared kwargs are factored into one `common` dict across the two stages.
Dropped the now-dead CU_NUM/get_cu_num plumbing (a16w4 persist=False never reads
cu_num). Behavior unchanged: kimik3_fp4 AOT 128 ok / 0 failed, run-only 32/32 with
0 cache misses (keys still match runtime).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…cache

Behavior-identical; keeps CI ruff green.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The a16w4 port baked situ_beta/linear_beta at 1.0 (epilog called
_situ_mul_batch with defaults), so fused_moe hard-blocked non-default
beta with NotImplementedError. Thread beta/linear_beta as runtime f32
scalars (host passes value + reciprocal, mirroring swiglu_limit); the
compile cache key stays beta-free so one kernel serves any beta (incl.
AOT). Unblocks kimi k3 (beta=4.0, linear_beta=25.0).

_situ_mul_batch clamp_limit now defaults to +inf (no clamp), matching
the situv2() reference and the a8w4 situv2 path (swiglu_limit=inf).
The old hardcoded 7.0 was invisible at beta=1.0 (tanh/sigmoid saturate)
but diverged badly for large linear_beta (tanh(u/25) is near-linear).

Rename _is_a16w4 -> _is_a16w4_situv2 and document why the SiTUv2 gate
is load-bearing: gpt-oss is also bf16 A x MXFP4 W but Swiglu and takes
swiglu_mxfp4_bf16_cktile; the activation is what disambiguates them.

Validated gfx950: k3 4.0/25.0 cos_diff 1.5e-5; beta=1.0 unchanged;
N betas -> 1 compile; gptoss still routes to cktile.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@coderfeli
coderfeli force-pushed the flydsl-a16w4-mxfp-moe branch from c044d33 to 05bd519 Compare August 3, 2026 11:03
# Conflicts:
#	aiter/ops/flydsl/kernels/mixed_moe_gemm_2stage_common.py
#	aiter/ops/flydsl/moe_kernels.py
@zufayu
zufayu requested a review from yadaish August 4, 2026 01:27
coderfeli and others added 3 commits August 4, 2026 09:34
… a8w4)

The prior runtime-beta commit changed the a16w4 SiTUv2 clamp to +inf (no
clamp), which was wrong: the model (and the a8w4 mixed_moe path) clamp
SiTUv2 at swiglu_limit. At large linear_beta (k3 uses 25) tanh(u/25) is
near-linear, so no-clamp vs clamp-at-7 diverges massively — a16w4 vs a8w4
cos_diff was 0.46. (It passed earlier only because the torch reference used
for validation also didn't clamp.)

Thread swiglu_limit into the a16w4 port as a runtime f32 scalar (same
pattern as beta) and clamp via maximumf(-limit), identical to the
a8w4/mixed_moe situv2 clamp: +inf -> no clamp, finite -> clamp. Host passes
runtime_swiglu_limit(swiglu_limit, act) so a16w4 == a8w4 for any limit; the
scalar is out of the compile cache key (one kernel serves any limit/beta).

Validated gfx950: a16w4(lim=7,b=4,lb=25) vs a8w4 cos_diff 0.46 -> 5.5e-3
(residual is a8w4 fp8 quant); no-limit default unchanged (1.5e-5); a16wfp4
pytest 6/6.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
coderfeli and others added 3 commits August 5, 2026 14:52
Align the a16w-mix a16w4 (bf16 A x mxfp4 W, SiTUv2) kernel to main's GGUU
(standard, separated gate/up) W1 layout, and add real tuner support so the
kimik3_fp4 CSV is produced by gemm_moe_tune (was hand-written).

- Layout: a16w4 dispatch uses w_layout="standard" (GGUU) in the runtime
  (moe_kernels) and AOT (aot/flydsl/moe) paths; tuner W1 prep and both tests
  set to GGUU. a8w4 (fp8) stays gate/up-interleaved (GUGU).
- gemm_moe_tune a16w4 support: skip the unsupported CK b16xfp4 enumeration;
  pass situ_beta to ref+kernel; map the SORTED stage1 output <-> unsorted ref
  by expert-match (_a16w_sorted_row_map) for both the stage1 compare and the
  stage2 sorted intermediate input; guard out invalid/degenerate candidates
  (stage2 sbm, tile_n=256 LDS overflow, tile_k not dividing inter_dim, and
  stage1 num_acc_n<1 e.g. tile_n=32/kw1 which emits NaN).
- Retune kimik3_fp4_tuned_fmoe.csv (all 32 shapes).

Validated on gfx950: a16w4 pytest 6/6 (tok 1/16/128 x inter {512,384}) and
tok 1024/4096 x {512,384} 4/4; CSV has no degenerate/invalid tile picks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Merge the three separate a16w-mix stage2 skip checks (sbm tile_m, tile_n=256
LDS overflow, tile_k not dividing inter_dim) into a single guard. Logically
identical (same candidates skipped) -> no retune needed; a16w4 pytest 6/6.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli and others added 3 commits August 7, 2026 09:51
The stage1 gather / stage2 scatter I added ran inside the tuner's timed region
(run_perftest times run_func), so the recorded us was kernel + a full re-sort
over the sorted buffer (~28k rows) -- 2-3x inflated, and for small tokens the
re-sort dominated, corrupting kernel selection too.

Make run_func kernel-only; move the sort into the untimed ref/data-prep:
- stage1: run_flydsl_stage1_out returns the kernel's SORTED output as-is;
  accuracy compares in sorted space against run_a16w_stage1_sorted_ref (computes
  the torch ref then gathers to sorted order) via _a16w_sorted_cos.
- stage2: the sorted intermediate input (a2_a16w_sorted) is built once in
  generate_data and fed via a conditional run-arg; run_flydsl_stage2_out no
  longer re-sorts (just a2_scale=None for bf16).
- drop the timed-path _a16w_sorted_row_map / topk_ids plumbing.

Retune kimik3_fp4_tuned_fmoe.csv: us now tracks e2e (e.g. 512/1 333->27,
384/4096 1893->944), not inflated. Validated: a16w4 pytest 6/6 (tok 1/16/128)
and tok 1024/4096 4/4; no degenerate/invalid tile picks.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…cnt (ISA byte-exact)

Cleanup of the a16w4/a16wi4 package leaf helpers (no codegen change; gemm1+gemm2
AMDGCN verified byte-exact A/B via FLYDSL_DUMP_IR):

- Raw ptr casts -> flydsl llvm-ptr API: _lds_ptr3/_global_base_ptr1 now go through
  a single _int_to_llvm_ptr() using fx.to_llvm_ptr(fx.inttoptr(..., AS.Shared/Global)),
  dropping the hardcoded ir.Type.parse("!llvm.ptr<3|1>") (backend resolves the AS).
- Drop the hand-encoded CDNA3 s_waitcnt shim: flydsl 0.3.0 has the arch-dispatched
  keyword form fx.rocdl.s_waitcnt(lgkmcnt=0); same encoding (ISA identical).
- b_cache_mod: use the CSV-driven b_nt directly (constant default when unset); remove
  the dead/inconsistent per-token U-shape heuristic (moe_kernels always passes b_nt).
- Extract the 18 shared leaf helpers gemm2 imported from gemm1 into a package-local
  utils.py; gemm1 shrinks ~185 LOC. gemm2/__init__/moe_kernels repoint to .utils.

Validated: a16w4 SiTUv2 pytest 6/6 cold (cache cleared), pyflakes + black clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
CI Black check flagged gemm_moe_tune.py; apply black to the a16w4
stage2 ternary run-arg lists.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

_act = "situv2" if act in ("situv2", "situ") else act
sorted_size = int(sorted_expert_ids.shape[0]) * int(tile_m)
_alloc = torch.zeros if inter_dim_pad > 0 else torch.empty

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

zeros ?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

ported from old code, https://github.com/ROCm/aiter/pull/4502/changes#diff-d5d52b4d0fede962b29252c6940ba9aa70479833fc6b767058574c607e140004L1496. seems for bug fixes of gptoss. No sure if needed for other models. will not be trigger for inter_dim_pad=0

_alloc = torch.zeros if inter_dim_pad > 0 else torch.empty
inter_sorted = _alloc(sorted_size, inter_dim, dtype=torch.bfloat16, device=dev)
flydsl_a16w4_gemm1(
a_bf16=a.to(torch.bfloat16).contiguous(),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

contiguous?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

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

no real call. just make sure no error. can be deleted.

…a!=1

- gemm1/gemm2 builders: assert TILE_N//(4//k_wave) >= 16 (gemm1) and
  TILE_N//4 >= 16 (gemm2). num_acc_n==0 (e.g. tile_n=32,k_wave=1) empties
  every accumulate/store loop -> silent all-zero output that times fast.
  The shipped CSV avoids it and the tuner skips it; this closes the sharp
  edge for future tuning.
- fused_moe_2stages: skip the throwaway (token_num, topk, inter_dim) `a2`
  allocation for the a16w4 port. Its stage1 allocates its own sorted
  [sorted_size, inter_dim] bf16 intermediate and ignores `out`, so the
  pre-alloc (up to tens of MB at prefill) was pure waste. Tightly gated to
  the a16w4 flydsl path; a8w4/a4w4/mxfp8/non-flydsl unaffected.
- test: parametrize situ_beta/linear_beta with (4,25) (kimi-k3 betas) to
  cover the runtime SiTUv2 beta path directly (beta is not a compile key).

Validated gfx950: a16w4 pytest 12/12 (incl. beta!=1), a8w4 15/15.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@coderfeli
coderfeli force-pushed the flydsl-a16w4-mxfp-moe branch from d6df09f to e2fb8da Compare August 8, 2026 05:02
… types

Cleanup only; ISA byte-exact (gemm1+gemm2 21_final_isa.s identical A/B),
a16w4 pytest 12/12 cold.

- Remove a16wmix_use_k16(arch); inline `"gfx95" not in str(get_rocm_arch())`
  at the 4 call sites (2 __init__ launchers, 2 moe_kernels compile entries).
- Remove kmchunks_for/lds_acc_bytes_for from the package utils; inline
  BM//16 and BM*TILE_N*4 (the shared mxfp4_gemm_common copies are untouched).
- Drop the _get_compiled_gemm{1,2}_a16w4 @functools.cache wrappers from
  __init__; move @functools.cache directly onto compile_gemm{1,2}_a16w4_port
  and call them all-keyword from the launchers, matching moe_kernels.py's
  direct calls so the compile cache key is shared (fixes the previously
  uncached moe_kernels compile-entry path).
- ir dialect -> fx types: ir.Type.parse("vector<2xbf16>") -> T.vec(2, T.bf16)
  (gemm1/gemm2), ir.IntegerType.get_signless(32) -> T.i32 (utils cvt_pk);
  drop the now-unused `from flydsl._mlir import ir` imports. The load-bearing
  raw llvm.load/StoreOp/AtomicRMWOp/inline_asm (raw !llvm.ptr consumers,
  side-effecting cvt_pk) are kept.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants