Phase 8.4.4 — fp16 / bf16 matmul on apple_gpu (mirror of CPU BNNS bf16) by gstoner · Pull Request #8 · gstoner/tessera

gstoner · 2026-05-08T03:56:28Z

Summary

Adds fp16 and bf16 dtype variants to the apple_gpu matmul runtime path. Mirrors the Phase 8.2 BNNS bf16 follow-up on the CPU side:

fp16: native MPSDataTypeFloat16 (Apple Silicon GPUs run fp16 at higher throughput than fp32)
bf16: fp32-conversion path inside the runtime shim (MPS does NOT natively support bf16 matrix descriptors as of macOS 14, so the shim decodes bf16 bit-patterns to fp32, runs MPSDataTypeFloat32, and encodes back with round-to-nearest-even — same shape as the CPU bf16 cblas_sgemm fallback)

Stacked on top of Phase 8.4.3 PR #7; base set to `claude/phase-8-4-3-msl-fusion` so this can land cleanly after #7 merges.

Scope

✅ matmul fp16 + bf16 only. Other custom MSL kernels (rope, softmax, gelu, flash_attn, matmul_softmax_fusion) remain f32-only this phase. Their dtype variants are 8.4.4.x followups.

What changed

MLIR / runtime

Two new C symbols in `apple_gpu_runtime.mm`:
- `tessera_apple_gpu_mps_matmul_f16` — native MPSDataTypeFloat16. ABI is `uint16_t*` for fp16 bit-pattern transmission (no `_Float16` dependency).
- `tessera_apple_gpu_mps_matmul_bf16` — fp32-conversion path. Decodes bf16 bit-pattern via bit-shift, runs MPSDataTypeFloat32, encodes back with round-to-nearest-even.
`apple_gpu_runtime_stub.cpp`: matching reference fallbacks for non-Darwin (both fp16 and bf16 use the same fp32-conversion path).
`MatmulToAppleGPU.cpp` picks the runtime symbol by input element type. Same `i64×3 + i32×3` ABI shape across all three dtypes — the element type is encoded in the symbol name only, not the signature.

Python

`driver.py`: `_apple_gpu_matmul_dtype_suffix` extracts dtype from Graph IR operand types and routes the backend artifact's symbol selection.
`schedule_ir.py`: `_base_attrs` now surfaces `dtype` on every schedule op by parsing `IROp.operand_types`. The attr propagates through Schedule → Tile → Target IR so `target_ir.py`'s `mps_matmul` emission carries the right `dtype` attribute ("f32" / "f16" / "bf16").
`runtime.py`: `_apple_gpu_dispatch_matmul` detects input array dtype at launch time (call-site dtypes are runtime-only since `@jit` function signatures are type-polymorphic) and routes to the matching ctypes wrapper. fp16 + bf16 paths both use `uint16_t*` ABI via `numpy.view(np.uint16)`; `ml_dtypes.bfloat16` is byte-compatible with the C ABI.
New ctypes wrappers `_apple_gpu_mps_matmul_f16` / `_bf16`. Loader gate now requires both new symbols (forces rebuild after Phase 8.4.4).

Tests

New lit fixture `apple_gpu_matmul_dtypes.mlir` — three positive cases (f32, f16, bf16 matmul lower to the right runtime symbol with the shared `i64×3 + i32×3` ABI) and one negative case (mixed-dtype operands fall back to the artifact path).
4 new unit tests in `test_apple_backend_roadmap.py`:
- f32 default artifact contract — compile-time symbol selection
- fp16 end-to-end — matches fp32-converted reference at fp16 tolerance
- bf16 end-to-end — matches fp32-converted reference at bf16 tolerance (gated on `ml_dtypes` presence, mirrors CPU bf16 soft-dep)
- fp16 + bf16 ABI shim correctness — direct ctypes invocation against a freshly-compiled shim

Why this design

Compile-time vs runtime dtype dispatch

The Graph IR is type-polymorphic in this codebase — `@jit` decorators don't know call-site dtypes. So:

Lit / MLIR pass-level tests exercise the compile-time symbol selection (when the Graph IR has explicit type tensors like `tensor<*xf16>`)
Python end-to-end tests exercise the runtime dispatcher (which detects dtype from input arrays at launch time)

