DSL-ify raw arith float ops in kernels (flash/mla/pa/moe), fastmath v… - #930
Open
xudoyuan wants to merge 12 commits into
Open
DSL-ify raw arith float ops in kernels (flash/mla/pa/moe), fastmath v…#930xudoyuan wants to merge 12 commits into
xudoyuan wants to merge 12 commits into
Conversation
…ia fast_fp_math Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace raw MLIR integer arith in kernels with DSL operators: - arith.cmpi(pred, a, b) -> a == / != / < / <= / > / >= - arith.divsi/divui -> //, remsi/remui -> % - arith.andi/ori/xori -> & / | / ^, shli -> <<, shrui/shrsi -> >> Operands stay DSL values (drop _raw/as_mlir_value/.ir_value unwraps). At raw-i1 consumers (scf.IfOp / llvm.intr_expect / arith.select cond), use as_ir_value(...), which handles both ArithValue (comparison over ArithValue operands) and Boolean (comparison over Numeric operands). Deliberately left as raw arith, with reasons: - unsigned compare where an operand can be negative (signed compare would misclassify) — must stay ult - pow2 strength-reduced shift/and standing in for divide/remainder - unsigned-division helpers whose contract is unsigned - integer min/max (no DSL integer min/max), ceil-div (no operator) - operands that are genuinely raw ir.Value (ballot / readfirstlane / scf result / dpp) with no DSL wrapper Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
De-index kernel code by replacing raw MLIR index builders with DSL: - arith.index(n) / arith.constant(v, index=True) -> fx.Index(v) - arith.index_cast(T.i32/i64, x) -> fx.Int32(x) / fx.Int64(x) At raw low-level consumers (get_llvm_ptr / llvm.* / rocdl.* / buffer DMA offset), the DSL value is unwrapped with as_ir_value(...). scf.for loop-carried counters that were arith.constant(0, index=True) become fx.Index(0) (type-preserving). Left as-is (with reasons): - index values used as layout coordinates (vec_load/vec_store/ linear_offset/cs_off/memref) and anything combined or compared with them: the layout shim casts coordinates from index, so they must stay index-typed - arith.index_cast(T.index, x) producing grid/launch or coordinate index values - raw index operands of raw shift/and builders (strength-reduced pow2 paths) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The last raw arith.cmpi here needed unsigned ult (local_expert_id can be negative for non-local experts; a signed compare would misclassify them). fx.Uint32 reinterprets the same i32 bits as unsigned, so the DSL `<` emits ult — behavior-preserving. Result feeds a dynamic if, so wrap with as_ir_value. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
arith.select unwraps DSL Numeric/Boolean conditions itself (via _to_raw), so wrapping the condition in as_ir_value was unnecessary. Pass the DSL comparison directly. as_ir_value stays only where the consumer is a raw MLIR builder that requires an ir.Value (scf.IfOp, llvm.*, rocdl.*, get_llvm_ptr, vector.*, buffer DMA offsets). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The as_ir_value wrapper was redundant: is_local feeds a dynamic `if` (the AST rewriter converts the condition to i1 itself) and is_local.select(...) (a DSL method). Keep it as the DSL Boolean from fx.Uint32 comparison. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The paged-decode accumulation/reduction ops originally used the narrow
`contract` fastmath flag (FMA fusion only, preserves accumulation order).
DSL-ifying them under a `fast_fp_math` hint widened the flag to `fast`,
whose reassoc reorders fp accumulation and loses precision on fp8 (bf16/
fp16 tolerate it). Switch the pa launchers' compile hint from
{"fast_fp_math": True} to {"fastmath": contract} so the DSL operators emit
`contract`, matching the original per-op flag.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
xudoyuan
marked this pull request as ready for review
July 31, 2026 05:09
coderfeli
reviewed
Jul 31, 2026
| m_new_raw = (m_running).maximumf(row_max) | ||
| if const_expr(traits.CAUSAL): | ||
| m_new_raw = _fmax(m_new_raw, ctx.c_neg_floor, fm_fast) | ||
| m_new_raw = (m_new_raw).maximumf(ctx.c_neg_floor) |
coderfeli
reviewed
Jul 31, 2026
| fmath.log(as_mlir_value(l_row), fastmath=self.fm_fast), | ||
| self.fm_fast, | ||
| ) | ||
| lse_val = ((m_row) * (self.c_ln2_f)) + (fmath.log(as_mlir_value(l_row), fastmath=self.fm_fast)) |
Collaborator
There was a problem hiding this comment.
also clean duplicated "()" ?
Collaborator
Author
There was a problem hiding this comment.
Got it. I will remove it.
Behavior-preserving cleanup of the DSL-ified attention kernels: - drop redundant parentheses left by the operand-wrapping conversion (e.g. `(a) - ((b) * (c))` -> `a - b * c`), keeping only precedence- required grouping - replace the `X.maximumf(Y)` method form with the `fx.maxnumf(X, Y)` free function (and `.minimumf` -> `fx.minnumf`); maxnumf matches the original flash reduction op and differs from maximumf only on NaN, which softmax max operands never produce Verified behavior-identical by AST equivalence (redundant parens do not appear in the AST; only the exact maximumf->maxnumf transform differs), plus flash gfx942 numerics. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
fmath is just `flydsl.expr.math`, and fx.rsqrt/fx.log/fx.fma/... are the same functions re-exported at top level (fmath is fx.math; fx.rsqrt is fmath.rsqrt). Use the fx.* names for consistency with the rest of the DSL and drop the redundant as_mlir_value on fx.log args (fx.log unwraps its argument internally). Behavior-identical (same function objects); flash gfx942 numerics unchanged. Note: rocdl.exp2/rcp are intentionally kept -- they lower to the bare v_exp/v_rcp hardware instructions, whereas fx.exp2 lowers to an __ocml_exp2_f32 call (slower, range-reduced), so they are not equivalent. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The amax exponent extraction used raw arith.bitcast/shrui/shli on _raw-unwrapped DSL operands. Rewrite with DSL ops (bit-exact): - arith.bitcast(T.i32, amax_f) + arith.shrui(.., 16) -> amax_f.bitcast(Uint32) >> 16 (Uint32 so >> emits the unsigned shift, matching shrui) - arith.shli(amax_dpp, 16) -> amax_dpp << 16 - arith.bitcast(T.f32, f32b) -> f32b.bitcast(Float32) - .maximumf() -> fx.maxnumf() for consistency Pure bit manipulation, no fastmath involved; moe a8w4/fp4 numerics pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
_fabs_f32 wrapped a raw llvm.call_intrinsic("llvm.fabs.f32"). fx.absf
(math.absf) lowers to the same llvm.fabs.f32 (via llvm.intr.fabs), so
replace the helper with fx.absf at all call sites and delete it. Also
`.maximumf()` -> fx.maxnumf() in gemm1 for consistency. Pure bit-op abs,
behavior-identical; moe a8w4/fp4 numerics pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Full-scan follow-up. Replace raw ops whose operands are already DSL and
whose DSL form lowers to the same instruction (bit-exact / same intrinsic):
- mfma_preshuffle: arith.bitcast(T.i32/f32, x) -> x.bitcast(fx.Int32/Float32)
- silu_and_mul_fq: llvm.call_intrinsic("llvm.fabs.f32") -> fx.absf; and
arith.maximumf -> fx.maxnumf
- pa_decode_swa: arith.bitcast(T.i32, weight_local) -> weight_local.bitcast(fx.Int32)
Left as-is (raw operands / no equivalent / unvalidatable): rocdl.exp2/rcp
(bare v_exp/v_rcp), arith.minsi/maxsi (no DSL int min/max), the shrui on a
raw index param, cmpf-result andi, ds_bpermute-result bitcast, and the
multi-gpu dispatch bitcasts (raw i64 operand, no local validation path).
Validated: test_preshuffle_gemm 58 passed; imports clean.
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
2 participants
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.
…ia fast_fp_math
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist