feat(rocm): compiler-generated flash_attn fwd+bwd + IR-stack matmul glue + int/FA perf ladders - #90
Merged
Merged
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1fe77ef81a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+1219
to
+1220
| TargetOp("tessera_rocm.wmma_gemm", {**base, "name": "gemm", | ||
| "m": 16, "n": 16, "k": 16, "dtype": dtype}), |
There was a problem hiding this comment.
When a ROCm function contains more than one matmul, this lowering emits every executable directive with name = "gemm". The generate-wmma-gemm-kernel pass uses that name to create gpu.module @<name>_mod and gpu.func @<name>, so consuming the Target IR for a multi-matmul graph produces duplicate gemm_mod/gemm symbols and the compiled pipeline cannot materialize the kernels. Include the ordinal/result (or another stable unique suffix) in the directive name instead of using a constant.
Useful? React with 👍 / 👎.
…piled op Brings flash_attn into the compiler-generated lane (was hand-written HIPRTC only). The Stage L machinery (directive -> generated WMMA kernel -> in-process hsaco) now covers a second op. - New tessera_rocm.flash_attn directive (head_dim, dtype) + the generate-wmma-flash-attn-kernel pass: a faithful MLIR re-emission of the hardware_verified hand-written FA-2 forward kernel — one wave per (16-query tile, b*h), LDS-staged Q (gpu.func workgroup attributions), S = scale*Q@K^T on WMMA over head-dim chunks, causal/ragged mask, online softmax (running max/sum, rescale), O += P@V on WMMA. Scores are staged in LDS so the QK^T accumulator layout is reread in the P@V A-fragment layout (the layout bridge). head_dim (mult of 16) is compile-time; Sq/Sk/scale/causal are runtime args. - tessera-opt: registered the math ConvertToLLVM external model (so convert-gpu-to-rocdl lowers the softmax math.exp -> llvm exp) + the math dialect; gpu.barrier + workgroup LDS already lower. The flash_attn pipeline is the same in-process chain as the GEMM lane (no mlir-opt). - Test: the compiler-generated FA-2 forward executes on gfx1151 matching a numpy attention reference (maxerr < 2e-2) across head_dim 16/64, causal/non-causal, and ragged Sq/Sk. test_rocm_flash_attn_compiled.py. Honest scope: forward only; the runtime.launch() executor-table lane (a flash_attn op-metadata contract + executor + matrix row) is the remaining glue, same additive step matmul took at L4 — not yet wired, so no execution-matrix row claimed. backward + perf ladder remain (audit item 10). drift in sync, ruff clean, rocm/wmma/flash_attn regression green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Item #2: close the Decision #19 gap. The Graph tessera.matmul -> Tile -> Target-IR lowering (_lower_rocm_op on tile.mma) now EMITS the executable tessera_rocm.wmma_gemm directive (m=n=k=16 WMMA tile + dtype) alongside the abstract tessera_rocm.mfma marker. So a @jit(target="rocm") matmul's target_ir contains the directive the generate-wmma-gemm-kernel pass consumes — the directive is produced by the IR stack, not only synthesized by the runtime. - target_ir.py: tile.mma -> [mfma (abstract marker, kept for the hardware-free contract + lit), wmma_gemm (concrete RDNA executable directive), async_copy, wait]. dtype threaded from the tile op (f16 default). - Test (GPU-free, CI-runnable): test_rocm_matmul_front_end_glue.py — the directive appears in target_ir with the right attrs AND the extracted directive feeds the generate pass into a gpu.func + WMMA op (directive consumed). - The abstract mfma marker stays (target_ir_contract / lit assertions unchanged). - The runtime lane still synthesizes a clean directive at launch for the per-shape mt/nt perf choice; the canonical lowering now owns directive production. Docs updated (op .td + ROCM_AUDIT). drift in sync, ruff + mypy clean, target_ir + rocm/wmma regression green (pre-existing test_apple_value_target_ir failures are Apple-backend-off, unrelated). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…t4 deferred) Item #3. Measured the compiler-generated int paths rather than assuming. benchmark_rocm_compiled_gemm_dtype.py (new): kernel-only dtype sweep of the compiled WMMA GEMM (f16/bf16/int8/int4, best macro-tile), honest-gated, JSON schema. On gfx1151 at 2048^3: f16 ~23.2 TFLOP/s, bf16 ~23.1, int8 ~21.0 TOP/s, int4 ~23.8 TOP/s (within ~10%) Finding: RDNA 3.5 WMMA runs iu8/iu4 at the SAME matrix-op rate as f16 (no low-precision FLOP-rate multiplier), so the compiled int paths are already compute-competitive and the int4 in-kernel nibble-pack is amortized. Consequence (measured, not assumed — Decision #25): packed-memory int4 (2 int4/byte) would buy memory footprint (1/2) + bandwidth, NOT compute on this arch. Its large sub-byte-strided-B layout is therefore deliberately DEFERRED — unjustified by a compute speedup that doesn't exist on RDNA 3.5. Documented in ROCM_AUDIT with the numbers. drift in sync, ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Item #4 (perf-ladder half). Moves the compiled FA-2 forward from "rung-0 correctness-only, no perf data" to a measured ladder. benchmark_rocm_flash_attn_compiled.py (new): kernel-only hipEvent ladder of the compiler-generated FA forward across (head_dim, seqlen), honest-gated, JSON schema. On gfx1151: ~4.0 TFLOP/s at head_dim 64, ~2.4 at 128 (FA-2 fwd FLOPs = 4*B*H*Sq*Sk*D). Modest by design — the kernel is correctness-first (one wave per query tile, LDS round-trips, online-softmax barriers, no KV pipelining / double buffering / multi-wave query tiles); the ladder quantifies the headroom. Audit item 10 updated: forward + forward-ladder done; flash_attn BACKWARD is the largest remaining attention piece (no hand-written oracle — a new kernel validated vs a numpy attention-backward reference; a focused standalone effort comparable to the forward), plus the runtime.launch() executor-table lane. drift in sync, ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Expand a single tessera_rocm.flash_attn_bwd directive into the textbook FA-2 backward as three fragment-materialized RDNA WMMA kernels (no stored attention matrix; S/P recomputed per tile): _pre scalar logsumexp L + D=rowsum(O*dO) _dkdv per key-tile: dP=dO@V^T, dS=P*(dP-D), dV+=P^T@dO, dK+=scale*dS^T@Q _dq per query-tile: dQ+=scale*dS@K All use the same C[m,n]=sum_k A[m,k]B[n,k] WMMA primitive + Stage J->I lowering as the forward; P/dS staged in LDS and reread transposed (the layout bridge). Executes on gfx1151 vs a numpy attention-backward reference (itself checked against finite differences): rel-err ~2-4e-4 (f16 storage, f32 accumulate) across head_dim 16/64, causal+non-causal, ragged. flash_attn (fwd+bwd) is now the third compiler-generated op on ROCm after matmul. Backward perf ladder measured: ~1.1-1.3 TFLOP/s @ D=64 (correctness- first; scalar logsumexp pre-pass + 5 matmuls dominate — WMMA logsumexp / causal tile-skip / pipelining are the next rung). - ROCM_FlashAttnBwdOp ODS + GenerateWMMAFlashAttnBwdKernel pass (registered) - tests/unit/test_rocm_flash_attn_bwd_compiled.py (on-device, skip-clean) - benchmarks/rocm/benchmark_rocm_flash_attn_bwd_compiled.py (ladder) - ROCM_AUDIT item 10 updated Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
force-pushed
the
rocm/flash-attn-compiled-lane
branch
from
June 24, 2026 00:52
1fe77ef to
f646bd6
Compare
gstoner
changed the title
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.
Closes the flash_attn compiler-generated lane (forward and backward) plus the supporting glue and measured perf ladders. All work validated on gfx1151 (Strix Halo / RDNA 3.5) where a GPU is present; skip-clean elsewhere.
What landed
flash_attn forward — compiler-generated.
tessera_rocm.flash_attndirective →generate-wmma-flash-attn-kernelpass → real FA-2 forward kernel (LDS-staged Q, QK^T on WMMA, online softmax viamath.exp, P@V on WMMA), lowered through Stage J (rocdl.wmma) + Stage I (→hsaco, in-process). Matches numpy (maxerr <2e-2) across head_dim 16/64, causal/non-causal, ragged.Front-end glue (Decision Apple GPU Tier-2/3: reductions, native GQA, fused batched attention #19).
_lower_rocm_op(tile.mma)now emits the executabletessera_rocm.wmma_gemmdirective, so a@jit(target="rocm")matmul's Target IR contains it — the IR stack produces it, not just the runtime. GPU-free test.int8/int4 perf sweep — measured.
benchmark_rocm_compiled_gemm_dtype.py: f16/bf16/int8/int4 all within ~10% because RDNA 3.5 runs iu8/iu4 WMMA at the same matrix-op rate as f16. Packed-memory int4 deliberately deferred (footprint-only on this arch — measured, not assumed).flash_attn backward — compiler-generated (NEW).
tessera_rocm.flash_attn_bwd→generate-wmma-flash-attn-bwd-kernelexpands to the textbook FA-2 backward as three fragment-materialized WMMA kernels (no stored attention matrix; recompute S/P per tile):_pre: scalar logsumexpL+D=rowsum(O·dO)_dkdv(per key-tile):dP=dO@Vᵀ,dS=P·(dP−D),dV+=Pᵀ@dO,dK+=scale·dSᵀ@Q(P/dS staged in LDS, reread transposed — the layout bridge)_dq(per query-tile):dQ+=scale·dS@KExecutes on gfx1151 vs a numpy attention-backward reference (itself checked against finite differences): rel-err ~2–4e-4 (f16 storage, f32 accumulate) across head_dim 16/64, causal+non-causal, ragged. Backward perf ladder: ~1.1–1.3 TFLOP/s @ D=64 (correctness-first; scalar logsumexp pre-pass + 5 matmuls dominate — WMMA logsumexp / causal tile-skip / pipelining are the documented next rung).
flash_attn (fwd+bwd) is now the third compiler-generated op on ROCm after matmul.
Validation
test_apple_value_target_irfailures confirmed unrelated (Apple backend OFF in the ROCm-only build).Still open (additive, documented in ROCM_AUDIT item 10)
runtime.launch()executor-table lane (the same step matmul took at L4).🤖 Generated with Claude Code