[FlyDSL] clean attn-aux kernels: fx.* modernization cleanup - #5113
Draft
coderfeli wants to merge 9 commits into
Draft
[FlyDSL] clean attn-aux kernels: fx.* modernization cleanup#5113coderfeli wants to merge 9 commits into
coderfeli wants to merge 9 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
PR title tags: |
github-actions
Bot
changed the title
coderfeli
marked this pull request as draft
August 29, 2026 14:51
Surgical fx.* migration of four FlyDSL attention-auxiliary kernels, behavior-preserving with byte-exact / ISA-parity gating: - fused_compress_attn.py: vector.*->fx.Vector, ArithValue index math ->fx typed-int operators, dead consts/banner removed (net -25 LOC). - fused_compress_attn_hca.py: LDS pointer index math ->fx operators (ISA byte-identical on all 3 compiled variants). - qk_norm_rope_quant.py: copy_atom_call->fx.copy (8), fx.Index->fx.Int64 (7), llvm rcp intrinsic + bare rocdl.*->fx.rocdl.*, dead import dropped (ISA byte-identical on all 5 gfx950 variants). - flash_attn_func_gfx1201.py: ArithValue.select/_bitcast_i32->fx ops, dead ArithValue import removed (LLVM-IR parity verified on 3 configs). Gate (gfx950 local + IR dump for gfx1201): - compress_attn combined sweep: all main paths err_pct=0; the 2 legacy(NW=1) failures confirmed pre-existing (identical on base). - qk_norm: 465 passed / 0 failed; perf equal-or-better. - gfx1201: RDNA target, not executable here; gated via byte-identical fully-lowered LLVM-dialect MLIR (stage-18) diff. - Perf medians within noise on all gfx950 kernels. Left legacy on purpose (per skill): buffer_ops scalar-base+per-thread loads with no layout form, non-default-fastmath arith.*FOp float compute, raw WMMA/MFMA plumbing, and the gfx1201 to_llvm_ptr AS-change (drifts ISA, needs real hardware to land). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…ndex Second cleanup pass, byte-exact ISA gated (net -83 LOC on top of the first pass): - fused_compress_attn.py / fused_compress_attn_hca.py: collapse the explicit arith.AddFOp/MulFOp/subf(fastmath=fm_fast) float layer + its feeder vector.extract/arith.extf/ArithValue to plain fx operators. Proven redundant: the module already compiles with fast_fp_math=True + unsafe_fp_math=True, so plain fx `+`/`*`/maximumf emit byte-identical ISA. Verified on 6 kernel variants (hca_main, csa_main ksplit, csa_indexer fp8; hca_compress_forward, hca_norm_rope_scatter bf16/fp8) -- all 21_final_isa.s byte-identical. Prove-then-apply caught 3 real ISA-drift cases, kept raw + documented: GPT-J RoPE subtract (default non-fastmath flags -> fx fuses differently, 123-line drift), the softmax-padding cmpf(OEQ)+select (unordered-compare+cndmask drift), and the fp8/fp4 ue8m0 scale bit-trick. - flash_attn_func_gfx1201.py: fx.Index -> fx.Int64 where byte-identical (12 of 23). Kept 11 that drift or fail to compile: signed<->unsigned divide (index udiv/urem vs i64 sdiv/srem), index-typed loop-induction / cmpi / select peers. Gated via stage-18 lowered LLVM-dialect MLIR diff (RDNA WMMA not executable on gfx950). Gate: gfx950 correctness err_pct=0 on all non-preexisting cases (the 3 legacy(NW=1)/prefill-bs32 cache_scale failures confirmed pre-existing via git-stash A/B); per-variant ISA byte-identical; black+ruff clean. qk_norm fx.copy sites reviewed: the 8 single-atom loads are the skill §7b degenerate case (kept as fx.copy, not promoted); the tileable ones were already make_tiled_copy via main's #4606. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…yc.jit Replicate the ASM-identical cleanup blueprint (ref 1e0a4da) on the current shared MoE epilogue helper: - Drop the injected MLIR dialect-module parameters (arith/vector/gpu/scf/ range_constexpr) from default_epilog/c_shuffle_epilog/mfma_epilog; import them at module scope instead. Update the sole caller (mixed_moe_gemm_2stage_common.py): remove 32 pass-through kwargs across the default_epilog + 6 c_shuffle_epilog sites. - arith.constant(index)/cmpi/select → fx.Index/fx comparison/fx.Boolean .select; vector.load_op → Vec.load. - scf.IfOp → @flyc.jit if/else (4 conversions: group-write branch + two row-guards). The one value-yielding load-select scf.IfOp is kept (documented): @flyc.jit can't merge a both-branch result to the caller. Left legacy on purpose (matching the reference): the now-unused _if_then SCF-then compat shim and the always-true `scf is not None` guards. Gate: 20/20 distinct mixed_moe kernel ISAs byte-identical (21_final_isa.s, git-stash BEFORE/AFTER), incl. sbm128 split-LDS variants exercising the group-write @flyc.jit + kept scf.IfOp + row-guards. moe_2stage correctness green (logits_diff unchanged); black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
coderfeli
force-pushed
the
coderfeli/flydsl-attn-aux-fx-cleanup
branch
from
August 29, 2026 15:34
01b5b8f to
d6c63b4
Compare
Third pass (net -72 LOC), targeting _to_raw/ArithValue/.ir_value/raw arith·vector·buffer_ops. Every change byte-exact ISA gated; anything that drifted the ISA was reverted (not forced). - hca: ArithValue 34->0, _to_raw 23->3, arith 71->31, vector 18->8. The integer offset math is now fx-typed throughout (buffer_ops natively accepts fx Numeric offsets); buffer_ops primitives kept because create_buffer_resource's base_byte_offset (64-bit descriptor rebase, guards a documented 4 GiB 32-bit-offset overflow) has no make_buffer_tensor equivalent. Caught: fx Int32 >> is arithmetic (ashr) vs the original logical shift -> fixed via fx.Uint32; no fx int max/min -> kept arith.maxsi/minsi. - fused_compress_attn: vector.*->0, arith.extf->0, arith.trunci->0 (all bf16->f32 / fp8 byte-store casts). Proven NOT byte-exact: converting a single arith.remui->fx % inside the loop-carried softmax scf.for drifts the ISA 533 lines (identical opcode histogram, pure SSA-order/regalloc reshuffle) -> reverted. Its offset+buffer_ops layer stays raw by proof. - qk_norm: dropped 5 redundant .ir_value() bridges into from_elements. The 28 buffer_ops are the layer the FlyDSL maintainers' own #4606 full-fx-surface refactor deliberately retained; kept. - gfx1201: dropped 7 .ir_value() bridges, rocdl.exp2->fx.rocdl.exp2 (IR-identical on 3 configs; the high-level fx.exp2 form drifts 35 lines). - mfma_epilogues: _is_group_b now a plain fx Boolean predicate (@flyc.jit if + .select, no .ir_value re-wrap). Remaining .ir_value bridges feed raw memref.load/vector.extract in caller callbacks (empirically load-bearing); the value-yielding load-select scf.IfOp is irreducible (fx.select would double LDS traffic). What stays raw is now the proven floor: ops that drift the ISA if converted (loop-carried offset math, RoPE subtract, cmpf(OEQ)+select, fp8/fp4 bit-tricks, gfx1201 signed-cmp/divide) and genuine no-fx-form boundaries (buffer_ops base_byte_offset, cvt_pk_fp8, rcp/rsqrt/exp2, emit_mx_e8m0_scale, raw WMMA, caller raw-memref callbacks). Gate: gfx950 compress_attn err_pct=0 (3 pre-existing cache_scale failures unchanged, A/B-confirmed) + qk_norm 465/0; per-variant ISA byte-identical (compress_attn x4, hca x3, qk_norm x5, mfma_epilogues 20/20); gfx1201 stage-18 IR byte-identical x3 configs; black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Relaxed gate this pass from byte-exact-ISA to correctness + perf-neutral (median-of-5/CUDA-graph, revert >2%), which unlocks the buffer_ops migration and the raw-op conversions that only rescheduled the ISA. buffer_ops -> fx.rocdl.make_buffer_tensor via a moe_reduce-style helper that folds the 64-bit descriptor base into the pointer (inttoptr(ptrtoint(get_iter)+base_i64)) BEFORE make_buffer_tensor, so the 4 GiB / 32-bit-offset overflow guard holds by construction: - fused_compress_attn_hca.py: 44 -> 2. Also finished the raw-dialect collapse: ArithValue 34->0, vector 8->0, and the RoPE subtract + softmax cmpf/select + int max/min -> fx (all correctness+perf verified; RoPE even faster). Net large LOC drop. - fused_compress_attn.py: 76 -> 4. (The _build_kernel offset math and the RoPE/softmax float ops stay raw here -- PROVEN gate failures: the former regresses hca_main +8.9%, the latter breaks the fp8/fp4 e8m0 cache_scale bit-exact check by <=1 ULP. Kept raw, bisected.) - moe_sorting_kernel.py: 82 -> 4; ArithValue 33->0, fx.Index 13->0, index_cast 30->0, trunci 2->0. GPU-probed that make_buffer_tensor's V# descriptor preserves the .select(idx,0x7FFFFFFF) OOB-drop semantics. The survivors are proven boundaries, not oversights: - 2/4 per file feed the off-limits _common fp8 emitter (raw V# descriptor + scf.IfOp) that can't take a fx buffer-tensor without editing siblings. - moe_sorting's 4: uint8 byte-addressed scatter -- an i8 byte-mode store at the sentinel FAULTS (hipErrorIllegalAddress) instead of dropping, so it stays guarded on buffer_ops (probe-confirmed). Critical units trap found+fixed (hca first, applied everywhere): the fx buffer-tensor offset scales by the tensor's elem_ty, but buffer_load scaled by the LOAD dtype -- a bf16 tensor read as i32 dwords got a 2x offset (err 0.57). Fix: pass elem_ty = load dtype per site (bf16-dword paths->Int32, state->Float32, cos/sin->BFloat16, byte-scatter->Int8). Tiled-copy / layout-algebra evaluated for the regular loads and correctly NOT applied: the f32 state loads are descriptor-rebased dynamic per-K row gathers (a TV layout can't express a dynamic gather), the bf16 loads are dword-packed bit-extracts, and the one affine site already lowers to a single buffer_load (partition machinery = net LOC up, skill 7b). Gate: compress_attn combined err_pct=0 (3 pre-existing cache_scale failures A/B-confirmed unchanged) + moe_sorting 48/48 bit-exact vs torch ref + GPU functional smoke (all tokens correctly routed); perf neutral on all variants; black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The 4 survivors were the uint8 byte-scatter stores in p0_scatter/p0v2, already wrapped in `if valid:` / `if is_mine:` guards -> they only ever fire in-bounds, so the V# descriptor's OOB-drop was never used (the sentinel path is dead under the guard, which is why an unguarded byte store faulted at the sentinel). A plain guarded LLVM store suffices and needs neither buffer_ops nor make_buffer_tensor. Added `_i8_global_ptr(tensor)` (inttoptr(ptrtoint(get_iter(t))) as i8 Global) and store `ws_i8[byte_off] = val_i8` inside the existing guard (i8 element == byte, so the byte offset is the element index directly). buffer_ops: 4 -> 0, import removed. Gate: 48/48 cold bit-exact vs torch ref (all paths incl. the mask+padding cases that drive the byte scatter); perf neutral (CUDA-graph interleaved A/B; an apparent +1-2% was measurement-window drift, present on the unchanged oneshot path too); black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The shared nm-asm fp8 scatter emitter (emit_group_fp8_nm_asm_scatter in fused_compress_attn_common.py) used no OOB — it writes plan-bounded, in-bounds kv_cache / k_rope slots — so the V# descriptor was unnecessary. Switch it to direct global pointer stores built from an i64 block base, and drop the 6 create_buffer_resource descriptors the callers fed it. - fused_compress_attn_common.py: emitter takes out_base_i64 / krope_base_i64 (i64 block base) + cache_base / krope_base (i32 within-block byte off) instead of raw V# resources; _global_ptr(base_i64, byte_off, elem, align) = fx.inttoptr over base+off, .store(...). 6 buffer_store -> 0. Rest of the emitter migrated to fx (cvt_pk_fp8, fx.Vector, typed ints); 3 guarded scf.IfOp kept (§5: plain `if` doesn't rewrite inside an imported helper), non-default-fastmath arith.MulFOp kept. - callers pass out_base_i64=fx.Int64(fx.ptrtoint(fx.get_iter(kv_cache))) + fx.Int64(block_base): fused_compress_attn.py (2 sites) + hca (1) — now 0 buffer_ops (dead import dropped); fused_compress_attn_gfx1250.py + hca_gfx1250.py (1 each). Gate: gfx950 compress_attn err_pct=0 on all emitter-driven cases (hca_main + csa_main nm_asm_fp8 nope/scale/rope; cache_scale bit-exact) — the 3 pre-existing legacy(NW=1)/prefill-bs32 cache_scale failures unchanged; gfx1250 CSA + HCA scatter cross-COMPILE clean (execution not validatable here). black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The _buf_tensor/_buf_load/_buf_store V#-helper trio was byte-identical in fused_compress_attn.py and fused_compress_attn_hca.py. Move it to the shared fused_compress_attn_common.py (as buf_tensor/buf_load/buf_store), import via alias in both callers, drop the now-unused bare `ptrtoint` imports. Net -~55 LOC, single source of truth. Gate: compress_attn err_pct=0 (3 pre-existing cache_scale failures unchanged); black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The emitter builds its store vector directly from rotated_lane; the vecVf32 (T.vec(VEC,f32)) parameter was unused. Remove it from the signature and all 5 call sites, and the 3 now-dead local vecVf32 assignments in the wave64 callers (gfx1250 callers still use their local vecVf32 for vector.from_elements, kept). Gate: gfx950 compress_attn err_pct=0 (3 pre-existing unchanged); gfx1250 csa+hca cross-compile clean; black+ruff clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
coderfeli
marked this pull request as ready for review
August 30, 2026 11:49
coderfeli
marked this pull request as draft
August 30, 2026 11:51
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.
Surgical fx.* migration of four FlyDSL attention-auxiliary kernels, behavior-preserving with byte-exact / ISA-parity gating:
Gate (gfx950 local + IR dump for gfx1201):
Left legacy on purpose (per skill): buffer_ops scalar-base+per-thread loads with no layout form, non-default-fastmath arith.*FOp float compute, raw WMMA/MFMA plumbing, and the gfx1201 to_llvm_ptr AS-change (drifts ISA, needs real hardware to land).
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist