rocm/gfx1151: FA-2 backward runtime lane (MHA+GQA+attn_bias), KV-cache paged lane, lstm_cell + inventory reconciliation - #354
Merged
Conversation
…e paged lane, lstm_cell + inventory reconciliation Closes several gfx1151 gaps, all execute-compare verified on real RDNA 3.5 silicon. Lanes added / promoted (runtime.launch, execution_mode=hip_runtime): - rocm_flash_attn_bwd_compiled — FA-2 backward (dQ/dK/dV) as a first-class lane. The compiler-generated backward kernel existed in C++ but was never wired into runtime.launch(). Added _execute_rocm_compiled_flash_attn_bwd (fa_pre -> fa_dkdv -> fa_dq; O recomputed via the forward lane, nothing saved from forward). Covers MHA + GQA/MQA (grouped fa_dkdv atomic-accumulates dK/dV) + additive attn_bias (S = scale*Q@K^T + bias in the recompute). The attn_bias variant is a new C++ path in GenerateWMMAFlashAttnBwdKernel.cpp (attn_bias BoolAttr -> trailing f32 [bh,Sq,Sk] memref added to S in emitPre + recomputeScoreTile). f16/bf16, f32 accumulate; matches autodiff vjp_flash_attn. - rocm_kv_cache_compiled — KV-cache paged-movement core (append/read/prune) over a resident cache buffer, composing the existing device scatter (append row write) + masked-gather (read/prune) kernels with host page-index math. Matches the KVCacheHandle reference. quantize_kv rides the intquant lane. - lstm_cell — added F.lstm_cell reference (gate order i,f,g,o) and registered it in the structured-compute lane (x86 + rocm), parity with gru/simple_rnn. Docs / automation: - Reconciled docs/rocm_mfma_kernel_inventory.md with the generated runtime execution matrix (it had drifted badly: claimed only matmul + flash_attn executed while dozens of compiled HIP lanes had landed). Points at the generated matrix as status truth (Decision #26). - Hardened the inventory guard test: TestH3RocmInventoryExecutionStatus now cross-checks every rocm_*_compiled lane named in the doc against the drift-gated matrix, so the execution-status prose cannot silently rot again. Verification: 147 unit tests pass (FA fwd/bwd runtime+direct, gqa-direct, structured x86+rocm, matrix, inventory guard); generated-doc drift gate in sync. ROCm runtime lanes 65 -> 67. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…scatter-add)
Promote the MoE transport ops off the reference oracle to real gfx1151 execution:
- moe_dispatch → the device gather kernel: token_of_slot = sort_perm // top_k
row gather (reuses _rocm_gather_rows). Bit-exact vs the numpy oracle.
- moe_combine → the device scatter (add) kernel: host pre-scales each packed row
by its route weight (exact f32), then atomic scatter-add to token order.
f32 accumulate — matches the f64 oracle within ~1.2e-7.
Both report native_gpu via the per-op (output, execution_kind) launch override;
they fall back to the stdlib oracle + reference_cpu off-box. grouped_swiglu (the
expert GEMM) stays on the oracle / reference_cpu — a native f32-exact grouped
GEMM is a separate follow-up (WMMA is f16). This was exactly the promotion the
matrix note called for ("requires HIP gather/scatter transport kernels").
Result: every ROCm row in the generated runtime execution matrix is now
native_gpu — zero reference_cpu ROCm lanes remain.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
![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: a3de4ff0e1
ℹ️ 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".
…ouped_swiglu) Add a plain f32 GEMM device kernel and use it to promote grouped_swiglu, so the whole rocm_moe_transport_compiled row is native (the last non-transport op). - New C++ pass generate-rocm-gemm-f32-kernel (GenerateROCMGemmF32Kernel.cpp, tessera_rocm.gemm_f32 ODS op): C[m,n] = sum_k A[m,k]*B[k,n], one thread per output element, scalar f32 k-loop (f32 accumulate). RDNA WMMA is f16/bf16 only, so this is the f32-exact VALU/FMA fallback the expert GEMM needs; correctness- first (no LDS tiling — the blocked perf ladder is a follow-up). - Runtime _rocm_f32_gemm (build hsaco + HIP launch) + _moe_grouped_swiglu_native: per expert group, g=x@Wg, u=x@Wu on the device GEMM, silu*mul host-side, then h@Wd on the device GEMM. grouped_swiglu now reports native_gpu; matches the f64 oracle within ~2.4e-7 (f32). Falls back to the oracle + reference_cpu off-box. Every ROCm row in the generated runtime execution matrix is now native_gpu. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Review feedback (P2): the structured-compute lane dispatched tessera.lstm_cell to F.lstm_cell, which has a DIFFERENT contract than the canonical tessera.ops.lstm_cell (the public op uses (4H,In)/(4H,H) weights via W_ih.T/W_hh.T and returns a packed concat([h_t, c_t]) for lstm_state_h/c extraction; F.lstm_cell took transposed weights and returned a tuple). Normal LSTM artifacts would fail shape validation or produce output the state extractors can't consume on the newly-advertised x86/ROCm/Apple structured lanes. - Dispatch tessera.lstm_cell -> ops.lstm_cell (canonical packed contract). - Remove the misleading F.lstm_cell duplicate + its __all__ entry. - Update the x86/ROCm structured fixtures to the packed contract ((4H,In) weights). Also fix the mypy ratchet failure: assert bias_c is not None in the FA backward attn_bias device-copy (mypy could not narrow the has_bias flag). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Complete the flash_attn backward variant surface — the runtime lane now covers the full set of forward variants (MHA + GQA/MQA + attn_bias + sliding-window + logit-softcap), all execute-compare verified vs autodiff vjp_flash_attn on gfx1151. C++ (GenerateWMMAFlashAttnBwdKernel.cpp): - ScoreCtx gains W (window width) + cap (soft-cap). recomputeScoreTile applies the soft-cap S=cap*tanh(sRaw/cap) before the softmax and scales dS by the chain-rule factor 1-tanh^2(sRaw/cap); the window mask (implicitly causal, keys older than W masked) is applied in both emitPre (for L) and recomputeScoreTile (for P/dS). emitDkDv/emitDq use useCausal for the causal tile-skip bounds. - Trailing runtime args appended in a FIXED order: gqa(heads,kv_ratio) | window(W:index) | softcap(cap:f32) | attn_bias([bh,Sq,Sk] f32, LAST), driven by the sliding_window / logit_softcap / attn_bias BoolAttrs. Runtime: thread window / logit_softcap kwargs (hsaco flags + trailing launch args in the same order); the forward-O recompute uses the same window/softcap so O matches the backward's score recompute. Verified: 22/22 gfx1151 tests (incl. window, softcap, and window+softcap+bias compose). mypy clean, drift gate in sync, 144-test regression sweep green. 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
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 several gfx1151 (Strix Halo / RDNA 3.5) gaps from the ROCm MFMA kernel inventory §5/§9 — all execute-compare verified on real silicon.
Lanes added / promoted (
runtime.launch,hip_runtime)rocm_flash_attn_bwd_compiled— FA-2 backward, first-classThe compiler-generated backward kernel (
generate-wmma-flash-attn-bwd-kernel→fa_pre/fa_dkdv/fa_dq) existed in C++ and passed direct-launch fixtures, but was never wired intoruntime.launch(). Added_execute_rocm_compiled_flash_attn_bwd— a self-contained VJP over(dO, Q, K, V[, bias])that recomputesOvia the forward lane (nothing saved from forward) and launches the three kernels in sequence.fa_dkdvatomic-accumulates dK/dV across thekv_ratioquery heads sharing each KV head; host pre-zeros dK/dV)attn_bias— a new C++ path inGenerateWMMAFlashAttnBwdKernel.cpp: anattn_biasBoolAttr adds a trailing f32[bh,Sq,Sk]memref toSinemitPre(for L) and the sharedrecomputeScoreTile(for P/dS), soS = scale·Q@Kᵀ + biasbefore the softmax.dbiasis not produced (constant-bias VJP contract).vjp_flash_attn.rocm_kv_cache_compiled— KV-cache paged movementkv_cache_append/read/pruneover a resident cache buffer, composing the existing device scatter (append row write) + masked-gather (read/prune) kernels with host page-index math. Matches theKVCacheHandlereference;quantize_kvrides the intquant lane.lstm_cellAdded
F.lstm_cell(gate order i,f,g,o) and registered it in the structured-compute lane on x86 + rocm, parity withgru_cell/simple_rnn_cell.Docs / automation
docs/rocm_mfma_kernel_inventory.mdwith the generated runtime execution matrix — it had drifted badly (claimed only matmul + flash_attn executed while dozens of compiled HIP lanes had landed). Now points at the generated matrix as status truth (Decision Apple GPU MLA: paged-cache decoder (production-serving wiring) #26).TestH3RocmInventoryExecutionStatuscross-checks everyrocm_*_compiledlane named in the doc against the drift-gated matrix, so the execution-status prose can't silently rot again.Verification
test_rocm_flash_attn_bwd_runtime_lane.py(15 cases: MHA/GQA/MQA/bias, f16+bf16, causal, ragged),test_rocm_kv_cache_compiled.py(7 cases).🤖 Generated with Claude Code