Skip to content

Apple GPU Tier-2: promote mla_decode to a real GPU kernel - #23

Merged
gstoner merged 1 commit into
mainfrom
apple-gpu-mla-decode-gpu
May 30, 2026
Merged

gstoner merged 1 commit into
mainfrom
apple-gpu-mla-decode-gpu

Conversation

@gstoner

@gstoner gstoner commented May 30, 2026

Copy link
Copy Markdown
Owner

Promotes tessera_apple_gpu_mla_decode_f32 from a host CPU reference to a real
on-GPU path (the compressed-KV phase) — the last, highest-effort Tier-2 item.

What

One cached MPSGraph fuses the whole MLA decode:

  1. latent down-projection c = X @ Wdkv[S_kv, D_lat]
  2. K/V up-projections K = c @ Wuk, V = c @ Wuv[S_kv, D_h]
  3. attention O = softmax((Q @ Kᵀ)·scale) @ V

The B·S_q query rows fold to a single matmul dimension (K/V are shared across
batch, 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_f32 stays as the fallback when Metal is
unavailable. ABI unchanged — purely a body swap.

Tests

  • tests/unit/test_apple_gpu_mla_decode_gpu.py6 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 — 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)

  • MLA suite (new + existing): 36/36 pass
  • apple_gpu + driver + runtime + mla sweep: 593 passed, 3 skipped
  • buffer-pool RAII gate + runtime ABI audit: pass (no new symbols)
  • mypy ratchet: clean (only the pre-existing environmental torch-import error)

CI on this repo is uniformly red on main (Python 3.8 to 3.11 matrix, missing
optional deps) — same state PRs #17 to #22 merged through. The local signal above
is green.

🤖 Generated with Claude Code

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>
@gstoner
gstoner merged commit 03bd626 into main May 30, 2026
8 of 26 checks passed
@gstoner
gstoner deleted the apple-gpu-mla-decode-gpu branch May 30, 2026 13:11
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>
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>
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant