Flydsl kernel cleanup - #4501
Merged
Merged
Conversation
Migrate the fused gate-activation+quant kernel off legacy constructs: - scf.IfOp + ir.InsertionPoint + scf.YieldOp -> Python if/else (runtime cond) - ArithValue wrappers dropped; index/address/condition math uses fx.Int32 with Python-int RHS (col masks, shifts, byte offsets, tiled-scale index) - arith.constant integer consts -> fx.Int32; arith.index_cast -> fx.Int64 - arith.TruncIOp on packed results -> fx .to(fx.Int8/Int16) The f32 activation core is left intact on purpose: llvm.amdgcn exp2/rcp/fabs intrinsics, the vendored vector.* shim, and buffer_ops manual-addressing (fp4/fp8 packing + tiled E8M0 scale writes) are the documented exceptions to the migration (no clean layout/copy-atom form). Verified byte-exact vs the pre-cleanup kernel on 48 configs (silu/swiglu/ situv2 x fp4/fp8/none x separated/interleaved x inter_dim 256/512, plus padding-row branch) on a COLD flydsl recompile. 637 -> 568 LOC. Also fix the op test: it called _run_compiled with raw tensors, but the kernel takes fx.Pointer args (production wraps every arg in ptr_arg). Wrap the test args in ptr_arg so it exercises the real calling convention. Note: enable_bias=True is broken on main (NameError: topk_rsrc at trace time) and remains identically broken here -- out of scope for this cleanup. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- scf.IfOp (row_valid, dw_valid) + InsertionPoint/YieldOp -> Python if
- ArithValue dropped; index/address/bit-twiddle math uses fx.Int32 with
Python-int RHS (col/lane/block, byte->dword offsets, hi/lo half select)
- arith.constant integer consts inlined; arith.index_cast -> fx.Int64
f32 swiglu core kept intact (arith.minimumf/maximumf/constant + llvm.amdgcn
exp2/rcp intrinsics); the bf16<->f32 arith.bitcast stays at the raw boundary
(fed .ir_value()); buffer_ops manual dword addressing kept (2-cols/dword
pack, no clean layout form).
Verified byte-exact vs pre-cleanup kernel on inter_dim {256,512,1024,1536}
x rows {1,4,17}, COLD flydsl recompile (12/12). 207 -> ~165 LOC.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The kernel is already on the modern fx.* surface (fx.SharedAllocator, @fx.struct, Python if/else, .select(), fx.math.exp2). Only the launch grid still used raw arith: - arith.index_cast(T.index, num_programs/grid_y_dim) -> fx.Int64(...) - remove the now-unused `from flydsl.expr import arith` import buffer_ops loads/stores are manual per-thread addressing (halo/LDS staging, conv_state blend) with no clean layout form -- kept per the migration guide. Verified: op_tests/test_causal_conv1d_prefill_split_qkv.py all cases OK (torch + flydsl backends), COLD flydsl recompile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…e_quant.py Surgical cleanup of the already-mostly-modern kernel (it uses make_buffer_tensor, fx.copy_atom_call, fx.slice, Python if/else): - ArithValue(x) * arith.constant(N,i32) -> x * N (Python-int RHS) for the q/kv per-token index math (q_my_off, row_off/row_base, scale_base_off, kv_off) and the two store helpers -- operands are already fx (bid_x/bid_t/ tid/head_idx/strides) - drop the now-dead c8 local; shift/const literals inlined Deliberately kept (documented keepers, not legacy to migrate): - arith.AddFOp(..., fastmath=fast) shuffle-xor reductions (non-default fastmath) - arith.cmpf(OLT/OGT) fnuz flush-to-zero clamp; arith.MulIOp in T.index for per-token byte offsets (index-typed on purpose, see inline note); the SWA scatter divsi/remsi/index_cast; vendored vector.bitcast; buffer_ops. Verified: test_flydsl_qk_norm_rope_quant.py (Q/KV checkAllclose passed, 0 err) and test_fused_qk_norm_rope_group_quant.py (0 err) on COLD recompile; perf TFLOPS unchanged (~47.5). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…refill}.py Both kernels are already on the modern surface (make_buffer_tensor, fx.copy_atom_call, fx.slice, range(init=...) loop-carry). Cleaned the remaining raw index/offset helpers: - arith.extsi/muli/constant i64 base-byte-offset -> fx.Int64(row)*fx.Int64( stride)*4 (.ir_value() at the create_buffer_resource boundary) - fx.Index(runtime loop bound) -> fx.Int64(...) (explicit width; keeps the range() init= loop-carry path per the migration guide) - arith.index_cast(T.index, gx) -> fx.Int64(gx) for the grid dim - ArithValue(packed) >> fx.Int32(8*nt) & fx.Int32(0xFF) -> fx.Int32 + int RHS (kv-scale unpack, fp4 only) Kept as intrinsic keepers: rocdl.mfma_scale_f32_16x16x128_f8f6f4 (hand-packed scaled MFMA) and the ds_bpermute xor-add lane reduce. Verified: test_flydsl_pa_mqa_logits_fp4.py and _prefill.py both PASS on COLD recompile (gfx950); perf unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…+ qk_norm The vendored kernels/vector.py shim's own docstring says it is legacy and to "prefer using Vector methods"; migrate the call sites in the two cleaned kernels that still used it: - vector.from_elements(ty, elems) -> fx.Vector.from_elements(elems, dtype=...) - vector.bitcast(ty, v) -> fx.Vector(v).bitcast(elem_dtype) - vector.extract(v, [i]) -> v[i] (fx.Vector indexing) - drop the now-unused `vector` import and dead vec_t/vec2_i32/vec_bf16xV locals silu: scalars feeding the raw f32 activation core (llvm.amdgcn exp2/rcp/fabs) are extracted as v[i].ir_value() so the intrinsic operands stay raw and the numeric core is byte-identical. qk_norm: also refresh a stale "scf.if" comment (the code uses a Python predicated branch now). Verified: silu 36/36 byte-exact vs prior commit (all act/quant/layout/inter_dim combos, COLD recompile) + pytest; qk_norm both tests pass (411 checkAllclose, 0 err), perf unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… chunk_gated_delta_h Infra (backward-compatible, index/i32 paths unchanged): - buffer_ops: new _to_i32_offset() normalizes any offset to i32 (index-> index_cast, i32->as-is, i64->trunc, narrower->sign-ext). Replaces the 4 IndexCastOp(i32, x) spots that raised "index_cast i64<->i32 incompatible" on fx.Int64 offsets. - tensor_shim GTensor.get_llvm_ptr: accept index / i32 / i64 byte offsets (keep i64 as-is; index->index_cast; narrower->extsi) so static_bytes_offset can be fx.Int64 math. - vector.py _as_index_ir_value: index_cast integer offsets to index for vector/memref indexing. chunk_gated_delta_h.py (now unblocked): - fx.Index(...) -> fx.Int64(...) everywhere (27 sites); the double index_cast(index)->index_cast(i64) LDS-ptr helper -> fx.Int64(x).ir_value() - vector.from_elements(T.f32x4, e) -> fx.Vector.from_elements(e, dtype=Float32) - drop now-unused arith + vector imports (rocdl.mfma + _to_raw are keepers) Verified: chunk 26/26 correctness (COLD); silu pytest + causal_conv1d 40/40 unaffected by the shared-infra change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…elect
Enabled by the tensor_shim/buffer_ops fx.Int64 support:
- fx.Index(...) -> fx.Int64(...) everywhere (GTensor static_bytes_offset,
vec_load/vec_store offsets, LDS indices, loop-carry scalars)
- softplus value-returning scf.IfOp -> Boolean.select: for beta_x > threshold
softplus(x)==x, so selecting between the (overflow-safe-to-discard) log1p arm
and x is bit-identical to the old branch. select is a plain op so it composes
inside the raw valid-guard region.
Kept as raw scf.IfOp (documented): the outer pool_idx>=0 whole-body guard and
the write_cond store guard -- flydsl's Python dynamic-if rejects branches whose
"state" includes GTensor stores ("state variable out_tensor is GTensor"), and
the write guard is nested inside that region.
Verified: test_flydsl_linear_attention.py 4/4 (gdr_decode) on COLD recompile.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…n if The earlier blocker (dynamic-if rejects GTensor stores as scf "state") is sidestepped by the skill's branch-helper idiom: wrap the guarded body in a local closure and call it under a Python `if`. The rewriter then sees an opaque call (no out_tensor assignment to thread through an scf.if yield) and lowers it to scf.if itself -- no raw region needed. - outer pool_idx>=0 whole-body guard: scf.IfOp -> `def _do_decode(): ...; if pool_idx >= 0: _do_decode()` - write_cond q-store guard: scf.IfOp -> `def _write_q(...): ...; if warp_k_vec_start == 0: _write_q()` (loop vars bound as defaults) - drop now-unused scf + ir imports gdr_decode is now scf-free (softplus was already select). Verified test_flydsl_linear_attention.py 4/4 on COLD recompile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… -> closure Prove the closure/@flyc.jit idiom on the compress path: the sentinel-skip scf.IfOp(position>=0) that wrapped the entire CSA kernel body becomes a local `def _body(): ...` invoked under `if fx.Int32(position) >= 0: _body()`. The rewriter sees an opaque call (no state to yield) and lowers it to scf.if, guarding every store -- byte-identical to the manual IfOp. Verified: op_tests/test_flydsl_compress_attn.py 204/204 (incl bit-exact e8m0 scale checks) on COLD recompile. Inner guards/value-return tail still scf; converted in follow-ups. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Both the CSA and ksplit kernels are now scf-free: - whole-body sentinel-skip scf.IfOp -> `def _body(): ...; if position>=0: _body()` - wave-0 block scf.IfOp -> `def _wave0(): ...; if wid==0: _wave0()` - lane/group store guards (is_even/is_lane0/is_grp_rep) -> inline Python `if` (buffer_store has no assignment target, so no scf "state" to thread) - value-returning tail scf.IfOp (phase2 online-softmax vs passthrough) -> per-lane Boolean.select (both arms pure arithmetic; overflow arm discarded -> bit-identical), extracted via .ir_value() for the downstream raw intrinsics - drop dead _if_then helper + contextmanager/scf/ir imports; refresh docstring Verified: test_flydsl_compress_attn.py 204/204 (incl bit-exact e8m0 scale checks across CSA / nm-asm / ksplit) on COLD recompile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Both HCA kernels are now scf-free using the closure-under-runtime-if idiom: - compress_forward: whole-body position>=0 guard -> `def _body()`; wave-0 reduce/scatter block -> `def _wave0()`, each called under a Python `if` - norm_rope_scatter: whole-body (position>=0 & pid<cap) guard -> `def _body()` under a Python `if` (the fp8 emitter it calls now runs in a traced region) - drop dead _if_then helper + contextmanager/scf/ir imports fused_compress_attn_common.py's emit_group_fp8_nm_asm_scatter KEEPS its scf.IfOp: it's an external IR-emitting helper, and the @flyc.kernel AST rewriter only rewrites `if` lexically inside the kernel (and its nested closures), not in a called module function -- a plain Python `if` there hits "cannot evaluate dynamic Boolean". Documented as load-bearing. Verified: test_flydsl_compress_attn.py 204/204 (bit-exact) on COLD recompile. 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):
|
# Conflicts: # aiter/ops/flydsl/kernels/tensor_shim.py
fx.thread_idx.x / fx.block_idx.x (and fx.Int32 kernel params, arith results) already carry an fx type and overload operators, so re-wrapping them is dead: - silu/swiglu: `bid = fx.Int32(fx.block_idx.x)` -> `bid = fx.block_idx.x` (same for thread_idx) - fused_compress_attn: `fx.Int32(tid)`/`fx.Int32(lid)` guards -> bare tid/lid - fused_compress_attn_hca: `fx.Int32(pid) < fx.Int32(plan_capacity)` -> `pid < plan_capacity` - gdr_decode: `pool_idx >= fx.Int32(0)` -> `pool_idx >= 0` Kept the wraps that actually introduce/convert a type: raw vector.extract results (fx.Int32(position) for the runtime `if`), buffer_load/GTensor loads, i32->i64 widening for grid dims / byte offsets, and raw rocdl/arith results feeding fx arithmetic. Verified byte-exact: silu pytest, gdr 4/4, swiglu 9/9 old-vs-new, compress 204/204, all COLD recompile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ast->fx.Int64
- gdr_decode: mlir_math.{exp,log1p,rsqrt} -> fx.math.* (drop the mlir_math
import); rsqrt result fed to raw vector.BroadcastOp via _to_raw.
- fused_compress_attn / _hca / qk_norm launchers: the host-side grid casts
arith.index_cast(T.index, plan_capacity/NUM_SPLIT/nblocks/num_tokens) ->
fx.Int64(...), matching the silu/pa_mqa launcher idiom.
Deliberately kept (no clean fx surface / documented):
- gdr_decode mlir_gpu.ShuffleOp (mode="idx" warp lane-broadcast has no fx
equivalent) + mlir_vector.ReductionOp + vector.FMAOp/BroadcastOp -- the
cohesive raw warp-reduction/MMA cluster; migrating around FMA/idx-shuffle
would add more _to_raw than it removes.
- k_static index_cast inside the raw arith score-index chains (raw MMA core).
- qk_norm bid_t_idx (index-typed on purpose for MulIOp byte offsets).
Verified: compress 204/204, qk_norm 411 checkAllclose, gdr 4/4 -- COLD recompile.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…cdl/fx.math
fx exposes the same AMD instructions, so the verbose llvm.call_intrinsic calls
are unnecessary:
- llvm.call_intrinsic("llvm.amdgcn.exp2.f32", [x]) -> fx.rocdl.exp2(f32, x)
- llvm.call_intrinsic("llvm.amdgcn.rcp.f32", [x]) -> fx.rocdl.rcp(f32, x)
- llvm.call_intrinsic("llvm.fabs.f32", [x]) -> fx.math.absf(x) (silu)
- silu fp8 cvt: raw rocdl.cvt_pk_fp8_f32 -> fx.rocdl.cvt_pk_fp8_f32 (drop the
local raw rocdl import + the now-empty `if _need_fp8:` block)
- drop the now-unused `from flydsl._mlir.dialects import llvm` in all 5 files
Same instructions -> byte-identical. Verified: silu 36/36 old-vs-new + pytest,
swiglu 9/9 old-vs-new, qk_norm 411 checkAllclose, compress 204/204 -- all COLD
recompile.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…le->shuffle_xor - mlir_vector.ReductionOp(T.f32, ADD, v).dest -> fx.Vector(v).reduce( fx.ReductionOp.ADD) (drop the mlir_vector import) - mlir_gpu.ShuffleOp(..., mode="xor") -> value.shuffle_xor(offset, WARP_SIZE) - reduced scalars are now fx, so: sum_hq.truncf(dtype_) -> sum_hq.to(fx_dtype_) (scalar cast), and the idx-shuffle input v_new wrapped via _to_raw. Kept (no clean fx surface): mlir_gpu.ShuffleOp(mode="idx") -- reads a per-lane source lane (K-group leader broadcast); readlane needs a uniform lane and ds_bpermute needs a manual bitcast dance, neither a byte-exact drop-in. Also vector.FMAOp/BroadcastOp -- fx has no vector FMA (a*b+c would be 2-rounding). Verified: test_flydsl_linear_attention.py 4/4 (gdr_decode) on COLD recompile. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
CI pins ruff 0.16.0 whose default set flags two things 0.15.19 did not: - SIM102 in silu_and_mul_fq padding branch: the outer `if const_expr(_need_quant):` wrapped only `if (col0 & 31) == 0:` -> combined into one `and` condition (const_expr short-circuits at trace time; byte-identical). - I001 import order in the silu op test (ptr_arg import). Verified: silu padding path byte-exact old-vs-new (fp4/fp8 x num_valid<rows, 6/6), ruff 0.16.0 check + format --check clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The pre-checks CI runs a `psf/black@stable` job in addition to `ruff check`.
My earlier `ruff format` runs rewrote black-compliant `assert (cond), "msg"`
into ruff's `assert cond, ("msg")` style (the two formatters diverge here),
which failed the black job. Re-run black on all touched files.
black --check clean + ruff 0.16.0 check clean on all 14 changed files.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 task
JH-Leon-KIM-AMD
added a commit
that referenced
this pull request
Aug 5, 2026
Merging main brought in #4501, which trims flydsl's stable interface, and the CI image's flydsl no longer exports buffer_ops from flydsl.expr. The two decode top-k kernels are the last files still importing it from there, so setup.py died collecting AOT jobs before it reached any kernel: aiter/ops/flydsl/kernels/topk_per_row_decode_tiered.py, line 51 ImportError: cannot import name 'buffer_ops' from 'flydsl.expr' Point both at aiter/ops/flydsl/kernels/{buffer_ops,vector}.py, which is where every kernel main touched in #4501 now gets them. The shims cover all five entry points these kernels use -- buffer_load, buffer_store, create_buffer_resource, create_llvm_ptr, extract_base_index -- plus vector.extract, with the same signatures, so no call site changes. arith, const_expr, gpu, range_constexpr and rocdl stay on flydsl.expr; kernels main updated still import those from there. Verified by running the step that failed: AOT job collection now succeeds for all six kinds (MOE 2659, MXFP4_MOE 33, GEMM 1852, GROUPED_MOE 142, CHUNK_GDN_H 12, TOPK 224). Kernel suite 99/99 and the gate smoke 12/12 are unchanged.
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