Apple GPU Tier-2: projection ops + batched MHA composition - #18
Merged
Merged
Conversation
Projections (linear_general, qkv_projection) now run on Apple GPU: - runtime.py: _apple_gpu_dispatch_linear_general routes the standard last-axis x @ W (+ bias) through the matmul/bmm lane (general tensordot -> numpy); _apple_gpu_dispatch_qkv_projection does x @ W_qkv then split-3 (multi-output tuple). _APPLE_GPU_PROJECTION_OPS gating (driver + runtime) -> metal_runtime. - tests/unit/test_apple_gpu_projections.py (9): rank-2/3, bias, general-axis fallback, qkv split, @jit metal_runtime gates. Batched multi-head attention via bmm composition (no per-head loop): - A full MHA block composes from the Tier-2 ops — qkv_projection -> bmm(Q,K^T) * scale -> softmax -> bmm(_,V) -> linear_general — with batch = B*H and an unbounded head dim (no flash_attn head_dim<=256 limit). - tests/unit/test_apple_gpu_batched_mha.py (6): 3 shapes vs float64 numpy, runtime/metal_runtime gates, and a head_dim=320 case. GQA/MQA Phase-1 (repeat-KV -> flash_attn) already lives in nn.functional. Docs: apple_gpu_tier2_tier3_plan.md (items 1-3 marked done), CLAUDE.md. Verified on Apple Silicon: full suite 5,810 passed / 0 failures under numpy 2.4.6; mypy 0; tessera-ir lit 61 PASS / 0 FAIL. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
May 30, 2026
Apple GPU: Gumbel-max inference sampler (#18-safe, reproducible)
gstoner
added a commit
that referenced
this pull request
May 31, 2026
1. SIMD device-caps probe ✅ tessera_apple_gpu_simd_caps() / Python apple_gpu_simd_caps() — an MTLGPUFamily-based bitmask (reduction / shuffle / shuffle-and-fill / simdgroup-barrier). Reports 0xF on this M-series Mac. Honest introspection, since these are universally true on Apple Silicon. 2. SIMD-reduction rowop optimization → tried, measured slower, reverted⚠️ Implemented simd_max/simd_sum in the tiled fused matmul_softmax_{f32,f16} kernels (T==32 == one SIMD-group, so the ~10-barrier tg_max/tg_sum tree collapses to two ops). Bit-correct, but A/B-benchmarked 0.72–0.93× (slower) — these kernels are matmul-K-loop + global-write bound, not reduction-bound, and simd_sum/simd_max (a 5-deep shuffle chain) cost more than the on-chip 32-lane tree. Reverted the kernels (documented in-code), kept the probe. Same discipline as the ThunderMittens result — the "obvious" optimization regressed, so it doesn't ship. 3. MPSMatrixRandomPhilox GPU RNG (opt-in) ✅ apple_gpu_random_uniform / apple_gpu_random_normal (f32). Philox-family, deterministic by seed, validated on range/mean/std. Deliberately a separate opt-in stream — NOT wired into tessera.rng (its Philox stream isn't bit-identical to the CPU reference, which would break check_determinism/Decision #18). 4. f16 linalg + QR + SVD ✅ f16/bf16 linalg: a _linalg_dtype_policy runs low-precision inputs on the GPU in f32 and casts back; f64 now routes to numpy in full precision (it was silently downcasting — a latent bug, fixed). QR (apple_gpu_qr): Cholesky-QR reusing the existing GPU Cholesky + tri-solve (no new MSL), with a compute-then-verify ‖QᵀQ−I‖ guard → numpy-Householder fallback. Caught the κ² instability — a barely-PD Gram factored but gave a non-orthonormal Q; the verify rejects it so the returned Q is always orthonormal. SVD (apple_gpu_svd): honest numpy deferral — MPS has no SVD/eigensolver, and a robust Jacobi-MSL kernel is a substantial standalone effort. Documented the design rather than shipping something fragile.
gstoner
added a commit
that referenced
this pull request
Jul 29, 2026
Extends the two survey documents with six more ROCm projects. Documentation only. Compiler survey gains section 4.7 on rocisa, TensileLite's nanobind assembly generator — the same Python-driving-C++ shape we have. Three findings worth copying: IR nodes carry a mandatory clone() deep-copy contract; exporting a vector to Python is a copy, so elements are mutable through their shared_ptr but cannot be assigned or replaced; and import raises if any C++ source is newer than the built extension. That last one is added to the take list — we lost time this session to a tessera-opt binary that silently did not match its sources. Patterns doc gains four project briefs and a rocWMMA re-read: rocFFT has the best cache design in the ecosystem. The kernel name is the cache key, with every differentiating parameter encoded into it, so profiler output and cache identity are the same string and the cache needs no schema update when a new parameter appears. Three further key fields guard staleness — architecture, HIP version, and generator version. A read-only system cache ships with the library alongside a read-write user cache, the shipped one populated at build time by a helper that shares the generator but is not installed. AOT and JIT are one path with a policy knob rather than two lanes. Also records that hipRTC holds process-wide locks, so parallel compilation needs a helper process. rocPRIM turns tuning output into generated headers, and its fallback_config is a typed fallback ladder: an untuned type inherits the config of a representative matched on size range and floating-pointness. That is dtype bucketing, the same move Decision #28 makes for shapes. rocRAND is the one with a direct bearing on us. Under dynamic ordering it picks launch geometry per device, and AMD states plainly that the number of generators and the sequence of generated numbers can vary as a result. So reproducibility versus performance is a named opt-in mode, not an emergent property. Worth confirming the same holds for Decision #18: if a tuned launch configuration ever fed an RNG offset scheme, autotuning would silently change numerical output. rocALUTION is included as a contrast, not a pattern. It selects execution location at run time via RTTI and silently migrates an object back to the host when a routine is unavailable on the accelerator. That is the opposite of Decision #21, which requires a diagnostic naming the op and target. Both are defensible for their audience; the contrast is worth recording because silent host migration is how a performance cliff hides. rocWMMA re-read adds that collaborative fragments are a movement concept and are explicitly unsupported in MMA functions, that partial and oversized tiles became the library's problem in 2.0.0, and that the wavefront-centric contract is undefined behaviour rather than a hint. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jul 29, 2026
Extends the two survey documents with six more ROCm projects. Documentation only. Compiler survey gains section 4.7 on rocisa, TensileLite's nanobind assembly generator — the same Python-driving-C++ shape we have. Three findings worth copying: IR nodes carry a mandatory clone() deep-copy contract; exporting a vector to Python is a copy, so elements are mutable through their shared_ptr but cannot be assigned or replaced; and import raises if any C++ source is newer than the built extension. That last one is added to the take list — we lost time this session to a tessera-opt binary that silently did not match its sources. Patterns doc gains four project briefs and a rocWMMA re-read: rocFFT has the best cache design in the ecosystem. The kernel name is the cache key, with every differentiating parameter encoded into it, so profiler output and cache identity are the same string and the cache needs no schema update when a new parameter appears. Three further key fields guard staleness — architecture, HIP version, and generator version. A read-only system cache ships with the library alongside a read-write user cache, the shipped one populated at build time by a helper that shares the generator but is not installed. AOT and JIT are one path with a policy knob rather than two lanes. Also records that hipRTC holds process-wide locks, so parallel compilation needs a helper process. rocPRIM turns tuning output into generated headers, and its fallback_config is a typed fallback ladder: an untuned type inherits the config of a representative matched on size range and floating-pointness. That is dtype bucketing, the same move Decision #28 makes for shapes. rocRAND is the one with a direct bearing on us. Under dynamic ordering it picks launch geometry per device, and AMD states plainly that the number of generators and the sequence of generated numbers can vary as a result. So reproducibility versus performance is a named opt-in mode, not an emergent property. Worth confirming the same holds for Decision #18: if a tuned launch configuration ever fed an RNG offset scheme, autotuning would silently change numerical output. rocALUTION is included as a contrast, not a pattern. It selects execution location at run time via RTTI and silently migrates an object back to the host when a routine is unavailable on the accelerator. That is the opposite of Decision #21, which requires a diagnostic naming the op and target. Both are defensible for their audience; the contrast is worth recording because silent host migration is how a performance cliff hides. rocWMMA re-read adds that collaborative fragments are a movement concept and are explicitly unsupported in MMA functions, that partial and oversized tiles became the library's problem in 2.0.0, and that the wavefront-centric contract is undefined behaviour rather than a hint. Co-Authored-By: Claude Opus 5 <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.
Projections (linear_general, qkv_projection) now run on Apple GPU:
Batched multi-head attention via bmm composition (no per-head loop):
GQA/MQA Phase-1 (repeat-KV -> flash_attn) already lives in nn.functional.
Docs: apple_gpu_tier2_tier3_plan.md (items 1-3 marked done), CLAUDE.md.
Verified on Apple Silicon: full suite 5,810 passed / 0 failures under numpy 2.4.6; mypy 0; tessera-ir lit 61 PASS / 0 FAIL.