Both paths use the same C ABI; they differ only in when the symbol is chosen.

Why bf16 doesn't get a native MPS path

MPS exposes `MPSDataTypeFloat16` and `MPSDataTypeFloat32` for matrix descriptors but not `MPSDataTypeBFloat16` as of macOS 14. Native bf16 matmul would need either custom MSL (a different phase) or upstream MPS support. The fp32-conversion path is correct, portable, and matches the CPU pattern.

Test plan

All 1,994 unit tests pass (1,991 → 1,994, +3 net new fp16/bf16 tests)
12/12 Phase 8 lit fixtures pass against the in-tree `tessera-opt` rebuilt with `-DTESSERA_BUILD_APPLE_BACKEND=ON`
End-to-end on Apple Silicon: fp16 matmul matches fp32-converted reference at `rtol=5e-2`; bf16 at `rtol=2e-2`
CMake builds with `TESSERA_BUILD_APPLE_BACKEND=ON` on macOS (LLVM/MLIR 21)
Linux CI: confirm `apple_gpu_runtime_stub.cpp` exports the new `f16` + `bf16` symbols (left for CI)

Followups

Phase 8.4.4.1 — fp16/bf16 for the simpler MSL kernels (rope, softmax, gelu) using MSL `half` type
Phase 8.4.4.2 — fp16/bf16 for the fused chain (matmul_softmax) and flash_attn
Phase 8.4.5 — more fusion patterns: matmul → softmax → matmul (full attention block)

🤖 Generated with Claude Code

@jit

Extends the apple_gpu matmul runtime path with fp16 and bf16 dtype variants. Mirrors the Phase 8.2 BNNS bf16 follow-up on the CPU side: - fp16: native MPSDataTypeFloat16 (Apple Silicon GPUs run fp16 at higher throughput than fp32 on most ops). - bf16: fp32-conversion path inside the runtime shim because MPS does NOT natively support bf16 matrix descriptors as of macOS 14. Same pattern as the CPU bf16 cblas_sgemm fallback. Scope is intentionally narrow — only the matmul kernel gets dtype variants this phase. The other custom MSL kernels (rope, softmax, gelu, flash_attn, matmul_softmax_fusion) remain f32-only; their dtype variants are 8.4.4.x followups. MLIR / runtime - Two new C symbols in apple_gpu_runtime.mm: * tessera_apple_gpu_mps_matmul_f16 — native MPSDataTypeFloat16. ABI is uint16_t* for fp16 bit-pattern transmission (no _Float16 dep). * tessera_apple_gpu_mps_matmul_bf16 — fp32 conversion path. Decodes bf16 bit-pattern via shift, runs MPSDataTypeFloat32 matmul, encodes back with round-to-nearest-even. - apple_gpu_runtime_stub.cpp gets matching reference fallbacks (fp32-via-conversion) for non-Darwin builds. - MatmulToAppleGPU.cpp picks the runtime symbol by input element type: f32 / f16 / bf16. Same i64×3 + i32×3 ABI shape across all three — the element type is encoded in the symbol name only. Python - driver.py: _apple_gpu_matmul_dtype_suffix extracts the dtype from the Graph IR operand types (tensor<*xf16>, tensor<*xbf16>) and routes the backend artifact's runtime symbol selection accordingly. - schedule_ir.py: _base_attrs now surfaces dtype on every schedule op by parsing the IROp's operand_types. The attr propagates through Schedule -> Tile -> Target IR layers so target_ir's mps_matmul emission carries the right dtype attr. - runtime.py: _apple_gpu_dispatch_matmul detects input array dtype at launch time (call-site dtypes are runtime-only since the @jit function signatures are type-polymorphic) and routes to the matching ctypes wrapper. fp16 and bf16 paths both use uint16_t* ABI via numpy's .view(np.uint16); ml_dtypes.bfloat16 is byte-compatible. - New ctypes wrappers _apple_gpu_mps_matmul_f16 / _bf16. Loader gate now requires both new symbols (forces rebuild after Phase 8.4.4). Tests - New lit fixture apple_gpu_matmul_dtypes.mlir — three positive cases (f32, f16, bf16 matmul lower to the right runtime symbol with the shared i64×3 + i32×3 ABI) and one negative case (mixed-dtype operands fall back to the artifact path). - 4 new unit tests in test_apple_backend_roadmap.py: * f32 default artifact contract (compile-time symbol selection) * fp16 end-to-end matches fp32-converted reference at fp16 tolerance * bf16 end-to-end matches fp32-converted reference at bf16 tolerance (gated on ml_dtypes presence, mirrors the CPU bf16 soft-dep) * fp16 + bf16 ABI shim correctness (direct ctypes invocation against a freshly-compiled shim) Verified on Apple Silicon (LLVM/MLIR 21, Metal active): 1994 unit tests passing (1991 + 3 net new fp16/bf16 tests); 12/12 Phase 8 lit fixtures passing against the in-tree tessera-opt. fp16 matmul matches fp32-converted reference at rtol=5e-2 (MPS does fp16 internal accumulation; minor drift from the per-element reference is expected). bf16 matches at rtol=2e-2. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8fa5e551c2

