Phase 8.3 + 8.4.0 — Apple GPU runtime: MPS matmul + custom MSL kernels (rope first) - #3
Merged
Merged
Conversation
Single rank-2 f32 matmul on @jit(target="apple_gpu") now executes through MPSMatrixMultiplication, mirroring the Phase 8.2 Apple CPU pattern. The runtime envelope is intentionally narrow — multi-op programs keep the existing metal_artifact contract until Phase 8.4 adds custom MSL kernels. MLIR - ODS: tessera_apple.gpu.mps_matmul / mps_softmax / mps_dispatch - Pass: MatmulToAppleGPU lowers tessera.matmul (rank-2, f32) to a func.call into tessera_apple_gpu_mps_matmul_f32 - Pipeline: tessera-lower-to-apple_gpu-runtime - target_ir verifier accepts the new ops; module-level execution_mode flips to "metal_runtime" only for single-matmul programs Runtime - apple_gpu_runtime.mm: Objective-C++ MetalDeviceContext singleton (MTLDevice + MTLCommandQueue) wrapping MPSMatrixMultiplication with shared-storage MTLBuffers - apple_gpu_runtime_stub.cpp: portable reference path on non-Darwin so the static lib + Python ctypes layer remain platform-agnostic - CMake links -framework Metal -framework MetalPerformanceShaders -framework Foundation on Darwin only Python - driver.py: _is_apple_gpu_mps_executable gate, "metal_runtime" mode - jit.py: _apple_gpu_fast_call cached-metadata fast path, new metadata branch with compiler_path="apple_gpu_mps" - runtime.py: _execute_apple_gpu_mps_metadata, _apple_gpu_dispatch_matmul, _load_apple_gpu_runtime with on-the-fly compilation; launch() gets an apple_gpu_mps branch with telemetry Tests - New lit fixture apple_gpu_runtime.mlir (positive + negative paths) - Three unit tests in test_apple_backend_roadmap.py covering the runtime contract, ABI shim correctness across multiple GEMM shapes, and multi-op metal_artifact preservation - test_target_ir_contract.py split into single-matmul (mps_runtime) and multi-op (metal_artifact) cases Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 1947 unit tests passing; 7/7 Phase 8 lit fixtures passing against the in-tree tessera-opt. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Establishes the custom MSL kernel emission and dispatch path for apple_gpu, broadening the runtime envelope beyond Phase 8.3's MPS-only matmul. First concrete kernel: rope (rank-2 f32, x.shape == theta.shape). Flash-attention follows in 8.4.1 — the infrastructure here is what makes that incremental. MLIR - ODS: tessera_apple.gpu.msl_kernel carrying entry_point + msl_source + cache_key as StringAttrs (the IR is the self-contained, replayable record of the kernel — runtime caches by (msl_source, entry_point)) - Pass: RopeToAppleGPU lowers rank-2 f32 tessera.rope to a func.call into tessera_apple_gpu_rope_f32 (mirrors MatmulToAppleGPU.cpp structurally) - Pipeline tessera-lower-to-apple_gpu-runtime now composes matmul + rope patterns in one go - target_ir verifier accepts the new op; runtime-mode lowering picks msl_kernel for rope sources and de-duplicates tile.rope/tile.rotary_pair decompositions to a single emission Runtime - apple_gpu_runtime.mm: kernel cache (std::unordered_map keyed by (msl_source, entry_point)) + compile_msl_kernel helper + tessera_apple_gpu_rope_f32 C symbol with embedded rope MSL source. Encoded via MTLComputeCommandEncoder. New capability probe tessera_apple_gpu_runtime_msl_cache_size for cache-hit assertions. - apple_gpu_runtime_stub.cpp: portable rope reference + stub cache size symbol so non-Darwin builds expose the same C ABI Python - driver.py: split runtime envelope into _APPLE_GPU_MPS_OPS + _APPLE_GPU_MSL_OPS; backend artifact picks the right symbol/framework per op - target_ir.py: broadened _apple_gpu_module_is_mps_runtime to "all compute ops share one source in the envelope"; added duplicate suppression for runtime mode; added _APPLE_GPU_ROPE_MSL_SOURCE constant + sha256-derived cache_key - runtime.py: _apple_gpu_dispatch_rope, _apple_gpu_rope_f32 ctypes wrapper; loader now requires both matmul + rope symbols (forces rebuild after Phase 8.4) Tests - New lit fixture apple_gpu_msl.mlir round-trips the new msl_kernel op through tessera-opt - Three new unit tests in test_apple_backend_roadmap.py: MSL artifact contract (IR carries kernel source), end-to-end rope numerical correctness vs numpy, MSL pipeline cache reuse (cache size 0->1 on first call, stays 1 on second) - Pass-level lit lowering test for tessera.rope is intentionally out-of-scope (tessera.rope is not a registered Tessera dialect op); Python unit tests cover that path end-to-end Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 1950 unit tests passing; 8/8 Phase 8 lit fixtures passing against the in-tree tessera-opt. End-to-end rope rtol=1e-5 vs numpy. 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: 94bc23c917
ℹ️ 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".
gstoner
added a commit
that referenced
this pull request
May 7, 2026
Merge pull request #3 from gstoner/claude/fervent-varahamihira-cd1ab3
gstoner
added a commit
that referenced
this pull request
May 18, 2026
Summary — all four issues fixed P1 #1 — Kernel widening + sweep envelope (FIXED) The MSL kernel had a per-thread float y_local[256] register buffer that hard-capped D at 256. With fixed gradient the T-step recurrence is exactly closed-form y_T = y0 − T·η·grad, so I rewrote the kernel to stream through D once accumulating squared-norm energy, then a second pass for the winner. No register-vector spill; D is now unbounded; K is still bounded at 256 by the threadgroup-size budget for the argmin reduction. Both Python and runtime guards updated. P1 #2 — Proof of native dispatch (FIXED) Added tessera.ebm.ebt_tiny_dispatched_on_gpu() / ebt_tiny_last_route() probes. The workload now returns a 5-tuple including dispatched_on_gpu, and run_workload_apple_gpu: If dispatched_on_gpu=False (silent fallback): row is degraded to backend="python_ref", mode="reference_chain_fallback", ok=False. Caller cannot mistake it for a native win. If True: row is backend="apple_gpu", mode="fused_chain", includes dispatched_on_gpu=True field. Sweep summary gains status per shape (native_dispatched vs degraded_fallback) and degraded_count in the envelope. speedup only computed when the native attempt actually fired. P2 #3 — Bridge breadth (FIXED, with honest narrowing) Migrated 10 fast paths through jit_bridge.dispatch_via_manifest — every op the workloads actually use: GA: reverse, grade_involution, conjugate, hodge_star, exp_mv, log_mv, geometric_product, wedge, left_contraction, rotor_sandwich, norm, inner (12) EBM: inner_step, ebt_tiny (2) Docs now explicitly say 14 of 26 fast paths route through the bridge — the rest still call _apple_gpu_dispatch.bind_symbol directly and are correctness-equivalent but invisible to the trace. Migrating the long tail is an open follow-up; the breadth is no longer overstated. P2 #4 — Stale docs (FIXED) benchmark_ga_ebm.py:26 header now lists the 8 native EBM kernels + ebt_tiny correctly (was "EBM ops do not yet ship Apple GPU MSL kernels"). README.md workload section: workload now described as fused single-dispatch ebm.ebt_tiny, with the new sweep numbers (~55× peak, not the old 116×). docs/status/ga_ebm_milestone.md: non-claim #1 rewritten — was "first_native_win_shape is None" and "ebm_refinement now wins"; now coherently describes the streaming closed-form rewrite + the proof bit + headline numbers will drift. docs/audit/apple_ga_ebm_native_execution_gap.md:232: first_native_win_shape=None replaced with the current numbers + dispatch proof description; "only ebm_inner_step routes through public API dispatch" replaced with the 14-of-26 breakdown. Roadmap top status line refreshed. Sweep run on M-series with fresh sample artifact: every sweep point dispatches on GPU (degraded_count=0); first native win at B=16,K=32,D=128,T=8 (~1.1×); peak ~55× at B=64,K=128,D=1024,T=256. Final: 3607 / 1 skipped / 0 failed.
gstoner
added a commit
that referenced
this pull request
May 18, 2026
Summary P2 — Stale "14 of 26" bridge coverage claims (4 spots fixed) docs/status/ga_ebm_milestone.md:77 — "Known non-claims" #3 — replaced "JIT-bridge coverage is partial (14 of 26 fast paths)... the remaining 12 fast paths still call the shared loader directly" with the now-correct statement: all 26 fast paths route through the bridge and the benchmark uses the trace as proof-of-dispatch. docs/audit/ga_ebm_roadmap.md:2 — top status banner — "wires 14 of 26 fast paths" replaced with "wires all 26 fast paths (17 GA + 9 native EBM), so every public-API GPU dispatch flows Python → manifest_for(op) → shared loader → result and produces a JitBridgeRoute row". docs/audit/apple_ga_ebm_native_execution_gap.md:250 — "What is missing" bullet — replaced the 12-fast-paths-still-direct claim with "All 26 fast paths route through the jit_bridge... bivector_langevin reuses the langevin_step kernel but tags itself as ebm_bivector_langevin in the trace via the helper's bridge_op_name kwarg." benchmarks/apple_gpu/README.md:30 — bullet was still listing "10 ops (bridge-migrated set)" + the legacy hand-written GA/EBM splits — replaced with "All 26 fast paths (17 GA + 9 native EBM) route through dispatch_via_manifest..." (matches the milestone TLDR now). P3 — Narrow proof-bit explanation (1 spot fixed) docs/status/ga_ebm_milestone.md:54 — non-claim #1 — previously said "Every native row in the report carries a dispatched_on_gpu proof bit sourced from tessera.ebm.ebt_tiny_dispatched_on_gpu()". Rewrote as a two-bullet split: Native EBM primitive rows + JIT-bridge benchmark rows → proof bit from the jit_bridge route trace (one-shot probe span) EBT-tiny workload + --ebt-sweep rows → proof bit from ebm.ebt_tiny_dispatched_on_gpu() + per-shape status in the sweep summary This matches the benchmark module header which already had the correct breakdown. Verification: grep -rn "14 of 26 | remaining 12 fast paths | JIT-bridge coverage is partial | bridge-migrated set" across docs/, benchmarks/, and python/ returns no stale claims (the only 14 of 26-shaped hit is in primitive_coverage_state.md, which is about a different metric — backend_kernel 227 partial / 147 planned). 138 benchmark + bridge tests pass.
gstoner
added a commit
that referenced
this pull request
May 18, 2026
Summary
Task 1 — Compiler-integrated GA vertical slice ✅
New module python/tessera/compiler/clifford_jit.py:
@clifford_jit(target="apple_gpu") decorator
First call traces the function under jit_bridge.jit_context("apple_gpu") with tracing on
Captures the bridge route trace as the canonical op plan
Verifies every op has an apple_gpu=fused manifest entry (raises CliffordJitError at compile-time otherwise)
Builds a frozen CliffordCompiledArtifact carrying the plan + plan hash + Apple target metadata
Subsequent calls execute under a fresh jit_context so the per-call route trace can be checked against the plan via plan_matches_routes()
Demo: point_cloud_rotor_invariant(rotor, points) = ga.norm(ga.rotor_sandwich(rotor, points)) → plan (clifford_rotor_sandwich, clifford_norm), plan_hash 9bb1cd4d73ae1bc5. Benchmark emits a namespace="vertical_slice" row with the full compiled_artifact JSON embedded.
11 / 11 tests in tests/unit/test_clifford_jit.py — trace + plan capture, manifest verification, hash determinism, per-call route trace match, JSON round-trip, error semantics, frozen-dataclass.
Task 2 — Fused GA + EBM workload ✅
rotor_conditioned_ebt — ga.exp_mv → ga.rotor_sandwich → ebm.ebt_tiny through public APIs. Routes all three ops through the bridge with a dispatched_on_gpu proof bit. Native ~21× speedup vs the equivalent numpy chain (~0.8 ms vs ~16 ms on M-series). Tests verify all 3 symbols appear in row["symbols"] and the proof bit fires.
Task 3 (partial) — Metal buffer pool ✅
New MetalDeviceContext::buffer_pool — 19-bucket size-class pool (16B → 4MB) with metal_buffer_acquire / metal_buffer_release / metal_buffer_acquire_with_bytes helpers. Wired into dispatch_clifford_unary_8x8_f32_msl + dispatch_ebm_ebt_tiny_refinement_argmin_f32_msl (the two heaviest dispatchers in the workloads). Remaining ~25 dispatchers still call newBufferWithBytes directly — migration is mechanical and explicitly listed as the #1 next target.
Task 4 — On-device RNG + energy lowering (next sprint)
Sized as #2 + #3 in the Next Targets list — these are a sprint of their own (Philox-in-MSL kernel + AST → MSL visitor for restricted energy_fn shapes) so I scoped them out of this turn and documented the path.
Task 5 — Docs honest + pruned ✅
docs/audit/apple_ga_ebm_native_execution_gap.md — front-matter changed to status: Historical (superseded by docs/status/ga_ebm_milestone.md); a banner block at the top redirects readers to the canonical milestone page. The roadmap remains the canonical sprint sequence.
docs/status/ga_ebm_milestone.md gains rows for Compiler vertical slice, Fused GA + EBM workload, and Metal buffer pool in the TL;DR table; the "Next targets" section reshuffled with finish-the-buffer-pool-sweep at #1, Philox-in-MSL at #2, energy lowering at #3, broaden @clifford_jit at #4.
Sample artifact regenerated
benchmarks/apple_gpu/sample_ga_ebm_report.json — now includes the vertical_slice row (with embedded compiled_artifact JSON), the rotor_conditioned_ebt workload row, and the updated pool-accelerated workload timings.
Sweep: 3623 passed, 1 skipped, 0 failures.
gstoner
added a commit
that referenced
this pull request
May 23, 2026
… Triage of all 7 findings While that runs, let me give you the triage summary so you can react. Triage of all 7 findings # Finding My verdict Action taken 1 mypy ratchet 7 errors STALE / FALSE None — mypy python/tessera/ reports 0 errors across 185 files; test_static_analysis_baseline.py 3/3 passes. The reviewer was looking at an older state 2 MockRankGroup repeated-collective race REAL ✅ Fixed — added per-rank generation counters; (kind, gen) tuple keys make every invocation unique; deposit and withdraw can no longer race 3 bind_symbol ctypes signature mutation race REAL ✅ Fixed — replaced getattr(handle, symbol) + .argtypes mutation with ctypes.CFUNCTYPE(restype, *argtypes)((symbol, handle)) so each signature gets a genuinely independent prototype; the whole bind now runs under _lock with a double-check after acquire 4 Graph IR cache no LRU/lock REAL but explicitly intentional Deferred — graph_ir_cache.py:18 documents the no-LRU design; turning it into a bounded LRU is a feature decision, not a bug fix 5 Monolith files REAL but architectural Deferred — apple_gpu_runtime.mm 6,988 LOC, vjp.py 4,262, etc. Splitting requires its own sprint 6 Queue verifier no-op stubs + shape_inference XFAIL REAL Deferred — confirmed QueueVerifiers.cpp is 3 trivial return success(); stubs and tests/tessera-ir/phase6/shape_inference.mlir is XFAIL: *. Next verifier-quality pass should cover them 7 trust_treedef=True default contradicts docstring REAL ✅ Fixed — flipped default to sentinel None and emit DeprecationWarning on implicit use; explicit True/False stays silent. Sets the trajectory to flip the implicit default to False in a future release without breaking current callers today Why I didn't fix #4, #5, #6 in this session #4 is documented intent (process-local cache for hot loops). Adding LRU + locks would change observable behavior; needs a design decision, not a drive-by fix. #5 is real but every file is structurally coherent — slicing them apart without a target architecture for each (Apple runtime by kernel family, vjp.py by op family, etc.) risks creating worse seams than the current single-file layout. #6 queue verifier stubs and shape_inference XFAIL are pre-existing surface area, not new debt. Worth a focused sprint along with the verifier-stub audit the reviewer flagged ("187 return success(); occurrences across src/"). Once the full unit sweep finishes I'll confirm green. Ran Full sweep result 5254 passed, 0 failures — same count as before the fixes, no regression. The 3 actionable findings (#2 MockRankGroup race, #3 bind_symbol mutation, #7 trust_treedef default) are landed and verified. Final verification mypy: 0 errors / 185 source files Full unit sweep: 5,254 passed, 7 skipped, 0 failures Static-analysis baseline: 3/3 passing Apple GPU MLA E2E proof: 6/6 passing MLIR verifier sprint guards: 64/64 passing Apple GPU buffer pool + JIT no-leak: 14/14 passing
This was referenced May 30, 2026
gstoner
added a commit
that referenced
this pull request
Jun 2, 2026
#1 — MTL4 fp16 matmul routing (the headline) ✅ honest flip Measured fp16 MTL4 matmul2d vs MPS end-to-end on this Mac (macOS 26.5, Metal 4 live). Finding: MPS's fp16 GEMM is well-tuned and wins on square shapes (0.6–0.8×) and most decode shapes — except the M==1 GEMV decode step, where MTL4 wins a rock-solid 3.2–3.4× (3 trials, tight variance; MPS has a slow fp16 M==1 path). Added _mtl4_route_matmul2d_f16 size-gated to M==1 by default (TESSERA_APPLE_GPU_MTL4_F16: auto=M==1 only / all=force / 0=off). A blanket flip would regress square matmuls, so I deliberately didn't take it — the honest win is the narrow one. Re-runnable benchmark benchmarks/apple_gpu/benchmark_mtl4_matmul_routing.py + 6 gate/numerics tests + P7 doc entry with the measured table. #2 — TPP embedded-MLIR driver ✅ subprocess killed New libtessera_tpp_capi shared lib (TPPCApi.cpp) runs the tpp-space-time pipeline in-process via PassManager; tpp.solve(mlir) loads it via ctypes (mirrors the Apple-runtime pattern) — no tessera-opt subprocess, no PATH dependency. Built clean, verified in-process, python_driver_wired now reports True. 10 tests. #3 — Optimizer batching rules ✅ Proved vmap composes with all functional optimizers (elementwise: vmap==whole-batch; lamb/muon: vmap==per-row loop) and flipped optimizer/functional_optimizer_step batching_rule partial→complete. 9 tests. #4 — Lit counts + G4 spot-check ✅ CLAUDE.md refreshed to the live 61 PASS / 19 UNSUPPORTED / 18 XFAIL / 0 FAIL across 98. Confirmed G4's matrix-as-dispatcher invariant holds — launch() resolves the executor via _exec_row_for_metadata with an explicit "no new branch in launch()" contract (the prior codex P1 was closed by G6/G6.1). Plus caught + fixed two registry-drift dashboards from #3 (s_series_status, regenerated) and added s_series_status to the #12 pre-commit drift gate so that class can't recur. ruff clean, mypy ratchet 0.
gstoner
added a commit
that referenced
this pull request
Jun 3, 2026
… where the three items landed. All green — drift gate passes, 71 guard tests pass, everything consistent. Here's where the three items landed. ✅ #4 — Coverage prose stops copying stale counts (complete) COVERAGE_AUDIT.md, MASTER_AUDIT.md, and BACKEND_AUDIT.md now link the generated dashboards instead of copying numbers (the entry-point had already drifted — it claimed "218 C ABI symbols" vs the live 226+). Added a banner stating the rule. Drift gate owns every count. ✅ #3 — Manifold Langevin ops off planned (complete, honestly scoped) Found a real name mismatch (_EBM_PRIMITIVES listed ebm_sphere_langevin but the registry/runtime use ..._step). Fixed it and added manifest entries: ebm_sphere_langevin_step → partial, fused (dedicated MSL kernel) ebm_bivector_langevin_step → partial, fused (reuses the affine ebm_langevin_step kernel on grade-2 coeffs) ebm_sphere/bivector_langevin_sample → partial, reference (real numpy chain wrappers) 4 of 5 flipped honestly (the 5th, rng_langevin_sample, is a different namespace). Backed by 2 guard tests + 38 existing geo-sampling tests. The universal complete gate stays intact (nvidia/rocm still planned); their distributed-mesh axis stays correctly Phase-G-gated. 🟡 #2 — Long-tail axes (batching done; transpose/sharding remain) Batching closed 27 → 4 by promoting the textbook-batchable categories (collective, recurrent, state_space, linalg decomposition+solver, sparse, segment_reduce) with documented justification — leaving only the genuinely mesh-aware ones (moe / moe_transport / kv-cache state) honestly partial. Hit and fixed a real duplicate-key shadowing bug in the category table along the way. Locked by test_batching_rule_closure.py. Transpose (40) and sharding (30) remain — I stopped here rather than rush them, because they need careful per-category judgment, not flag-flipping: transpose: closes mostly via not_applicable for non-linear families (optimizers, ebm, recurrent, linalg-decomp, moe) + complete for the genuinely-linear ones (spmm, segment_reduce, avg_pool). sharding: largely Phase-G-mesh-pending (distributed attention/spectral/moe) — the axis where staying partial is most honest; only trivially-replicated/not_applicable ones close cleanly. Verification: drift gate green; 130 EBM-benchmark + 71 guard/manifest/roadmap tests pass; dashboards (s_series_status, apple_target_map, standalone_primitive_coverage) regenerated and consistent; lint clean.
gstoner
added a commit
that referenced
this pull request
Jun 3, 2026
#4 — Coverage prose stops copying stale counts Audit docs link the drift-gated dashboards instead of copying numbers. ✅ #3 — Manifold Langevin ops off planned 4 ops flipped planned → partial with real Apple GPU kernels (fixed a real _EBM_PRIMITIVES name mismatch along the way). Distributed-mesh axis correctly stays Phase-G-gated. ✅ #2 — Long-tail axes closed to their honest limit batching_rule: 27 → 4 — textbook-batchable families complete; mesh-aware ones (moe/transport/kv-cache) stay partial. transpose_rule: 40 → 0 — closed on the linear-vs-nonlinear principle: linear maps (sparse spmm/sddmm/bsmm, moe_transport gather/scatter adjoints, segment_reduce, tri_solve, avg_pool) → complete; nonlinear families (optimizers, recurrent cells, linalg decomposition, ebm energy/sampling, moe routing, max/min/adaptive pool) → not_applicable (their backward is the registered VJP, not a linear transpose). avg_pool handled per-name since pooling is a mixed category. sharding_rule: 30 — left honestly partial. Every open entry is genuinely mesh-aware (distributed FFT transpose, seq/head-partitioned attention, all-to-all MoE, block-cyclic linalg, sequence-parallel SSM). There's no trivially-replicated subset to close — this is the same Phase-G-mesh gate as the universal backend_kernel axis, not closeable without real multi-GPU hardware. Verification: transpose 40→0, 432 complete; 5 drift gates green; 284 coverage/manifest/roadmap/memory/reasoning tests pass + the new test_batching_rule_closure.py transpose guards (zero-open + linear-complete + nonlinear-N/A, both directions locked).
gstoner
added a commit
that referenced
this pull request
Jun 6, 2026
…ph kernels) #2 bf16 in the whole-graph lane: author_graph gains an io_bf16 boundary flag (bf16 placeholders -> cast f32 -> f32 body -> cast bf16). The mlpkg reflection path (MTLTensorDataTypeFromMPSDataType) hard-asserts on bf16 bindings today, so GraphFn.run_mlpkg(elem="bf16") authors an f32 package and converts at the Python boundary (bf16 in/out, f32 internal compute = ABI f32-accumulate). io_bf16 C path retained for when bf16 bindings become reflectable. #3 native bf16 MPSGraph kernels: tessera_apple_gpu_{mpsgraph_unary, mpsgraph_binary,rmsnorm_gpu,layer_norm}_bf16 — native via the dtype-parameterized mpsg_run_* helpers (bf16 supported per mpsgraph_bf16_supported()), with an upcast->f32-extern->round host fallback; stub parity. The back-half's rmsnorm/layer_norm/silu+unary/elementwise bf16 paths now run native bf16 instead of host-upcast. _SENTINEL_SYMBOL bumped; runtime ABI docs regenerated. bf16 is now uniform across the GPU back-half AND both GraphFn engines (run + run_mlpkg). 256/256 production-lane + 1166 apple_gpu tests green; mypy clean on host and --platform linux. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 8, 2026
Locks the decision from the Thrust-#3 cross-platform discussion: the Mamba-2 SSD fusion is matmul-dominant by construction, so it belongs at Tile IR as a tiled GEMM schedule with the matmul intrinsic selected per backend (simdgroup_matrix / WGMMA / MFMA) — NOT a one-off Apple Metal kernel. * Records why SSD = matmul (the gram C@B^T is shared across channels; a naive per-(b,d) Apple fused kernel recomputes it D times and is slower than the current 3-bmm path — an explicit anti-pattern). * The cross-platform tiled schedule (chunk tiling, gram-tile in fast memory, masked-matmul + register decay, chunk-state carry) is identical across Apple/ NVIDIA/AMD; only the fast-memory keyword, matmul intrinsic, and barrier differ. * Apple is the executable validation backend (only one that runs today); the NVIDIA/AMD lowerings inherit the schedule and slot in WGMMA/MFMA when silicon lights up. Sequenced to the P0 backend-kernel hardware-proof timeline. * Current Apple selective_ssm (chunked-parallel, 3 MPS-bmm + host) stays as the functional reference until the tiled schedule lands. docs/architecture/proposals/tiled_ssd_tile_ir_schedule.md + a deferred-work pointer in ROADMAP_AUDIT.md. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 8, 2026
The canonical rotor-invariant norm(rotor_sandwich(R, x)) previously dispatched
two GA ops with an intermediate multivector round-trip through global memory.
This fuses it into a single kernel.
Kernel — tessera_apple_gpu_clifford_rotor_sandwich_norm_cl30_f32 in
apple_gpu_runtime.mm: reuses the Cl(3,0) double-geometric-product expansion,
keeps the 8-vector sandwich result in registers, and writes only the scalar
norm per batch element (one dispatch, no intermediate spill). C++ reference
fallback for non-Metal hosts.
Surface — ga.rotor_sandwich_norm(R, x) is the direct entry (numpy fallback =
norm(rotor_sandwich(...))); routes Cl(3,0) f32 through the bridge manifest.
Fusion pass — _fuse_rotor_sandwich_norm rewrites the rotor_sandwich→norm
adjacency into one clifford_rotor_sandwich_norm op, but ONLY when the
intermediate sandwich is consumed exactly once and isn't the program return.
Applied by the @clifford_jit decorator (and the lazy compile path); the
structural lower_function_to_ir stays a faithful unfused 1:1 AST→IR projection.
So @clifford_jit(norm(rotor_sandwich(...))) now compiles to a single-op plan +
single dispatch.
Wiring — _CLIFFORD_APPLE_GPU_FUSED + _GA_ATTR_TO_OP_NAME + a new
_CLIFFORD_FUSION_OPS set (kept OUT of the 17 _CLIFFORD_PRIMITIVES so the
"seventeen primitives" audits/counts stay exact); clifford_manifest_for now
reports fusion ops (apple_gpu fused fp32 + CPU reference). _SENTINEL_SYMBOL
bumped.
Tests: test_clifford_rotor_sandwich_norm_fusion.py (8 — numerics, fusion pass,
non-fusable cases, manifest, @clifford_jit single-op plan + single route).
Updated the canonical-chain plan/route assertions across test_clifford_jit,
test_compile_report{,_auto_emission}, test_benchmark_{ga_ebm,row},
test_compiler_audit, test_ga_backend_manifest to reflect the fused plan.
mypy host+linux + ruff clean; generated docs in sync.
Closes GA/EBM close-out gaps #1/#2/#5/#6; remaining: Apple-CPU GA/EBM native
kernels (#3), Cl(1,3) kernels (#4, gated with a diagnostic), exp/log GA autodiff.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 10, 2026
The Apple GPU kernel C-ABI symbols are `void` and signal an internal failure (timeout / device lost / command-buffer error) by leaving the output buffer untouched — so Python read garbage as success and every numerical test still passed. Rather than break ~70 symbol signatures with an int return (a massive ABI change), generalize the existing g_mlpkg_last_error_kind errno pattern: - .mm: a thread-local last-error set at the shared command-buffer choke point commit_and_wait_with_timeout (~72 callers) on its two failure branches (timeout/hang = kind 1, cb.error = kind 2). New C ABI: tessera_apple_gpu_last_error_kind / _message / _clear_last_error, with non-Darwin stub parity (kind always 0). - runtime.py: _apple_gpu_arm_gpu_error / _apple_gpu_consume_gpu_error wrap the matmul lane's GEMM call (reference consumer) — a silent kernel failure now funnels through _note_dispatch_fallback (strict mode raises) and recomputes on host instead of returning the untouched buffer. No-channel builds (older dylib / stub) are a safe no-op. Validated on a Metal host: real matmul succeeds with no false-positive; new symbols present in the test dylib AND the canonical libTesseraAppleRuntime.a; new strict-dispatch tests cover the simulated-error funnel+raise and the no-channel no-op. Regenerated runtime_abi (+3 symbols) and test_coverage. Remaining follow-on: unary/binary/rowop/bmm adopting the same arm/consume. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 13, 2026
…odel (#4) #4 — tessera.dflash_reference.ReferenceDecoderLM: a small numpy causal decoder (pre-norm MHA + SwiGLU, rope, tied/untied LM head) with a multi-layer hidden tap (the DFlash conditioning signal) and a stateless forward() that is the greedy-AR ground truth. random_decoder_lm builds one with small random weights. #3 — stateful KV cache + rollback: step(tokens) does causal cached decoding and appends roped-K/V per layer; rollback(n) drops the over-speculated tail. Verified that incremental step() (in 3 chunks) reproduces the stateless full-sequence forward to 1e-3, and that rollback restores exact cache state. dflash_generate_cached ties it together: cached draft (#1) + stateful target with rollback (#3) + greedy or rejection sampling (#2). Verified the whole efficient loop reproduces greedy AR exactly, is block-size independent, and sampling is reproducible + in-vocab. tests/unit/test_dflash_reference_target.py (5). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 13, 2026
…odel (#4) #4 — tessera.dflash_reference.ReferenceDecoderLM: a small numpy causal decoder (pre-norm MHA + SwiGLU, rope, tied/untied LM head) with a multi-layer hidden tap (the DFlash conditioning signal) and a stateless forward() that is the greedy-AR ground truth. random_decoder_lm builds one with small random weights. #3 — stateful KV cache + rollback: step(tokens) does causal cached decoding and appends roped-K/V per layer; rollback(n) drops the over-speculated tail. Verified that incremental step() (in 3 chunks) reproduces the stateless full-sequence forward to 1e-3, and that rollback restores exact cache state. dflash_generate_cached ties it together: cached draft (#1) + stateful target with rollback (#3) + greedy or rejection sampling (#2). Verified the whole efficient loop reproduces greedy AR exactly, is block-size independent, and sampling is reproducible + in-vocab. tests/unit/test_dflash_reference_target.py (5). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 13, 2026
* feat(dflash): cached drafting (#1) + sampling & rejection acceptance (#2) #1 — Draft KV cache. block_diffusion_attention gains return_ctx_kv to expose this step's projected+roped context KV; DraftKVCache accumulates it per layer. dflash_decoder_layer_cached / dflash_draft_forward_cached thread the cache so the draft attends to the full accumulated context instead of recomputing it. Verified: cached(accumulated) == non-cached(full context) to 1e-3, and the cache accumulates/advances correctly across steps. #2 — Non-greedy sampling + distribution-preserving acceptance. make_sampler (temperature / top-k / top-p, rng-reproducible), sampler_probs (matching truncated distribution), and dflash_speculative_verify (Leviathan rule: accept d_i w.p. min(1, p_t/p_d); on reject draw from normalize(relu(p_t-p_d)); bonus from the target's next-position distribution). Verified: greedy == argmax, top-k restricts support, draft==target accepts all, and the speculative-sampling theorem — the emitted token's marginal equals the target distribution (Monte Carlo, 40k draws, max abs err < 0.02). tests/unit/test_dflash_cached_sampling.py (7). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(dflash): stateful target with rollback (#3) + reference target model (#4) #4 — tessera.dflash_reference.ReferenceDecoderLM: a small numpy causal decoder (pre-norm MHA + SwiGLU, rope, tied/untied LM head) with a multi-layer hidden tap (the DFlash conditioning signal) and a stateless forward() that is the greedy-AR ground truth. random_decoder_lm builds one with small random weights. #3 — stateful KV cache + rollback: step(tokens) does causal cached decoding and appends roped-K/V per layer; rollback(n) drops the over-speculated tail. Verified that incremental step() (in 3 chunks) reproduces the stateless full-sequence forward to 1e-3, and that rollback restores exact cache state. dflash_generate_cached ties it together: cached draft (#1) + stateful target with rollback (#3) + greedy or rejection sampling (#2). Verified the whole efficient loop reproduces greedy AR exactly, is block-size independent, and sampling is reproducible + in-vocab. tests/unit/test_dflash_reference_target.py (5). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(dflash): GPU draft attention (#5) + training loss (#9a) + checkpoint I/O (#7) #5 — attention_fn threaded through dflash_decoder_layer / dflash_draft_forward (+ cached variants) so the whole draft forward runs its attention on the Apple GPU metal_runtime lane via apple_gpu_attention_fn. Verified the whole draft (2 layers) matches the numpy reference on Metal (rtol/atol 1e-3). The matmul- heavy projections/MLP/LM-head stay host-side (GPU gather/embedding is the remaining blocker for a single fully-jitted artifact). #9a — position-weighted block training loss: dflash_position_weights (wₖ = exp(-k/γ), normalized), dflash_block_loss (mean/sum/none) and the explicit gradient dflash_block_loss_grad. Verified the gradient vs finite differences (<1e-7), that a grad step lowers the loss, and reduction consistency. #7 — checkpoint I/O (tessera.dflash_io): a dependency-free safetensors reader/writer + HF state-dict <-> DFlashWeights mapping (transposing the nn.Linear (out,in) weights to the x@W (in,out) convention; embedding/LM head supplied from the target). load_dflash_weights reads a z-lab/*-DFlash draft; verified safetensors round-trip, the (out,in) transpose, and that round-tripped weights produce identical draft logits. tests: test_dflash_train_io.py (7) + #5 GPU draft case. ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * feat(dflash): nn.Module (#6) + rotating cache (#9b) + tokenizer/scheduler (#9c/#9d) #6 — DFlashDraft(nn.Module): holds every draft tensor as a Parameter (so it participates in parameters()/state_dict/.to(dtype)), forwards through the functional draft (cached or not), from_weights()/to_weights() round-trip. Verified module forward == functional (<1e-5), 5 + 11*N params registered, weight round-trip. #9b — RotatingDraftKVCache: bounds the draft context cache to the last max_size tokens (the draft analogue of MLX RotatingKVCache for sliding layers). Verified it caps per-layer length and, when unbounded, is identical to DraftKVCache. #9c/#9d — tessera.dflash_serve: dflash_generate_text (string-in/out via any encode/decode tokenizer) and DFlashScheduler (holds draft + stateful target, serves generation requests, greedy == AR). Verified scheduler greedy == AR and generate_text round-trips through a tokenizer. tests/unit/test_dflash_module_serve.py (5). ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(dflash): GQA repeat note (#8) + MASTER_AUDIT integration 1–9 landed #8 — annotate block_diffusion_attention's GQA: repeat is numerically exact; the native flash_attn_gqa kernel doesn't support DFlash's concat-context+proposal KV with an additive bias, so the reference materializes the repeat (no code change — correctness is unaffected). MASTER_AUDIT records DFlash integration items 1–9 as landed, with the two remaining gates flagged as external (real-checkpoint numerical parity needs a network download; a single fully-jitted GPU draft artifact needs GPU gather). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs: record DFlash + attn_bias in README, add docs/dflash.md - README status table: new "Speculative decoding — attn_bias substrate + DFlash block-diffusion draft" row (honest status: Python reference + attention core on Apple GPU metal_runtime; greedy spec-decode == greedy AR proven vs the MLX reference; real-checkpoint parity + fully-jitted GPU draft are external gates). - README: refresh stale Apple C ABI counts to the generated truth (256→264 symbols, 109→112 kernel families). - New docs/dflash.md: user-facing overview — the attn_bias substrate, the module map (dflash / dflash_reference / dflash_io / dflash_serve), a quick start, what's proven (per-test), and the external gates. Linked from the README doc index. docs lint passes; all links resolve. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(api): fold DFlash + attn_bias public API into PYTHON_API_SPEC + CANONICAL_API PYTHON_API_SPEC.md: - flash_attn signature + parameter table gain attn_bias (additive (B,Sq,Sk) mask, Apple GPU flash_attn_bias_* / metal_runtime, causal+bias, broadcast fallback, positional-bias VJP). - Module hierarchy lists tessera.dflash / dflash_reference / dflash_io / dflash_serve. - New §18 "Speculative Decoding (DFlash)" documents the full public surface across the four modules + nn.functional.block_diffusion_attention / mask_token_block; TOC + Appendix A symbol index updated. CANONICAL_API.md: - flash_attn ops row gains attn_bias; functional table gains block_diffusion_attention + mask_token_block; new "tessera.dflash — Speculative Decoding (DFlash)" section with canonical names (one per concept) + a quick-start. Verified: check_spec_sync + docs lint pass; every documented symbol exists and every module __all__ symbol is documented (zero drift, confirmed programmatically). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * docs(api): close tessera.nn / tessera.ops doc-coverage gaps (full sweep) A programmatic sweep of the public surface vs both API docs found the surface ~99% documented with concentrated gaps; this closes them to zero. tessera.ops (312 ops): added the 2 missing — bmm (batched matmul + broadcast, Apple GPU metal_runtime) and fake_quantize (QAT STE) — to both ops tables. tessera.nn (77 public attrs): added the 10 missing functional layers (linear_general, lora_linear, spectral_norm, conv_transpose, avg/max/min/adaptive pool, gru_cell, simple_rnn_cell, bidirectional_scan) to CANONICAL's functional table, and the 10 missing Module classes (LinearGeneral, Einsum, LoRALinear, ConvTranspose1d/ConvTranspose, SpectralNorm, GRUCell/SimpleRNNCell, NativeSparseAttention, MixtureOfRecursions) to the stateful class table, with accurate constructor/forward signatures. Verified programmatically: tessera.ops, tessera.nn, and nn.functional.__all__ now have ZERO undocumented public symbols. check_spec_sync + docs lint pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 14, 2026
…dle (#3) AMA-Bench's arbitrary-horizon finding is the forcing function: memory must stay on-device across decode steps, not be re-uploaded each step. ResidentBank keeps the bank's keys resident (one upload) and scores each query against them via the encode-session bmm_enc lane — per-read traffic is O(query), not O(bank). Hardware-verified (tests/unit/test_resident_bank.py): resident reads == reference reads (metamorphic), and the bank uploads once vs recompute's re-upload-per-read — ~28× upload reduction on Metal for a 256-row bank. Falls back to a numpy reference when the encode-session runtime is absent (portable). This lands the READ side of resident_state_handle (the expensive scan-against-a- large-bank path). It moves from MEMORY_PRIMITIVE_GAPS to PARTIAL_MEMORY_PRIMITIVES; the remaining piece — incremental on-device append (offset-write, no full re-upload) — stays in GAPS as kv_cache_append_read. long_memory_core's run_core gains a resident_bank_read_residency proof row + the append gap row. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
pushed a commit
that referenced
this pull request
Jun 24, 2026
…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>
gstoner
added a commit
that referenced
this pull request
Jun 24, 2026
…lue + int/FA perf ladders (#90) * feat(rocm): compiler-generated flash_attn (FA-2 forward) — second compiled 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> * feat(rocm): front-end glue — IR stack emits the wmma_gemm directive 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> * perf(rocm): int8/int4 compiled-GEMM dtype sweep — measured (packed int4 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> * perf(rocm): compiler-generated flash_attn forward perf ladder (measured) 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> * feat(rocm): compiler-generated flash_attn BACKWARD (3 WMMA kernels) 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> --------- Co-authored-by: gstoner <angstroms01@gmail.com> Co-authored-by: Claude Opus 4.8 <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.
Summary
Two phases bundled in one PR — they share the same Apple GPU runtime contract (
MetalDeviceContext,tessera-lower-to-apple_gpu-runtimepipeline,metal_runtimeexecution mode) and split cleanly only at the kernel boundary.@jit(target=\"apple_gpu\")executes throughMPSMatrixMultiplicationMulti-op programs (
tiny_decode,simple_transformer, MoE) intentionally stay on the existingmetal_artifactcontract until Phase 8.4.1 (flash-attention) broadens the runtime envelope further.What changed
MLIR (
src/compiler/codegen/Tessera_Apple_Backend/)Phase 8.3:
tessera_apple.gpu.mps_matmul / mps_softmax / mps_dispatchMatmulToAppleGPU— rank-2 f32tessera.matmul→func.call @tessera_apple_gpu_mps_matmul_f32tessera-lower-to-apple_gpu-runtimePhase 8.4.0:
tessera_apple.gpu.msl_kernelcarryingentry_point+msl_source+cache_keyas StringAttrs (the IR is the self-contained, replayable kernel record)RopeToAppleGPU— rank-2 f32tessera.rope→func.call @tessera_apple_gpu_rope_f32target_ir.pyverifier accepts the new ops; runtime-mode lowering de-duplicates rope'stile.rope+tile.rotary_pairdecomposition to a single emissionRuntime (
runtime/apple_gpu_runtime.mm+ stub)Phase 8.3:
MetalDeviceContextsingleton (MTLDevice+MTLCommandQueue) wrappingMPSMatrixMultiplicationwith shared-storageMTLBuffersMTLCreateSystemDefaultDevicereturns nilapple_gpu_runtime_stub.cpp— portable C++ TU on non-Darwin so the static lib + Python ctypes layer stay platform-agnosticPhase 8.4.0:
kernel_cache: std::unordered_map<std::string, id<MTLComputePipelineState>>keyed by(msl_source, entry_point), mutex-guardedcompile_msl_kernelhelper —[device newLibraryWithSource:options:error:]+newComputePipelineStateWithFunctionwith cache-or-compile semanticstessera_apple_gpu_rope_f32C symbol with embedded rope MSL kernel string, encoded viaMTLComputeCommandEncodertessera_apple_gpu_runtime_msl_cache_sizefor cache-hit assertionsCMake links
-framework Metal -framework MetalPerformanceShaders -framework Foundationon Darwin only; non-Darwin compiles the stub instead.Python
driver.py—_is_apple_gpu_mps_executablegate covers both MPS (matmul/gemm) and MSL (rope) ops;metal_runtimeexecution mode for single-op runtime programstarget_ir.py—_apple_gpu_module_is_mps_runtimeaccepts any single-source program in the envelope;_lower_tile_opsde-duplicates runtime emissions;_APPLE_GPU_ROPE_MSL_SOURCEconstant + sha256 cache_keyjit.py—_apple_gpu_fast_callcached-metadata fast path; new metadata branch withcompiler_path=\"apple_gpu_mps\"runtime.py—_apple_gpu_dispatch_matmul/_apple_gpu_dispatch_rope, ctypes wrappers,_load_apple_gpu_runtime+ on-the-fly compilation;launch()gets anapple_gpu_mpsbranch with telemetryTests
apple_gpu_runtime.mlir,apple_gpu_msl.mlir)test_apple_backend_roadmap.py:test_target_ir_contract.pysplit into single-matmul (mps_runtime) and multi-op (metal_artifact) casesWhy the split between MPS and MSL?
Phase 8.4.0 is the on-ramp: rope is small (~10 lines of MSL) with a tight numpy reference, ideal for proving the infrastructure end-to-end. Phase 8.4.1 (flash-attention) is incremental on this foundation.
Architectural notes for reviewers
metal-cppvendoring. Per CLAUDE.md design lock, the.mmruntime uses the system Metal/MPS frameworks directly via Objective-C++.{matmul, gemm}. Phase 8.4.0 ={matmul, gemm, rope}. Multi-op programs stay onmetal_artifactuntil a kernel earns its way in.(msl_source, entry_point). Same source + same entry point → cache hit. Either changing forces a recompile. The IR carriescache_key = sha256(msl_source)[:16]so callers can pre-check without reading the full source.Compatibility
Existing
apple_gputests assertingexecution_mode == \"metal_artifact\"andnot uses_compiled_pathcontinue to pass — multi-op programs (tiny_decode,simple_transformer, MoE) are not in the runtime envelope.The one contract change is
test_jit_apple_gpu_target_emits_metal_artifact— split into single-matmul (mps_runtime) and multi-op (metal_artifact) cases. The split makes the gate explicit and reviewable; future phases broaden the runtime side without touching the artifact-only side.Test plan
tessera-optrebuilt with-DTESSERA_BUILD_APPLE_BACKEND=ONTESSERA_BUILD_APPLE_BACKEND=ONon macOS (LLVM/MLIR 21)apple_gpu_runtime_stub.cpppath links and exports the C ABI symbols correctly (left for CI)Followups
_runtime_flash_attn)🤖 Generated with Claude Code