[CK] [FlyDSL] Retune Kimi-K3 a16w4 MoE tile geometry - #5118
Conversation
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
amd-wsung102
changed the title
amd-wsung102
changed the title
valarLip
left a comment
valarLip
left a comment
There was a problem hiding this comment.
Review
The tile_m=16 decode work and the widened a16w4 tile space look like the right direction, and I confirmed two things up front so they don't get re-litigated: all 64 kernel names in the new CSV resolve as exact registry keys (stage-2 tile_m=16 was already registered for fp4, so the decode rows are safe), and the CSV rename does not break config loading — both loaders glob (jit/core.py:427, moe_recipes.py:129). The _persist/_w{n} suffix strips are genuinely behavior-preserving under the old code, and the CSV passes the duplicate-shape check against every other tuned CSV.
The problems below are mostly about blast radius outside a16w4 and about the provenance of the new CSV rows.
Blocking
1. aiter/ops/flydsl/moe_kernels.py:1529 — the new _g1_waves_per_eu gate silently switches the a16wi4 (int4) port from waves_per_eu=None to 3.
get_flydsl_stage1_kernels_int4_bf16._emit (~line 516) never writes a waves_per_eu key — verified: get_flydsl_kernel_params('flydsl_moe1_abf16_wint4_bf16_t32x128x256') returns no such key. aiter/fused_moe.py:1482 then does waves_per_eu=parsed.get("waves_per_eu", 3), so (3 is not None and int(3) > 1) is True and 3 is forwarded. gemm1.py:615 attaches rocdl.waves_per_eu=3 (a real minimum-occupancy / VGPR clamp) and gemm1.py:477 appends _w3 to the JIT cache key. All 88 rows of aiter/configs/model_configs/kimik2_i4_tuned_fmoe.csv were tuned with no such attribute and now run a differently-named, occupancy-capped kernel — a silent perf change plus a guaranteed cold JIT on a model that has nothing to do with this PR. Any of these would fix it: gate on b_dtype == 'fp4', register waves_per_eu: 1 in the int4 _emit, or default the stage-1 dispatchers to None the way stage 2 already does (fused_moe.py:1552).
2. aiter/configs/model_configs/kimik3_a16w4_tuned_fmoe.csv:14 — 8 rows change tile_n but keep byte-identical timings from the deleted file.
Joining old (merge-base d8a10bf) and new on (inter_dim, token): inter=512, tokens 2048/4096/8192/16384 move t32x128x128* → t32x256x128_bnt0_xcd1 (stage 1) and t32x256x128_atomic (stage 2), while us1 stays 472.3426 / 701.4285 / 1161.1315 / 2114.6148 to four decimals. inter=384, tokens 2048-16384 move to t32x192x128_bnt0_xcd1 / t32x256x128_atomic with us1 still 362.3881 / 534.9361 / 878.3311 / 1595.7209 (token 4096 additionally flips xcd4→xcd1). tile_n is a real constexpr — at inter=512 it halves NUM_N_BLOCKS from 4 to 2 — so identical microsecond timings are not possible. The 11 new tile_m=16 rows by contrast all carry fresh 2-decimal numbers, which is what a real tuning run looks like. As written, these 8 rows ship an unvalidated kernel choice (and copied err%) for the highest-throughput shapes, and jit/core.py's merge dedup picks the lowest us per shape using exactly these numbers. Could you re-run the sweep for those 8 shapes and paste the real numbers?
3. csrc/ck_gemm_moe_2stages_codegen/gemm_moe_tune.py:3475 — the a16w4 tuner hard-filters blockM=16, so the new tile_m=16 rows cannot be reproduced.
gen_flydsl_2stages_task — the only consumer of get_flydsl_stage1_kernels('bf16','fp4',...) — does if blockM not in [32, 64, 128] or not use_g1u1: continue, then if s1_tile_m != blockM: continue (3512); the v2 path at 3822 has the same gate. blockMs = [16, 32, 64, 128] at 4882 and the int4 sibling at 4301 already allows 16, so this looks like an oversight rather than a deliberate exclusion. Consequences: 384 of the 960 registered a16w4 names are unreachable, and any retune of kimik3 3584/384 silently regenerates block_m=32 rows, discarding the decode win this PR is built around. The 11 shipped block_m=16 rows have out-of-tree provenance that nobody can audit.
4. gemm_moe_tune.py:3676 — the 8 new flydsl_moe2_abf16_wfp4_bf16_t32x256x128_atomic rows name a stage-2 config the tuner explicitly refuses to compile.
if a_dtype_str == "bf16" and (s2_tile_m != blockM or kparams["tile_n"] == 256 or inter_dim % kparams["tile_k"] != 0): continue, with the comment directly above: "tile_n=256 over-allocates LDS at large tile_m (compile failure that takes the worker pool down); tile_n=128 covers the shape." No tuner run could have produced these rows. Either the comment/filter is stale and should be relaxed in this PR with evidence, or these rows point at a config that will hit the documented LDS over-allocation inside the JIT on the first token>=2048 request — at serving time instead of at tune time.
Should fix before merge
5. aiter/aot/flydsl/moe.py:881 — AOT still hardcodes "waves_per_eu": None, and its docstring's invariant is now false.
The docstring (~line 855) asserts "the cache key then matches runtime by construction (waves_per_eu=None ... are all applied inside the launcher)". After this PR the launcher no longer forces None — the decision moved to the caller (moe_kernels.py:1529) on the runtime side only. AOT builds ..._bm32_tn128 while runtime asks for ..._w3 (every a16wi4 row today) or ..._w{2,3,4}. That's a 100% AOT cache miss with the .pkl present on disk — precisely the failure mode aiter/aot/flydsl/moe.py:222-229 documents — i.e. a multi-second first-call JIT stall. Both the value and the docstring should change here.
6. moe_kernels.py:218 — tile_n=192/256 is added to a shape-independent registry, and the a16w port is the one path that bypasses the tile_n clamp.
tile_ns = [32, 64, 128, 192, 256] if is_a16w4 is keyed only on dtype, so t32x192 is offered to every a16w4 shape — including the inter_dim=512 rows in this same CSV (512 % 192 = 128). _flydsl_moe_stage1_impl returns from the _is_a16w_port branch at 1575, before the resolve_flydsl_stage1_tile_n(inter_dim, tile_n) clamp at 1581 — which the PR's own comment concedes ("that port takes tile_n as given"). The tuner's a16w4 filter (3479-3515) checks k_batch/k_wave divisibility and num_acc_n >= 1 but has no inter_dim % tile_n check, so tuning inter_dim=512 or any gptoss_a16w4 shape dispatches t*x192 candidates straight into assert _INTER % TILE_N == 0 (gemm1.py:436) / NotImplementedError, aborting the sweep. Also stale: the comment at 1576-1579 still says the OOB downgrade "Applies to both a16w4 and a8w4".
7. moe_kernels.py:216 — the % 16 filter was added to the int4 registry only, and this PR doubles the unbuildable a16w4 set.
For tn=32, kw=1 the kernel computes TILE_N // (4 // 1) = 8, tripping the pre-existing >= 16 assert at gemm1.py:441 — whose comment names this exact case ("silent all-zero output that times fast (e.g. TILE_N=32, k_wave=1)"). Adding tile_m=16 grows the a16w4 unbuildable set from 96 to 192 names. This PR fixed the class in get_flydsl_stage1_kernels_int4_bf16 by tightening the generator (541), but on the a16w4 side only added an assert. Any consumer that enumerates the registry without replicating the tuner's ad-hoc filter (AOT sweep, a new autotuner, a hand-written CSV) gets an AssertionError instead of a skip — and under python -O the assert vanishes, making num_acc_n=0 a fast-timing no-op the tuner would happily select.
8. gemm_moe_tune.py:3505 — third copy of the num_acc_n invariant, still < 16 while the kernel assert and the int4 registry moved to % 16.
if a_dtype_str == "bf16": _n_waves = max(1, 4 // _kw); if (kparams["tile_n"] // _n_waves) < 16: continue. The registered set {32,64,128,192,256} × kw{1,2,4} divides cleanly today, so the gap is latent — but the next tile_n like 96 or 160 passes this filter and then AssertionErrors inside a compile worker mid-sweep. Worth noting separately: the int4 filter change this PR advertises (< 16 → % 16) is a provable no-op — enumerated over tile_ns = [16,32,64,128], the name count is 76 before and 76 after, zero added or removed.
9. aiter/ops/flydsl/kernels/moe_2stage_a16wmix/gemm2.py:347 — the new assert's message doesn't describe its condition.
It says "TILE_N must be a multiple of 64" but tests (TILE_N // 4) % 16 == 0. 193 // 4 == 48, 48 % 16 == 0 → the assert passes for a TILE_N that is not a multiple of 64, nor even of 4. The condition should be TILE_N % 64 == 0. gemm1.py:445 has the same floor-division gap for k_wave > 1. No registered tile_n trips it today, but this assert exists specifically to be the backstop against exactly this, and a reader trusting the message either rejects 192 (legal) or accepts 193 (illegal).
10. moe_kernels.py:341 — _persist was stripped from the CSV but the registry still emits it for a16w4.
get_flydsl_stage2_kernels sets persist: True for every (a=bf16, b=fp4) config, but the a16w branch of _flydsl_moe_stage2_impl (~2059) returns via flydsl_a16w4_gemm2 without ever forwarding persist (first read is at 2154), so _persist=False always. get_flydsl_stage2_kernels_int4_bf16's docstring (560) documents exactly this hazard and refuses to register _persist for int4 — the fp4 twin was never given the same treatment. The tuner therefore times persist and non-persist as identical duplicates and picks one at random, recording a us2 under a name that never runs. Same for mode="reduce", which the a16w branch also ignores while the registry advertises it.
11. aiter/configs/model_configs/kimik3_fp4_untuned_fmoe.csv — the tuned file was renamed to kimik3_a16w4_* but its untuned twin wasn't.
All 30 sibling fmoe pairs share a prefix (gptoss_a16w4_tuned/untuned, kimik3_a4w4_tuned/untuned, kimik3_a8w4_tuned/untuned, …). The orphan still lists exactly the 32 shapes the new tuned file covers, so a retune driven off it writes results back to the name-derived kimik3_fp4_tuned_fmoe.csv, resurrecting the deleted file next to the new one with 32 identical shape keys — the next AITER_CONFIG_FMOE merge then raises "duplicate shape entries during merge" (jit/core.py:365) and refuses to start. Runtime today is unaffected (both loaders glob and exclude untuned), and op_tests/tuning_tests/test_csv_validation.py checks only a fixed 7-file list, so CI won't catch it.
Minor
12. moe_kernels.py:249 — the a16w4 k_wave gate is computed against a wave count the kernel doesn't use.
The registry computes num_n_waves = min(4, tn // 32) and rejects on num_n_waves * kw > 8, but compile_gemm1_a16w4_port always partitions exactly 4 waves as (4 // k_wave) N-waves (gemm1.py:85, known_block_size=[256,1,1]) — never 8, never a function of tn. The "≤8 total waves (≤512 threads)" comment this diff edits describes a different kernel's threading model. Meanwhile get_flydsl_stage1_kernels_int4_bf16 filters on _gemm1_lds_bytes(...) > _MAX_LDS_BYTES (544) and the a16w4 branch doesn't, though both feed the same port. Max LDS over the widened a16w4 set is 128 KiB < 160 KiB so nothing breaks today, but the next tile widening will register unbuildable names that AOT silently drops.
13. moe_kernels.py:200 — dead pre-loop tile_ns, and the inventory comment is missing the family this PR is about.
tile_ns = [32, 64, 128] if is_fp4_b else [128] is unconditionally overwritten on the first iteration of for tm in tile_ms: (both branches at 216-220 assign it), so this diff turned a 2-place duplication into a 3-place one with the pre-loop copy now inconsistent with the new a16w4 set. Separately, the comment at 195-197 was edited to (t32x{64,128,192,256}x128 / _xcd1) in the same hunk that adds the entire t16x* family the new CSV mostly relies on (t16x32x256_xcd4_kw2, t16x32x128_kw4, t16x64x256_kw2, t16x128x256_xcd1, t16x128x128_xcd4) — so the widest part of the new search space is missing from its own inventory comment, while t32x64x128 is listed and appears in no CSV.
14. moe_kernels.py:202 and :247 — two comments state rules the code doesn't implement.
(a) 202 claims "the decode win at E=896 (1.17-1.38x, token≤512)", but the PR's own CSV gives 1.02x-1.65x on the us column (token 1: 26.7689 → 17.57 = 1.52x; token 512: 381.3451 → 374.05 = 1.02x), and inter_dim=512 gets no tile_m=16 rows at all — so the claim isn't true of "E=896" generally, and someone benchmarking at token=512 sees 1.02x and suspects a regression. (b) 247 says k_wave applies to "small-M tiles (tile_m in {16,32})", but _small_m = tm == 32 or (tm == 16 and is_a16w4) excludes tile_m=16 for fp8xfp4, which line 205 also registers — someone adding a8w4 k_wave tuning will assume t16x*_kw* names exist and hit "Invalid FlyDSL kernel name".
15. moe_kernels.py:534 — the int4 registry comment was half-updated.
It now reads "each N-wave covers tn/(4/kw) cols, a multiple of 16 for the 16x16 MMA (kw=1 therefore requires tn >= 64)". Under the new (tn // num_n_waves) % 16 condition, kw=1 requires tn to be a multiple of 64, not merely ≥ 64 — tn=80 and tn=112 are ≥ 64 and silently skipped. Someone extending tile_ns beyond [16,32,64,128] will read "≥ 64" as the rule, add a tn the loop drops, and then debug a kernel name the registry never emitted.
The int4 stage 1 registry does not emit a waves_per_eu key, so _flydsl_stage1_wrapper's numeric default of 3 was the value that reached the a16w port. That was harmless while the port hardcoded waves_per_eu=None, but now that the value is forwarded it clamps occupancy on a path this work does not touch: all 88 tuned kimik2_i4 rows were measured with no such attribute, and rocdl.waves_per_eu=3 also appends _w3 to the JIT cache key, so every one of them would take a cold rebuild and run a differently scheduled kernel. Default to None instead, which is what stage 2 already does. Every other stage 1 family sets waves_per_eu explicitly in the registry, so the default only applies to int4 and the fp4 families keep the axis live.
…ith measured ones The eight token>=2048 rows changed tile_n but kept the us1/us2/us/tflops/bw values from the pre-retile file, so they described a kernel that is no longer the one named. That is worse than cosmetic: update_config_files sorts by us and keeps the lowest per shape when merging tuned CSVs, so those numbers can decide which row wins. Re-measured each row at its own parsed config, block_m, tile_n, tile_k, b_nt, xcd_swizzle, k_wave and waves_per_eu taken from its own kernelName, median of 3 repeats at 200 iterations. us is us1+us2 and tflops/bw are recomputed with gemm_moe_tune.py's own formulas, which reproduce every unchanged row in this file exactly. The new values run 15 to 19 percent above the replaced ones because this host is capped at 1000 W against a 1400 W TBP, the same host the block_m=16 rows were measured on. No other tuned CSV covers these eight shape keys, so the merge dedup has nothing to compare them against. err1 and err2 are left as they are. tile_n does not change gemm1's K-reduction order, so its output is bitwise identical across 64, 128, 192 and 384, and err1 genuinely does not move. gemm2 changes only the order of its bf16 atomic scatter, which cannot shift a value printed to one decimal place. Dispatch is untouched: block_m, ksplit, kernelName1, kernelName2, run_1stage, xbf16 and flat are byte-identical for all 32 shape keys.
Both asserts checked the per-wave column count after a floor division, so a TILE_N of 64k+1 to 64k+3 slipped through: 193 // 4 is 48 and 48 % 16 == 0, so the gemm2 assert passed for a TILE_N that is not a multiple of 64, which is exactly what its message claimed to enforce. A reader trusting the message would either reject 192, which is legal, or accept 193, which is not. Test the divisibility directly instead. gemm2 wants TILE_N % 64 == 0; gemm1 partitions 4 // k_wave N-waves, so it wants TILE_N % (16 * (4 // k_wave)) == 0, which is 64 at k_wave=1, 32 at 2 and 16 at 4. No registered tile_n changes verdict: of every a16w4 and a16wi4 stage 1 name that clears the pre-existing >= 16 assert, zero pass the old condition and fail the new one, and the stage 2 set of 128 and 256 is unaffected. Verified reachable with D_INTER = 193 * 128, where the divisibility and 256-alignment asserts both pass and TILE_N=193 now fails the wave-partition check.
The tuned config became kimik3_a16w4_tuned_fmoe.csv but its untuned input kept the old kimik3_fp4_ prefix, leaving the only unpaired fmoe config in the directory; every other pair shares a prefix. That is a trap rather than a cosmetic issue. The orphan lists exactly the 32 shape keys the tuned file covers, and the tuner derives its output name from the input name, so a retune driven off it writes kimik3_fp4_tuned_fmoe.csv and recreates the deleted file beside the new one with all 32 keys duplicated. The next config merge then hits the duplicate-shape check and refuses to start. No runtime effect either way: both loaders glob and exclude any name containing untuned. test_csv_validation.py checks a fixed seven-file list under configs/, not model_configs/, so CI would not have caught this.
Three filters in gen_flydsl_2stages_task made the shipped kimik3 a16w4 rows unreproducible, so a retune would silently regenerate different ones. blockM was gated to 32, 64 and 128, so no block_m=16 candidate was ever timed even though the a16w4 and a8w4 stage-1 registries emit t16 names and the int4 sibling already allows 16. Any retune of 3584/384 would have discarded the decode configs this work is built around. Added 16; the s1_tile_m != blockM filter below still drops the dtypes with no t16 names, so a4w4 gains nothing and a8w4 gains the 32 candidates its registry already advertises. Stage 2 refused tile_n=256 for bf16 outright, which is why no run could have produced the eight t32x256x128_atomic rows. The comment was right that the hazard is LDS over-allocation at large tile_m, but the code did not implement that. The port requests tile_m*(tile_k*2 + tile_n*4), so check the bound directly. Confirmed at both sides of it by forcing real codegen, since compile_gemm2_a16w4_port only builds the launcher: tile_m=128/tile_n=256/ tile_k=128 is 163840 bytes and runs, while tile_k=256 is 196608 and fails with "local memory (196608) exceeds limit (163840)" -- the predicted figure exactly. The shipped tile_m=32 config is 40 KiB. The num_acc_n filter used a floor divide, so it would pass a tile_n like 96 or 160 to a compile worker that then hits the kernel assert mid-sweep. Made it the same divisibility test the kernel now uses. The tuner would now generate the config named by all 32 rows of the CSV, up from 21. The v2 path's identical blockM gate is left alone: it returns early unless the activation is fp4 or fp8, so a16w4 never reaches it.
The k_wave comment said the option set applies to "tile_m in {16,32}", but the
guard is tm == 32 or (tm == 16 and is_a16w4), so fp8xfp4 gets k_wave only at
tile_m=32 even though its registry does emit 32 t16 names, none of them with a
_kw suffix. Someone adding a8w4 k_wave tuning would assume t16 _kw names exist
and hit "Invalid FlyDSL kernel name".
The int4 comment still said kw=1 "requires tn >= 64" after the filter moved to
a divisibility test. Under that test tn of 80, 96, 112 and 160 are all >= 64
and silently dropped, so anyone extending tile_ns past 16, 32, 64, 128 would
read the wrong rule and debug a name the registry never emitted.
Comments only; the parsed AST is unchanged.
The comment attributed the win to E=896 generally, which reads as covering both tuned inter_dim families. Only inter=384 has block_m=16 rows; inter=512 was never swept for it and all 16 of its rows stay at 32. It also gave no basis for the figure, so a reader comparing it against the us column of this CSV lands on 1.02x to 1.65x instead, because those rows were measured on two different hosts. State the measurement instead: 1.18-1.35x on the isolated GEMM pair at E=896 inter=384 for token<=512, which is the same-host A/B both this branch and the original sweep produced, and note that inter=512 is excluded. Comment only; the parsed AST is unchanged.
github-actions
Bot
changed the title
Eight of the comments this branch adds ran to two, three or four lines. Each is now a single line, and the two that sat inside pre-existing multi-line comments are back to the original wording with only the one clause that the code change invalidated rewritten, so the diff shows one added line rather than the whole paragraph. The whole branch now adds 14 comment lines across five files, none adjacent, longest 88 characters. Comments only; the parsed AST of every file is unchanged.
|
Hi @valarLip, thank you for reviewing this PR! I have addressed each point with a commit. Addressed
Point 2 — measured numbers
Point 4 — why the filter was relaxed rather than removedYou're right about the mechanism, but the code was over-broad by 4×. Forcing real codegen (
which is exactly what I left the v2 path's identical Deferred with reasons
|
amd-wsung102
changed the title
3 participants
Summary
This PR retunes the tile geometry of the Kimi-K3 a16w4 (bf16 activation × MXFP4 weight) SiTUv2 expert GEMMs, on a different axis for each end of the token range:
token <= 512): halve the M-block, 32 → 16. With 896 experts and top-16 routing a decode batch spreads ~2k routed rows over ~800 distinct experts — about 2.5 real rows per expert.moe_sortingpads every expert up to a fullblock_m, so atblock_m=32roughly 92% of the rows the two GEMMs process are padding. Halving the M quantum halves that waste.token >= 2048): widen the N-tile. Pick the largest gemm1tile_nthat dividesinter_dimexactly, up to 4 accumulator groups — 192 forinter_dim=384, 256 for 512 — collapsingNUM_N_BLOCKS3→2 and 4→2 and cutting both the A-gather traffic and the CTA count. gemm2 moves 128 → 256.Performance improvement:
Relevant Files
aiter/configs/model_configs/kimik3_fp4_tuned_fmoe.csv—block_m32→16 on the 11inter_dim=384,token<=512rows; N-tiles widened on the 8token>=2048rows of bothinter_dimfamilies; 25 inert_w{n}and 21 inert_persistsuffixes droppedaiter/ops/flydsl/moe_kernels.py— registertile_m=16andtile_n192/256 stage-1 names for a16w4; lettile_m=16sharetile_m=32's tile-N and k-wave sets on the a16w port; forwardwaves_per_euinstead of hardcodingNone; tighten the int4 registry's wave-partition filteraiter/ops/flydsl/kernels/moe_2stage_a16wmix/gemm1.py— assert the per-wave column count is a multiple of 16 (+8 lines)aiter/ops/flydsl/kernels/moe_2stage_a16wmix/gemm2.py— assertTILE_Nis a multiple of 64 (+6 lines)Optimizations Applied
Decode — M-block 32 → 16 on
token <= 512block_m=16wins for every bucket from 1 to 512 tokens (1.17–1.38× on the isolated GEMMs) and loses from 1024 up (0.85–0.98×). Becauseblock_mis a per-bucket CSV field there is no runtime coupling between the decode and prefill configurations.t16stage-1 names for a16w4.tile_msgains 16 when the weight is fp4 and the activation is bf16, matching the existing fp8 case. Strictly additive: the a16w4 stage-1 registry grows,fp4×fp4/fp8×fp8/fp8×fp4are bit-identical, and every kernel name in the shipped CSVs still resolves.tile_m=16sharestile_m=32's tile-N and k-wave sets on the a16w port. The tiny-batch decode configs use narrow N-tiles (t*x32/t*x64) andk_wave2/4; both are legal atBM=16(m_repeat <= 2either way, and the port takestile_nas given) and they are needed to name buckets 1–4. Guarded tois_a16w4.Prefill — N-tile widening on
token >= 2048tile_nthat dividesinter_dim, capped at 4 accumulator groups.pick_flydsl_stage1_tile_nonly ever returns 256 or 128, and for a non-256-alignedinter_dimit silently forces 128 while warning that 256 "is NOT tunable for such shapes" — so the values that divide exactly were never considered. Legal only on the a16w port, which takestile_nas given; the other fp4 dtypes route throughresolve_flydsl_stage1_tile_n, which would force the tile back to 128 and make the kernel name a lie. Registered for a16w4 only.tile_n=384(6 accumulator groups) and 512 (8) both lose to register pressure — 512 is 0.56× (2448.1 vs 1361.6 µs at token 8192). The occupancy loss swamps the A-traffic saving.tile_nchanges:xcd_swizzle=1beats 0 and 4 (1057.8 vs 1182.4 / 1066.8 µs),b_nt=0beats 2 (1057.8 vs 1221.8 µs), andblock_m=32remains optimal at prefill.Performance on ATOM
Measured on MI355X (gfx950, 256 CU). Kimi-K3 routed-expert geometry at TP=8:
model_dim=3584,inter_dim=384, E=896, top-16, SiTUv2, bf16 × MXFP4.kv_cache_dtype=fp8,block-size=128, prefix caching off. 8k/1k, 128 prompts, concurrency 128.Kernel-level
MoE layer
Kimi-K3 end-to-end benchmark
Performance on vLLM
Measured on MI355X (gfx950, 256 CU). Kimi-K3 routed-expert geometry at TP=8: model_dim=3584, inter_dim=383, E=896, top-16, SiTUv2, bf16 × MXFP4. kv_cache_dtype=fp8, block-size=128, prefix caching off. 8k/1k, 128 prompts, concurrency 128
Kernel-level
MoE layer
Kimi-K3 end-to-end benchmark
Accuracy
The N-tile half is bitwise identical.
tile_nonly changes how the N axis is blocked across workgroups and waves; it does not change the K-reduction order. Cross-checking every legaltile_nagainst 128 atinter_dim=384:End-to-end through the production
get_2stage_cfgs→fused_moe_2stagesdispatch against a bf16 SiTUv2 torch reference, cos/logits_diff, baseline vs this branch:Identical to four significant figures at every prefill point.
op_tests/flydsl_tests/test_flydsl_moe_a16wfp4.pypass (2 shapes × 3 token counts × 2 SiTUv2 beta pairs) atcos_diff < 1e-2.kernelName1/kernelName2entries parse against the registry; no residual_w{n}or_persistnames remain.get_2stage_cfgsverified to select the intendedblock_mand tiles at every bucket, including the unchangedt32x128x128at 1024.TILE_Nguards verified to reject 96 (gemm1) and 112 (gemm2), accept 192/256.Submission Checklist