Moe reduce extract - #4596
Merged
Merged
Conversation
The topk-reduction kernel (compile_moe_reduction) lived in moe_gemm_2stage.py alongside the int4 (a16wi4) stage1/stage2 GEMM builders, but it is dtype-agnostic: it is the epilogue for stage2 mode="reduce" for every dtype (fp8/a8w4/mxfp4/int4), reached via moe_kernels._run_moe_reduction. Move it to its own module so it lives independently of the int4 GEMM builders. - New self-contained kernels/moe_reduce.py (own _if_then/_if_else SCF helpers, no back-reference to moe_gemm_2stage). - moe_kernels._run_moe_reduction now imports from .kernels.moe_reduce (the live production path). - moe_gemm_2stage.py imports compile_moe_reduction back so the (unused) compile_moe_gemm2_ex still resolves; int4 gemm1/gemm2 untouched. Cleanup on the moved kernel: drop the `if True:` scaffold, the unused i8_type() helper, and two dead no-op statements (get_hip_arch(), ir.ShapedType.get_dynamic_size()); trim step-narration comments while keeping the load-bearing rationale (64-bit base-offset fold, unsigned Index->ult, EP gather). Behavior-preserving: fx.Index, buffer_ops, and the scf.IfOp control flow are kept as-is. Verified on gfx950 (FLYDSL_RUNTIME_ENABLE_CACHE=0): _run_moe_reduction matches a torch reference for f16/bf16/f32 x masked/unmasked, over both the full-vector and tail (model_dim=4100) paths (max err <= 1.9e-6). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Apply the flydsl-kernel-cleanup skill to the extracted reduction kernel: replace the raw scf.IfOp + _if_then/_if_else context-manager boilerplate with idiomatic Python if/else on typed fx.Index compares (the current FlyDSL idiom for guarded side-effects), and drop the now-unused contextlib / _mlir.ir / _mlir.dialects.scf imports and the two helpers. Also inline the trivial compute_type()/i32_type() shims (-> T.f32/T.i32). Kept as-is (load-bearing, per the skill's "behavior-preserving" rule): - fx.Index for the bounds compares -- unsigned, lowers to ult; the guards rely on it. - buffer_ops.create_buffer_resource_from_addr with the 64-bit base-offset fold -- a raw-addr + i64 base form with no make_buffer_tensor equivalent. moe_reduce.py: 305 -> 222 lines. Verified on gfx950: byte-identical results to the scf.IfOp version across f16/bf16/f32 x masked/unmasked over full-vector and tail paths; and perf-aligned (min-of-5, bf16 md=7168 topk=8): tok=32768 862.8->866.0us, tok=8192 211.8->211.9us, tok=1024 19.8->19.6us -- all within noise, as expected since Python-if lowers to the same scf.if. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…gned Replace the buffer_ops + fx.Index kernel with the FlyDSL layout API, and close the perf gap that made the first layout attempt ~10% slower. Layout API: - fx.inttoptr per-token base + make_view + make_buffer_tensor (drops buffer_ops.create_buffer_resource_from_addr + manual buffer_load/store). - TV tiling via logical_divide/slice; BufferCopy128b copy atoms into register fragments. No fx.Index (explicit fx.Int32/Int64). - Descriptor sized to the token slab -> OOB-checked model_dim tail drops the scalar tail path and the token/col bounds guards. Perf alignment (this was the crux; the naive layout port lost ~10% on this latency-bound kernel at only ~55% of HBM peak, where instruction overhead is not hidden): - The topk rows are reached with a UNIFORM SCALAR soffset=k*model_dim via fx.copy(..., soffset=), so all topk loads share ONE per-thread voffset and issue back-to-back (scalar offsets land on the SALU). This removes the 7 per-k v_add_u32 the naive port emitted before each load. ISA now: 8 buffer_load_dwordx4 sharing v54 + s_mov soffsets, 0 v_add. - Guard threads whose column group starts past model_dim (their loads would read the next row -- in-descriptor, wasted BW), matching base's tail skip. Emitted only when TILE does not divide model_dim (const_expr), so the common divisible case pays no branch. Verified on gfx950 (FLYDSL_RUNTIME_ENABLE_CACHE=0): byte-exact vs the buffer_ops version for f16/bf16/f32 x masked/unmasked over divisible, whole-tile-tail (md=7168) and partial-vector-tail (md=4100) paths, plus >4 GiB (70000 tokens, past the i32-voffset boundary). Perf min-of-6 vs the buffer_ops baseline (bf16, topk=8): md=7168 aligned across tok=1024/8192/ 32768 (within ~0.5%, faster at low tok); md=4096 within ~2% (a compiler load-drain-granularity nuance, leaner ISA + faster at tok=1024). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Condense the module/function docstrings and inline comments to the load-bearing "why" only (base fold, uniform-scalar soffset, OOB guard), compact the compile_moe_reduction signature, drop the _view layout arg (both callers pass make_layout(model_dim,1)), and inline one-shot temporaries. Trace-time-only changes -> identical IR; 183->131 LOC. Verified byte-exact via _run_moe_reduction (f16/bf16/f32 x masked/unmasked over divisible / md=7168 / md=4100 tails + >4 GiB). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Drop the unused keyword defaults (both callers -- _run_moe_reduction and compile_moe_gemm2_ex -- pass all five args explicitly) so the keyword-only signature fits on one line under black's 88-col limit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the compile_moe_reduction(...) closure factory + functools.lru_cache with the idiomatic FlyDSL surface: a module-level @flyc.kernel (moe_reduction_kernel) and @flyc.jit (moe_reduction) whose shape/dtype params are Constexpr, so flyc specializes per shape. moe_reduce.py no longer has a Python compile wrapper. Dispatch: a direct @flyc.jit call re-runs ~34us of jit dispatch per call (4x slower at tok=1/128 -- the decode path). So _run_moe_reduction caches the compiled function per Constexpr set (_MOE_REDUCTION_CF) and fast-dispatches via cf(*args), same mechanism _run_compiled used. Perf matches the old factory (tok=1 ~11us, tok=32768 ~856us; byte-exact across f16/bf16/f32 x masked/unmasked over the divisible / md=7168 / md=4100 tails + >4 GiB). Also removes the now-dead reduce machinery orphaned by the extraction from moe_gemm_2stage.py: MoeGemm2Mode, _MoeGemm2ReduceWrapper, compile_moe_gemm2_ex (all unreferenced), plus the then-unused compile_moe_reduction / _run_compiled imports there (-246 LOC). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
…te-out Resolves the #4596 conflict with main and ports main's new MXFP8 route-out reduction feature onto the layout-API reduce kernel. - moe_gemm_2stage.py: take the PR side (reduction lives in moe_reduce.py); keeps main's T.f8 -> default_f8_type() migration outside the conflict. - moe_reduce.py: add dtype_str="fp8" branch + out_dtype_str param. Reduces the flat uint8 route-out rows [N fp8 | N/8 e8m0] -> bf16/f16, ported verbatim from main's _compile_moe_reduction_fp8 (same atoms/address math) but inlined into the cache-safe module-level @flyc.kernel (Constexpr-specialized, auto-name), not the factory form. Dense path byte-identical to the PR. - moe_kernels._run_moe_reduction: fold main's is_fp8 handling (flat X, dtype_str=fp8, out_dtype_str) through the _MOE_REDUCTION_CF dispatch; thread out_dtype_str into the args tuple + cache key. Validated on gfx950 (cache ON): fp8 output byte-exact vs main across fast/tail paths, mask on/off, bf16/f16; dense path byte-exact vs torch; isolated dispatch perf at parity with main (ratio 0.99-1.03). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Collapse the two separate reduce branches into a single make_tiled_copy / partition_S/D body parametrized by a compile-time is_fp8 flag. Dense (f16/bf16/f32) and fp8 route-out share buffer setup, the model_dim tiled copy, the uniform-soffset topk load loop, the masked f32 accumulate and the truncating store; only the input element/atom, row_stride and the decode step (extf vs cvt_pk_f32_fp8 * e8m0) differ. - Port both paths to fx.make_layout_tv + fx.make_tiled_copy + partition_S/D (layout-API tiled copy) instead of manual per-thread views. - Drop the fp8 per-lane scalar tail: with model_dim % 8 == 0 (MX) and 8-aligned col_base, in-range threads always have full windows, so the col_base guard + descriptor store-drop cover the tail (fp8 ISA 153 -> 17 buffer insns). - Use T.f32x2 for the cvt result; drop the _NUMERIC module global; trim comments. Validated on gfx950 (cache ON): fp8 byte-exact vs main across fast/tail, mask on/off, bf16/f16; dense byte-exact vs main all dtypes. Dense ISA byte-identical to the pre-unification kernel (perf parity; measured deltas are shared-GPU noise). 327 -> 212 LOC. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…duction factory Restore the compile_moe_reduction(**constexprs) factory + _run_compiled dispatch (the codebase convention) instead of the module-level _MOE_REDUCTION_CF dict. The factory is @functools.lru_cache'd and returns a distinct per-shape @flyc.jit launcher, so _run_compiled's per-exe _cf cache stays correct across shapes. It wraps the auto-named Constexpr moe_reduction_kernel, so it is cache-safe under the production disk cache (main's old bug was the explicit-name @flyc.kernel form, not the factory itself). Verified on gfx950: dense byte-exact vs torch (bf16/f16/f32, mask on/off), fp8 deterministic, cache-safe across a cold second process (fresh + warm disk cache, reversed shape order). Decode dispatch 3.6us (>= the dict's 4.5us); a bare jit call is 28us. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
coderfeli
pushed a commit
to ROCm/FlyDSL
that referenced
this pull request
Aug 7, 2026
…ion.py Ports the stage2 topk-reduce (aiter ROCm/aiter#4596) onto the tiled-copy layout API, replacing the legacy buffer_ops/scf.IfOp reduction.py. New kernels/moe/moe_gemm_2stage/moe_reduce.py: - make_layout_tv / make_tiled_copy / partition_S/D / zipped_divide / memref_load_vec, uniform soffset=k*row_stride, 64-bit base-ptr fold (voffsets stay i32-safe for X > 4 GiB). - Unifies dense (f16/bf16/f32) + fp8 MXFP8 route-out ([N fp8 | N/8 e8m0]) via fx.Constexpr params; fuses the EP gather valid = expert_mask[topk_ids[t,k]] != 0 (expert_mask + topk_ids, not a precomputed mask). New launcher takes fx.Pointer args. gemm2.py: _MoeGemm2ReduceWrapper.__call__ now takes expert_mask/topk_ids and dispatches the pointer launcher via tensor_shim._run_compiled + from_c_void_p; compile_moe_gemm2_ex switches valid_mask=None -> use_mask: bool and passes num_experts=experts. test_moe_reduce.py rewired to the new interface. Verified: 8/8 non-large test_moe_reduce cases pass on GPU (dense + masked EP-gather); fp8 path traces/compiles/runs. Bandwidth on par with the legacy kernel (~5 TB/s on prefill shapes; ~15% lower latency on small decode). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist