Apple GPU Tier-2: promote mla_decode to a real GPU kernel - #23
Merged
Merged
Conversation
Promotes tessera_apple_gpu_mla_decode_f32 from a host CPU reference to a real on-GPU path (compressed-KV phase). - apple_gpu_runtime.mm: mpsg_run_mla_decode_f32 — one cached MPSGraph fuses the whole decode: latent down-projection c = X@Wdkv, the K/V up-projections K = c@Wuk / V = c@Wuv, then attention O = softmax((Q@Kᵀ)·scale)@v. The B·S_q query rows fold to a single matmul dim (K/V are shared across batch, so no batched-broadcast is needed); all matmuls accumulate in fp32. RAII buffer-pool acquires; graph cached by (B,S_kv,D_x,D_lat,S_q,D_h). The host reference_mla_decode_f32 stays as the fallback when Metal is unavailable. (mpsg_cache_get/put forward-declared since they are defined later in the TU.) ABI unchanged. - tests/unit/test_apple_gpu_mla_decode_gpu.py: 6 tests — batched (B>1) KV broadcast across 4 shapes vs numpy, decode-step (S_q=1), KV-shared-across-batch independence, and MPSGraph graph-cache reuse. The existing 30 MLA tests (test_mla_decode_fusion / _apple_gpu_mla_e2e / _mla_primitives / _mla_example) continue to pass against the GPU path. - docs: plan item 5 marked done (compressed-KV phase); decoupled-rope noted as the remaining follow-on. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This was referenced May 30, 2026
gstoner
added a commit
that referenced
this pull request
May 30, 2026
) Addresses the validate.yml CI redness — the parts that are real, in-scope, and clearly correct. lint lane (required check) -> green: - pyproject.toml: add torch.* to the no-stubs mypy overrides. torch is a soft import only (a hardware-smoke oracle behind try/except; Decision #23 keeps Tessera torch-free), so "missing stubs for torch" is not a real error. mypy ratchet now reports 0 errors (was 1). - cache/resident_decode.py: drop the unused `import math` (ruff E F401). docs lint (Python Quality lane) -> removes false positives: - scripts/lint_docs.py: a trailing location annotation on a path is a pointer, not part of the filename — strip `:N`, `:N+`, `:N-M`, `:N–M`, `:~N` before the existence check (the path:line convention is documented). Skip pytest node ids (`test_x.py::TestY`). Resolve extension-less module references (`python/tessera/ops`) against `.py` / `.pyi` / package `__init__`. (~10 pre-existing false positives cleared; the remaining ~36 are genuine doc-debt — planned files, literal placeholders, moved files — unrelated to this work.) unit lane -> Apple-GPU tests skip on non-Darwin: - tests/unit/conftest.py: Apple GPU is a macOS-only backend (Metal/MPSGraph); its runtime-execution tests need a real Metal device and cannot pass on the Linux runner (no Metal; the stub returns degenerate half-precision values). Like the lit/sanitizer lanes are opt-in, these 39 execution test files now SKIP on non-Darwin instead of failing. Platform-agnostic Apple tests (target-IR text, backend manifest, pass-order, buffer-pool source scan) are excluded from the skip set and keep running on Linux. Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 14, 2026
…oise marker + MSA layout hardening (#74) * cgg What landed: score_combine is now compiler-executable: verifier: [TesseraOps.cpp (line 579)](/Users/gregorystoner/dev_project/tessera/src/compiler/ir/TesseraOps.cpp:579) ODS contract: [TesseraOps.td (line 1061)](/Users/gregorystoner/dev_project/tessera/src/compiler/ir/TesseraOps.td:1061) tessera-to-linalg lowering: [TesseraToLinalgPass.cpp (line 135)](/Users/gregorystoner/dev_project/tessera/src/transforms/lib/TesseraToLinalgPass.cpp:135) lit tests for lowering and invalid shapes: [score_combine_lowering.mlir (line 6)](/Users/gregorystoner/dev_project/tessera/tests/tessera-ir/phase2/score_combine_lowering.mlir:6) Added CGG benchmark harness: [cgg_benchmark.py (line 48)](/Users/gregorystoner/dev_project/tessera/examples/diffusion_guidance/cgg_benchmark.py:48) supports smoke/full sweeps, gamma grid, adapter strengths, aggregate JSON, fixed seed. Added guided_denoise_region Graph IR marker: ODS: [TesseraOps.td (line 1075)](/Users/gregorystoner/dev_project/tessera/src/compiler/ir/TesseraOps.td:1075) verifier: [TesseraOps.cpp (line 598)](/Users/gregorystoner/dev_project/tessera/src/compiler/ir/TesseraOps.cpp:598) visibility/negative lit fixtures added. Hardened MSA Phase 2.5 selected-block layout: msa_select_blocks now verifies block_ids: (B,Hkv,Sq,top_k) i64. negative fixture: [msa_verifier.mlir (line 68)](/Users/gregorystoner/dev_project/tessera/tests/tessera-ir/phase3/msa_verifier.mlir:68) docs updated: [docs/msa.md (line 61)](/Users/gregorystoner/dev_project/tessera/docs/msa.md:61) Added MSA Phase 3 CUDA/NVIDIA plan: [docs/msa_cuda_phase3_plan.md (line 9)](/Users/gregorystoner/dev_project/tessera/docs/msa_cuda_phase3_plan.md:9) CUDA13 KV-outer sparse attention contract fixture: [msa_kv_outer_sparse_attention.mlir (line 15)](/Users/gregorystoner/dev_project/tessera/tests/tessera-ir/phase3/cuda13/msa_kv_outer_sparse_attention.mlir:15) * review(cgg): fix spec-sync blocker + mypy ratchet + manifest CGG examples Code-review fixes on top of the "cgg" batch before merge: - BLOCKER: score_combine was in op_catalog OP_SPECS but missing from docs/spec/PYTHON_API_SPEC.md → check_spec_sync.py failed (2 unit tests + pre-push gate). Added the spec row (matches the C++ lowering: base + gamma·delta). - mypy ratchet: diffusion_guidance.py had 3 errors — DiffusionSchedule._alpha_bar is a derived cache set via object.__setattr__ in a frozen dataclass, invisible to mypy. Declared it as field(init=False, repr=False, compare=False). Runtime unchanged; full ratchet back to 0 (256 files). - Consistency: registered both CGG examples (cgg_benchmark.py --smoke, cgg_diffusion_gemma.py) in examples_manifest.py as runnable entries; regenerated surface_status / docs_freshness dashboards. Verified: spec-sync + drift + docs-lint + examples-audit + mypy all green; score_combine / guided_denoise / msa_verifier lit fixtures pass; CGG + MSA + operator-registry + verifier-sprint tests green. score_combine numpy ref ≡ C++ lowering; CGG diffusion math validated; no torch/jax (Decision #23 OK). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 18, 2026
…e + MLX reference Read the Apple Metal Feature Set Tables (rev 2026-05-21) + the MLX docs. Added a grounded hardware-capability reference to APPLE_AUDIT.md: - GPU-family map: M1=Apple7, M2=Apple8, M3/M4=Apple9, M5=Apple10 (all M-series Metal 3 & 4). - This dev Mac (M1 Max = Apple7) ALREADY has the full Metal-4 ML compute surface: bf16 (Apple6), simdgroup_matrix (Apple7), SIMD reductions (Apple7), MTLTensor + ML encoding (Metal 4, Apple7), float atomics. What's gated past Apple7 is mostly graphics/sparse/texture, not compute -> Apple ML work is not blocked on newer silicon. - The one ML gap is toolchain, not hardware: FP8/FP4/MX MTLTensor dtypes + multi-plane scale machinery are macOS-27.0 SDK-gated (this machine is 26.5.1); the M1 gets them on a 27.0 SDK. Bridge: compiler/microscaling.py. - MLX = production Apple-Silicon reference (unified memory + custom MSL kernels + simdgroup_matrix), reference/oracle only, never a runtime dep (Decision #23). Re-confirmed the existing apple7 feature-set memory is current vs the 2026-05-21 table (no drift). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 18, 2026
… (autotune seeds) Acting on the user's MLX deep-dive: turn MLX's GEMM tile selection into Tessera schedule candidates "rather than hard-coded one-offs". Grounded by reading mlx/backend/metal/matmul.cpp (GEMM_TPARAM_MACRO + the align/axpby/swizzle function constants). python/tessera/compiler/apple_gemm_schedules.py: - MLX_SEED_TILES: the distinct (bm,bn,bk,wm,wn) tiles MLX selects -> the autotuner sweep set. Every tile's bm/bn/bk is an 8x8-fragment multiple, so it feeds emit_steel_gemm_msl directly (the integration test proves a 64x64 seed -> 64 output fragments, validated). - select_seed_tile(device_class, dtype, transpose_b, large): MLX's heuristic seed (the autotuner warm start), matching the grounded matmul.cpp branches. - GemmScheduleAxes + schedule_axes_for: the orthogonal knobs MLX expresses as Metal function constants (align_M/N/K, do_axpby alpha*AB+beta*C epilogue, swizzle_log). Recorded the broader 5-area MLX mining checklist (GEMM/attn seeds [this], feature probes, allocator/residency, microscaling oracle, custom-MSL hook) in APPLE_AUDIT.md + memory. MLX = production policy source, not kernels to copy (Decision #23). 11 tests; mypy + drift clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 19, 2026
…chet Multimodal model-class workstream: reference contracts + IR lowering for the media/vision and JEPA latent-prediction families, wired into MiniMax-M3. - models/multimodal.py — media processor / span / patch-grid / projected-embedding contracts + HF processor-config parsing. - models/vision_transformer.py — numpy reference vision tower + projector with explicit patch/merge/project shape contracts. - models/jepa.py — context/target masking, stop-gradient target encoder, EMA update, continuous-latent prediction + selective decode. - IR lowering: MEDIA_OPS / JEPA_OPS flow Graph -> Schedule (schedule.media.* / schedule.jepa.*) -> Tile (resource estimates) -> Target (per-arch contracts), all status="artifact_only" (no execution claim); ops.* Graph-IR contracts in __init__.py. - minimax_m3: processor_config / vision_transformer_config helpers + multimodal importer extension; moe_transformer_runtime updates. - fix: multimodal.processor_config_from_hf narrows the nested image_processor mapping before dict() so mypy no longer flags an Any|None arg (ratchet 0->0). Full unit suite green (10176 passed, 23 skipped) after the fix; ruff + doc-lint + generated-doc drift all clean. No torch/jax imports (Decision #23). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 19, 2026
Extends the multimodal model-class workstream (reviewed + full-suite-tested). - minimax_m3.py: build_multimodal_graph / verify_multimodal_graph + MiniMaxM3MultimodalGraph — a shape-only image/video tower + text/media splice contract (reuses diffusion_gemma.GraphNode), width/length validated. - minimax_m3_importer.py: HFTokenizerAdapter + optional HF-tokenizer import, apply_chat_template (jinja2 soft import), expected_hf_text_tensor_shapes_all_layers, load_text_runtime_weights[_from_safetensors] with tensor-name aliasing. - IR: JEPA selective_decode + train_step ops threaded through schedule/tile/target lowering (artifact_only); target_ir preserves JEPA contract attrs (decay / stateful / gating / branches / ema_update / latent_loss / gradient). - tests: +78 assert/raises across jepa / minimax contract / importer / multimodal. - build: mypy override ignore_missing_imports for the optional tokenizers / jinja2 soft imports (ratchet 0->0). Full unit suite green (10188 passed, 23 skipped); ruff + doc-lint + drift clean; no torch/jax imports (Decision #23). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 19, 2026
Adds a speculative-decode / prompt-cache rewind complement to evict_oldest: trim(n) removes unaccepted tail tokens and preserves the prefix, across all cache types (reviewed + full-suite-tested). - KVCacheHandle.trim / is_trimmable — zeros freed key/value/scale slots. - LatentKVCacheHandle.trim — zeros freed latent slots. - MLAPagedDecoder.trim — delegates to latent+rope caches; leaves _abs_base unchanged (surviving prefix keeps the same absolute RoPE positions). - MLABlockPagedCache.trim — zeros the partial-block tail, returns fully-emptied physical blocks to the shared pool, and truncates the block table. - tests: +23 assert/raises (test_kv_cache_handle, test_mla_block_paged_cache). - docs: apple_backend_capability_roadmap, apple_gpu_mlx_ecosystem_survey, tessera_inference_serving_plan. Full unit suite green (10194 passed, 23 skipped); ruff + mypy ratchet (0/0) + doc-lint + drift clean; no torch/jax imports (Decision #23). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jun 21, 2026
…e + profiler trace Combined landing of two in-progress efforts plus mypy/lint cleanup. DiffusionBlocks ideas (arXiv:2506.14202) — hardware-free, pure-numpy, test-gated: - compiler/diffusion_schedule.py: equi-probability (CDF) noise-band partition (+γ overlap), EDM preconditioning (c_skip/c_out/c_in/c_noise), σ-weighted loss w(σ), Karras ρ=7 inference schedule. Standalone inverse-normal-CDF (Acklam + Halley), no SciPy (Decision #23). - compiler/denoise_reference.py: classification-as-denoising reference model (band dispatch + EDM precond + Euler ODE sampler) with metamorphic self-consistency invariants — Evaluator conformance fixture. - pipeline_planner: decoupled-stage (local-objective) schedule — zero bubble, no cross-stage activation dependency. - checkpoint: CheckpointPolicy.DECOUPLED_BLOCK + peak-activation ≈ 1/num_blocks (structural lever, orthogonal to recompute). - primitive_coverage: edm_precondition / edm_loss_weight / equiprob_band_partition / karras_sigma_schedule rows (+ fp32 numeric_policy); generated dashboards regenerated (drift-gate green). - docs/audit/roadmap/decoupled_stage_pipeline.md design note. CDNA3/MI300X attention cost-model spine + ROCm backend (see BACKEND_AUDIT.md): MFMA accumulator-footprint cost model, LDS XOR-swizzle IR emit, decoupled vmcnt/lgkmcnt waits, target-parametric wave specialization, tail-KV split (flash-decoding) plan, chiplet/XCD-aware grid mapping, canonical rounding.py (RTNE/RTNA/RTZ sweep), LDS-budget-aware attn tile sizing. Profiler provider-trace/status spine + tools/profiler updates; new fixtures. Cleanup: fixed 2 mypy errors in profiler_provider_status.py / profiler_trace_merge.py (explicit dict[str, Any] annotation; narrow provider.get summary before .get). mypy clean across 336 files; ruff clean; full Python test delta green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
pushed a commit
that referenced
this pull request
Jun 25, 2026
#23) The shipped runtime is standalone — torch/transformers are reference vocabularies only (test/benchmark numerical oracles + a soft try/except GPU smoke import in gpu_smoke.py), never imported by python/tessera/ at runtime. Declaring them in [project.dependencies] forced pip to try resolving heavy, unused wheels and emitted a dependency-resolver conflict in CI ("tessera 0.1.0 requires torch>=2.0.0, which is not installed"). Move both to a new opt-in `reference` extra (pip install tessera[reference]); required deps are now numpy/scipy/matplotlib/pyyaml/click/rich/tqdm. test_standalone_compiler_roadmap.py + test_static_analysis_baseline.py stay green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This was referenced Jun 27, 2026
gstoner
added a commit
that referenced
this pull request
Jul 2, 2026
Fold the AOCL-DLP reference (amd/aocl-dlp — AMD's BLIS-family DL primitives: low-precision GEMM/batch GEMM, pre/post-ops, INT4/FP16, symmetric quant, OpenMP) into the north-star plan: - Theory Tier-3 list: add it to the x86 line (CPU analog of cuBLAS/rocWMMA). - Refactor Plan C1: the x86 TargetPlugin registers AOCL-DLP as a Zen-family Tier-3 candidate — AVX512-based (fits the Zen 5 fleet box, no AMX), fills the x86 OpenMP + INT4/FP16 gaps, opt-in behind a build flag (BLAS-family lib like Accelerate, Decision #23-clean), arbiter-selected only where measured faster; license check before a shipped lane. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Jul 2, 2026
* Add compiler north-star plan pair + fix sm_120 capability drift Compiler direction (docs): - New paired plan + theory: COMPILER_THEORY_OF_OPERATION.md (three-tier kernel model, accuracy-budgeted measured arbiter, three-system fleet, W1-W8 scope register) and COMPILER_REFACTOR_PLAN.md (workstreams A-E spine + F-K world-class, coordination + §9 source-verified seam verdicts). - Reassess OPTIMIZING_COMPILER_PLAN.md: F0-F5 landed on Apple; rewrite F6 (the backend-build seam) since its "Mac can't run CUDA/ROCm" premise is dead (Strix Halo gfx1151 + NR2 Pro sm_120 now execute) and scope the anti-goal. - Dynamic-shapes decision pulled into the spine: symbolic-dim-aware KernelEmitter/TargetPlugin API + shape-bucket arbiter key, bucket-specialize first. - Wire the north star into MASTER_AUDIT.md, docs/audit/README.md, README.md, docs/README.md, and CLAUDE.md (new Decision #28 + reference row). sm_120 capability fix (code): - gpu_target.py: route all coarse capability properties (supports_wgmma / tcgen05 / tmem / cta_pairs / mbarrier / tma / block_scaled_mma) through the authoritative _CUDA_13_3_FEATURES matrix via cuda_feature_status instead of isa >= SM_x. Fixes consumer Blackwell sm_120 wrongly reporting Hopper wgmma + datacenter tcgen05/TMEM/CTA-pairs as supported (it is NOT a superset of sm_100; its matrix path is mma.sync). Also fixes sm_120 wrongly inheriting the SM_90 FA-4 attn default in jit.py. - test_gpu_target.py: bug-encoding test_sm120_runtime_arch becomes a test_sm120_consumer_blackwell_capabilities regression guard; drop stale "rubin_placeholder" naming. Doc drift cleanups: - CLAUDE.md: fix stale "Execution reality" (gfx1151 + sm_120 now execute). - docs/README.md: ROCm row artifact-only -> gfx1151 hardware-runtime. - CANONICAL_API.md: sm_120 "Rubin placeholder" -> Blackwell consumer; correct the WGMMA column + footnote to match the fixed properties. Gates: mypy clean, generated-doc drift in sync, doc lint passed, test_gpu_target + test_audit_docs green. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Sync string-alias capability path with sm_120 feature matrix Address PR review: GPUTargetProfile(SM_120) was fixed, but the string-alias capability path (get_target_capability / backend_capabilities) still advertised invalid consumer-Blackwell features. - capabilities.py: drop wgmma / wgmma_sparse / tcgen05 / tcgen05_pair / tmem from nvidia_sm120.features (consumer Blackwell is NOT a superset of datacenter sm_100; FP4 goes through mma.sync.block_scale). Now mirrors the cuda_feature_set(SM_120) "ready" flags. - test_compiler_capabilities.py: add test_nvidia_features_match_cuda_matrix — a single-source-of-truth guard asserting no NVIDIA capability entry advertises a feature the CUDA matrix marks not_supported, plus a positive lock that sm_120 excludes the datacenter/Hopper flags. Prevents this drift from recurring. - Regenerate test_coverage dashboards (deterministic negative_refs count shift from the added guard test). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> * Note AOCL-DLP as the x86 Tier-3 candidate for Zen Fold the AOCL-DLP reference (amd/aocl-dlp — AMD's BLIS-family DL primitives: low-precision GEMM/batch GEMM, pre/post-ops, INT4/FP16, symmetric quant, OpenMP) into the north-star plan: - Theory Tier-3 list: add it to the x86 line (CPU analog of cuBLAS/rocWMMA). - Refactor Plan C1: the x86 TargetPlugin registers AOCL-DLP as a Zen-family Tier-3 candidate — AVX512-based (fits the Zen 5 fleet box, no AMX), fills the x86 OpenMP + INT4/FP16 gaps, opt-in behind a build flag (BLAS-family lib like Accelerate, Decision #23-clean), arbiter-selected only where measured faster; license check before a shipped lane. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This was referenced Jul 11, 2026
gstoner
added a commit
that referenced
this pull request
Jul 14, 2026
* docs: fix drift vs source — 33 stale/minor docs corrected
Formal drift review (145 hand-written prose docs vs. source at HEAD; the 786-file
auto-generated ISA archive and script-owned dashboards excluded) found 6 STALE +
27 MINOR docs. The existing docs_freshness dashboard reported "0 stale" because it
only tracks last_updated dates, not whether content matches code — this pass
checked every concrete claim (API/symbol names, paths, signatures, status) against
source. All 33 corrected here; every fix re-verified against current source.
Highlights (STALE):
- architecture/Compiler/Tessera_Compiler_ScheduleIR_Design.md — fictional `sched.*`
dialect → real `schedule.*` (mesh/pipeline/stage ops); inspection examples →
JitFn.schedule_ir/tile_ir/target_ir (verified to exist).
- architecture/Compiler/tessera_tile_ir_documentation.md — `tessera_tile.*` →
real `tile.*` dialect; invented pass names → real TileIRLowering/WarpSpec/etc.
- architecture/compiler_gaps_1_3_5_plan.md — Gaps 1 & 5 marked "deferred" but
LANDED (StencilLoopMaterializePass, HaloTransportLowerPass); never-created
planned files annotated as such.
- architecture/workloads/attention-family.md — NVIDIA flash_attn "absent/matmul-
only" → device-verified (manifest + committed baseline); drifted line cites fixed.
- audit/roadmap/S_SERIES_GAP_CLOSURE_PLAN.md — undercounts (5/33 ODS ops, JVP
scatter-only) → real 11/33 + full structural JVP family.
- audit/stub_surface.md — regenerated (stale embedded counts).
Recurring MINOR themes fixed: execution-status lag ("Apple-only executes" →
sm_120 + gfx1151 now execute), effect lattice (4/5 levels → real 8-level lattice
in GRAPH_IR_SPEC + LANGUAGE_SPEC), version string (0.2.0 → 0.1.0 per tsr_version.h),
dead `TesseraPrivilegeError` → TesseraConstraintError, PyTorch-interop claims vs
Decision #23, and drifted file:line citations.
Also regenerates two script-owned dashboards for the pre-push gate:
docs_freshness.md (last_updated bumps) and test_coverage.csv (a pre-existing
selective_ssm count drift, 80→81, unrelated to the prose fixes). No source files
changed.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* NVIDIA sm_120 SSM replay-decode device-state lane
Adds the NVIDIA replay-decode state handle for selective-SSM streaming decode
(REPLAYSSM_PLAN): runtime.nvidia_ssm_replay_state_handle exposes append /
read_output / flush / step_block / submit_block_async over a persistent device
state, backed by the synthesized CUDA lanes in emit/nvidia_cuda.py
(_synthesize_ssm_replay_decode_cuda / _synthesize_ssm_replay_device_cuda,
NvidiaReplayDeviceState with async CudaEvent / CudaDeviceBuffer).
Verified live on RTX 5070 Ti (sm_120, CUDA 13.3): tests/unit/test_ssm_nvidia_replay.py
plus test_ssm_replay_route.py + test_delta_state_handle.py all green (38 passed).
ruff + mypy clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
* Address PR review on the SSM replay lane (4× P2)
CUDA-backed _Handle (nvidia_ssm_replay_state_handle) hardening + test gating:
- reset(): override to recreate the CUDA device context from the freshly-zeroed
S0 (was inheriting the base reset, which only zeroed the host mirror — a new
sequence could decode against the stale device checkpoint after a flush()).
- read_output(): apply the base handle's (B, N) reshape/size check on c_t before
NvidiaReplayDeviceState.decode(), which otherwise C-copies B*N floats from the
pointer and could over-read a mis-sized c_t instead of raising ValueError.
- submit_block_async(): validate [T,B,D] / [T,B,N] block shapes before the CUDA
enqueue (parity with step_block) — the async memcpy would otherwise over-read.
- test_ssm_nvidia_replay.py: gate the CUDA-only assertions on _cuda_host_ready()
— the factory test now asserts _device is not None on a CUDA host and the
reference-fallback (None) off-device; the async test skips off-device. Fixes
the failing `unit` check (unguarded _device assertion + async hard-require).
Verified on RTX 5070 Ti (sm_120): full test_ssm_nvidia_replay.py green (8),
reset/validation behaviors checked directly; off-device path simulated (factory
declines, read_output/reset/step_block fall back, async skips). ruff + mypy clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---------
Co-authored-by: angst <angstroms01@gmail.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Aug 7, 2026
Implements finding T3 from DIFFERENTIABLE_PROGRAMMING_REVIEW.md. Functions
defined implicitly as the root of F(x, θ)=0 don't decompose into elementary
ops, so ordinary autodiff can't reach through them; the implicit function
theorem gives their derivatives directly (Blondel & Roulet Ch. 10), built only
from the residual's JVP/VJP and a linear solve — constant memory, no solver
unrolling.
New autodiff/implicit.py:
* cg_solve — matrix-free conjugate gradient (Algorithm 8.1), fails closed on
non-convergence.
* ihvp — inverse-Hessian vector product via CG on the existing hvp, the map
Newton / natural-gradient needs without materializing the Hessian.
* root_vjp / root_jvp / custom_root — differentiate x*(θ): VJP solves Aᵀr=u
then -Bᵀr; JVP solves At=-Bv, with A=∂_xF, B=∂_θF (§10.4).
* adjoint_state_grad — the constrained-objective specialization (§10.5):
∂₁cᵀr=-∇₁L then ∇₂L+∂₂cᵀr.
Reuses the CG/GMRES vocabulary compiler/solver_config.py already names; no
SciPy dep (Decision #23). Exported from tessera.autodiff. Unblocks second-order
optimizers, energy/Langevin samplers, and the manifold/OT work — none of which
had a differentiable fixed point.
Every piece verified against closed-form or finite-difference truth
(test_implicit_diff.py): CG vs np.linalg.solve, IHVP vs (2·diag)⁻¹, sqrt/linear
roots, adjoint state vs the reduced objective's gradient, and honest raises on
singular ∂_xF.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
gstoner
added a commit
that referenced
this pull request
Aug 7, 2026
Implements finding T3 from DIFFERENTIABLE_PROGRAMMING_REVIEW.md. Functions
defined implicitly as the root of F(x, θ)=0 don't decompose into elementary
ops, so ordinary autodiff can't reach through them; the implicit function
theorem gives their derivatives directly (Blondel & Roulet Ch. 10), built only
from the residual's JVP/VJP and a linear solve — constant memory, no solver
unrolling.
New autodiff/implicit.py:
* cg_solve — matrix-free conjugate gradient (Algorithm 8.1), fails closed on
non-convergence.
* ihvp — inverse-Hessian vector product via CG on the existing hvp, the map
Newton / natural-gradient needs without materializing the Hessian.
* root_vjp / root_jvp / custom_root — differentiate x*(θ): VJP solves Aᵀr=u
then -Bᵀr; JVP solves At=-Bv, with A=∂_xF, B=∂_θF (§10.4).
* adjoint_state_grad — the constrained-objective specialization (§10.5):
∂₁cᵀr=-∇₁L then ∇₂L+∂₂cᵀr.
Reuses the CG/GMRES vocabulary compiler/solver_config.py already names; no
SciPy dep (Decision #23). Exported from tessera.autodiff. Unblocks second-order
optimizers, energy/Langevin samplers, and the manifold/OT work — none of which
had a differentiable fixed point.
Every piece verified against closed-form or finite-difference truth
(test_implicit_diff.py): CG vs np.linalg.solve, IHVP vs (2·diag)⁻¹, sqrt/linear
roots, adjoint state vs the reduced objective's gradient, and honest raises on
singular ∂_xF.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Promotes
tessera_apple_gpu_mla_decode_f32from a host CPU reference to a realon-GPU path (the compressed-KV phase) — the last, highest-effort Tier-2 item.
What
One cached MPSGraph fuses the whole MLA decode:
c = X @ Wdkv→[S_kv, D_lat]K = c @ Wuk,V = c @ Wuv→[S_kv, D_h]O = softmax((Q @ Kᵀ)·scale) @ VThe
B·S_qquery rows fold to a single matmul dimension (K/V are shared acrossbatch, so no batched-broadcast is needed); all matmuls accumulate in fp32.
RAII buffer-pool acquires; the graph is cached by
(B,S_kv,D_x,D_lat,S_q,D_h).The host
reference_mla_decode_f32stays as the fallback when Metal isunavailable. ABI unchanged — purely a body swap.
Tests
tests/unit/test_apple_gpu_mla_decode_gpu.py— 6 tests: batched (B>1) KVbroadcast across 4 shapes vs numpy, decode-step (S_q=1), KV-shared-across-batch
independence, and MPSGraph graph-cache reuse.
test_mla_decode_fusion,_apple_gpu_mla_e2e,_mla_primitives,_mla_example) continue to pass — now against the GPU path.Scope
This is the compressed-KV phase. Decoupled RoPE (RoPE-carrying dims +
per-head split) is the remaining follow-on, noted in the plan.
Verification (local, Apple Silicon)
torch-import error)🤖 Generated with Claude Code