ℹ️ 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".

chatgpt-codex-connector · 2026-05-08T03:59:07Z

-    auto outMemTy = MemRefType::get({M, N}, f32Ty);
+    auto lhsMemTy = MemRefType::get({M, K}, lhsElem);
+    auto rhsMemTy = MemRefType::get({K, N}, rhsElem);
+    auto outMemTy = MemRefType::get({M, N}, lhsElem);


Preserve f32 matmul result types

For a valid tessera.matmul with bf16 operands and an f32 result (the Graph IR spec uses this form in docs/spec/GRAPH_IR_SPEC.md), this lowering now allocates and returns the output tensor with lhsElem instead of the op's declared result element type. That means tensor<...xbf16> -> tensor<...xf32> matmuls either produce invalid IR when returned from an f32 function or silently lose the f32 accumulation contract; please read the op result type and only take this runtime path when the selected symbol can produce that result dtype.

Useful? React with 👍 / 👎.

…lanes CODE_AUDIT_2026_06_10 follow-ons #5/#8 + P2 refactor: - SwiGLU fusion-group derivation (finding #8 follow-on): canonical_compile ._match_swiglu_at matches the SwiGLU DAG (gate/up share %x, both feed silu_mul) inside the known-chain scan, tried first as the longest (4-op) fusion. The apple_gpu executor consumes fused_kernel == "swiglu", short-circuiting the per-invoke structural re-matcher. - Dtype-lane dispatch tables (P2 partial): _apple_gpu_dispatch_matmul and _apple_gpu_dispatch_unary become per-dtype (router, symbol) tables sharing _apple_gpu_gemm2d_call / _apple_gpu_mpsgraph_unary_call, replacing the copy-pasted f32/f16/bf16 branches. Per-dtype routing (MTL4 routers, bf16 capability gate) stays visible in the tables; missing-symbol paths now route through the strict-dispatch funnel instead of crashing. - Strict-dispatch CI lanes (finding #5): conftest _STRICT_DISPATCH_LANES forces TESSERA_STRICT_DISPATCH=1 (Darwin-only) on the two differential generators + the 10 manifest-declared Metal execute_compare_fixture modules, so a failure-class GPU->numpy fallback raises instead of silently passing against the numpy oracle. Verified zero fallbacks on a Metal host. - Tests: SwiGLU derivation + executor-consumption cases; generated test_coverage regenerated for the new test references. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Update CODE_AUDIT_2026_06_10 findings #5/#8/#10 + P2 to reflect the landed follow-ons: numeric_policy propagation (C++), SwiGLU fusion-group derivation + executor consumption, strict-dispatch CI lane wiring, and the matmul/unary dtype-table refactor (with the consciously-skipped symbol-getter memoization noted). Sync the matching COMPILER_AUDIT line. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

#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>

* 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>

@jit

