Conversation
Single rank-3 f32 tessera.flash_attn programs on @jit(target="apple_gpu") now execute through a purpose-built MSL kernel. Same online-softmax algorithm as flash-attention's algorithm 1, fused into a single kernel — avoids materializing the (B, Sq, Sk) score matrix entirely. Builds directly on the Phase 8.4.0 MSL infrastructure (kernel cache, MTLComputePipelineState dispatch). MLIR - Pass: FlashAttnToAppleGPU lowers rank-3 f32 tessera.flash_attn (head_dim <= 256) to func.call @tessera_apple_gpu_flash_attn_f32. Reads `causal` BoolAttr and optional `scale` FloatAttr; defaults scale = 1/sqrt(D). - Pipeline tessera-lower-to-apple_gpu-runtime extended to compose matmul + rope + flash_attn patterns - target_ir runtime-mode lowering accepts tessera.flash_attn as a recognized envelope source; emits msl_kernel + mps_dispatch with the embedded MSL source as a StringAttr - Bug fix: dedup-key tracking in _lower_tile_ops now only consumes a key when the op actually produces a runtime emission. Previously, filtered ops (tile.async_copy carrying source="tessera.flash_attn") would prematurely consume the slot and prevent the real compute op from emitting. Runtime - apple_gpu_runtime.mm: embedded flash_attn_f32 MSL kernel (online softmax with running max + denominator + per-thread output accumulator), one thread per (batch, q_row), grid (Sq, B, 1). Causal mask via i32 flag; head_dim <= 256 via per-thread stack array. Edge case: l == 0 (fully-masked row) returns zeros instead of NaN. - apple_gpu_runtime_stub.cpp: portable C++ flash_attn reference for non-Darwin builds (matches the MSL algorithm in plain C++) Python - driver.py: _APPLE_GPU_MSL_OPS extended with tessera.flash_attn; backend artifact picks the new symbol/framework/abi - target_ir.py: _APPLE_GPU_FLASH_ATTN_MSL_SOURCE constant + sha256 cache_key; emits msl_kernel for single-source flash_attn programs - runtime.py: _apple_gpu_dispatch_flash_attn dispatcher reads scale/causal kwargs and routes through the C ABI; loader gate now requires the flash_attn symbol (forces rebuild after 8.4.1) Tests - New lit fixture apple_gpu_flash_attn.mlir — positive (rank-3 static) and 2 negative cases (dynamic shapes, head_dim > 256). Uses the registered Tessera_FlashAttnOp directly (head_dim attribute is required by the verifier). - Three new unit tests in test_apple_backend_roadmap.py: MSL artifact contract (IR carries kernel source), end-to-end execution across 3 shapes with both causal and non-causal masks, runtime shim ABI correctness with direct ctypes invocation - Updated test_flash_attention_apple_gpu_gets_metal_kernel_contract -> test_flash_attention_apple_gpu_gets_msl_runtime_contract reflecting the contract change for single-flash_attn programs - Updated test_lower_tile_to_apple_gpu_target_ir_maps_fa4_to_metal_contract -> ..._maps_fa4_to_msl_runtime_contract - Updated compiler_examples.py: flash_attn_contract foundation example now claims runtime-executable on apple_gpu; manifest test reads claimed stages from the manifest instead of hardcoding artifact_only Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 1953 unit tests passing; 9/9 Phase 8 lit fixtures passing against the in-tree tessera-opt. End-to-end flash_attn rtol=1e-4 vs numpy across multiple shapes (B/Sq/Sk/D = 1/4/4/8, 2/8/8/16, 1/16/32/64) for both causal and non-causal masks. 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: ca9e43d2fb
ℹ️ 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".
|
|
||
| for (int k_row = 0; k_row < Sk; ++k_row) { | ||
| // Causal mask: skip keys above the query row. | ||
| if (causal != 0 && k_row > q_row) break; |
There was a problem hiding this comment.
For causal flash-attn with prefix/cached keys (Sk > Sq), the existing numpy reference masks with an offset (1 + max(Sk - Sq, 0)), so query row 0 can still attend to the prefix keys. This kernel stops as soon as k_row > q_row (the stub uses the same condition), so the apple_gpu runtime only attends to q_row + 1 keys and returns incorrect results for valid non-square causal inputs; for example B=1, Sq=2, Sk=4, D=3, causal=True diverges from _runtime_flash_attn.
Useful? React with 👍 / 👎.
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
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
gstoner
added a commit
that referenced
this pull request
May 30, 2026
Item #4 — two compute-efficiency features for the Gumiho draft/verify. Half-precision draft (--mode precision --dtype f16|bf16): - AppleBackend gains a compute_dtype that threads through every dense op (matmul/rmsnorm/silu_mul/relu/softmax). Native half MSL on the GPU keeps fp32 accumulation, so logits barely move and the serial argmax tokens match f32: f16 logit err ~0.004, bf16 ~0.022, serial_tokens_match=True. Off Metal it falls back to f32. run_precision_demo + PrecisionSummary. Paged-KV prefix sharing (--mode prefix): - PrefixSharedVerifier serves the context K/V from a tessera.cache.KVCacheHandle: prefill the context once, append only accepted tokens, recompute K/V for just the tree nodes — N rows/step instead of C+N. verify_tree reproduces build_draft's target log-probs bit-for-bit (err ~1.7e-7). run_prefix_sharing_demo over the full growing context: ~31% fewer K/V rows at tiny scale; the saving compounds as the context grows. DraftBundle.path_node_ids (already exposed) lets the verifier reconstruct the trie. demo.py gains --mode {precision,prefix} + --dtype. +5 tests in test_example_gumiho.py; ruff + lint_docs clean; 38 adjacent pass. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
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 7, 2026
Strategic reframe before building test infra: the compiler is NOT stub-riddled. Reading the existing audit dashboards shows the software-actionable gap surface is small and specific; most incompleteness is hardware-gated (expected) or thin-test (better closed by generative differential testing than hand-written tests). - scripts/stub_surface_report.py reads the checked-in dashboards (verifier_coverage, op_target_conformance, test_coverage) and emits a single ranked rollup that SEPARATES software-actionable gaps from hardware-gated rows. Rerunnable, derives every number from the CSVs (no hand-copied counts that drift). -> docs/audit/generated/stub_surface.md. - Findings: ~9 trivial-stub verifiers (Arch* NAS + KVCacheCreate/RingCreate); ~9 software conformance cells (conv2d/kv_cache_read -> cpu stop at numerical; conv2d/flash_attn/kv_cache_read -> metalium/nvidia/rocm stop at codegen); the rest of the conformance failures are hardware-gated (honest); ~123 needs-direct- test ops are the differential-generator target. - MASTER_AUDIT: a "Compiler-Completeness & Testing Program" section linking the report + the #1/#2/#4 workstream, and the key reframe (stubs are an oracle/conformance problem, not a fuzz problem; fuzz layers on the oracle). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 7, 2026
…safety fix Generative round-trip test (gen op-chain -> render DSL -> parse via lower_text_to_graph_ir -> assert recovered op names) plus a malformed-input fuzz that asserts the parser only ever raises a NAMED diagnostic (FrontendSyntaxError/FrontendSemanticError), never an uncaught crash. The fuzz immediately found a real crash-safety bug: parse_module()'s loop guard 'while _peek_text() != "}"' is also true at EOF, so a module missing its closing '}' reached an 'assert token is not None' (a bare AssertionError crash) instead of a named parser error. Fixed to emit FrontendSyntaxError 'unexpected end of input: expected }' (Decision #21 / crash-safety contract). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 7, 2026
Synthesizes random well-typed programs over the executable lane
(tessera.ops elementwise/matmul/norm subset + tessera.control.fori_loop/cond)
and diffs two evaluators of the same function:
* oracle — eager fn(*arrays): tessera.ops numpy reference + eager control
* candidate — run_traced(...): the real trace -> GraphFn / execute_traced
fused run_graph_* Apple GPU (Metal) path
A mismatch is a miscompile in the production lane — the bug class hand-written
per-op tests (item #1's ~123 needs_direct_test ops) miss. 51 numerical cases
green on Apple GPU across straight-line, fused run_graph_loop, and fused
run_graph_cond; runtime-free trace/op-name + vocab-subset properties run
everywhere. Deterministic stdlib random over fixed seeds; square NxN tensors
so every op composes; unstable generated programs (non-finite/overflow eager
output) are skipped so generator noise can't masquerade as a miscompile.
Closes #2 + #4 of the compiler-testing program; MASTER_AUDIT updated.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 7, 2026
…age parser fix Thrust #4 (IR contracts/verifiers), first batch. Adds real MLIR verifiers to four heavily-used ops that had none — tessera.rmsnorm, tessera.rmsnorm_safe, tessera.softmax_safe, tessera.log_softmax — mirroring the proven SoftmaxOp/ LayerNormOp contracts: rank + per-axis static dim + element-type preservation, eps>0 (norms), axis-in-range (log_softmax). Factored a shared verifyShapeDtypePreserving() helper. hasVerifier=1 in TesseraOps.td + verify() bodies in TesseraOps.cpp; tessera-opt builds clean (MLIR 22). Lit: tests/tessera-ir/phase2/sprint_v8_norm_softmax_verifiers.mlir (8 cases, positive twins + shape/dtype/eps/axis violations) passes -verify-diagnostics + FileCheck. verifier_coverage: real 73 -> 77. Also fixes a latent verifier_coverage.py parser bug this change surfaced: _extract_op_block's `[^{]*` greedily ran a block-less one-line def (`def X : Base<...>;`) into the NEXT op's `{...}` block, mis-attributing that op's hasVerifier. Two fixes: (1) stop the scan at `;` so a one-line def yields an empty block; (2) resolve base-class hasVerifier (one-line ops like the Clifford family inherit hasVerifier=1 from Tessera_CliffordBinaryOp/UnaryOp, which the parser previously only saw by accidental brace-bleed). Net: the optimizer one-liners (Tessera_OptimizerTreeOp, no base verifier) correctly read no_verifier; the Clifford one-liners correctly read real; absent = 0. 142 tests green (verifier_coverage 50 + mlir_verifier_sprint + tessera_opt build smoke); ruff/mypy clean; spec-sync + drift gates pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 7, 2026
…7 -> 100) Thrust #4, second batch — 23 new MLIR verifiers; verifier_coverage real 77->100, trivial_stub 9->1, absent 0. Control flow (Phase-G ops): ControlForOp/ControlIfOp/ControlWhileOp. A factored verifyControlPayload() enforces the run_graph op-list ABI invariant — the serialized body/branch arrays (opcodes/in0/in1/iattr/fattr) must be mutually length-consistent with out_id set, else the runtime reads past the op-list. Plus: step!=0 + loop-carried #result/type match (for); max_iters>0 + carry/result type match (while); flag_arg_index in range + then/else both-or-neither payload (if); carry/flag index bounds against operands. Closed 8 of 9 trivial stubs with real contracts: KVCacheCreate (max_seq/head_dim/ page_size > 0), RingCreate (capacity>0), ArchParameter (size>0), ArchGumbelSoftmax /ArchHardConcrete (temperature>0), ArchWeightedSum/ArchSwitch (non-empty candidates, each shape-matching the result), ArchMixed (non-empty candidates). ArchSTEOneHot stays a structural success() — an opaque ArchParam->ArchGate with no scalar/shape contract to check. MoR: MorRouter (max_depth>0), MorPartition (step>=0), MorScatter (out preserves full's shape/dtype). Quant: QuantizeFP8/DequantizeFP8 (format in {e4m3,e5m2} + shape-preserving), QuantizeFP4/DequantizeFP4 (format in {e2m1,nvfp4}). FFT family (FFT/IFFT/RFFT/IRFFT/DCT, via a hasVerifier on the Tessera_SpectralUnaryOp base): axis-in-range against input rank. tessera-opt builds clean (MLIR 22). Lit: tests/tessera-ir/phase2/sprint_v9_control_stub_misc_verifiers.mlir (positive + negative cases per family) passes -verify-diagnostics + FileCheck; no regression on v1/v4b/v8. 142 verifier/build tests green; mypy + spec-sync + drift gates pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 8, 2026
The Multivector front-end allow-list is {Cl(3,0), Cl(1,3)} but v1 ships Apple-GPU
kernels only for Cl(3,0) (cl30). A @clifford_jit callable invoked with a
spacetime Cl(1,3) Multivector previously routed silently to the numpy reference
inside tessera.ga.* — the silent-fallback anti-pattern (Decision #21).
Add a call-time signature gate (mirroring the existing decoration-time
dtype!='f32' gate): _require_cl30_args refuses any non-Cl(3,0) Multivector arg
with CLIFFORD_UNSUPPORTED_SIGNATURE, naming the offending signature and pointing
to the plain tessera.ga.* numpy lane for non-Cl(3,0) algebras. Wired into both
the IR-compiled and lazy (first-call-compile) call boundaries — the lazy gate
fires before the numpy trace. New diagnostic code
ConstrainedDiagnosticCode.CLIFFORD_UNSUPPORTED_SIGNATURE.
The plain tessera.ga.* lane is unaffected (numpy Cl(1,3) still works); only the
GPU-plan decorator gates.
Tests: tests/unit/test_clifford_jit_signature_gate.py (Cl(3,0) runs; Cl(1,3)
single-op + multi-op plans gated; plain ga lane unaffected). 165 clifford_jit/
dialect regression tests green; mypy host+linux + ruff clean.
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 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
… CompilerGym) CompilerGym exposes compiler tasks as RL environments (phase ordering being the canonical one); Tessera already has the agents (magellan/alphaevolve, gated perf-behind-correctness search) and the environment authority (the evaluator), but lacked the env/agent split for pass ordering. PassOrderEnv is that environment. An action is an ordering of the lowering- pipeline passes (the real apple_gpu/NVIDIA spine: EffectAnnotation → Canonicalize → fusion passes → TileIRLowering → WarpSpec; CollectiveInsertion after EffectAnnotation per CLAUDE.md). Reward is gated by correctness = dependency validity — an ordering that runs a pass before its prerequisite scores INF and is rejected (the Sakana invariant) — and among valid orderings the cost is fusion effectiveness: a fusion pass only eliminates ops when it runs *before* lowering freezes the graph, so order matters. magellan.search / magellan.evolve drive it unchanged (search_best_order / evolve_best_order); the env exposes fitness/reward, the agent optimizes it. The cost model is an honest fuse-before-lower proxy, not a cycle-accurate simulator. Guards: tests/unit/test_pass_order_env.py (8) — validity gate, INF-rejection of dependency violations, fuse-before-lower beats fuse-after, gated search finds the min-cost valid order, evolve never accepts an invalid order. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
pushed a commit
that referenced
this pull request
Jun 22, 2026
…ough the bridge The real RDNA WMMA matrix instruction now runs on the gfx1100/gfx1151 device and produces a numerically correct GEMM, routed through Tessera's C-ABI launch bridge — the first on-hardware execute-and-compare of a Tessera matmul on non-Apple silicon. - tests/unit/test_rocm_wmma_execute_compare.py: a hipcc harness whose launched kernel uses __builtin_amdgcn_wmma_f32_16x16x16_f16_w32 (the same v_wmma_f32_16x16x16_f16 rocdl_emit.py emits) to compute a 16x16x16 f32<-f16 GEMM, routed through tsrLaunchKernel, compared to a host reference. The operand/accumulator fragment layout matches rocdl_emit.py's grounded mapping (col = lane&15, row = 2*e + lane>>4). maxerr ~3e-8 standalone, <1e-2 through the bridge (f16 rounding). f32<-f16 first (bf16 has documented gfx115x bugs). Honest status — NOT promoted to hardware_verified / backend_kernel complete. This clears the *numerical-proof* half of the backend_manifest hardware_verified contract (execute_compare_fixture), but that status also requires a *shipped* runtime_symbol (an auto-registered ROCm runtime launcher); today the kernel + launcher live in the test harness (like the Apple G7 proof), so flipping the status would be Decision #25 inflation. backend_kernel stays 474/0. The flip is gated on shipping the launcher (ROCM_AUDIT.md Next Work #4) and becomes mechanical once the symbol ships. Docs: STRIX_HALO_EXECUTION_PLAN.md + ROCM_AUDIT.md (Stage D proof + honest gate). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
gstoner
pushed a commit
that referenced
this pull request
Jun 24, 2026
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>
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.
Single rank-3 f32 tessera.flash_attn programs on @jit(target="apple_gpu") now execute through a purpose-built MSL kernel. Same online-softmax algorithm as flash-attention's algorithm 1, fused into a single kernel — avoids materializing the (B, Sq, Sk) score matrix entirely. Builds directly on the Phase 8.4.0 MSL infrastructure (kernel cache, MTLComputePipelineState dispatch).
MLIR
causalBoolAttr and optionalscaleFloatAttr; defaults scale = 1/sqrt(D).Runtime
Python
Tests
Verified on Apple Silicon (LLVM/MLIR 21, Metal active):
1953 unit tests passing; 9/9 Phase 8 lit fixtures passing against the
in-tree tessera-opt. End-to-end flash_attn rtol=1e-4 vs numpy across
multiple shapes (B/Sq/Sk/D = 1/4/4/8, 2/8/8/16, 1/16/32/64) for
both causal and non-causal masks.