Apple GPU MLA: multi-sequence block-paged cache (vLLM-style paged attention) - #27
Conversation
…ention) Adds tessera.cache.MLABlockPagedCache — the production-serving follow-on that manages many concurrent sequences over a single physical block pool, the core idea behind vLLM's PagedAttention. - python/tessera/cache/mla_block_paged.py: physical pool of fixed-size blocks (latent [num_blocks, block_size, Dl] + rope [num_blocks, block_size, dr]); per-sequence block tables; on-demand block allocation from a free list; free-on-finish page reuse (no per-sequence contiguous reservation, no external fragmentation); logical->physical indirection so non-contiguous blocks serve a sequence; add_sequence / append / decode / decode_batch / free_sequence; pool accounting (num_free_blocks / utilization / block_table / cache_bytes_per_token). Ragged decode loops per sequence (different lengths) over the shared absorbed decode kernel. - mla_paged.py: factored a module-level absorb_decode_one() (single-sequence weighted-absorbed decode with rope-table generation + GPU dispatch / numpy fallback) shared by MLAPagedDecoder and the block cache; removed the now-dead _rope_tables/_gpu_decode methods. One decode implementation. - cache/__init__.py: export MLABlockPagedCache + MLABlockPagedCacheError. - tests/unit/test_mla_block_paged_cache.py: 8 tests — block growth, free+reuse, non-contiguous block tables, concurrent ragged decode_batch, incremental two-sequence loop, pool exhaustion + recovery, utilization/footprint, lifecycle errors. Decode cross-checked against the single-sequence MLAPagedDecoder. - docs: plan item 5 (e) marked done; same-length B>1 batching noted as the remaining compute-side optimization. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9e4371292a
ℹ️ 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".
| while len(st.block_table) < needed_blocks: | ||
| if not self._free: | ||
| raise MLABlockPagedCacheError( | ||
| f"block pool exhausted: {self.num_blocks} blocks all in use " | ||
| f"(free a finished sequence to reclaim pages)") | ||
| st.block_table.append(self._free.pop()) |
There was a problem hiding this comment.
Preflight block capacity before mutating the free list
When an append needs multiple new blocks but the pool has only some of them free, this loop pops the available blocks into st.block_table and then raises on the next iteration. The failed append leaves the sequence at the old length while still consuming pages, so a rejected oversized prefill can strand capacity and cause unrelated requests to see false pool exhaustion until that sequence is explicitly freed. Check needed_blocks - len(st.block_table) against len(self._free) before popping anything, or roll back the blocks on failure.
Useful? React with 👍 / 👎.
Tessera authors its own production .mtlpackage from the MPSGraph lane — no coremltools, no DXIL, no offline toolchain: build MPSGraph matmul → compileWithDevice: → MPSGraphExecutable → serializeToMPSGraphPackageAtURL: → wrap manifest.json → .mtlpackage The authored package flows through the existing PK1–PK7 lifecycle unchanged (load → reflect → prepare → dispatch) and the GPU output is bitwise-exact vs numpy A @ B (max abs err 0.0 at 4×4×4; passes fp32 tol across 3 shapes). The one real discovery along the way MPSGraph-serialized packages expose positionally-indexed, unnamed bindings (MPSGraphTensor has no name property — the Apple sample's inputA/inputB/output names came from its CoreML origin). That's an Apple-side fact, not a Tessera bug. So I added index-addressing (fill_input_at/read_output_at + a tensorsByIndex map) rather than forcing names that don't exist. Changes (all uncommitted, on main) apple_gpu_runtime.mm — 4 new C ABI symbols: tessera_apple_gpu_mlpkg_author_matmul, _first_function_name, _fill_input_at, _read_output_at; tensorsByIndex in prepare_tensors. apple_gpu_runtime_stub.cpp — matching non-Darwin stubs. apple_mlpkg.py — author_matmul_package, first_function_name, Pipeline.fill_input_at/read_output_at. test_apple_mlpkg_pk8.py (new) — 7 tests; full author→dispatch→numpy proof. APPLE_AUDIT.md — PK8 marked landed; item 5 closed. runtime_abi.md — regenerated (218 → 222 symbols, drift gate green). Verification: all 76 test_apple_mlpkg_* pass (incl. 7 PK8), 80 broader Apple tests pass, ruff + mypy clean, runtime-ABI drift gate green. The arc this closes Three times this conversation I wrote ".mtlpackage is blocked" from memory. Each was wrong: MTLBinaryArchive/dynamic-library serialization exists, metal-package-builder is on-machine, and MPSGraphExecutable.serializeToMPSGraphPackageAtURL: makes the whole authoring path real — which PK8 now proves end-to-end. The process fix (CLAUDE.md Decision #27 / skills.md "READ FIRST": ground Metal claims in SDK headers, not memory) is exactly what turned the third "blocked" into a working, bitwise-exact feature.
L5 — Apple CPU runtime (Accelerate LAPACK) Added tessera_apple_cpu_cholesky_f32(A, L, N) → info to apple_cpu_runtime.cpp, backed by LAPACK spotrf (grounded in the SDK header per Decision #27), with a portable Cholesky–Banachiewicz fallback for non-Apple CI. Handled the column-major/row-major subtlety correctly: spotrf(UPLO='U') on the col-major view of a row-major symmetric buffer yields the row-major lower factor with no transpose; strict upper triangle zeroed to match numpy. Validated: exact match (max|A−LLᵀ|=0, L=[[2,0,0],[6,1,0],[-8,5,3]]); non-SPD → info>0; both LAPACK and fallback paths exact. Guard: test_apple_cpu_runtime_cholesky_f32_correctness (n=1..16 vs numpy.linalg.cholesky + non-SPD). Apple backend suite: 70 passed. L6 — Seam-closure execution The harness test_apple_cholesky_seam_closure.py proves the executed result comes from the IR-named symbol, not a hardcoded dispatcher: Lowers tile.cholesky through the real tessera-opt Tile→Apple pass. Parses symbol straight out of the emitted Target IR. Compiles the runtime, resolves that exact symbol via getattr(runtime, symbol) (never hardcoded), runs it on random SPD matrices → matches numpy.linalg.cholesky (n=3,8,32). Plus: GPU IR names tessera_apple_gpu_cholesky_f32 + tags metal_runtime, and that symbol genuinely exists in the GPU runtime source. 3/3 pass. This is the key guarantee you asked for: a compiler change that emitted the wrong symbol would make the harness execute the wrong thing and fail — so every layer from tessera.cholesky down to the executed kernel is now on the correctness critical path for the CPU lane (LAPACK-vs-numpy), with the GPU lane structurally verified. Where the pilot stands Layers L1→L6 are all green: Graph IR → Schedule IR → Tile IR → Apple Target IR → CPU runtime execution, seam-closed. Only L7 remains: a single tessera-lower-to-apple_{cpu,gpu}-full pipeline alias chaining the whole spine (distribution-lowering → tiling → tile-to-apple) so one invocation does end-to-end, plus its drift/lit gates — turning the manual multi-pass chain into the reusable template for tri_solve/svd and the other ops.
…orical
First LDT / lattice-reasoning family PR (from the model-family gap analysis).
Four self-contained primitives, each wired across every layer the new compiler
path cares about:
* op_catalog: 4 OpSpec entries (count_nonzero=reduction, popcount=elementwise,
masked_categorical=random/indexing, asymmetric_bce=loss) -> Graph IR identity.
* numpy reference (python/tessera/__init__.py ops namespace):
- count_nonzero(x, axis, keepdims): candidate-cardinality reduction.
- popcount(x): per-element set-bit count (numpy>=2 bitwise_count fast path,
uint64 masking fallback for numpy<2).
- masked_categorical(logits, mask, key=None): masked greedy argmax (key=None,
deterministic/testable) or Gumbel-max sample (key given). Returns indices.
- asymmetric_bce (losses.py): pos/neg-weighted BCE-with-logits in the stable
softplus form; reduces to binary_cross_entropy_loss at weights=1.
* autodiff: asymmetric_bce VJP + JVP (finite-diff verified ~1e-10/1e-6); the
three index/integer ops marked vjp/jvp not_applicable via
_NONDIFFERENTIABLE_PER_NAME.
* registry: 435 -> 439 entries; asymmetric_bce vjp+jvp complete, the trio N/A.
* Apple GPU dispatch: all four execute correctly under @jit(target="apple_gpu")
via the numpy-fallback chain (functional, per PR scope). Metal-4 kernel path
for each documented (grounded in MPSGraph SDK headers, Decision #27) in
docs/ldt_primitives_metal4_mapping.md: 3/4 compose from already-linked
MPSGraph nodes + the shipped gumbel_argmax/argreduce kernels with zero new
MSL; popcount is the one genuine MSL intrinsic follow-on.
Tests: test_ldt_primitives.py (20 - catalog/registry/numpy/VJP+JVP/apple_gpu) +
the differential-generator harness extended (stdlib + hypothesis) with an LDT
case set in _diff_lane.ldt_cases (oracle vs @jit(apple_gpu), exact for the
integer/index ops, f32 tol for the loss). Generated dashboards regenerated.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Promotes the router z-loss and asymmetric BCE from the numpy-fallback (metal_artifact) lane to metal_runtime via MPSGraph subgraphs (mirroring the PPO loss precedent), grounded in the SDK headers (Decision #27). * apple_gpu_runtime.mm: mpsg_run_z_loss_f32 (reductionMaximum/exp/reductionSum/ log → logsumexp²→mean) + mpsg_run_asymmetric_bce_f32 (stable softplus = relu(u)+log(1+exp(-|u|)) via relu/abs/exp/log) + extern "C" wrappers with reference fallbacks; graph-cached by shape. Stub parity. * _apple_gpu_backend gpu_z_loss / gpu_asymmetric_bce ctypes wrappers + argtypes; _SENTINEL_SYMBOL -> asymmetric_bce (forces recompile); driver + runtime envelopes + dispatch branches. Only reduction="mean" runs on GPU; other reductions fall back to the numpy reference. Gotcha fixed: the loss ops' Graph IR name is tessera.loss.{z_loss,asymmetric_bce} (not tessera.*), so the envelope + dispatch branch use the .loss. prefix. GPU matches the reference at ~1e-8 (incl. large-logit stability); both flip execution_mode artifact -> metal_runtime. test_apple_gpu_ldt_loss_ops.py (18) + buffer-pool RAII guard green; runtime_abi regenerated; ruff/mypy/gates pass. Lattice benchmark: 4 -> 2 remaining apple_gpu artifact rows (load_balance_loss + masked_categorical next). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ll LDT artifact rows Promotes the last two LDT/MoE-aux ops from numpy-fallback (metal_artifact) to metal_runtime via MPSGraph subgraphs (argMax / oneHot / select), grounded in the SDK headers (Decision #27). * apple_gpu_runtime.mm: mpsg_run_load_balance_loss_f32 (argMax->oneHot->mean f_e, mean P_e, E*sum(f*P)) + mpsg_run_masked_categorical_f32 (mask>0 select(-inf) -> argMax -> int32 indices) + extern "C" wrappers with reference fallbacks; graph-cached. Stub parity. * _apple_gpu_backend gpu_load_balance_loss / gpu_masked_categorical wrappers + argtypes; _SENTINEL_SYMBOL -> masked_categorical; driver + runtime envelopes (tessera.loss.load_balance_loss, tessera.masked_categorical) + dispatch branches. load_balance: mean + default top-1 argmax run on GPU; explicit assignment / non-mean reduction fall back. masked_categorical: greedy (no rng key) runs on GPU; keyed sample falls back. GPU matches the reference at ~1e-8 (load_balance bounds uniform->1, conc->E) and exact greedy argmax for masked_categorical; both flip artifact -> metal_runtime. Lattice benchmark: ALL 8 apple_gpu rows are now optimized_native — 0 artifact_only (was 6). Integrated-step row (lattice_reasoning_compiler_artifact) stays artifact-only by nature (host-orchestrated branching composition); its gap list + registry-mismatch notes updated to reflect the 6 GPU-backed primitives. test_apple_gpu_ldt_loss_ops.py extended (28 total) + buffer-pool RAII guard + benchmark schema/core tests green; runtime_abi regenerated; ruff/mypy/gates pass. Relaxed two sentinel-hardcoding tests (the sentinel moves as kernels land). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…metal_runtime) Routes the float-output elementwise math + comparison primitives through the existing MPSGraph unary/binary opcode lane so a @jit(target="apple_gpu") call runs them natively (execution_mode == "metal_runtime") instead of the numpy artifact. No new C ABI symbols — extends the existing opcode switches. C (apple_gpu_runtime.mm) — mpsg_unary_node cases 12-29: sin/cos/tan/asin/acos/ atan/sinh/cosh/erf/erfc/expm1/log1p/reciprocal/sign/floor/ceil/round/trunc; mpsg_binary_node cases 7-10 (pow/atan2/mod[floor-mod]/floor_div) + 11-16 (eq/ne/lt/le/gt/ge → f32 0/1 mask via predicate cast). Matching host-reference fallbacks for the non-Darwin/CI path. All nodes grounded in the MPSGraph SDK headers (Decision #27). Python (runtime.py) — _APPLE_GPU_UNARY_OPCODES extended; new _APPLE_GPU_BINARY_OPCODES table (add/sub/mul/div reuse C nodes 0-5; scalar second-operand form supported) + _apple_gpu_dispatch_mpsgraph_binary + _apple_gpu_binary_numpy + an erf reference. Envelope union + dispatch branch. driver.py envelope mirror + the C++ TileToApple kRuntimeOps mirror (rebuilt tessera-opt) updated together so the G3 single-source enforcer stays green (its non-envelope sentinel moved tessera.add → tessera.gather). Accelerator-proof map auto-updates: proven 75 → 109, eligible 204 → 170. Tests: test_apple_gpu_elementwise_opcodes.py (41 — kernel numerics via the GPU dispatcher for all 34 ops + literal-@jit metal_runtime classification + floor-mod + scalar form). Dashboards regenerated; mypy host+linux + ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Completes the MPSGraph reduce/scan opcode switches: C (apple_gpu_runtime.mm) — mpsg_reduce_node case 7 = logsumexp (stable log(Σ exp(x−max))+max); mpsg_run_scan ops 2/3 = cumulativeMaximum/Minimum. Matching host-reference fallbacks (reference_reduce op 7; scan ops 2/3 via a running max/min). Nodes grounded in the MPSGraph SDK headers (Decision #27). Python (runtime.py) — _APPLE_GPU_REDUCE_OPS gains logsumexp→("reduce",7), cummax→("scan",2), cummin→("scan",3); the dispatcher's numpy fallback handles op 7 + np.maximum/minimum.accumulate. op_catalog OpSpecs for cummax/cummin (logsumexp already had one). driver.py reduction-envelope mirror + the C++ TileToApple kRuntimeOps mirror (rebuilt tessera-opt) updated together so the G3 enforcer stays green. Accelerator-proof map: proven 109 → 112, eligible 170 → 167. Tests: test_apple_gpu_reduce_scan_opcodes.py (12 — kernel numerics across axes + literal-@jit metal_runtime). PYTHON_API_SPEC reductions list extended; dashboards regenerated; reductions regression 95 passed; mypy host+linux + ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…+16 metal_runtime) Closes the rest of the batch-2 eligible set (numeric_helper / logical / reduction) on Apple GPU: C (apple_gpu_runtime.mm) — mpsg_unary_node cases 30-34: isfinite/isinf/isnan (predicate → f32 mask), logical_not (x==0), bitwise_not (int32 cast); mpsg_binary_node cases 17-22: logical and/or/xor (on a!=0,b!=0 → f32 mask) + bitwise and/or/xor (int32). Matching host-reference fallbacks. Nodes grounded in the MPSGraph SDK headers (Decision #27). Python (runtime.py) — unary/binary opcode tables + numpy fallbacks extended; max/min added to the reduce lane (reduce-max/min, ops 2/3) + op_catalog OpSpecs; clamp/clip and where routed via NEW compose dispatchers that chain the GPU binary lane (clamp = max(min(x,hi),lo); where = c*a + (1-c)*b) — no new C symbol. _APPLE_GPU_COMPOSE_OPS envelope + dispatch branches. driver.py mirror + the C++ TileToApple kRuntimeOps mirror (rebuilt tessera-opt) updated together (G3 enforcer green). Accelerator-proof map: proven 112 → 128, eligible 167 → 151. The opcode lanes (unary/binary/comparison/reduce/scan) + predicate/logical/bitwise/compose are now fully closed — every remaining eligible op is a bespoke kernel (losses, attention, norms, pooling, quant), not an opcode. Tests: test_apple_gpu_predicate_logical_opcodes.py (18). Dashboards regenerated; batch-1/2 + enforcer + map regression 132 passed; mypy host+linux + ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ph FFT) Builds the FFT/spectral kernel class flagged 'special' in s_series_accelerator_proof.md. MPSGraph ships native FFT (macOS 14+, MPSGraphFourierTransformOps.h, grounded in the SDK headers per Decision #27): fastFourierTransform / realToHermiteanFFT / HermiteanToRealFFT. - C++ FFT lane in apple_gpu_runtime.mm: tessera_apple_gpu_fft_f32(mode,in,out, batch,n) — mode 0=fft 1=ifft 2=rfft 3=irfft, 1-D over the last axis, batched. Complex crosses the C ABI as INTERLEAVED real/imag f32 (identical layout to MPSDataTypeComplexFloat32 AND numpy complex64 — a plain reinterpret on both sides). Normalization matches numpy (forward unscaled; inverse 1/n via ScalingModeSize; roundToOddHermitean for odd irfft). Naive-DFT CPU fallback for ctx-not-ok / macOS<14; non-Darwin stub parity. - Python: _apple_gpu_backend.gpu_fft1d + ctypes signature; sentinel bumped. - runtime._apple_gpu_dispatch_spectral routes all 9: fft/ifft/rfft/irfft direct; dct/stft/istft/spectral_conv compose over them (host glue + GPU FFT, the heavy-FLOPs-on-GPU pattern); spectral_filter is elementwise. - Envelope: _APPLE_GPU_SPECTRAL_OPS on the 'spectral' lane; runtime handler; C++ kRuntimeOps (.inc) regenerated (166->175 ops); tessera-opt rebuilt; G3 enforcer confirms the pass tags all 9 metal_runtime. - accelerator_proof: spectral category default special->eligible; per-primitive all 9 flip to proven (special 29->20; the remaining 20 are device RNG). Numerically validated vs numpy.fft at fp32 tol (~1e-6) incl. odd-length irfft, off-last-axis, and the 4 composites vs the host reference. New test_apple_gpu_spectral.py (envelope membership, accelerator-proof flip, fft/ifft/rfft/irfft x shapes, composites, @jit metal_runtime). mypy/ruff clean; 16 generated docs + .inc all in sync. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… gate Closes the numerical-proof discipline gap for Apple GPU `fused` manifest rows that had genuine dedicated GPU execute-compare tests but no wired fixture. 21 ops (GA/Clifford x17, complex x2, EBM x2) now carry their execute_compare_fixture — each verified (Decision #27) to run the op's kernel and assert_allclose vs a numpy/GA reference, re-run green on Metal. Fixes a latent bug: manifest_for's clifford_/ebm_/complex_ early-return paths bypassed _attach_numerical_fixtures, so those domains could never have received a fixture; the early returns now attach like the main path. (test_ga_backend_manifest updated — its "manifest_for == clifford_manifest_for" invariant encoded that very bypass.) New manifest-level gate test_apple_gpu_numerical_proof_discipline.py asserts every Apple GPU fused/hardware_verified row has a fixture or is on an explicit allowlist (ebm_self_verify / ebm_langevin_step / kv_cache_read — no dedicated GPU execute-compare), plus a stale-allowlist guard and a hardware_verified- implies-fixture lock. Regenerated test_coverage. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The macOS 27.0 SDK ships the low-precision tensor types M3's contract anticipated (doc dump, Decision #27 authority tier 2): MTLTensorDataType.{float8e4m3,float8e5m2,float4e2m1,float8ue8m0,int2,uint2}, plus the multi-plane tensor machinery (MTLTensorAuxiliaryPlaneDescriptor. blockFactors, auxiliaryPlanes, MTLTensorBufferAttachments, per-plane getBytes/replace) that is the runtime image of a ScaleLayout: one data plane (element dtype) + one auxiliary scale plane (float8ue8m0 for MX / float8e4m3 for NVFP4) whose blockFactors encode the block size. Confirmed on-machine (Decision #27): this Mac is macOS 26.5.1 / SDK 26.5, whose MTLTensorDataType tops out at Int4/UInt4 (@26.4) with no float8/float4/ e8m0 and no auxiliaryPlanes -- so the *execution* path stays gated on a 27.0 SDK; no new silicon needed. Add the hardware-free bridge in microscaling.py: mtl_tensor_data_type(dtype) -> MetalTensorType(swift_case, mtl_symbol, min_macos) and metal_plane_plan( fmt, shape) -> MetalPlanePlan(element, aux_planes, min_macos). Availability is honest per type (fp8/fp4/e8m0/int2/uint2 -> 27.0; int4/uint4 -> 26.4; int8/f32 -> 26.0); per-tensor int8 needs no aux plane; MX/NVFP4 emit one scale plane whose scale_shape round-trips ScaleLayout.scale_shape (one source of truth). Validates the contract maps 1:1 onto the concrete API and gives the future runtime its per-plane dtype + blockFactors target. Guard: tests/unit/test_microscaling_metal_bridge.py (7). 259-test low-precision sweep green; generated-doc drift clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ne displacement tessera.transpose now runs on Metal via transposeTensor:permutation: (SDK-header grounded, Decision #27) instead of the numpy-reference fallback. The first "real_gap_structural" displacement from the Phase 2 disposition map. - .mm: mpsg_run_transpose helper + tessera_apple_gpu_mpsgraph_transpose_{f32,f16}. Transpose is value-preserving N-D data movement, so f32 goes native and f16/bf16 share the 2-byte raw path. Generic N-D permuted-stride host fallback for non-Darwin / GPU-miss; stub parity. - envelope: _APPLE_GPU_TRANSPOSE_OPS (first-class runtime op) + "transpose" lane. - runtime: _apple_gpu_dispatch_transpose (axes default reversed; ctypes dims/perm; _apple_gpu_run_checked host fallback). - C++: apple_runtime_ops.inc regenerated + tessera-opt rebuilt (enforcer + drift gate pass). A single-op @jit(apple_gpu) transpose now reports native_gpu / metal_runtime (was fallback_eager). Caveat: a transpose mid-program still demotes to artifact_only until _apple_gpu_chain_kind learns a general "all ops GPU-capable → per-op metal" recognizer (the next structural step). Guards: tests/unit/test_apple_gpu_transpose.py (7 — 2D/3D/4D + explicit permute, f16, jit, no-fallback-on-Metal). Verified: 7 dedicated + 2346 broad apple_gpu tests pass; mypy clean; .inc enforcer + drift gates pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…y), flag bf16 bugs User confirmed ROCm 7.2.2 supports gfx1151. Grounded against community evidence (Decision #27): ROCm 7.2.x runs on the Strix Halo iGPU in practice (ollama, llama.cpp known-good stack, TinyComputers 7.0->7.2 guide), generally enumerating without HSA_OVERRIDE on 7.2 — so the previous "partial/unofficial, may need HSA_OVERRIDE" caveat was too pessimistic. Corrected, but kept honest: gfx1151 on 7.2.x is community/nightly support, not AMD's official matrix (official = 6.4.4), and Tessera's 7.2.3 pin (>= 7.2.2) comfortably covers it. Added a load-bearing caveat: documented gfx1151 bf16 correctness bugs (ROCm#6034, "5 critical bf16 bugs") directly affect Stage D's first proof (a bf16 WMMA GEMM) — mitigation is to bring up f16/fp32 WMMA combos first and cross-check bf16 mismatches against the upstream bug list, which the execute-and-compare oracle catches by construction. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…t Tables) Two grounded corrections to apple_target.py (Decision #27 — ground every Apple capability against the Feature Set Tables, not memory): 1. Apple7 (M1-series) DOES have native MTLDataType.bfloat (Apple6+) and SIMD-scoped matrix multiply / simdgroup_matrix (Apple7+). The table previously marked both not_supported on Apple7, which wrongly forced M1 down the bf16 host-upcast path and excluded it from the simdgroup_matrix GEMM lane. 2. Chip→family map: M3 AND M4 are both Apple9; M5 is Apple10; there is no Apple11 family. The enum/comments previously had M4=Apple10 / M5=Apple11. runtime.py docstrings on _apple_gpu_supports_native_bf16 + the matmul dispatch bf16 lane updated to match (bfloat native on every modeled arch incl. M1). 82 apple_target/feature-limit tests green; ruff + mypy clean; generated-doc drift in sync. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ck (llc-verified) Completes the AMD Stage-A emit track (A1–A4). emit_wmma_rdna4_llvmir + wmma_intrinsic_rdna4 + validate_wmma_rdna4_structure establish the RDNA 4 (gfx1200/1201) WMMA path, grounded by `llc` on this host across f16/bf16/fp8_e4m3/fp8_e5m2. Decision #27 correction (grounded against LLVM 22 IntrinsicsAMDGPU.td + llc on this host): the earlier "gfx12 v2 ABI = extra format/reuse operands" assumption was WRONG. The mods/reuse ABI (i1 A_mod / i16 C_mod / reuse flags, wmma.f32.16x16x32.f16) is gfx1250/1251 — a LATER arch, not RDNA 4. RDNA 4 actually: - keeps the plain 3-arg wmma(A,B,C) ABI; - uses DENSER <8 x elem> fragments (gfx11 is <16 x elem> — RDNA 4 drops the wave32 lane 0-15 → 16-31 duplication); - adds native FP8/BF8: fp8_e4m3→fp8.fp8, fp8_e5m2→bf8.bf8 → v_wmma_f32_16x16x16_{fp8_fp8,bf8_bf8}. Cross-checked: the FP8 intrinsic "Cannot select" on gfx1151, confirming the unlock is genuinely RDNA-4-only. Corrected the misleading gfx1200 comment in the constants block. 7 new tests (5 rung-3, run here on gfx1200 via Homebrew LLVM 22). Honesty ceiling: single-intrinsic ABI proof (the gfx11 path's starting point); RDNA 4's GEMM/operand-layout/threadgroup generalizations are follow-ons, and its denser VGPR layout makes the D→C mapping its own grounding job. Numbers wait for silicon (rungs 6-7). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Adds the gfx1250/1251 WMMA path — the "v2" mods/reuse ABI, grounded by
`llc` (LLVM 22 AMDGPU) + LLVM IntrinsicsAMDGPU.td, distinct from RDNA 4.
Emitter (rocdl_emit.py): emit_wmma_gfx1250_llvmir + wmma_intrinsic_gfx1250
+ validate_wmma_gfx1250_structure. Grounded differences from gfx11/RDNA4:
- K is DOUBLED — f16/bf16 are 16x16x32 (A/B = <16 x elem>).
- 5 extra immediate operands — wmma(i1 A_mod, A, i1 B_mod, B, i16 C_mod,
C, i1 a_reuse, i1 b_reuse) (all ImmArg; passed 0 = the plain product).
- native bfloat — bf16 is <16 x bfloat>, NOT the <_ x i16> bit-pattern
gfx11/RDNA4 require.
llc-verified on gfx1250 AND gfx1251 for f16/bf16; cross-checked the K=32
intrinsic "Cannot select" on gfx1200 (proves it's a distinct arch class).
FP8 (16x16x64/128, the ModsC ABI — (A,B,i16 C_mod,C,reuse,reuse), no A/B
negate) scoped out as a documented follow-on like iu4 was for RDNA4.
Target profile (rocm_target.py): GFX_1250/GFX_1251 in the AMDArch enum +
WMMA variants {16x16x32, 16x16x64, 16x16x128} + wave32 + gfx125x arch
strings (all grounded), empty MFMA (mutual-exclusion holds across all 9
arches). Per Decision #27, the fields with no gfx1250 ISA source —
LDS bytes, occupancy, and non-WMMA features — are marked PROVISIONAL /
"tba" rather than fabricated; the grounded executable surface is the
llc-verified emitter.
16 new tests (4 rung-3 run here on gfx1250/1251 via Homebrew LLVM 22).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ogy) Two drifts surfaced while reviewing what's open for the Apple compiler: 1. MASTER_AUDIT Apple row was stale — its "still open" listed items that have landed (binding specs, feature-limit-guided lowering, production packaged kernels, one-command-buffer JIT — all closed per APPLE_AUDIT.md's now-empty "Open Work"). Rewrote the row: Apple CPU+GPU execute natively and Apple is now the reference impl of the shared KernelEmitter/Runner/F4 framework (Workstream B); the real open frontier is performance + precision — a native simdgroup_matrix "steel-like" GEMM lane (clear-MPS), FP8/FP4/MX execution (macOS-27.0-SDK-gated, not hardware), and the world-class dims. 2. The `_APPLE_FEATURES` `metal4` key reads like a bug (`not_supported` on Apple7, whose SDK ships the MTL4 headers and whose machine runs Metal 4) but is NOT: it gates the MTL4 cooperative-tensor-op *runtime* (command model + packaged ML), deliberately M5-gated and test-pinned, distinct from "Metal 4 the API / MSL 4.0" (Apple7+) and from the simdgroup_matrix/bfloat/MTLTensor compute surface (separate keys, already "ready"). Added a precise note to the status legend + fixed the misleading APPLE7 "No Metal 4" comment so the terminology collision stops reading as a bug. Did NOT flip the value — whether M1 hardware-accelerates MTL4 cooperative tensor ops is an open device-probe question (task_fbb4d13b) and Decision #27 forbids flipping a capability claim without grounding. Confirmed task_fbb4d13b's original items (simdgroup_matrix/bfloat wrongly not_supported; M3/M4/M5 family map) are already fixed in the current code. Comment/doc-only — no capability values or tests changed. feature-limits 29/29, frontmatter 8/8, drift gate in sync (19), docs lint passed; ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ogy) (#304) * docs: fix Apple status drifts (MASTER_AUDIT row + metal4-key terminology) Two drifts surfaced while reviewing what's open for the Apple compiler: 1. MASTER_AUDIT Apple row was stale — its "still open" listed items that have landed (binding specs, feature-limit-guided lowering, production packaged kernels, one-command-buffer JIT — all closed per APPLE_AUDIT.md's now-empty "Open Work"). Rewrote the row: Apple CPU+GPU execute natively and Apple is now the reference impl of the shared KernelEmitter/Runner/F4 framework (Workstream B); the real open frontier is performance + precision — a native simdgroup_matrix "steel-like" GEMM lane (clear-MPS), FP8/FP4/MX execution (macOS-27.0-SDK-gated, not hardware), and the world-class dims. 2. The `_APPLE_FEATURES` `metal4` key reads like a bug (`not_supported` on Apple7, whose SDK ships the MTL4 headers and whose machine runs Metal 4) but is NOT: it gates the MTL4 cooperative-tensor-op *runtime* (command model + packaged ML), deliberately M5-gated and test-pinned, distinct from "Metal 4 the API / MSL 4.0" (Apple7+) and from the simdgroup_matrix/bfloat/MTLTensor compute surface (separate keys, already "ready"). Added a precise note to the status legend + fixed the misleading APPLE7 "No Metal 4" comment so the terminology collision stops reading as a bug. Did NOT flip the value — whether M1 hardware-accelerates MTL4 cooperative tensor ops is an open device-probe question (task_fbb4d13b) and Decision #27 forbids flipping a capability claim without grounding. Confirmed task_fbb4d13b's original items (simdgroup_matrix/bfloat wrongly not_supported; M3/M4/M5 family map) are already fixed in the current code. Comment/doc-only — no capability values or tests changed. feature-limits 29/29, frontmatter 8/8, drift gate in sync (19), docs lint passed; ruff clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * fix: annotate JEPA preserved dict[str, Any] for mypy 2.2.0 (unblock lint ratchet) CI's unpinned `mypy>=1.11.0` resolved to the freshly-released mypy 2.2.0, which narrows the `preserved` comprehension (over string-literal keys) to dict[Literal[...], Any] and then rejects `**preserved` into TargetOp's dict[str, Any] attrs — target_ir.py:1379, `[dict-item]`. Latent on main; my docs PR was just the first to run under 2.2.0 (local mypy 2.1.0 doesn't flag it). Fix: explicit `preserved: dict[str, Any]` annotation. Verified errors=0 under mypy 2.2.0 + numpy<2.0 (CI's exact stack) via a throwaway venv, and under local mypy 2.1.0. No behavior change. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
* J: roofline attainment — % of peak as the hot-path bar (W7) The E2 latency ratchet is a relative bar (did it get slower?). J adds the absolute bar the plan asks for: % of peak. - benchmarks/roofline.py: a grounded per-device peak table (rocm:gfx1151 = 29.7 TF fp32 / 59.4 TF fp16 / 256 GB/s, each with a `source` string deriving it from rocminfo CU/SIMD/clock + documented RDNA3 rates — Decision #27, auditable), FLOP/byte models (matmul 2*MNK, flash_attn 4*B*H*S^2*D), and achieved_tflops / pct_peak / evaluate_attainment. - The committed gfx1151 ratchet rows gain pct_peak + achieved_tflops + an attainment_floor (= pct_peak / margin, symmetric with the latency cap), computed from the EXISTING medians (no re-timing — ratchet caps unchanged). - perf_gate --attainment: gates a row that regresses below its floor (the absolute analog of the latency ratchet). record_hot_path_baseline annotates future baselines automatically. Honest scope: the ratchet median is end-to-end wall-clock (H2D/launch/D2H + tessera-opt shell-out), so pct_peak is an END-TO-END attainment — a lower bound on kernel efficiency. The current gfx1151 lanes sit at ~0.3-2.9%, so the metric's immediate value is making the headroom visible and giving it a ratchet floor. Proof: test_roofline_attainment.py (12) — FLOP/peak/attainment model, gate pass/fail/coverage, and the committed baseline self-passes `perf_gate --attainment`. ruff clean; drift gate green. Still open (J): kernel-isolated attainment (strip host overhead); NV sm_120 + Apple peak rows; floors that ratchet upward as lanes optimize. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> * PR #313 review: read latency_ms in attainment gate + package-safe roofline import - P1: measured ratchet-report rows carry latency_ms (as evaluate_ratchet reads), not median_ms — evaluate_attainment fell back to 0.0, so pct_peak returned None and every measured row false-failed on coverage. Read latency_ms with a median_ms fallback (baseline/self-check rows use median_ms). - P2: `from roofline import` broke under `python -m benchmarks.perf_gate` and `from benchmarks import perf_gate` (package context — roofline not on sys.path). Try `from benchmarks.roofline import` first, fall back to the script-dir import. Tests: a measured-row latency_ms gate case (pass at baseline latency, fail 10x slower) and a package-import case (from benchmarks import perf_gate + main --attainment on a latency_ms report). 14 pass; ruff clean; both script and `-m benchmarks.perf_gate --attainment` invocations verified. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> --------- Co-authored-by: gstoner <angstroms01@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ified_jit prose (#404) Three related pieces of cleanup on top of the Apple GPU C-ABI registry (#403): 1. Reconcile the ABI registry with the runtime (the registry net now PASSES on-device). The registry listed two device-limit probe symbols the runtime never exported, so the strict net (test_dylib_exports_resolve_every_registry _symbol) failed. Add them to apple_gpu_runtime.mm (+ the non-Darwin stub): - tessera_apple_gpu_max_threadgroup_memory_length ([device maxThreadgroupMemoryLength]; 0 = "use static floor") - tessera_apple_gpu_family_integer (raw MTLGPUFamilyApple* value, e.g. Apple7 == 1007; -1 sentinel) Enum values + selectors grounded in the on-machine SDK headers (Decision #27). apple_target.probe_apple_runtime_limits already binds these defensively. 2. Fix 4 pre-existing stale tests (all test-side, not product regressions): - test_apple_gpu_tiny_decode..._kv_cache: diagnostic wording unified to "KV-cache mutation ..." in c53010c; Decision #21 contract still holds. - test_apple_gpu_simple_moe...: tessera.moe gained a native Apple GPU compute lane (_APPLE_GPU_MOE_COMPUTE_OPS) → now metal_runtime; assert that + numerical proof. Renamed to ..._runs_metal_runtime. - test_apple_gpu_multi_op_with_non_gpu_op_stays_metal_artifact: moe is no longer a non-lane op; swap to tessera.flip (+ lane_for self-guard) to keep the conservative-residency-gate guard meaningful. - test_compile_loads_real_metal_package (pk1): c53010c's blanket "compiled"→"device_verified_jit" rename clobbered a repr assertion + docstrings; revert to "compiled". 3. Repo-wide follow-on to (2): c53010c intentionally renamed the *status token* compiled→device_verified_jit but over-reached into English prose. Revert every word-usage back to "compiled" (diagnostic message strings, NVRTC-/ HIPRTC-compiled, emit/* local variable names, docstrings/comments) while preserving all genuine status-token references (quoted "device_verified_jit" values, backtick doc-refs, and status-name prose like `native/ device_verified_jit`, `= device_verified_jit`, `fused (x86) / device_verified_jit (rocm)`). Each reversion git-verified against the pre-c53010c image. Generated dashboards regenerated (new ABI symbols, reworded notes, flip test count); drift gate clean (22 in sync). Apple suite: 2060 passed, 3 skipped. The emit/* variable renames are covered by test_kernel_cache / test_dynamic_shape_emit / test_spectral_candidates / test_tpp_candidates. Retest target: CUDA (sm_120) + ROCm (gfx1151) — this branch touches the nvidia/rocm emit + manifest prose and needs the ROCm-enabled tessera-opt those boxes have (the local Mac build is CPU+Apple only). Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
The eight quantize/dequantize ops were declared `same_as_first` -- a
single-tensor claim for ops that return `(codes, scale)`. Parking them in
DELIBERATELY_UNDECLARED recorded the problem; this fixes it.
Why it survived: the differential probe did `np.asarray(fn(...))`, which raises
on a tuple, so those ops were silently SKIPPED. A gate that skips what it
cannot express is indistinguishable from a passing gate -- the same failure
mode as the substring assertions this registry replaced. The multi-result gate
added here would have caught it.
Vocabulary:
* rules may now return a tuple of IRType;
* `_infer_result_types` exposes the full contract, while
`_infer_result_type` returns only the primary tensor and cannot express
multi-result ops at all -- pretending it could is how the false
declaration happened;
* `quantize_per_tensor` for fp8/fp6/fp4 (rank-0 scale);
* `quantize_per_block` for nvfp4, whose scale is one per 16-element block
along the last axis. Folding nvfp4 into the per-tensor rule would misstate
the micro-scaled format Blackwell actually implements -- wrong for exactly
the architecture that motivates the format.
* dequantize is single-result and declared `same_as_first`.
Verified against measured behavior: quantize_fp8 -> (4,16)/f32 + ()/f32;
quantize_nvfp4 -> (4,16)/f32 + (4,2)/f32 for a 32-wide last axis.
The remaining sub-byte question is recorded where it belongs -- as a BACKEND
PATH item, not a shape rule. The reference returns f32 codes (fake-quant), but
fp8_e4m3 / fp8_e5m2 / fp4_e2m1 / nvfp4 are canonical dtypes the Graph IR type
system can already carry; no lowering produces them. Filed per backend under
SUBBYTE-STORAGE-PATH-2026-08-03: NVIDIA owns it first (Blackwell has native FP8
and NVFP4, so "the backend upcasts anyway" is not the answer there); gfx1151 is
not-applicable by hardware (RDNA 3.5 has no FP8 WMMA, which the dtype contract
already records as `unsupported`); x86 is `emulated` and can carry storage
without native arithmetic; Apple's Metal capability is unverified and per
Decision #27 must come from on-machine SDK headers.
declared 184 | deliberately undeclared 19 | not yet examined 106.
Unit 14213 passed / 0 failed; ruff clean; 24 generated docs in sync.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…s, cross-backend sync Three review findings on PR #634, all confirmed against the repo: 1. Slot key was not unique under wave specialization: sm90_attention_plan places four consumer waves in one role, giving four writers per slot. The index gains a wave-in-role coordinate (compile-time known, so the size formula stays closed-form); single-writer election was rejected because intra-role imbalance is signal. 2. L2 aggregate stall counters cannot feed the §5 producer-attribution join — register accumulation destroys the endpoints the join needs. Slot contents are now level-dependent by contract: L2 = phase endpoints + aggregate stall (fraction only), L3 = regions one-to-one with wait intervals and publish endpoints; the stall join and realized critical path are L3 analyses by definition, and IKF-P4's gate says so. 3. AGENTS.md cross-backend rule: the plan proposes a shared schema, Tile IR ops, and a runtime buffer contract, so all four architecture queues now carry a disposition under sync key IKF-INTRA-KERNEL-CONTRACT-2026-08-27 — ROCm follow-up (owning lane, P0/P3), NVIDIA follow-up at P6 (%globaltimer; no gfx1151 evidence transfers), Apple follow-up at P6 (synthesizer/MLIR seam; in-kernel MSL timestamp primitive unverified — ground per Decision #27), x86 deferred with reason (TPROF-X86 owns CPU visibility; no consumer yet, #29). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Closes the last decline, and corrects the plan that produced it. The queue listed "masked loads" and "threadgroup staging" as two items. Metal's simdgroup_load has NO bounds predicate -- masking happens when the tile is copied in (`As[e] = (gr < M && gk < K) ? A[...] : 0`) and the MMA then reads a tile that is always in range. Ragged support REQUIRES staging; doing masked loads first was not possible. Adds threadgroup_alloc, landed with its producer rather than ahead of one. Its budget is carried IN THE IR as `budget_bytes` rather than compiled into the verifier: the limit is a device property (queried, [MTLDevice maxThreadgroupMemoryLength] = 32768 on this Apple7, family 1007 -- not recalled, per Decision #27), and baking it into C++ would make the verifier's answer depend on the host it runs on. That would break the toolchain-free property Decision #19 actually buys. As an attribute the check stays exact and the compiler stays host-free. The pass stages unconditionally rather than only on the ragged path -- one code path is easier to trust than two -- and pads the accumulator to whole tiles so a ragged edge writes into the pad instead of out of bounds. The epilogue copies back only the valid region. Mathematics checked three ways rather than asserted: * Zero padding is EXACT, not approximate: verified against a reference matmul for (17,13,23), (8,8,8), (1,1,1) and (31,9,7). A zero operand contributes nothing to the dot product, so valid outputs are unaffected and the padded tail is never copied out. * The load sits INSIDE the guard. Computing the address and selecting afterwards would still have read out of bounds, so the scf.if yields the value rather than selecting on one. * The budget is bytes, not elements: 16384 elements fits exactly in f16 and is double the limit in f32. Counting elements would accept both. Two fixture failures on the way, both correct consequences I had missed: staging moves the row stride from the global buffer (K, N) onto the tile (8), and threadgroup_alloc hoists out of the loop nest so it precedes simdgroup_fill in the output. Evidence: lit 443/443 on the M1 Max, 17 Python tests, 11 rejection cases including the two new budget ones, ruff clean, 29 generated docs in sync.
The production-serving follow-on to
MLAPagedDecoder:tessera.cache.MLABlockPagedCachemanages many concurrent sequences over a single physical block pool — the
core idea behind vLLM's PagedAttention.
What
Physical storage is a pool of fixed-size blocks:
Each sequence owns a block table (ordered physical block ids). Logical token
ilives at(block_table[i // block_size], i % block_size). Blocks areallocated from a free list on demand and returned on
free_sequence, so afinished request's pages are immediately reusable — no per-sequence contiguous
reservation, no external fragmentation.
API:
add_sequence/append(prefill + per-step, allocates blocks as needed) /decode/decode_batch(ragged, concurrent) /free_sequence, plus poolaccounting (
num_free_blocks,utilization,block_table,cache_bytes_per_token).Correctness
MLAPagedDecoderfor the same tokens — so the block-table gather is provenequivalent to a contiguous window.
test_non_contiguous_block_tablesinterleaves appends across two sequences soeach gets a non-contiguous set of physical blocks, then decodes correctly
through the indirection.
test_concurrent_ragged_decode_batchruns three sequences of lengths 3/9/16in one
decode_batch, each matching an independent decoder.Refactor
Factored a module-level
absorb_decode_one()(single-sequence absorbed decode +RoPE tables + GPU dispatch / numpy fallback) shared by
MLAPagedDecoderand theblock cache — one decode implementation, no duplication.
Tests
tests/unit/test_mla_block_paged_cache.py— 8 tests: block growth, free +page reuse, non-contiguous block tables, concurrent ragged
decode_batch,incremental two-sequence loop, pool exhaustion + recovery, utilization/footprint
(8.9× cache win at DeepSeek-V2 dims), lifecycle errors.
Verification (local, Apple Silicon)
torch-import error)Scope
Compute loops per sequence (lengths are ragged); the contribution is the
block-table memory manager. Batching same-length sequences through the kernel's
B>1path is a noted compute-side optimization.🤖 Generated with Claude Code