…d-KV native rung Close the recurring meta-gap (typed contracts with no consuming pass) across the 8-item KV/attention/quant/TP/multimodal audit. Each contract now has a consuming pass AND a semantics-preserving oracle, tracked by a drift-gated dashboard. Phase 0 — compiler/contract_consumers.py: live-probed meta-gap dashboard (docs/audit/generated/contract_consumers.{md,csv}); reports 6 live / 0 declared. A — cache/paged_kv.py: PagedKVState ABI unifying contiguous + tiered KV; ops.paged_attention + flash_attn(kv_state=); evaluator.paged_kv_equivalence. A follow-on (#8) — native Apple-GPU (Metal 4.0) paged attention via the shipped fused matmul→softmax→matmul kernel with provenance gating; LATENT (MLA expand) + QUANTIZED_TAIL kinds; evaluator.paged_kv_native_equivalence (native rung). B — compiler/phase_specialization.py: prefill/decode as distinct schedules + CacheHandoff ABI; @jit(phase=,slo=); verify_phase_split oracle. C — fusion.select_attention_lowering: IO-byte selector replaces the hard Nk<=SYNTH_MAX_N branch; paged_stage_bytes feeds page-gather cost. D — compiler/smoothquant.py: W8A8 activation-scale migration producer pass + anti-fallback oracle (operands stay int8). E — compiler/tensor_parallel.py: column/row/sequence-parallel rewrite_linear + cross-rank gradient-equivalence oracle (MockRankGroup). F — compiler/model_walk.py: named multimodal walks (partition_walks) + first-class encoder-free ops (patch/coordinate/audio projection) + walk-parity oracle. Also: numpy-2 uint8-safety fix in quantization.py int4/fp4 packers (pre-existing working-tree change, ridden along). Tests: ~120 new unit tests across the new modules; mypy clean; 17 generated docs in sync (drift-gated). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

@jit

…ckend_kernel (#106) * fix(nvidia): finish CUDA 13.2.1→13.3 propagation + verify sm_120 smem carve-out on silicon (#8, #14) The 13.2.1→13.3 pin bump (landed 2026-06-18) had only partially propagated. Source-of-truth pins were correct, but the surrounding surface had drifted. #8 — complete the bump: - validate_nvcc_compile.py: MIN_NVCC_VERSION (13,2,1)→(13,3,0); fix the user-facing error message and usage example (real drift, not prose). - capabilities.py: rename feature marker cuda_13_2_u1 → cuda_13_3 (4 entries) to match the existing C++ pin-test naming; update the test assertion + name accordingly. - Refresh stale "CUDA 13.2 U1" strings in diagnostics (2 fix-hints), backend_manifest notes, probe_collective_libs, and docs (CLAUDE.md, README, PROJECT_STRUCTURE, kernel inventory incl. the stale _CUDA_13_2_FEATURES identifier ref, execution plan). - Regenerate nvidia_sm90_target_map.{csv,md} (drift-gated). #14 — resolve the 100 KB vs 128 KB shared-memory carve-out on silicon: Measured on RTX 5070 Ti (driver 610.62 / CUDA UMD 13.3, nvcc 13.3.33, -arch=sm_120): sharedMemPerMultiprocessor = 102400 B (100 KiB) <- exact match cudaDevAttrMaxSharedMemoryPerBlockOptin = 101376 B (99 KiB) cudaDevAttrMaxSharedMemoryPerBlock = 49152 B (48 KiB) 100 KB wins; the release-note 128 KB is the unified data cache (Table 32), not the shared-memory carve-out. _SMEM_BYTES[SM_120]=102400 is correct. Recorded the on-silicon confirmation + the per-block 99 KiB opt-in ceiling (a lowering must cap dynamic smem below 101376 on sm_120) in gpu_target.py and BLACKWELL_SM120_EXECUTION_PLAN.md. Verification: drift gate clean (17 docs in sync); toolchain-pin and NVIDIA lane tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * spike(nvidia): sm_120 mma.sync.m16n8k16 PTX — emit→assemble→launch→compare on RTX 5070 Ti (#6) NVIDIA action item #6, proven end-to-end on real consumer-Blackwell silicon (RTX 5070 Ti, CC 12.0, CUDA 13.3, nvcc/ptxas 13.3.33). A hand-emitted, Tessera-style sm_120 bf16 mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 single-tile GEMM (D[16x8] = A[16x16]·B[16x8], f32 accumulate; warp-level mma.sync, NOT the Hopper warpgroup wgmma path) clears the full rung ladder: rung 2.5 (emit): raw PTX .version 9.3 / .target sm_120a (matches pin) rung 3 (assemble): ptxas --gpu-name=sm_120a -> 6168-byte cubin rung 5 (execute): Driver-API cuModuleLoadDataEx + cuLaunchKernel, max abs err 4.8e-7 vs CPU ref, 0/128 elements off Artifacts (committed under spikes/sm120_mma_sync/, with README + reproduce steps): the hand-written PTX, a CUDA inline-asm oracle that nails the m16n8k16 fragment layout, a driver-API run harness, and the smem device-query used to resolve #14. The .ptx is force-added past the global *.ptx ignore because here it is source, not a build artifact. Gotchas captured for productization: PTX must be ASCII (driver JIT ptxas rejects non-ASCII that standalone ptxas tolerates); f32 accumulator regs must be zero-initialized before the mma; CUDA 13.3 cuCtxCreate is v4 (use cuDevicePrimaryCtxRetain). Updated BLACKWELL_SM120_EXECUTION_PLAN.md sequencing: Stage A spike done; remaining work is wiring the path into ptx_emit.py + tsrRegisterGpuLauncher + an execute-compare oracle test. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(nvidia): productize sm_120 mma.sync — emit_mma_sync_matmul_ptx + CUDA launch bridge + execute-and-compare (#6) Turns the spike (#6) into a proven in-tree path. All three pieces run on real consumer-Blackwell silicon (RTX 5070 Ti, CC 12.0, CUDA 13.3). 1. ptx_emit.emit_mma_sync_matmul_ptx (+ mma_sync_mnemonic, is_valid_mma_sync_bf16_shape, validate_mma_sync_ptx_structure): emits a COMPLETE, assemblable, launchable sm_120 mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 kernel (warp-level, register operands — no smem descriptors/TMA, unlike the WGMMA skeleton). ASCII-only (driver JIT ptxas rejects non-ASCII); accumulator explicitly zeroed. Output is the validated spike kernel. 2. CUDA launch bridge: test_conformance_execute_compare_nvidia.py registers an nvidia_mma_launcher via tsrRegisterGpuLauncher that loads the EMITTED PTX through the Driver API (cuModuleLoadDataEx -> cuLaunchKernel) and runs it through tsrLaunchKernel — mirrors the ROCm WMMA bridge template but exercises Tessera's own emitter output. Skip-clean (no nvcc / no runtime lib / no GPU); also asserts the negative UNIMPLEMENTED path. 3. CMake fix: src/runtime/CMakeLists.txt wired CUDA include/link for TESSERA_ENABLE_CUDA (mirror of the existing HIP block) — it was enabled but unbuildable (cuda_backend.cpp couldn't find cuda_runtime.h). Proof on box: emitter output -> ptxas sm_120a (rung 3) -> driver launch -> execute-and-compare vs CPU ref at max abs err 4.8e-7, 0/128 off (rung 5). Tests: test_ptx_emit (incl. ptxas-gated rung-3 assemble) + the NVIDIA execute-compare both pass; drift gate clean (17 docs in sync); mypy clean. Not flipped: the backend_manifest NVIDIA row stays artifact_only. ROCm's hardware_verified rows ship a real auto-built C-ABI runtime_symbol .so; the NVIDIA kernel here runs via the emitted PTX through a harness launcher, not a shipped libtessera_nvidia_*.so. Flipping to hardware_verified honestly needs that shipped runtime lib (the well-scoped next step). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * spike(nvidia): general mma.sync GEMM across shapes for bf16/f16/tf32/fp8 (e4m3,e5m2) Extends the #6 spike from a single 16x8x16 tile to a GENERAL tiled/K-looped GEMM (arbitrary M/N/K, ragged zero-padded) and sweeps every CC 12.0 Tensor-Core input dtype, all NVRTC-compiled (compute_120) and execute-and-compared on the RTX 5070 Ti across 7 shapes: bf16 m16n8k16 worst maxerr 9.5e-6 f16 m16n8k16 worst maxerr 1.7e-5 tf32 m16n8k8 worst maxerr 2.5e-5 e4m3 m16n8k32 bit-exact (0) e5m2 m16n8k32 bit-exact (0) Each dtype uses its own MMA shape + fragment layout (16-bit: 2 elems/reg; tf32: one tf32/reg, K=8; fp8: 4 bytes/reg, K=32). The host reference quantizes inputs to the same dtype the hardware consumes (OCP fp8 decode), so fp8 matches bit-for-bit. nvgemm_proto.cpp is the bf16/f16 general GEMM that backs the shipped libtessera_nvidia_gemm.so; nvgemm_dtypes.cpp is the full dtype sweep. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * spike(nvidia): NVFP4 m16n8k64 block-scale MMA — assembles + executes on sm_120a (#9, partial) The consumer-Blackwell headline (action item #9): warp-level block-scaled NVFP4 (e2m1 data + ue4m3 per-16-block scale), m16n8k64. Grounded on the RTX 5070 Ti (ptxas/runtime = source of truth): - mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X. f32.e2m1.e2m1.f32.ue4m3 ASSEMBLES and EXECUTES on sm_120a. Accepted operand grammar: {d4},{a4},{b2},{c4},{sfa},{byteid_a,tid_a},{sfb},{byteid_b,tid_b}. - Arch-specific: only sm_120a SASS accepts it; base compute_120/sm_120 PTX rejects .kind::mxf4nvf4 / .block_scale / .scale_vec::4X. Build with -gencode arch=compute_120a,code=sm_120a. NOT yet numerically verified (honest, per the grounding rule): the on-box CUDA 13.3 headers expose only the datacenter tcgen05 block-scale variant, not the warp mma.sync scale-distribution + ue4m3 encoding semantics. A scale-byte sweep shows the scale is multiplicative (0x00 -> ~0) but standard e4m3 1.0 (0x38) yields ~3e9, so a guessed scale layout does not match a reference. Recorded as the precise next step; not claimed working. Artifacts: nvfp4_probe.cu (assemble probe), nvfp4_gemm.cu (per-lane fragment harness + scale sweep), README NVFP4 section. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(nvidia): ship libtessera_nvidia_gemm.so + flip sm_120 matmul to hardware_verified (#6) Completes execution-plan step 3: NVIDIA's first hardware-verified backend_kernel row, mirroring the ROCm Strix Halo shipped-symbol pattern. Proven on the RTX 5070 Ti (CC 12.0, CUDA 13.3). Shipped runtime symbol: - src/.../tessera_gpu_backend_NVIDIA/runtime/cuda/tessera_nvidia_gemm.cpp + CMake target `tessera_nvidia_gemm` -> libtessera_nvidia_gemm.so. Exports tessera_nvidia_mma_gemm_{bf16,f16,tf32,e4m3,e5m2}: a general tiled/K-looped warp-level mma.sync GEMM (ragged M/N/K zero-padded, f32 accumulate), NVRTC-compiled (compute_XX) for the device arch at first call, launched via the CUDA driver API. Per-dtype MMA shape: bf16/fp16 m16n8k16, tf32 m16n8k8, fp8 e4m3/e5m2 m16n8k32. Built only when TESSERA_ENABLE_CUDA (self-gated subdir). Numerical proof: - tests/unit/test_nvidia_mma_runtime_symbol.py dlopens the shipped .so and validates all 5 dtypes vs numpy/ml_dtypes references across aligned + ragged shapes. Skip-clean (no GPU/NVRTC -> rc=2). Manifest flip: - _NVIDIA_HARDWARE_VERIFIED["matmul"] + the builder injects a hardware_verified row for nvidia_sm120 (runtime_symbol + execute_compare_fixture), replacing the artifact_only row for that arch only; sm_80/90/100 stay artifact_only (proven only on sm_120). _NUMERICAL_FIXTURES[("matmul","nvidia_sm120")] added. dtypes=(bf16,fp16,fp32,fp8_e4m3,fp8_e5m2); fp32 = tf32-math (Decision #15a); feature_flags=(mma_sync,) — NOT tcgen05/tmem (consumer Blackwell lacks those). - Regenerated op_target_conformance + runtime_abi dashboards (drift gate clean). - Updated test_numerical_check_via_manifest to expect the sm_120 fixture (and none on sibling arches). Not flipped: the @jit default lane (execution_matrix executable row + runtime.launch dispatch) is the documented follow-up, cf. ROCm rocm_wmma symbol vs rocm_compiled lane. Verification: shipped-symbol (5 dtypes) + emitted-PTX execute-compare + ptx_emit pass on the box; 1236 affected unit tests pass (clean PATH); mypy + drift clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * feat(nvidia): wire @jit(target=nvidia_sm120) matmul to the shipped mma.sync lane Flips the execution_matrix row to executable and dispatches a real matmul through libtessera_nvidia_gemm.so end-to-end on the RTX 5070 Ti. Mirrors the ROCm rocm_wmma lane. - execution_matrix: new executable ("nvidia_sm120","nvidia_mma") native_gpu row (execution_mode=cuda_runtime) + KNOWN_EXECUTORS["nvidia_mma"]; nvidia_sm120 removed from _UNIMPLEMENTED_TARGETS (sm_80/90/100 stay — only sm_120 proven). - runtime.py: _execute_nvidia_mma_artifact + _load_nvidia_gemm_runtime + _nvidia_mma_runtime_available probe + _NVIDIA_GEMM_SYMBOLS; registered in _executor_table. Loads the shipped .so (preloading libcuda + libnvrtc), picks the dtype symbol (f16/bf16/fp32→tf32-math), runs, returns the f32 result. Never raises on a missing lib/device — the probe gates executability. - jit.py: _nvidia_mma_lane_available + _uses_nvidia_mma_default + a stamping branch -> executable/compiler_path=nvidia_mma/native_gpu when the host probe passes; off-device the artifact stays artifact_only (no behavior change). execution_kind also honors the nvidia lane so is_executable agrees with launch. - test_nvidia_launch_execute.py: host-independent matrix-row + executor-registry tests, plus hardware-gated execute-and-compare (f16/bf16 hand-built artifact) and the @jit(target="nvidia_sm120") default-dispatch test. - Regenerated runtime_execution_matrix + test_coverage dashboards. Proof on box: @jit(target="nvidia_sm120") matmul launches via the shipped symbol at f32 epsilon (maxerr ~1e-6). 552 affected tests pass (clean PATH); 13 GPU lane tests pass; mypy + drift clean. Follow-up: a compiler-GENERATED nvidia lane (the rocm_compiled analog) and NVFP4. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * test(nvidia): allow sm_120 mma.sync matmul as hardware_verified in honesty guard The branch landed a real on-silicon NVIDIA proof (ecd825f): matmul/nvidia_sm120 flipped to hardware_verified, backed by the shipped tessera_nvidia_mma_gemm_* C-ABI symbols (libtessera_nvidia_gemm.so) and a skip-clean execute-compare fixture validating a warp-level mma.sync GEMM on the RTX 5070 Ti (sm_120, CC 12.0, CUDA 13.3). The manifest moved forward but the honesty guard in test_backend_capability_extension.py still hard-coded "NVIDIA not allowed" — the "first proof landed -> update the guard" event its own docstring anticipates. Mirror the ROCm block: add _NVIDIA_HARDWARE_VERIFIED_OPS = {"matmul"}, extend _hardware_verified_claim_is_allowed to accept nvidia_sm120 matmul (still requiring both evidence fields), and refresh the docstrings + assertion message. The guard stays strict for every other NVIDIA op/target. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: angst <angstroms01@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

gstoner merged commit f886f68 into claude/phase-8-4-3-msl-fusion May 8, 2026
0 of 5 checks passed

chatgpt-codex-connector Bot reviewed May 8, 2026

View reviewed changes

gstoner mentioned this pull request Jun 25, 2026

NVIDIA sm_120: CUDA 13.3 verify + first hardware-verified mma.sync backend_kernel #106

Merged

gstoner mentioned this pull request Jul 10, 2026

docs: unstale three ROCm items that already landed (mfma_table, tessera-to-linalg, ROCm MSA IR mirror) #364

Merged

Navigation Menu

Phase 8.4.4 — fp16 / bf16 matmul on apple_gpu (mirror of CPU BNNS bf16) - #8

Phase 8.4.4 — fp16 / bf16 matmul on apple_gpu (mirror of CPU BNNS bf16)#8
gstoner merged 1 commit into
claude/phase-8-4-3-msl-fusionfrom
claude/phase-8-4-4-fp16-bf16

gstoner commented May 8, 2026

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot May 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

gstoner commented May 8, 2026

Summary

Scope

What changed

MLIR / runtime

Python

Tests

Why this design

Compile-time vs runtime dtype dispatch

Why bf16 doesn't get a native MPS path

Test plan

Followups

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot May 8, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant