rocm/gfx1151: register-block the f32 GEMM (grouped_swiglu) — MFU-ladder rung - #356
Merged
Merged
Conversation
…er rung Apply the Stage-F-proven lever (register-budget tiling, NOT LDS — LDS is wash-to-regression on this unified-memory APU) to the one GEMM kernel that had no blocking: the plain-VALU f32 GEMM (generate-rocm-gemm-f32-kernel) that grouped-SwiGLU rides. - Each thread now computes a TM×TN=4×4 output tile: per k-step it loads TM A-values + TN B-values and reuses them across 16 FMAs (each A elt reused TN times, each B elt TM times), cutting global loads ~2× and raising ILP. Bounds-guarded for arbitrary M/N/K (masked loads + guarded stores). - Runtime launch: one thread per tile (grid ceil(ceil(M/4)*ceil(N/4)/256)). - Measured on gfx1151 (end-to-end incl H2D/D2H): ~1.6× at 1024³ (12.9 -> 8.1 ms, 0.166 -> 0.264 TFLOP/s); the win scales with the working set (small shapes are launch/transfer-bound). Still correctness-first — not an MFU claim. - New test_rocm_gemm_f32_compiled.py locks correctness across shapes (esp. non-multiples of the 4×4 tile — the bounds guards). - Extended the hot-path perf ratchet with gemm_f32 rows + a live re-timing check so a regression in the blocked kernel fails CI. Verified: 144-test sweep green; drift gate in sync; ruff + mypy clean. 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: 0084822895
ℹ️ 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".
The gemm_f32 recorder rows + live ratchet test were gated on the FLASH probe (_rocm_compiled_flash_attn_available), but the f32 GEMM uses a DIFFERENT compiler pass (generate-rocm-gemm-f32-kernel). On a host where the f32 GEMM works but the flash lane is missing/broken, the gemm_f32 rows were silently omitted and the live test skipped — so blocked-GEMM regressions stopped being measured. Add _rocm_compiled_gemm_f32_available() (runs a tiny _rocm_f32_gemm, cached) and gate both the recorder rows and test_live_gemm_f32_within_ratchet on it instead. 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.
The perf-tuning follow-up. The WMMA GEMM MFU ladder is already thoroughly explored (STRIX_HALO Stage F/H: rung-1 register blocking is production; LDS/pipelining are wash-to-regression on this unified-memory APU — global bandwidth isn't the bottleneck). So this applies the proven lever — register-budget tiling, not LDS — to the one GEMM kernel that had none: the plain-VALU f32 GEMM (
generate-rocm-gemm-f32-kernel) that grouped-SwiGLU rides.Kernel change
Each thread now computes a TM×TN = 4×4 output tile instead of one element: per k-step it loads TM A-values + TN B-values and reuses them across 16 FMAs (each A elt reused TN times, each B elt TM times) — cutting global loads ~2× and raising ILP. Bounds-guarded for arbitrary M/N/K (masked loads + guarded stores). Launch is one thread per tile.
Measured (gfx1151, end-to-end incl H2D/D2H)
The win scales with the working set, as expected for register blocking (small shapes are launch/transfer-bound). Still correctness-first — not an MFU claim.
Tests
test_rocm_gemm_f32_compiled.pylocks correctness across shapes — especially non-multiples of the 4×4 tile (the bounds guards): 7×5×3, 13×17×9, 3×8×2, 65×63×31, …gemm_f32rows + a live re-timing check, so a regression in the blocked kernel (e.g. a lost register-blocking win) fails CI.Verification
144-test sweep green (f32 gemm, moe grouped_swiglu, grad_clip, perf ratchet, matrix, inventory guard); drift gate in sync; ruff + mypy clean.
The honest remaining lever (noted in inventory §9) is the deeper WMMA occupancy / dual-issue work (Stage F's named "next lever") — that's where ROCm's lead-performance status (Decision #28) gets fully earned.
🤖 Generated with Claude Code