Phase 8.4.3 — First multi-op MSL fusion (matmul → softmax) - #7
Conversation
Architectural step: the apple_gpu runtime gate moves from "single op in envelope" to "single op OR recognized op-chain in envelope." First chain: matmul -> softmax (axis=-1, rank-2, f32, N <= 256) — a common attention primitive. Both ops collapse into a single fused MSL kernel that avoids materializing the (M, N) intermediate score matrix on host. MLIR - Pass: MatmulSoftmaxFusionToAppleGPU pattern-matches a 2-op SSA chain (matmul -> softmax with single-use intermediate) and replaces both with one func.call into tessera_apple_gpu_matmul_softmax_f32. Pattern benefit=2 ensures the fusion wins the rewrite race against per-op patterns. The fusion pass runs first in the pipeline. - Pipeline tessera-lower-to-apple_gpu-runtime now composes 6 patterns: fusion + matmul + rope + flash_attn + softmax + gelu. - target_ir runtime-mode lowering accepts the fusion chain at the module level; emits a single fused msl_kernel + mps_dispatch pair carrying the embedded MSL source as a StringAttr (with a fusion="matmul_softmax" attribute for downstream introspection). Runtime - apple_gpu_runtime.mm: embedded matmul_softmax_f32 MSL kernel — one thread per output row, computes the row of A@B into a stack array (cap N <= 256), then numerically-stable row-wise softmax in place. Single MTLComputeCommandEncoder dispatch; reuses the Phase 8.4.0 MSL cache machinery. - apple_gpu_runtime_stub.cpp: portable C++ reference for non-Darwin builds (heap-allocated row buffer; correct for any N). Python - driver.py: _is_apple_gpu_mps_executable now accepts a 2-op plan that matches the matmul -> softmax chain (single-use intermediate, axis=-1 for the softmax). _apple_gpu_chain_kind classifies the chain. - target_ir.py: _APPLE_GPU_MATMUL_SOFTMAX_MSL_SOURCE constant + sha256 cache_key; _apple_gpu_module_fusion_kind detects the chain; _lower_apple_gpu_fusion emits a single fused msl_kernel for the chain (skipping the per-op walk to avoid double-emission). - runtime.py: _apple_gpu_dispatch_matmul_softmax dispatcher and _apple_gpu_matmul_softmax_f32 ctypes wrapper. The metadata dispatcher detects the 2-op chain at runtime and routes through the fused kernel; loader gate now requires the new symbol. Tests - New lit fixture apple_gpu_matmul_softmax_fusion.mlir — positive (rank-2 static, single-use intermediate), and 2 negative cases: matmul result with multiple uses (chain must NOT fire), N > 256 (out of GPU kernel envelope). - 4 new unit tests in test_apple_backend_roadmap.py: fused MSL artifact contract, end-to-end execution across 3 shapes, runtime shim ABI correctness, and a negative test for the multi-use case. - Updated test_jit_apple_gpu_multi_op_keeps_metal_artifact_contract -> split into matmul_softmax_chain_emits_fused_msl_runtime_contract (the now-fused pattern) and unrecognized_multi_op_keeps_metal_ artifact_contract (matmul -> gelu, still artifact-only). - compiler_examples.py: attention_matmul_softmax now claims runtime-executable on apple_gpu (Phase 8.4.3 contract change). Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 1991 unit tests passing (1979 + 12 picked up via origin/main merge); 11/11 Phase 8 lit fixtures passing against the in-tree tessera-opt. End-to-end fused chain rtol=1e-4 vs the per-op numpy reference across multiple shapes (M/K/N = 4/8/8, 8/16/32, 16/16/64). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5e6fc6c22b
ℹ️ 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".
| if (sources & matmul_sources) and (sources & softmax_sources) and ( | ||
| sources <= (matmul_sources | softmax_sources) | ||
| ): | ||
| return "matmul_softmax" |
There was a problem hiding this comment.
Keep Target IR fusion tied to the actual chain
When a module contains both a matmul/gemm source and a softmax/softmax_safe source but they are not the exact %matmul -> softmax dataflow chain, this set-based check still emits the fused matmul_softmax_f32 Target IR and marks the module metal_runtime. For example, an unused matmul(A, B) followed by return softmax(X) or a softmax with a non-fusible operand is not accepted by the stricter driver/runtime chain check, so the bundle becomes artifact-only while target_ir advertises a fused runtime kernel that does not implement the program.
Useful? React with 👍 / 👎.
| if first.op_name not in {"tessera.matmul", "tessera.gemm"}: | ||
| return None | ||
| if second.op_name not in {"tessera.softmax", "tessera.softmax_safe"}: |
There was a problem hiding this comment.
Do not advertise fusion variants the MLIR pass cannot match
This makes gemm -> softmax and matmul/gemm -> softmax_safe executable fused chains in the Python driver, but the new MLIR fusion pass is registered only on tessera.softmax and also requires the producer name to be exactly tessera.matmul. In those accepted variants, tessera-lower-to-apple_gpu-runtime will not emit the fused call (and can leave softmax_safe unlowered), while the metadata/backend report the fused native GPU path.
Useful? React with 👍 / 👎.
… 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
CodeQL #7 fixed. The stub recompiles (with __APPLE__ undefined) and the compute tests pass (4 passed) — the (double) cast is semantics-preserving (it computes the product in double precision, equal-or-better than the old float multiply). The fix: cv += (double)ckvb[…] * wuvb[…] — casting the first operand to double promotes the whole multiplication to double, so ckvb[…] * wuvb[…] no longer rounds/overflows in float before being accumulated into the double cv. This is the value-multiplication analogue of the earlier index fix (#5 cast j*Dl to size_t; #7 casts the float product to double). Note: I confirmed the sibling line 1282 (qabs[l] * ckvb[…]) is already safe — qabs is std::vector<double>, so that multiply is already double×float. Only line 1296 had the float×float→double pattern, so this is the complete fix for that alert, not a partial one.
…oint 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>
…oint 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): 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>
Summary
First multi-op MSL fusion on the apple_gpu runtime path. The runtime gate moves from "single op in envelope" to "single op OR recognized op-chain in envelope". First chain: matmul → softmax (axis=-1, rank-2 f32, N ≤ 256) — a common attention primitive (`softmax(QK^T)`).
Both ops collapse into a single fused MSL kernel that avoids materializing the (M, N) intermediate score matrix on host. One MTLCommandBuffer + one MTLComputePipelineState for the whole chain.
What changed
MLIR
Runtime
Python
Tests
Why this matters
Compatibility
The fusion changes one existing contract: `softmax(matmul(x, w))` programs flip from `metal_artifact` to `metal_runtime`. Two existing tests that pinned that contract were renamed/updated to use `gelu(matmul(x, w))` instead — still artifact-only since gelu chains aren't yet recognized. The change is explicit and reviewable in the diff.
The `flash_attn_contract` foundation example was also updated to claim runtime-executable on apple_gpu (Phase 8.4.1 oversight) — surfaced via a manifest-driven test now reading `stages_by_target` instead of hardcoding statuses.
Test plan
Followups
🤖 Generated with Claude Code