[FlyDSL] Raw-dialect cleanup (others): gdr_decode / chunk_gated_delta_h / pa_mqa_logits_fp4 - #4609
Open
coderfeli wants to merge 11 commits into
Open
[FlyDSL] Raw-dialect cleanup (others): gdr_decode / chunk_gated_delta_h / pa_mqa_logits_fp4#4609coderfeli wants to merge 11 commits into
coderfeli wants to merge 11 commits into
Conversation
vector.BroadcastOp/from_elements/FMAOp -> fx.Vector.filled/from_elements/fx.math.fma (keeps fused v_fma); arith.constant -> fx literals, drop arith import. rocdl.exp2 + gpu.shuffle kept raw (fx I/O). gfx950 byte-exact (16/16), perf within noise. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
llvm.inttoptr<3> -> fx.to_llvm_ptr; drop redundant _to_raw at scf.for range(init=) boundary. rocdl.mfma/ds_read kept raw. gfx950 byte-exact (15/15); lowered IR md5-identical (perf-neutral). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
create_buffer_resource+buffer_load/store -> make_buffer_tensor + copy atoms / indexed fx pointers; 1-writer scatter -> guarded fx.ptr_store; ArithValue -> fx operators. rocdl.mfma_scale/ds_bpermute + mlir_undef (poison pad) kept raw. gfx950 op tests pass; byte-exact (ISA+LLVM-IR identical); prefill ~11% faster. 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):
|
- _i32_buffer: replace inttoptr(Int64(ptrtoint(get_iter))) retype round-trip with recast_iter; drop redundant _I32_MAX_RECORDS (make_buffer_tensor defaults to max_size=True = descriptor 0xFFFFFFFF). - prefill: collapse the duplicated cta_info buffer (cta_info_bt + cta_info_flat over the same ptr) into one width-4 buffer; read fields 4,5 via 2D indexing. Fold the per-CTA base with add_offset (element offset) instead of int byte math. - port copy_atom_call -> fx.copy for the vec4 KV load and the Q/QS/W loads. Validated byte-exact (cos=1.0) on decode + prefill op-tests, perf on par. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_i32_buffer / _load_vec4_i32 / _pack_i32_pair_to_i64 / _pack_lo_i64x2_to_i32x8 were duplicated verbatim in the decode and prefill kernels. Move them to one pa_mqa_logits_fp4_common module and import from both (-88/+63 net, single source of truth). Drop the now-unused _llvm import from both kernels. Validated byte-exact (cos=1.0) on decode + prefill op-tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…kage Move fp8_mqa_logits, pa_mqa_logits_fp4, _prefill and _common into a dedicated kernels/mqa_logits/ package. Update the package __init__, the two op-test deep imports, and fp8_mqa_logits's sibling import (.tensor_shim -> ..tensor_shim). Validated: fp4 decode + prefill op-tests PASS (cos=1.0); fp8 module imports (kernel is gfx1250-only, skipped at runtime on gfx950). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…style) Replace the raw rocdl.mfma_scale_f32_16x16x128_f8f6f4 calls in the decode and prefill MQA-logits kernels with fx.make_mma_atom(cdna4.MFMA_Scale(...)) + fx.gemm(atom, c, a, b, c, scale_a=, scale_b=). Q/KV operands are now i32<4:1> register fragments (drops the manual v8i32 hand-packing), accumulators bridge the software-pipelined SSA carry via c_frag store/load, and e8m0 scale words are passed as plain i32 values. Prefill keeps its per-nt opsel_b via one atom per nt. Not ported to a tiled MMA on purpose: per-block e8m0 scales + prefill's per-nt opsel can't be expressed by a single make_tiled_mma (one atom / one opsel per tile); the only in-tree scaled-MMA precedent, mxmoe_gemm_v2, uses per-call fx.gemm for the same reason. The now-unused v8i32 pack helpers are removed from pa_mqa_logits_fp4_common. Validated byte-exact (cos=1.0) on decode (all shapes) + prefill (all shapes); decode perf on par (within run-to-run noise). ruff F401/F811/F821 clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Resolves the fp8_mqa_logits.py conflict: main's #4606 tiled-copy cleanup rewrote the file (dropping _to_raw) while this branch moved it into the kernels/mqa_logits/ subpackage. Kept main's body, re-pathed the shim import to '..tensor_shim'.
- black: the refactor left single blank lines before top-level defs in
pa_mqa_logits_fp4{,_prefill}.py, and the longer mqa_logits module path
pushed an import in the op-test past the line limit. Base was clean, so
these would have failed the format check.
- gdr_decode: fx.math.fma already returns an fx.Vector, so the
fx.Vector(...) re-wrap before .reduce() is a no-op. Dropped in 3 places.
The hoisted 'from .pa_mqa_logits_fp4_common import ...' sits directly against the third-party import block; isort wants a blank line between the third-party and local-folder groups. Both lines are PR-introduced, so reviewdog's diff_context filter failed the ruff job. Verified with the CI-pinned ruff 0.16.0 and black: both clean.
coderfeli
force-pushed
the
flydsl-cleanup-others
branch
from
August 8, 2026 08:51
e726adf to
9635490
Compare
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.
Summary
Migrates four gfx950 attention/delta FlyDSL kernels off raw MLIR dialects onto the
fx.*surface (part of a broader raw-dialect cleanup, split by area). Op-surface only — no logic/tiling/math/offset changes.gdr_decodevector.*,arith.*chunk_gated_delta_hllvm.inttoptr<3>, redundant_to_rawpa_mqa_logits_fp4buffer_ops.*,ArithValuepa_mqa_logits_fp4_prefillbuffer_ops.*,ArithValueWhat changed
vector.*→fx.Vector(.filled/.from_elements/fx.math.fmawhich keeps the fusedv_fma_f32).buffer_ops.create_buffer_resource+buffer_load/store→make_buffer_tensor+ copy atoms / indexed fx pointers; 1-writer scatter → guardedfx.ptr_store.arith.constant/ArithValue→ fx literals/operators;llvm.inttoptr<3>→fx.to_llvm_ptr(backend-resolved address space).scf.forloop-carry alreadyrange(init=); removed the redundant boundary_to_raw(the rewriter converts fx atinit=/yield).Kept raw (evidenced boundaries)
rocdl.mfma*,rocdl.ds_read_tr16_b64,rocdl.ds_bpermute,rocdl.exp2,gpu.shuffle(hand-packed MFMA operands / warp shuffles — no fx wrapper), and_llvm.mlir_undef(fx has no undef/poison). Tiled-copy (make_tiled_copy/partition_S/D) does not apply to any of these — they're per-lane MFMA / swizzled-LDS / paged-gather patterns, not contiguous bulk copies.Validation
gfx950, cache ON. Byte-exact vs the original per kernel (several proven via md5-identical lowered IR / ISA diff — stronger than
torch.equal);pa_mqaop tests pass before+after; no perf regression. Ruff + Black clean.Note: the
_gfx1250/_gfx1201attention variants are intentionally not in this PR — they can't be byte-exact/perf validated on a gfx950 box and need a matching-arch machine.