Phase 8.4.6 — Threadgroup-tiled matmul_softmax_f32 + benchmark harness - #13
Merged
gstoner merged 1 commit intoMay 9, 2026
Conversation
Lifts the N <= 256 constraint on the Phase 8.4.3 fused matmul -> softmax
kernel by adding a threadgroup-tiled variant that allocates the score
buffer in dynamic threadgroup memory. One row per threadgroup; 32 threads
cooperate on the K reduction, threadgroup-reduce max + sum, and final
write. Caps at N <= 8192 (typical device threadgroup memory bound).
Also adds a benchmark harness comparing fused vs sequential paths across
representative shapes — first concrete perf-characterization scaffolding
for the apple_gpu MSL kernels.
Scope cuts (followups):
- Tiling for f16/bf16 variants — same pattern, just hasn't been replicated
- Tiling for the 3-op fusion (matmul -> softmax -> matmul) — needs both
`scores[N]` AND `out[P]` in threadgroup memory, slightly different layout
- MPSGraph baseline comparison in the benchmark — needs a separate setup
MLIR / runtime
- New native MSL kernel matmul_softmax_tiled_f32 with threadgroup memory:
* `tg_scores` allocated as dynamic threadgroup memory ([[threadgroup(0)]])
* Two `threadgroup float[32]` scratch arrays for max + sum reductions
* Cooperative loops over N step THREADS for compute, exp, write
* Tree reduction (stride/2) for both max and sum to keep latency O(log T)
- New C symbol tessera_apple_gpu_matmul_softmax_tiled_f32 — same i64×3 + i32×3
ABI as the per-thread variant. Reference fallback for non-Darwin.
- The existing tessera_apple_gpu_matmul_softmax_f32 entry point becomes a
router: per-thread for N <= 256 (no threadgroup overhead), tiled for
larger N. Reference fallback still works for N > 8192.
Python
- runtime.py: _apple_gpu_dispatch_matmul_softmax now lifts the f32
upper bound to 8192. f16/bf16 stay at 256 until tiled variants land.
Loader gate updated to require the new tiled symbol.
- apple_gpu_runtime_stub.cpp gets matching reference fallback so non-Darwin
builds export the same C ABI.
Benchmarks
- benchmarks/apple_gpu/benchmark_fusion.py — Python timing harness that
compares fused matmul_softmax vs sequential per-op pipeline across a
shape sweep. Outputs JSON in the same schema as benchmark_gemm.py so
tools/roofline_tools/ can ingest it. Skips on non-Darwin with a clear
message and exits 0.
Tests
- 3 new unit tests in test_apple_backend_roadmap.py:
* test_apple_gpu_matmul_softmax_tiled_path_executes_for_large_n —
end-to-end @jit through tiled at N=512 and N=1024
* test_apple_gpu_matmul_softmax_tiled_runtime_shim_exposes_symbol —
direct ctypes ABI test for the new symbol
* test_apple_gpu_matmul_softmax_small_n_still_uses_per_thread_path —
pins that the router doesn't accidentally use tiled for small N
Verified on Apple Silicon (LLVM/MLIR 21, Metal active):
2015 unit tests passing (2012 + 3 net new tiled tests);
15/15 Phase 8 lit fixtures passing against the in-tree tessera-opt.
Tiled path matches numpy reference at rtol=1e-4 for both N=512 and
N=1024. Benchmark harness runs end-to-end and emits JSON.
Co-Authored-By: Claude Opus 4.7 <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: bb304eabcc
ℹ️ 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".
Base automatically changed from
claude/phase-8-4-5-attn-fusion
to
claude/phase-8-4-4-2-fused-flash-attn
May 9, 2026 12:47
gstoner
merged commit May 9, 2026
3240682
into
claude/phase-8-4-4-2-fused-flash-attn
0 of 5 checks passed
This was referenced May 9, 2026
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
Lifts the
N <= 256constraint on the Phase 8.4.3 fused matmul→softmax kernel via a threadgroup-tiled variant that allocates the score buffer in dynamic threadgroup memory. Also adds a benchmark harness comparing fused vs sequential execution paths.Stacked on PR #12.
Scope cuts (deliberate, called out as followups)
scores[N]ANDout[P]in threadgroup memory; slightly different layoutThe point of this PR is to demonstrate the technique on one kernel cleanly. Replicating to the other dtypes/kernels is a 1-PR-each followup.
What changed
MLIR / runtime
matmul_softmax_tiled_f32:tg_scores[N]via[[threadgroup(0)]]threadgroup float[32]scratch arrays for max + sumO(log T)latencyn in [lid, N) step THREADStessera_apple_gpu_matmul_softmax_tiled_f32(samei64×3 + i32×3ABI as the per-thread variant)tessera_apple_gpu_matmul_softmax_f32becomes a router: per-thread forN ≤ 256(no threadgroup overhead), tiled for larger N. Reference fallback forN > 8192.Python
runtime.py:_apple_gpu_dispatch_matmul_softmaxlifts the f32 upper bound to 8192. f16/bf16 stay at 256 until tiled variants land.Benchmarks
benchmarks/apple_gpu/benchmark_fusion.py— JSON-schema-compatible harness comparing fused vs sequential paths across a shape sweep.Tests
Test plan
rtol=1e-4for both N=512 and N=1024PR stack
🤖 Generated with Claude Code