diff --git a/benchmarks/baselines/nvidia_sm120_solver_krylov_performance.json b/benchmarks/baselines/nvidia_sm120_solver_krylov_performance.json index 52c3ec84e..a1d58f9de 100644 --- a/benchmarks/baselines/nvidia_sm120_solver_krylov_performance.json +++ b/benchmarks/baselines/nvidia_sm120_solver_krylov_performance.json @@ -10,8 +10,8 @@ "mode": "cooperative_grid:end_to_end", "selected_route": "cooperative_grid", "timing_domain": "end_to_end", - "median_ms": 1.480445, - "max_latency_ms": 3.701112, + "median_ms": 1.57365, + "max_latency_ms": 3.934125, "reduction_ctas": 3, "iterations": 8, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -24,8 +24,8 @@ "mode": "cooperative_grid:device_event", "selected_route": "cooperative_grid", "timing_domain": "device_event", - "median_ms": 0.289926, - "max_latency_ms": 0.724816, + "median_ms": 0.292467, + "max_latency_ms": 0.731168, "reduction_ctas": 3, "iterations": 8, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -38,8 +38,8 @@ "mode": "cooperative_grid:end_to_end", "selected_route": "cooperative_grid", "timing_domain": "end_to_end", - "median_ms": 2.752441, - "max_latency_ms": 6.881103, + "median_ms": 2.642597, + "max_latency_ms": 6.606492, "reduction_ctas": 5, "iterations": 8, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -52,8 +52,8 @@ "mode": "cooperative_grid:device_event", "selected_route": "cooperative_grid", "timing_domain": "device_event", - "median_ms": 0.513587, - "max_latency_ms": 1.283968, + "median_ms": 0.538221, + "max_latency_ms": 1.345552, "reduction_ctas": 5, "iterations": 8, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -66,8 +66,8 @@ "mode": "cooperative_grid:end_to_end", "selected_route": "cooperative_grid", "timing_domain": "end_to_end", - "median_ms": 6.632047, - "max_latency_ms": 16.580118, + "median_ms": 7.182736, + "max_latency_ms": 17.95684, "reduction_ctas": 9, "iterations": 8, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -80,8 +80,8 @@ "mode": "cooperative_grid:device_event", "selected_route": "cooperative_grid", "timing_domain": "device_event", - "median_ms": 1.03271, - "max_latency_ms": 2.581776, + "median_ms": 1.015552, + "max_latency_ms": 2.53888, "reduction_ctas": 9, "iterations": 8, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -94,8 +94,8 @@ "mode": "cooperative_grid:end_to_end", "selected_route": "cooperative_grid", "timing_domain": "end_to_end", - "median_ms": 1.587581, - "max_latency_ms": 3.968952, + "median_ms": 1.60096, + "max_latency_ms": 4.0024, "reduction_ctas": 3, "iterations": 6, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -108,8 +108,8 @@ "mode": "cooperative_grid:device_event", "selected_route": "cooperative_grid", "timing_domain": "device_event", - "median_ms": 0.367206, - "max_latency_ms": 0.918016, + "median_ms": 0.296947, + "max_latency_ms": 0.742368, "reduction_ctas": 3, "iterations": 6, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -122,8 +122,8 @@ "mode": "cooperative_grid:end_to_end", "selected_route": "cooperative_grid", "timing_domain": "end_to_end", - "median_ms": 2.693219, - "max_latency_ms": 6.733047, + "median_ms": 2.542125, + "max_latency_ms": 6.355313, "reduction_ctas": 5, "iterations": 6, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -136,8 +136,8 @@ "mode": "cooperative_grid:device_event", "selected_route": "cooperative_grid", "timing_domain": "device_event", - "median_ms": 0.572154, - "max_latency_ms": 1.430384, + "median_ms": 0.394662, + "max_latency_ms": 0.986656, "reduction_ctas": 5, "iterations": 6, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -150,8 +150,8 @@ "mode": "cooperative_grid:end_to_end", "selected_route": "cooperative_grid", "timing_domain": "end_to_end", - "median_ms": 6.833705, - "max_latency_ms": 17.084263, + "median_ms": 7.147036, + "max_latency_ms": 17.86759, "reduction_ctas": 9, "iterations": 6, "correctness_gate": "known_solution_plus_fp32_true_residual", @@ -164,8 +164,8 @@ "mode": "cooperative_grid:device_event", "selected_route": "cooperative_grid", "timing_domain": "device_event", - "median_ms": 0.974483, - "max_latency_ms": 2.436208, + "median_ms": 0.611917, + "max_latency_ms": 1.529792, "reduction_ctas": 9, "iterations": 6, "correctness_gate": "known_solution_plus_fp32_true_residual", diff --git a/benchmarks/rocm/benchmark_rocm_adafactor.py b/benchmarks/rocm/benchmark_rocm_adafactor.py index 14fcac953..35a5b0d64 100644 --- a/benchmarks/rocm/benchmark_rocm_adafactor.py +++ b/benchmarks/rocm/benchmark_rocm_adafactor.py @@ -29,7 +29,7 @@ def artifact() -> rt.RuntimeArtifact: "op_name": "tessera.adafactor", "result": "output", "operands": ["parameter", "gradient", "row", "col"], - "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": 8}, }], }) @@ -47,7 +47,7 @@ def backward_artifact() -> rt.RuntimeArtifact: "result": "gradients", "operands": ["parameter", "gradient", "row", "col"], "out_cotangent": "dy", - "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": 8}, }], }) diff --git a/docs/audit/backend/apple/todo.md b/docs/audit/backend/apple/todo.md index 257d6d227..3d61f1e10 100644 --- a/docs/audit/backend/apple/todo.md +++ b/docs/audit/backend/apple/todo.md @@ -7,6 +7,32 @@ last_updated: 2026-08-29 --- # Apple compiler, exact-device, and performance plan +Cross-backend sync `ADAFACTOR-BIAS-CORRECTION-2026-08-30` — **shared numerical +policy changed; per-backend outcome below.** +`optim.adafactor_decay` makes the Adafactor second-moment decay step-dependent +(`b2_t = b2*(1 - b2^(t-1))/(1 - b2^t)`), removing an early-step update +inflation of 1/sqrt(1 - b2^t) — 31.6x at step 1, 10.0x at step 10, 1.26x at +step 1000 for the default beta2. The correction is applied HOST-SIDE as a +scalar decay, so **no kernel ABI moves**: every physical kernel already takes +`beta2` as a scalar and receives the effective value instead of the nominal +one. The flat op gained an optional `step` kwarg matching the `adam`/`adamw` +ABI beside it. + +Two contract details a backend owner needs to know. `state["v"]` now carries +the DEBIASED estimate rather than the raw EMA, so the state dict grew a +`v_representation` marker and a state without one is migrated on load rather +than misread. And an absent `step` is NOT treated as step 1 — `decay(b2, 1)` +is exactly 0, so defaulting would have made a stateful caller that never +passes one discard its own moments; such a caller keeps the legacy +uncorrected decay. + +*Apple outcome: not applicable.* The Apple backend exposes no Adafactor +kernel — neither the Accelerate CPU lane nor the MSL/MPS GPU lane registers +one, so there is no Apple code path this policy reaches. The reference +`optim.adafactor` runs as host numpy here like any other pure-Python lane and +is covered by `tests/unit/test_s10_optim.py`; nothing Apple-specific to +validate. + Cross-backend sync `P2-REVIEW-SHARED-PASSES-2026-08-29` — **15 shared MLIR passes changed; only the Mac's fixture set could be run.** The P2 code-review batch touched passes every backend lowers through: diff --git a/docs/audit/backend/nvidia/todo.md b/docs/audit/backend/nvidia/todo.md index f37c250e9..a16c21e06 100644 --- a/docs/audit/backend/nvidia/todo.md +++ b/docs/audit/backend/nvidia/todo.md @@ -7,6 +7,57 @@ last_updated: 2026-08-29 --- # NVIDIA compiler test-suite evaluation and rearchitecture +Cross-backend sync `ADAFACTOR-BIAS-CORRECTION-2026-08-30` — **shared numerical +policy changed; per-backend outcome below.** +`optim.adafactor_decay` makes the Adafactor second-moment decay step-dependent +(`b2_t = b2*(1 - b2^(t-1))/(1 - b2^t)`), removing an early-step update +inflation of 1/sqrt(1 - b2^t) — 31.6x at step 1, 10.0x at step 10, 1.26x at +step 1000 for the default beta2. The correction is applied HOST-SIDE as a +scalar decay, so **no kernel ABI moves**: every physical kernel already takes +`beta2` as a scalar and receives the effective value instead of the nominal +one. The flat op gained an optional `step` kwarg matching the `adam`/`adamw` +ABI beside it. + +Two contract details a backend owner needs to know. `state["v"]` now carries +the DEBIASED estimate rather than the raw EMA, so the state dict grew a +`v_representation` marker and a state without one is migrated on load rather +than misread. And an absent `step` is NOT treated as step 1 — `decay(b2, 1)` +is exactly 0, so defaulting would have made a stateful caller that never +passes one discard its own moments; such a caller keeps the legacy +uncorrected decay. + +*NVIDIA outcome: follow-up required.* `sm120_adafactor_*` receives the +effective decay through the existing scalar; `tests/device/nvidia/test_optimizer_reverse.py` +was migrated to pass `step`. Not executed on sm_120 in this batch — an exact-device +run of the optimizer lanes is owed. + +Cross-backend sync `P3-DEVICE-VERIFIED-2026-08-30` — **the two NVIDIA rows +owed by `P3-SOURCE-ONLY` are now measured, and one of them was a regression.** + +* `emit/nvidia_solver_krylov.py` `tsr_matvec` — the warp-per-row rewrite was + shipped on a reasoned access-pattern claim. Measured on an RTX 5070 + (sm_120), medians of 9 reps, device_event: **dense_cg 0.44-0.63x (a + REGRESSION of up to 2.3x)** and **dense_gmres 1.22-1.56x (a win)**. The + coalescing argument was correct and still lost, because a COOPERATIVE + launch caps the grid at what stays resident, so warp-per-row also buys 32x + fewer rows in flight. The solvers no longer share one matvec: + `tsr_matvec_scalar` for CG, `tsr_matvec_warp` for GMRES, with the table in + the source. Re-measured after the split: CG back to 1.00-1.06x of scalar, + GMRES keeps 1.24-1.56x. `benchmarks/baselines/nvidia_sm120_solver_krylov_performance.json` + was recorded with the OLD matvec and **passed throughout the regression** — + re-recorded at 15 reps / 5 warmup, and the ratchet now measures reality. +* `emit/nvidia_cuda.py` flash-backward cleanup — the 20 Krylov/solver device + tests and the flash-backward route tests pass on sm_120. An induced + allocation failure is still not exercised; that remains the honest gap. + +Also closed here: the `rc=5` invoke failure (the runtime dispatches scheduled +sm_120 matmuls by NAME PREFIX while the compiler named the kernel after the +caller's Graph function) and the sm_120 packager reading matmul epilogue +edges from `op.kwargs` when the verifier requires operands. Device suite: +**81 failed -> 5 failed / 844 passed.** The 5 remaining are 4 stale +shared-staging assertions (`__tessera_sm120_ab_stage_bf16`, pre-existing and +a routing question, not a test-editing one) and NCCL not being installed. + Cross-backend sync `P3-SOURCE-ONLY-2026-08-30` — **two rows are fixed in source and have never run on a GPU; they are this queue's to close.** The P3 batch changed two NVIDIA emitters with no CUDA host available: diff --git a/docs/audit/backend/rocm/todo.md b/docs/audit/backend/rocm/todo.md index 33551c0b4..bfcfa04d3 100644 --- a/docs/audit/backend/rocm/todo.md +++ b/docs/audit/backend/rocm/todo.md @@ -6,6 +6,47 @@ scope: ROCm backend implementation and exact-device proof --- # ROCm backend TODO +Cross-backend sync `ADAFACTOR-BIAS-CORRECTION-2026-08-30` — **shared numerical +policy changed; per-backend outcome below.** +`optim.adafactor_decay` makes the Adafactor second-moment decay step-dependent +(`b2_t = b2*(1 - b2^(t-1))/(1 - b2^t)`), removing an early-step update +inflation of 1/sqrt(1 - b2^t) — 31.6x at step 1, 10.0x at step 10, 1.26x at +step 1000 for the default beta2. The correction is applied HOST-SIDE as a +scalar decay, so **no kernel ABI moves**: every physical kernel already takes +`beta2` as a scalar and receives the effective value instead of the nominal +one. The flat op gained an optional `step` kwarg matching the `adam`/`adamw` +ABI beside it. + +Two contract details a backend owner needs to know. `state["v"]` now carries +the DEBIASED estimate rather than the raw EMA, so the state dict grew a +`v_representation` marker and a state without one is migrated on load rather +than misread. And an absent `step` is NOT treated as step 1 — `decay(b2, 1)` +is exactly 0, so defaulting would have made a stateful caller that never +passes one discard its own moments; such a caller keeps the legacy +uncorrected decay. + +*ROCm outcome: follow-up required.* The gfx1151 `adafactor_row|col|mean|update` +kernels take `beta2` as a scalar and are unchanged; `tests/unit/test_rocm_optimizer_compiled.py` +and `benchmarks/rocm/benchmark_rocm_adafactor.py` were migrated to pass `step`. +The host-free lanes pass here, but a gfx1151 optimizer execution comparing the +corrected trajectory against the numpy reference is owed. + +Cross-backend sync `P3-DEVICE-VERIFIED-2026-08-30` — **the batch FFT seam +written blind against the `.hip` signature is real, correct and faster.** +`RocmStockhamFFTCandidate.run_rows` was landed unexecuted, declining to the +per-row path when the image lacked the batch ABI — which is precisely the +state it was tested in. Measured here on gfx1151: the shipped image DOES +export `ts_fft_stockham_amd_hostptr_batch`, the candidate reports available, +and a (512, 256) transform runs on lane `rocm_stockham` with **max abs error +1.9e-05** against `numpy.fft` and **62.6 ms vs 237.7 ms per-row — 3.80x**. + +Also verified here: the rank-4 dropout `stream_offset` contract change. The +fixture that broke the PREVIOUS attempt at this row, +`phase3/streaming_attention_backward_rocm.mlir`, is `REQUIRES: +tessera-rocm-backend` and therefore unsupported on the Mac — it runs on this +box and passes. `lit` 439 discovered / 384 passed / 0 failed; +`check-tessera-rocm` 67 passed / 1 unsupported, identical to baseline. + Cross-backend sync `P3-SOURCE-ONLY-2026-08-30` — **one row is written from the `.hip` signature and has never executed.** The P3 batch added `RocmStockhamFFTCandidate.run_rows`, wired to diff --git a/docs/audit/backend/x86/todo.md b/docs/audit/backend/x86/todo.md index 3d2faaac2..eccf0a963 100644 --- a/docs/audit/backend/x86/todo.md +++ b/docs/audit/backend/x86/todo.md @@ -8,6 +8,30 @@ scope: x86 AVX-512 implementation/proof and AMX access planning --- # x86 backend TODO +Cross-backend sync `ADAFACTOR-BIAS-CORRECTION-2026-08-30` — **shared numerical +policy changed; per-backend outcome below.** +`optim.adafactor_decay` makes the Adafactor second-moment decay step-dependent +(`b2_t = b2*(1 - b2^(t-1))/(1 - b2^t)`), removing an early-step update +inflation of 1/sqrt(1 - b2^t) — 31.6x at step 1, 10.0x at step 10, 1.26x at +step 1000 for the default beta2. The correction is applied HOST-SIDE as a +scalar decay, so **no kernel ABI moves**: every physical kernel already takes +`beta2` as a scalar and receives the effective value instead of the nominal +one. The flat op gained an optional `step` kwarg matching the `adam`/`adamw` +ABI beside it. + +Two contract details a backend owner needs to know. `state["v"]` now carries +the DEBIASED estimate rather than the raw EMA, so the state dict grew a +`v_representation` marker and a state without one is migrated on load rather +than misread. And an absent `step` is NOT treated as step 1 — `decay(b2, 1)` +is exactly 0, so defaulting would have made a stateful caller that never +passes one discard its own moments; such a caller keeps the legacy +uncorrected decay. + +*x86 outcome: follow-up required.* `tessera_x86_avx512_adafactor_*` takes +`beta2` as a scalar and is unchanged; `tests/unit/test_x86_optimizer_compiled.py` +was migrated to pass `step`. AVX-512 execution lives on the Strix Halo box, so +an exact-device run of the corrected optimizer is owed there. + Cross-backend sync `P2-REVIEW-SHARED-PASSES-2026-08-29` — **15 shared MLIR passes changed; only the Mac's fixture set could be run.** The P2 code-review batch touched passes every backend lowers through: diff --git a/docs/audit/compiler/CODE_REVIEW_2026-08-29.md b/docs/audit/compiler/CODE_REVIEW_2026-08-29.md index dff2dfd43..ff829d9e0 100644 --- a/docs/audit/compiler/CODE_REVIEW_2026-08-29.md +++ b/docs/audit/compiler/CODE_REVIEW_2026-08-29.md @@ -2,7 +2,7 @@ title: Full Code Review — 2026-08-29 (logic · mathematical correctness · algorithms · performance) last_updated: 2026-08-29 scope: Python numeric core (autodiff · RNG/quantization · losses/optimizers · fusion + four backend emitters) · MLIR passes (linalg lowering · autodiff · tiling · analysis · legality · solver dialects) -status: 102 findings confirmed by an independent refutation pass; all four severity tiers have fixes committed, with 3 P3 rows still owing device verification +status: 102 findings confirmed by an independent refutation pass; all four severity tiers closed and device-verified audit_role: snapshot --- @@ -22,7 +22,7 @@ audit_role: snapshot > | Severity | Closed in | Outcome | > |---|---|---| > | P0 · 9 | PR #635 | all 9 fixed | -> | P1 · 36 | PRs #636, #637, #638, and the P3 batch | 34 fixed (the 3 performance/scope rows closed alongside P3); 2 remain deferred with recorded ABI reasoning: Adafactor bias correction and the rank-4 dropout per-instance seed | +> | P1 · 36 | PRs #636, #637, #638, the P3 batch, and the deferred-P1 batch | all 36 fixed. The last 2 — Adafactor bias correction and the rank-4 dropout per-instance seed — were deferred 2026-08-29 with ABI reasoning that **did not survive re-examination**: neither needed a kernel ABI change. Both fixed 2026-08-30; see their rows. The dropout row's severity was also corrected downward (the executing gfx1151 lane was already iid across batch/head; the defect was in the Tile-IR contract). The rank-4 dropout fix is **not yet lit-verified on the ROCm box** | > | P2 · 42 | PR #640 | 39 fixed; 3 needed no change — two had already been fixed by the P0 batch and the review quoted the pre-fix body, one (`AutodiffPass.cpp:199`) was closed by the P0 seed-type fix | > | P3 · 16 | the P3 batch | all 16 fixed in source; 13 host-measured, **3 device-unverified** (see the P3 section); 2 of the 16 were correctness defects, not improvements | > @@ -316,19 +316,27 @@ By dimension: logic 56 · math 24 · algorithm 7 · performance 15. ### `python/tessera/optim.py:427` — Adafactor missing second-moment bias correction -> **Deferred 2026-08-29 — needs an ABI change, not a one-line fix.** The -> finding is correct and the tree reference is easy to correct, but Adafactor -> has **three** implementations that a test deliberately keeps in agreement: -> `optim.adafactor` (tree), the analytic VJP/JVP in `autodiff/`, and the flat -> `ts.ops.adafactor(params, grads, row, col)` op — and the flat form, which is -> what the native x86/ROCm kernels bind to, **carries no step counter in its -> signature**, so it cannot compute `1 - beta2**t` at all. Correcting only the -> reachable halves makes the Python reference disagree with the certified -> native path, which is the exact defect shape reported for `nesterov` above. -> Fixing this properly means threading a step operand through the flat op and -> its native kernels, and re-proving them on ROCm/x86 hardware. Attempted and -> reverted; `test_flat_adafactor_full_and_factored_match_tree_reference` is the -> gate that catches a partial fix. +> **FIXED 2026-08-30 — the deferral's premise was wrong; no kernel ABI change +> was needed.** The deferral said the flat `ts.ops.adafactor` "carries no step +> counter in its signature, so it cannot compute `1 - beta2**t` at all". Two +> things falsify that. (a) The flat **`adam`/`adamw`** ABI right next to it +> already takes `step: int = 1` as a kwarg and the ROCm/x86 executors already +> read `int(kwargs.get("step", 1))` — so a step kwarg on the flat adafactor is +> the house pattern, not a new ABI. (b) The correction is exactly expressible +> as a step-dependent **decay rate**, +> `b2_t = b2*(1 - b2**(t-1))/(1 - b2**t)`, for which the recursion carries the +> debiased estimate directly (`v_t == EMA_t/(1 - b2**t)`, `b2_1 = 0`) — and +> every physical Adafactor kernel already takes `beta2` as a **scalar**, so the +> correction is applied host-side and `tessera_x86_avx512_adafactor_*` / +> `adafactor_row|col|mean|update` / `sm120_adafactor_*` are untouched. +> Landed as one shared `optim.adafactor_decay` consumed by the tree form, the +> flat op, the analytic VJP, the x86/ROCm/NVIDIA forward and backward +> executors, and the backward state contract (which now records nominal +> `beta2`, `step`, and `beta2_effective`). Pinned by six tests in +> `tests/unit/test_s10_optim.py`; +> `test_flat_adafactor_full_and_factored_match_tree_reference` was strengthened +> to step 3, since at step 1 the corrected decay is 0 and the carried state +> would not be exercised at all. *Losses, optimizers, RL, nn.functional · math* @@ -511,17 +519,47 @@ By dimension: logic 56 · math 24 · algorithm 7 · performance 15. ### `src/transforms/lib/TileIRLoweringPass.cpp:306` — Identical dropout mask replicated across batch and head -> **Deferred 2026-08-29 — needs an operand, not a local edit.** Refusing to -> distribute a rank-4 `flash_attn` carrying dropout was implemented and then -> **reverted after it regressed a committed ROCm lane**: -> `tests/tessera-ir/phase3/streaming_attention_backward_rocm.mlir` drives -> `dropout_p = 0.25` through exactly this path and passed before. The defect is -> real but statistical (masks correlated across batch/head, not a wrong value -> per element), and the batch/head coordinates at the distribution site are SSA -> loop induction variables — so a per-instance seed cannot be an attribute and -> must be threaded as an OPERAND, changing the op signature and its native -> consumers. Same shape as the Adafactor row. Caught only by running lit on the -> ROCm box; the Mac skips this fixture entirely (`REQUIRES: tessera-rocm-backend`). +> **FIXED 2026-08-30 — as an operand, but on `block_dropout`, not on +> `flash_attn`; and the severity is lower than reported.** The deferral was +> right that a per-instance seed cannot be an attribute (the batch/head +> coordinates are `scf.for` induction variables). It was wrong about which op +> has to change and about the blast radius. +> +> **Correction to the finding's evidence.** "No backend can recover the +> distinction" does not hold. The only backend that executes dropout attention +> is ROCm, and it does not consume `block_dropout`'s operands at all: +> `TileToROCM.cpp:1549` records a bare `dropout = true/false` on +> `tessera_rocm.flash_attn`, and the generated kernels rebuild the mask from +> launch geometry — `counter = ((bh*Sq)+q)*Sk + k` in *both* +> `GenerateWMMAFlashAttnKernel.cpp:418` and `...BwdKernel.cpp:482`, with `bh` +> the fused batch-head block id. So the **executed** gfx1151 masks are already +> iid across batch and head, and forward/backward already agree. Apple refuses +> dropout outright (`APPLE_STREAMING_ATTN_DROPOUT_UNSUPPORTED`). The defect is +> therefore a **Tile-IR contract defect** (Decisions #29/#32 — the shared +> boundary understated what the physical kernels do), not a wrong numerical +> result on any lane that runs today. +> +> **Fix.** `tessera_attn.block_dropout` gained an `Index:$stream_offset` +> operand; `LowerFlashAttnToTileIR` derives the instance index `b*H + h` from +> the `tessera.attention_distribution` batch/query_head loops this same pass +> already annotates (Decision #30 — derive, don't ask) and passes +> `(b*H + h) * Sq * Sk_padded`, the same disjoint counter block the WMMA +> kernels use. A rank-2 attention that was never distributed passes 0 and its +> mask is bit-identical to before; a distributed instance whose annotated loops +> are unreachable **fails the match with a diagnostic** rather than silently +> emitting stream 0 (Decision #21a). `tessera.flash_attn` is untouched, so no +> `operandSegmentSizes` churn. +> +> **Why the ROCm lane cannot regress this time** (stated, not assumed — this +> Mac skips `streaming_attention_backward_rocm.mlir` via +> `REQUIRES: tessera-rocm-backend`, so it is unverified here): no fixture in +> the repo parses `tessera_attn.block_dropout` textually; `ROCMWaveLdsPipeline` +> never mentions attention or dropout ops; and `TileToROCM` tests only +> `blockDropout != nullptr`. The only change in that fixture's output is one +> more operand printed on `block_dropout`, which no CHECK line matches. **A +> lit run on the ROCm box is still owed.** Host-free coverage added at +> `tests/tessera-ir/phase3/streaming_attention_dropout_stream.mlir` (positive +> rank-4 case + `stream = %c0` rank-2 negative). *C++ — tiling & Tile IR lowering · math* @@ -1168,11 +1206,13 @@ By dimension: logic 56 · math 24 · algorithm 7 · performance 15. **Independently verified:** AsyncPrefetch.cpp:33-40 + 77-84: producesOperandOf compares prev->getResults() against op->getOperands() only, and the hoist at line 81 fires whenever prev is non-terminator and that SSA test is false. A zero-result writer — memref.store, a linalg op on memrefs, a collective writing a buffer — always fails the test, so the prefetch moves above the write. Nothing restricts the prefetch source to value semantics: Schedule_PrefetchOp takes AnyType:$source (ScheduleMeshPipelineOps.td:556-566) and its verifier (ScheduleDialect.cpp:748-757) checks only type preservation, non-empty `into`, and the overlap enum — a memref source is legal. The pass also never consults MemoryEffectOpInterface for prev or for the prefetch, and the pass runs on a plain module walk with no bufferization ordering constraint, so the header's 'dependency-safe' claim rests on an SSA-only fact (Decision #30, told-not-derived, failing open). The existing fixture (async_prefetch_overlap.mlir) uses only tensor values and comp.matmul, so it does not pin the memref case as intentional. Proposed fix (refuse the hoist when prev has write effects on a buffer the prefetch touches) is the right shape. -## P3 — improvement opportunities (16) — FIXED IN SOURCE 2026-08-30, 3 rows device-unverified +## P3 — improvement opportunities (16) — CLOSED 2026-08-30, all rows device-verified -> All 16 have a fix committed, together with the 3 P1 rows that had been left -> open as performance/scope items. **"Fixed" is not "closed" for three of -> them.** `nvidia_cuda.py:626` (flash-backward cleanup), +> All 16 fixed, together with the 3 P1 rows that had been left open as +> performance/scope items. The three rows that were device-unverified on +> 2026-08-30 have since been measured on their own hardware, and **one of +> them was a regression the reasoning had missed** — see the Krylov row in +> `../backend/nvidia/todo.md`. Superseded text follows for the record. `nvidia_cuda.py:626` (flash-backward cleanup), > `nvidia_solver_krylov.py:79` (warp-per-row matvec) and the ROCm half of the > spectral per-frame row are verified only as generated text plus a clean > `clang++ -fsyntax-only` parse — no GPU has executed any of them, and this diff --git a/docs/audit/generated/docs_freshness.md b/docs/audit/generated/docs_freshness.md index efd061fc2..b71e644d2 100644 --- a/docs/audit/generated/docs_freshness.md +++ b/docs/audit/generated/docs_freshness.md @@ -2,13 +2,13 @@ Generated from `python/tessera/compiler/docs_manifest.py`. Don't edit by hand — regenerate via `python -c "from tessera.compiler.docs_manifest import render_dashboard; open('docs/audit/generated/docs_freshness.md', 'w').write(render_dashboard())"`. Drift gated by `tests/unit/test_docs_freshness.py`. -Reference date for staleness: **2026-08-29**. +Reference date for staleness: **2026-08-30**. ## Headline - **150** docs catalogued across the canonical doc tree. - **149** carry a `last_updated:` marker; **1** are undated (invisible to the freshness audit until tagged). -- **59** updated within the last 30 days. +- **58** updated within the last 30 days. - **22** older than 90 days; **0** older than 180 days. ## Undated docs (no parseable `last_updated`) @@ -23,183 +23,183 @@ These docs need either YAML frontmatter (`last_updated: YYYY-MM-DD`) or a body-f | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `AUTODIFF_SPEC.md` | - | 2026-07-14 | 46 | ✓ | -| `CITL_ROCM_TRACE_PROFILER_SPEC.md` | Draft | 2026-08-06 | 23 | ✓ | -| `CLIFFORD_SPEC.md` | - | 2026-05-17 | 104 | ✓ | -| `COMPILER_REFERENCE.md` | Normative | 2026-06-25 | 65 | ✓ | -| `CONFORMANCE.md` | Normative | 2026-06-11 | 79 | ✓ | -| `CONTROL_FLOW_CONTRACT.md` | - | 2026-08-11 | 18 | ✓ | -| `EBM_SPEC.md` | - | 2026-05-16 | 105 | ✓ | -| `GA_EBM_EXECUTION_STATUS.md` | - | 2026-07-18 | 42 | ✓ | -| `GRAPH_IR_SPEC.md` | Normative | 2026-07-14 | 46 | ✓ | -| `LANGUAGE_AND_IR_SPEC.md` | Normative | 2026-05-06 | 115 | ✓ | -| `LANGUAGE_SPEC.md` | Normative | 2026-07-14 | 46 | ✓ | -| `LOWERING_PIPELINE_SPEC.md` | Normative | 2026-07-13 | 47 | ✓ | -| `MEMORY_MODEL_SPEC.md` | Normative | 2026-05-22 | 99 | ✓ | -| `NATIVE_ARTIFACT_SPEC.md` | Normative | 2026-07-19 | 41 | ✓ | -| `PRODUCTION_COMPILER_PLAN.md` | Ratified | 2026-06-05 | 85 | ✓ | -| `PYTHON_API_SPEC.md` | Normative | 2026-07-23 | 37 | ✓ | -| `RUNTIME_ABI_SPEC.md` | Normative | 2026-07-18 | 42 | ✓ | -| `SHAPE_SYSTEM.md` | Normative | 2026-05-22 | 99 | ✓ | -| `TARGET_IR_SPEC.md` | Normative | 2026-08-24 | 5 | ✓ | -| `TILE_IR.md` | Normative | 2026-08-10 | 19 | ✓ | -| `VALIDATION_SPINE.md` | Normative | 2026-08-02 | 27 | ✓ | -| `VALUE_TARGET_IR_CONTRACT.md` | Normative | 2026-06-04 | 86 | ✓ | +| `AUTODIFF_SPEC.md` | - | 2026-07-14 | 47 | ✓ | +| `CITL_ROCM_TRACE_PROFILER_SPEC.md` | Draft | 2026-08-06 | 24 | ✓ | +| `CLIFFORD_SPEC.md` | - | 2026-05-17 | 105 | ✓ | +| `COMPILER_REFERENCE.md` | Normative | 2026-06-25 | 66 | ✓ | +| `CONFORMANCE.md` | Normative | 2026-06-11 | 80 | ✓ | +| `CONTROL_FLOW_CONTRACT.md` | - | 2026-08-11 | 19 | ✓ | +| `EBM_SPEC.md` | - | 2026-05-16 | 106 | ✓ | +| `GA_EBM_EXECUTION_STATUS.md` | - | 2026-07-18 | 43 | ✓ | +| `GRAPH_IR_SPEC.md` | Normative | 2026-07-14 | 47 | ✓ | +| `LANGUAGE_AND_IR_SPEC.md` | Normative | 2026-05-06 | 116 | ✓ | +| `LANGUAGE_SPEC.md` | Normative | 2026-07-14 | 47 | ✓ | +| `LOWERING_PIPELINE_SPEC.md` | Normative | 2026-07-13 | 48 | ✓ | +| `MEMORY_MODEL_SPEC.md` | Normative | 2026-05-22 | 100 | ✓ | +| `NATIVE_ARTIFACT_SPEC.md` | Normative | 2026-07-19 | 42 | ✓ | +| `PRODUCTION_COMPILER_PLAN.md` | Ratified | 2026-06-05 | 86 | ✓ | +| `PYTHON_API_SPEC.md` | Normative | 2026-07-23 | 38 | ✓ | +| `RUNTIME_ABI_SPEC.md` | Normative | 2026-07-18 | 43 | ✓ | +| `SHAPE_SYSTEM.md` | Normative | 2026-05-22 | 100 | ✓ | +| `TARGET_IR_SPEC.md` | Normative | 2026-08-24 | 6 | ✓ | +| `TILE_IR.md` | Normative | 2026-08-10 | 20 | ✓ | +| `VALIDATION_SPINE.md` | Normative | 2026-08-02 | 28 | ✓ | +| `VALUE_TARGET_IR_CONTRACT.md` | Normative | 2026-06-04 | 87 | ✓ | ### `docs/guides/` | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `Tessera_Debugging_Tools_Guide.md` | Informative | 2026-05-06 | 115 | ✓ | -| `Tessera_Developer_Frontend_End_To_End.md` | Informative | 2026-05-06 | 115 | ✓ | -| `Tessera_Differentiable_NAS_Guide.md` | Draft | 2026-04-28 | 123 | ✓ | -| `Tessera_Error_Handling_And_Diagnostics_Guide.md` | Normative | 2026-04-28 | 123 | ✓ | -| `Tessera_Fault_Tolerance_And_Elasticity_Guide.md` | Informative | 2026-04-28 | 123 | ✓ | -| `Tessera_Inference_Server_Guide.md` | Informative | 2026-06-11 | 79 | ✓ | -| `Tessera_Production_Reliability_And_Chaos_Guide.md` | Informative | 2026-04-28 | 123 | ✓ | -| `Tessera_Profiler_Release_Gates.md` | Informative | 2026-08-06 | 23 | ✓ | -| `Tessera_Profiling_And_Autotuning_Guide.md` | Informative | 2026-08-06 | 23 | ✓ | -| `Tessera_QA_Reliability_Guide.md` | Informative | 2026-04-28 | 123 | ✓ | -| `Tessera_Runtime_ABI_Guide.md` | Tutorial | 2026-07-14 | 46 | ✓ | -| `Tessera_Tensor_Layout_And_Data_Movement_Guide.md` | Normative | 2026-07-14 | 46 | ✓ | -| `porting_advanced_examples.md` | Informative | 2026-05-09 | 112 | ✓ | +| `Tessera_Debugging_Tools_Guide.md` | Informative | 2026-05-06 | 116 | ✓ | +| `Tessera_Developer_Frontend_End_To_End.md` | Informative | 2026-05-06 | 116 | ✓ | +| `Tessera_Differentiable_NAS_Guide.md` | Draft | 2026-04-28 | 124 | ✓ | +| `Tessera_Error_Handling_And_Diagnostics_Guide.md` | Normative | 2026-04-28 | 124 | ✓ | +| `Tessera_Fault_Tolerance_And_Elasticity_Guide.md` | Informative | 2026-04-28 | 124 | ✓ | +| `Tessera_Inference_Server_Guide.md` | Informative | 2026-06-11 | 80 | ✓ | +| `Tessera_Production_Reliability_And_Chaos_Guide.md` | Informative | 2026-04-28 | 124 | ✓ | +| `Tessera_Profiler_Release_Gates.md` | Informative | 2026-08-06 | 24 | ✓ | +| `Tessera_Profiling_And_Autotuning_Guide.md` | Informative | 2026-08-06 | 24 | ✓ | +| `Tessera_QA_Reliability_Guide.md` | Informative | 2026-04-28 | 124 | ✓ | +| `Tessera_Runtime_ABI_Guide.md` | Tutorial | 2026-07-14 | 47 | ✓ | +| `Tessera_Tensor_Layout_And_Data_Movement_Guide.md` | Normative | 2026-07-14 | 47 | ✓ | +| `porting_advanced_examples.md` | Informative | 2026-05-09 | 113 | ✓ | ### `docs/programming_guide/` | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `Tessera_Goals.md` | Tutorial | 2026-07-14 | 46 | ✓ | -| `Tessera_Programming_Guide_Appendix_NVL72.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter10_Portability.md` | Tutorial | 2026-07-13 | 47 | ✓ | -| `Tessera_Programming_Guide_Chapter11_Conclusion.md` | Tutorial | 2026-07-14 | 46 | ✓ | -| `Tessera_Programming_Guide_Chapter1_Introduction_Overview.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter2_Programming_Model.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter3_Memory_Model.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter4_Execution_Model.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter5_Kernel_Programming.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter6_Numerics_Model.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter7_Autodiff.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter8_Layouts_Data_Movement.md` | Tutorial | 2026-06-11 | 79 | ✓ | -| `Tessera_Programming_Guide_Chapter9_Libraries_Primitives.md` | Tutorial | 2026-06-11 | 79 | ✓ | +| `Tessera_Goals.md` | Tutorial | 2026-07-14 | 47 | ✓ | +| `Tessera_Programming_Guide_Appendix_NVL72.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter10_Portability.md` | Tutorial | 2026-07-13 | 48 | ✓ | +| `Tessera_Programming_Guide_Chapter11_Conclusion.md` | Tutorial | 2026-07-14 | 47 | ✓ | +| `Tessera_Programming_Guide_Chapter1_Introduction_Overview.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter2_Programming_Model.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter3_Memory_Model.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter4_Execution_Model.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter5_Kernel_Programming.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter6_Numerics_Model.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter7_Autodiff.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter8_Layouts_Data_Movement.md` | Tutorial | 2026-06-11 | 80 | ✓ | +| `Tessera_Programming_Guide_Chapter9_Libraries_Primitives.md` | Tutorial | 2026-06-11 | 80 | ✓ | ### `docs/operations/` | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `Tessera_Standard_Operations.md` | Normative | 2026-07-13 | 47 | ✓ | -| `backend_local_proofs.md` | - | 2026-07-15 | 45 | ✓ | -| `release_gates.md` | Normative | 2026-07-13 | 47 | ✓ | +| `Tessera_Standard_Operations.md` | Normative | 2026-07-13 | 48 | ✓ | +| `backend_local_proofs.md` | - | 2026-07-15 | 46 | ✓ | +| `release_gates.md` | Normative | 2026-07-13 | 48 | ✓ | ### `docs/architecture/` | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `Compiler/Tessera_Compiler_Architecture_Overview.md` | Informative | 2026-07-14 | 46 | ✓ | -| `Compiler/Tessera_Compiler_Frontend_Design_GraphIR.md` | Informative | 2026-07-14 | 46 | ✓ | -| `Compiler/Tessera_Compiler_ScheduleIR_Design.md` | Informative | 2026-07-14 | 46 | ✓ | -| `Compiler/Tessera_Compiler_TargetIR_Design.md` | Informative | 2026-07-14 | 46 | ✓ | -| `Compiler/Tessera_Compiler_TileIR_Design.md` | Informative | 2026-07-14 | 46 | ✓ | -| `Compiler/tessera_ir_layers.md` | Informative | 2026-07-13 | 47 | ✓ | -| `Compiler/tessera_tile_ir_documentation.md` | Informative | 2026-07-14 | 46 | ✓ | -| `README.md` | Informative | 2026-05-20 | 101 | ✓ | -| `Tessera_Kernel_Compilation_Stages_Overview.md` | Informative | 2026-05-06 | 115 | ✓ | -| `compiler_gaps_1_3_5_plan.md` | - | 2026-07-14 | 46 | ✓ | -| `compiler_test_architecture.md` | Normative | 2026-08-02 | 27 | ✓ | -| `distributed/megamoe.md` | - | 2026-06-09 | 81 | ✓ | -| `frontend_substrate_plan.md` | Active | 2026-05-20 | 101 | ✓ | -| `inference/serving.md` | - | 2026-07-13 | 47 | ✓ | -| `proposals/cute_tessera_enhancement.md` | Proposal | 2026-04-26 | 125 | ✓ | -| `proposals/tile_fragment_abi.md` | Proposal | 2026-07-19 | 41 | ✓ | -| `proposals/tiled_ssd_tile_ir_schedule.md` | - | 2026-07-14 | 46 | ✓ | -| `stencil_materialize_and_window_lowering.md` | Informative | 2026-05-20 | 101 | ✓ | -| `system_overview.md` | Informative | 2026-06-11 | 79 | ✓ | -| `tessera_target_ir_usage_guide.md` | Informative | 2026-04-30 | 121 | ✓ | -| `workloads/attention-family.md` | Planning | 2026-07-14 | 46 | ✓ | -| `workloads/dflash.md` | - | 2026-07-14 | 46 | ✓ | -| `workloads/msa-cuda-phase3.md` | - | 2026-07-13 | 47 | ✓ | -| `workloads/msa.md` | - | 2026-07-13 | 47 | ✓ | +| `Compiler/Tessera_Compiler_Architecture_Overview.md` | Informative | 2026-07-14 | 47 | ✓ | +| `Compiler/Tessera_Compiler_Frontend_Design_GraphIR.md` | Informative | 2026-07-14 | 47 | ✓ | +| `Compiler/Tessera_Compiler_ScheduleIR_Design.md` | Informative | 2026-07-14 | 47 | ✓ | +| `Compiler/Tessera_Compiler_TargetIR_Design.md` | Informative | 2026-07-14 | 47 | ✓ | +| `Compiler/Tessera_Compiler_TileIR_Design.md` | Informative | 2026-07-14 | 47 | ✓ | +| `Compiler/tessera_ir_layers.md` | Informative | 2026-07-13 | 48 | ✓ | +| `Compiler/tessera_tile_ir_documentation.md` | Informative | 2026-07-14 | 47 | ✓ | +| `README.md` | Informative | 2026-05-20 | 102 | ✓ | +| `Tessera_Kernel_Compilation_Stages_Overview.md` | Informative | 2026-05-06 | 116 | ✓ | +| `compiler_gaps_1_3_5_plan.md` | - | 2026-07-14 | 47 | ✓ | +| `compiler_test_architecture.md` | Normative | 2026-08-02 | 28 | ✓ | +| `distributed/megamoe.md` | - | 2026-06-09 | 82 | ✓ | +| `frontend_substrate_plan.md` | Active | 2026-05-20 | 102 | ✓ | +| `inference/serving.md` | - | 2026-07-13 | 48 | ✓ | +| `proposals/cute_tessera_enhancement.md` | Proposal | 2026-04-26 | 126 | ✓ | +| `proposals/tile_fragment_abi.md` | Proposal | 2026-07-19 | 42 | ✓ | +| `proposals/tiled_ssd_tile_ir_schedule.md` | - | 2026-07-14 | 47 | ✓ | +| `stencil_materialize_and_window_lowering.md` | Informative | 2026-05-20 | 102 | ✓ | +| `system_overview.md` | Informative | 2026-06-11 | 80 | ✓ | +| `tessera_target_ir_usage_guide.md` | Informative | 2026-04-30 | 122 | ✓ | +| `workloads/attention-family.md` | Planning | 2026-07-14 | 47 | ✓ | +| `workloads/dflash.md` | - | 2026-07-14 | 47 | ✓ | +| `workloads/msa-cuda-phase3.md` | - | 2026-07-13 | 48 | ✓ | +| `workloads/msa.md` | - | 2026-07-13 | 48 | ✓ | ### `docs/reference/` | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `tessera-api-reference.md` | Informative | 2026-07-13 | 47 | ✓ | +| `tessera-api-reference.md` | Informative | 2026-07-13 | 48 | ✓ | | `tessera_frontend_lanes.md` | - | _undated_ | - | _body_ | -| `tessera_migration_guide_part1.md` | Pre-canonical | 2026-05-20 | 101 | ✓ | -| `tessera_migration_guide_part2.md` | Informative | 2026-05-20 | 101 | ✓ | -| `tessera_tensor_attributes.md` | Normative | 2026-05-11 | 110 | ✓ | +| `tessera_migration_guide_part1.md` | Pre-canonical | 2026-05-20 | 102 | ✓ | +| `tessera_migration_guide_part2.md` | Informative | 2026-05-20 | 102 | ✓ | +| `tessera_tensor_attributes.md` | Normative | 2026-05-11 | 111 | ✓ | ### `docs/audit/` | Path | status | last_updated | days stale | frontmatter | |------|--------|--------------|-----------:|--| -| `MASTER_AUDIT.md` | - | 2026-08-11 | 18 | ✓ | -| `README.md` | - | 2026-08-08 | 21 | ✓ | -| `backend/BACKEND_AUDIT.md` | - | 2026-07-31 | 29 | ✓ | -| `backend/E2E_COMPILATION_AUDIT.md` | - | 2026-07-27 | 33 | ✓ | -| `backend/X86_AVX512_ABI_INVENTORY.md` | - | 2026-07-22 | 38 | ✓ | -| `backend/apple/APPLE_AUDIT.md` | - | 2026-07-28 | 32 | ✓ | -| `backend/apple/APPLE_GPU_CODEGEN_PLAN.md` | - | 2026-07-13 | 47 | ✓ | -| `backend/apple/MPSGRAPH_RUNTIME_GLASS_JAWS.md` | - | 2026-07-13 | 47 | ✓ | -| `backend/apple/README.md` | - | 2026-07-13 | 47 | ✓ | -| `backend/apple/todo.md` | - | 2026-08-29 | 0 | ✓ | -| `backend/nvidia/BLACKWELL_SM120_EXECUTION_PLAN.md` | - | 2026-06-24 | 66 | ✓ | -| `backend/nvidia/NVIDIA_AUDIT.md` | - | 2026-07-18 | 42 | ✓ | -| `backend/nvidia/SM120_DIFFERENTIATION_DASHBOARD.md` | - | 2026-07-19 | 41 | ✓ | -| `backend/nvidia/VERIFY_TARGET_IR_TAIL.md` | - | 2026-07-13 | 47 | ✓ | -| `backend/nvidia/spikes/sm120_mma_sync/README.md` | - | 2026-06-24 | 66 | ✓ | -| `backend/nvidia/todo.md` | - | 2026-08-29 | 0 | ✓ | -| `backend/rocm/GEMM_PERF_LADDER.md` | - | 2026-08-04 | 25 | ✓ | -| `backend/rocm/GFX125X_CDNA5_COMPILER_REFERENCE.md` | - | 2026-08-14 | 15 | ✓ | -| `backend/rocm/GIN_EXACT_DEVICE_RUNBOOK.md` | - | 2026-08-09 | 20 | ✓ | -| `backend/rocm/ROCM_AUDIT.md` | - | 2026-08-14 | 15 | ✓ | -| `backend/rocm/ROCM_LANE_MAP.md` | - | 2026-08-05 | 24 | ✓ | -| `backend/rocm/ROCM_PATTERNS_FROM_AMD_ECOSYSTEM.md` | - | 2026-07-28 | 32 | ✓ | -| `backend/rocm/STRIX_HALO_EXECUTION_PLAN.md` | - | 2026-06-23 | 67 | ✓ | -| `backend/rocm/todo.md` | - | 2026-08-29 | 0 | ✓ | -| `backend/x86/todo.md` | - | 2026-08-29 | 0 | ✓ | -| `compiler/AMD_KERNEL_COMPILER_SURVEY.md` | - | 2026-07-28 | 32 | ✓ | -| `compiler/AUTODIFF_ARCHITECTURE_REVIEW.md` | - | 2026-08-18 | 11 | ✓ | -| `compiler/AUTODIFF_NEXTGEN_PLAN.md` | - | 2026-08-20 | 9 | ✓ | -| `compiler/AUTODIFF_UNIFICATION_PLAN.md` | - | 2026-08-18 | 11 | ✓ | -| `compiler/BLOCK_ATTNRES_ROCM_PLAN.md` | - | 2026-08-13 | 16 | ✓ | -| `compiler/CODE_REVIEW_2026-08-29.md` | - | 2026-08-29 | 0 | ✓ | -| `compiler/COMPILER_ARCHITECTURE_SWEEP.md` | - | 2026-08-11 | 18 | ✓ | -| `compiler/COMPILER_AUDIT.md` | - | 2026-08-10 | 19 | ✓ | -| `compiler/COMPILER_REFACTOR_PLAN.md` | - | 2026-08-08 | 21 | ✓ | -| `compiler/COMPILER_THEORY_OF_OPERATION.md` | - | 2026-07-28 | 32 | ✓ | -| `compiler/CORE_SUBSTRATE_VIEW.md` | - | 2026-08-24 | 5 | ✓ | -| `compiler/CUTE_IR_ASSESSMENT.md` | - | 2026-08-24 | 5 | ✓ | -| `compiler/DIFFERENTIABLE_PROGRAMMING_REVIEW.md` | - | 2026-08-08 | 21 | ✓ | -| `compiler/EGGROLL_SUPPORT_PLAN.md` | - | 2026-08-09 | 20 | ✓ | -| `compiler/EVALUATOR_PLAN.md` | - | 2026-08-08 | 21 | ✓ | -| `compiler/FORGE_ASSESSMENT.md` | - | 2026-08-15 | 14 | ✓ | -| `compiler/FRONTEND_GRAPH_SCHEDULE_REVIEW.md` | - | 2026-08-02 | 27 | ✓ | -| `compiler/FUNCTIONAL_ANALYSIS_TSOL_PLAN.md` | - | 2026-08-22 | 7 | ✓ | -| `compiler/GAME_THEORY_PLAN.md` | - | 2026-08-15 | 14 | ✓ | -| `compiler/INTEGRATED_COMPILER_PLAN.md` | - | 2026-08-24 | 5 | ✓ | -| `compiler/INTRA_KERNEL_FEEDBACK_PLAN.md` | - | 2026-08-27 | 2 | ✓ | -| `compiler/IR_STACK_INTEGRATION_REVIEW.md` | - | 2026-08-02 | 27 | ✓ | -| `compiler/LSE_CHECKPOINT_CONTRACT.md` | - | 2026-07-27 | 33 | ✓ | -| `compiler/MATRIX_CALCULUS_REVIEW.md` | - | 2026-08-20 | 9 | ✓ | -| `compiler/OPTIMIZING_COMPILER_PLAN.md` | - | 2026-08-08 | 21 | ✓ | -| `compiler/PDE_STENCIL_CAPABILITY_PLAN.md` | - | 2026-08-14 | 15 | ✓ | -| `compiler/README.md` | - | 2026-08-18 | 11 | ✓ | -| `compiler/RIEMANNIAN_OT_PLAN.md` | - | 2026-08-08 | 21 | ✓ | -| `compiler/SCHEDULE_OBJECT_DESIGN.md` | - | 2026-08-16 | 13 | ✓ | -| `compiler/SEQUENCE_MIXER_ENGINEERING_PLAN.md` | - | 2026-08-08 | 21 | ✓ | -| `compiler/SEQUENCE_MIXER_THEORY.md` | - | 2026-07-17 | 43 | ✓ | -| `compiler/SPARDA_REVIEW.md` | - | 2026-08-12 | 17 | ✓ | -| `compiler/TARGET_IR_REVIEW.md` | - | 2026-08-02 | 27 | ✓ | -| `compiler/TILERT_ASSESSMENT.md` | - | 2026-08-14 | 15 | ✓ | -| `compiler/TILESIGHT_ASSESSMENT.md` | - | 2026-07-30 | 30 | ✓ | -| `compiler/W1_1_TYPING_DESIGN.md` | - | 2026-08-18 | 11 | ✓ | -| `compiler/W1_1_TYPING_INVENTORY.md` | - | 2026-08-02 | 27 | ✓ | -| `compiler/W4_ADMISSIBLE_EFFECTS_PLAN.md` | - | 2026-08-25 | 4 | ✓ | -| `compiler/compiler_enhancement.md` | - | 2026-08-18 | 11 | ✓ | -| `coverage/COVERAGE_AUDIT.md` | - | 2026-08-11 | 18 | ✓ | -| `domain/DOMAIN_AUDIT.md` | - | 2026-06-11 | 79 | ✓ | -| `domain/GA_EBM_ARCHITECTURE_REVIEW.md` | - | 2026-08-02 | 27 | ✓ | -| `roadmap/CF_CROSS_ELEMENT_PLAN.md` | - | 2026-06-30 | 60 | ✓ | -| `roadmap/MODEL_CLASS_ROADMAP.md` | - | 2026-08-12 | 17 | ✓ | -| `roadmap/ROADMAP_AUDIT.md` | - | 2026-08-11 | 18 | ✓ | +| `MASTER_AUDIT.md` | - | 2026-08-11 | 19 | ✓ | +| `README.md` | - | 2026-08-08 | 22 | ✓ | +| `backend/BACKEND_AUDIT.md` | - | 2026-07-31 | 30 | ✓ | +| `backend/E2E_COMPILATION_AUDIT.md` | - | 2026-07-27 | 34 | ✓ | +| `backend/X86_AVX512_ABI_INVENTORY.md` | - | 2026-07-22 | 39 | ✓ | +| `backend/apple/APPLE_AUDIT.md` | - | 2026-07-28 | 33 | ✓ | +| `backend/apple/APPLE_GPU_CODEGEN_PLAN.md` | - | 2026-07-13 | 48 | ✓ | +| `backend/apple/MPSGRAPH_RUNTIME_GLASS_JAWS.md` | - | 2026-07-13 | 48 | ✓ | +| `backend/apple/README.md` | - | 2026-07-13 | 48 | ✓ | +| `backend/apple/todo.md` | - | 2026-08-29 | 1 | ✓ | +| `backend/nvidia/BLACKWELL_SM120_EXECUTION_PLAN.md` | - | 2026-06-24 | 67 | ✓ | +| `backend/nvidia/NVIDIA_AUDIT.md` | - | 2026-07-18 | 43 | ✓ | +| `backend/nvidia/SM120_DIFFERENTIATION_DASHBOARD.md` | - | 2026-07-19 | 42 | ✓ | +| `backend/nvidia/VERIFY_TARGET_IR_TAIL.md` | - | 2026-07-13 | 48 | ✓ | +| `backend/nvidia/spikes/sm120_mma_sync/README.md` | - | 2026-06-24 | 67 | ✓ | +| `backend/nvidia/todo.md` | - | 2026-08-29 | 1 | ✓ | +| `backend/rocm/GEMM_PERF_LADDER.md` | - | 2026-08-04 | 26 | ✓ | +| `backend/rocm/GFX125X_CDNA5_COMPILER_REFERENCE.md` | - | 2026-08-14 | 16 | ✓ | +| `backend/rocm/GIN_EXACT_DEVICE_RUNBOOK.md` | - | 2026-08-09 | 21 | ✓ | +| `backend/rocm/ROCM_AUDIT.md` | - | 2026-08-14 | 16 | ✓ | +| `backend/rocm/ROCM_LANE_MAP.md` | - | 2026-08-05 | 25 | ✓ | +| `backend/rocm/ROCM_PATTERNS_FROM_AMD_ECOSYSTEM.md` | - | 2026-07-28 | 33 | ✓ | +| `backend/rocm/STRIX_HALO_EXECUTION_PLAN.md` | - | 2026-06-23 | 68 | ✓ | +| `backend/rocm/todo.md` | - | 2026-08-29 | 1 | ✓ | +| `backend/x86/todo.md` | - | 2026-08-29 | 1 | ✓ | +| `compiler/AMD_KERNEL_COMPILER_SURVEY.md` | - | 2026-07-28 | 33 | ✓ | +| `compiler/AUTODIFF_ARCHITECTURE_REVIEW.md` | - | 2026-08-18 | 12 | ✓ | +| `compiler/AUTODIFF_NEXTGEN_PLAN.md` | - | 2026-08-20 | 10 | ✓ | +| `compiler/AUTODIFF_UNIFICATION_PLAN.md` | - | 2026-08-18 | 12 | ✓ | +| `compiler/BLOCK_ATTNRES_ROCM_PLAN.md` | - | 2026-08-13 | 17 | ✓ | +| `compiler/CODE_REVIEW_2026-08-29.md` | - | 2026-08-29 | 1 | ✓ | +| `compiler/COMPILER_ARCHITECTURE_SWEEP.md` | - | 2026-08-11 | 19 | ✓ | +| `compiler/COMPILER_AUDIT.md` | - | 2026-08-10 | 20 | ✓ | +| `compiler/COMPILER_REFACTOR_PLAN.md` | - | 2026-08-08 | 22 | ✓ | +| `compiler/COMPILER_THEORY_OF_OPERATION.md` | - | 2026-07-28 | 33 | ✓ | +| `compiler/CORE_SUBSTRATE_VIEW.md` | - | 2026-08-24 | 6 | ✓ | +| `compiler/CUTE_IR_ASSESSMENT.md` | - | 2026-08-24 | 6 | ✓ | +| `compiler/DIFFERENTIABLE_PROGRAMMING_REVIEW.md` | - | 2026-08-08 | 22 | ✓ | +| `compiler/EGGROLL_SUPPORT_PLAN.md` | - | 2026-08-09 | 21 | ✓ | +| `compiler/EVALUATOR_PLAN.md` | - | 2026-08-08 | 22 | ✓ | +| `compiler/FORGE_ASSESSMENT.md` | - | 2026-08-15 | 15 | ✓ | +| `compiler/FRONTEND_GRAPH_SCHEDULE_REVIEW.md` | - | 2026-08-02 | 28 | ✓ | +| `compiler/FUNCTIONAL_ANALYSIS_TSOL_PLAN.md` | - | 2026-08-22 | 8 | ✓ | +| `compiler/GAME_THEORY_PLAN.md` | - | 2026-08-15 | 15 | ✓ | +| `compiler/INTEGRATED_COMPILER_PLAN.md` | - | 2026-08-24 | 6 | ✓ | +| `compiler/INTRA_KERNEL_FEEDBACK_PLAN.md` | - | 2026-08-27 | 3 | ✓ | +| `compiler/IR_STACK_INTEGRATION_REVIEW.md` | - | 2026-08-02 | 28 | ✓ | +| `compiler/LSE_CHECKPOINT_CONTRACT.md` | - | 2026-07-27 | 34 | ✓ | +| `compiler/MATRIX_CALCULUS_REVIEW.md` | - | 2026-08-20 | 10 | ✓ | +| `compiler/OPTIMIZING_COMPILER_PLAN.md` | - | 2026-08-08 | 22 | ✓ | +| `compiler/PDE_STENCIL_CAPABILITY_PLAN.md` | - | 2026-08-14 | 16 | ✓ | +| `compiler/README.md` | - | 2026-08-18 | 12 | ✓ | +| `compiler/RIEMANNIAN_OT_PLAN.md` | - | 2026-08-08 | 22 | ✓ | +| `compiler/SCHEDULE_OBJECT_DESIGN.md` | - | 2026-08-16 | 14 | ✓ | +| `compiler/SEQUENCE_MIXER_ENGINEERING_PLAN.md` | - | 2026-08-08 | 22 | ✓ | +| `compiler/SEQUENCE_MIXER_THEORY.md` | - | 2026-07-17 | 44 | ✓ | +| `compiler/SPARDA_REVIEW.md` | - | 2026-08-12 | 18 | ✓ | +| `compiler/TARGET_IR_REVIEW.md` | - | 2026-08-02 | 28 | ✓ | +| `compiler/TILERT_ASSESSMENT.md` | - | 2026-08-14 | 16 | ✓ | +| `compiler/TILESIGHT_ASSESSMENT.md` | - | 2026-07-30 | 31 | ✓ | +| `compiler/W1_1_TYPING_DESIGN.md` | - | 2026-08-18 | 12 | ✓ | +| `compiler/W1_1_TYPING_INVENTORY.md` | - | 2026-08-02 | 28 | ✓ | +| `compiler/W4_ADMISSIBLE_EFFECTS_PLAN.md` | - | 2026-08-25 | 5 | ✓ | +| `compiler/compiler_enhancement.md` | - | 2026-08-18 | 12 | ✓ | +| `coverage/COVERAGE_AUDIT.md` | - | 2026-08-11 | 19 | ✓ | +| `domain/DOMAIN_AUDIT.md` | - | 2026-06-11 | 80 | ✓ | +| `domain/GA_EBM_ARCHITECTURE_REVIEW.md` | - | 2026-08-02 | 28 | ✓ | +| `roadmap/CF_CROSS_ELEMENT_PLAN.md` | - | 2026-06-30 | 61 | ✓ | +| `roadmap/MODEL_CLASS_ROADMAP.md` | - | 2026-08-12 | 18 | ✓ | +| `roadmap/ROADMAP_AUDIT.md` | - | 2026-08-11 | 19 | ✓ | diff --git a/docs/audit/generated/test_coverage.csv b/docs/audit/generated/test_coverage.csv index d702b044a..bd1de8e37 100644 --- a/docs/audit/generated/test_coverage.csv +++ b/docs/audit/generated/test_coverage.csv @@ -2,12 +2,12 @@ op,python_refs,lit_refs,negative_refs,total_refs,is_thinly_tested,dtype_variants abs,6,0,1,6,0,bf16 f32,directly_tested,6 direct test references absolute,6,0,0,6,0,bf16 f32 fp16 fp32,directly_tested,6 direct test references acos,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references -adafactor,24,0,0,24,0,fp32 fp64,directly_tested,24 direct test references +adafactor,38,0,1,38,0,fp32 fp64,directly_tested,38 direct test references adam,22,5,1,27,0,bf16 f16 fp16 fp32 fp64 fp8_e4m3,directly_tested,27 direct test references adamw,16,5,0,21,0,fp32 fp64,directly_tested,21 direct test references adaptive_pool,7,0,0,7,0,fp32,directly_tested,7 direct test references add,160,68,11,228,0,bf16 f16 f32 f64 fp16 fp32 int8,directly_tested,228 direct test references -add_decoupled_weight_decay,2,0,0,2,0,,directly_tested,2 direct test references +add_decoupled_weight_decay,2,0,0,2,0,fp32,directly_tested,2 direct test references alibi,7,0,0,7,0,bf16 f16 f32 fp32 int32,directly_tested,7 direct test references all_gather,6,0,0,6,0,fp32,directly_tested,6 direct test references all_reduce,20,2,4,22,0,bf16 f32 f64 fp32 fp8_e4m3 int8,directly_tested,22 direct test references @@ -52,8 +52,8 @@ cast,25,61,1,86,0,bf16 fp16 fp32 int8,directly_tested,86 direct test references cat,27,0,1,27,0,fp32,directly_tested,27 direct test references ceil,4,0,0,4,0,bf16 f32 fp16 fp32,directly_tested,4 direct test references center_crop,10,0,1,10,0,fp16,directly_tested,10 direct test references -centralize_grad,1,0,0,1,1,,structural_only,category default for 'grad_transform' -chained_schedule,1,0,0,1,1,,structural_only,category default for 'schedule' +centralize_grad,1,0,0,1,1,fp32,structural_only,category default for 'grad_transform' +chained_schedule,1,0,0,1,1,fp32,structural_only,category default for 'schedule' check_cauchy_riemann,0,0,0,0,1,,covered_by_family,exercised by complex_jit / CR conformance tests checkpoint,0,0,0,0,1,,structural_only,category default for 'transform' cholesky,28,31,0,59,0,bf16 f16 f32 fp16 fp32 fp64 fp8_e4m3 int32 int64 int8 nvfp4,directly_tested,59 direct test references @@ -80,8 +80,8 @@ clifford_rotor_sandwich,2,0,0,2,0,,directly_tested,2 direct test references clifford_vec_deriv,0,0,0,0,1,,covered_by_family,category default for 'geometric_algebra' clifford_wedge,1,0,0,1,1,,covered_by_family,category default for 'geometric_algebra' clip,16,0,1,16,0,fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,directly_tested,16 direct test references -clip_grad_norm,14,0,1,14,0,,directly_tested,14 direct test references -clip_grad_value,2,0,0,2,0,,directly_tested,2 direct test references +clip_grad_norm,14,0,1,14,0,fp32,directly_tested,14 direct test references +clip_grad_value,2,0,0,2,0,fp32,directly_tested,2 direct test references collective_permute,4,0,0,4,0,fp32,directly_tested,4 direct test references compilation_cache,2,0,0,2,0,,directly_tested,2 direct test references complex_abs,7,0,1,7,0,fp32,directly_tested,7 direct test references @@ -96,7 +96,7 @@ complex_sqrt,3,0,0,3,0,,directly_tested,3 direct test references cond,22,0,3,22,0,f32 fp16,directly_tested,22 direct test references conformal_energy_on_sphere,5,0,0,5,0,,directly_tested,5 direct test references conformal_jacobian,0,0,0,0,1,,covered_by_family,exercised by complex/conformal lane tests -constant_lr,3,0,0,3,0,,directly_tested,3 direct test references +constant_lr,3,0,0,3,0,fp32,directly_tested,3 direct test references contrastive_divergence_loss,13,0,1,13,0,fp64,directly_tested,13 direct test references contrastive_loss,9,0,1,9,0,,directly_tested,9 direct test references conv1d,6,0,0,6,0,fp16,directly_tested,6 direct test references @@ -106,8 +106,8 @@ conv_transpose,5,0,0,5,0,fp32,directly_tested,5 direct test references cos,6,0,0,6,0,bf16 f32 fp32 int16 int32 int64 int8,directly_tested,6 direct test references cosh,3,0,0,3,0,bf16 f32,directly_tested,3 direct test references cosine_embedding_loss,6,0,1,6,0,,directly_tested,6 direct test references -cosine_lr,2,0,0,2,0,,directly_tested,2 direct test references -cosine_warmup_lr,1,0,0,1,1,,structural_only,category default for 'schedule' +cosine_lr,2,0,0,2,0,fp32,directly_tested,2 direct test references +cosine_warmup_lr,1,0,0,1,1,fp32,structural_only,category default for 'schedule' count_nonzero,12,0,0,12,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,12 direct test references cross_attention,14,0,0,14,0,fp16,directly_tested,14 direct test references cross_entropy_loss,18,0,1,18,0,,directly_tested,18 direct test references @@ -123,7 +123,7 @@ custom_jvp,0,0,0,0,1,,structural_only,category default for 'extension' custom_lowering,0,0,0,0,1,,structural_only,category default for 'extension' custom_primitive,2,0,0,2,0,,directly_tested,2 direct test references custom_vjp,0,0,0,0,1,,structural_only,category default for 'extension' -cyclical_lr,1,0,0,1,1,,structural_only,category default for 'schedule' +cyclical_lr,1,0,0,1,1,fp32,structural_only,category default for 'schedule' dataset_batch,0,0,0,0,1,,structural_only,category default for 'data' dataset_checkpoint,0,0,0,0,1,,structural_only,category default for 'data' dataset_filter,0,0,0,0,1,,structural_only,category default for 'data' @@ -173,7 +173,7 @@ edm_loss_weight,3,0,0,3,0,fp16,directly_tested,3 direct test references edm_precondition,3,0,0,3,0,fp16,directly_tested,3 direct test references eigh,4,0,0,4,0,fp32,directly_tested,4 direct test references einsum,13,0,2,13,0,fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,directly_tested,13 direct test references -ema_update,1,0,0,1,1,,structural_only,category default for 'grad_transform' +ema_update,1,0,0,1,1,fp32,structural_only,category default for 'grad_transform' empty_state_tree,2,0,0,2,0,,directly_tested,2 direct test references entmax15,1,0,0,1,1,fp64,directly_tested,public forward plus simplex and FD-VJP tests eq,5,1,1,6,0,bf16 f32 fp16 fp32 fp4_e2m1 fp8_e4m3 int32 int64 int8,directly_tested,6 direct test references @@ -188,7 +188,7 @@ factorized_matmul,12,0,1,12,0,fp32,directly_tested,12 direct test references factorized_pos_emb,12,0,1,12,0,fp16,directly_tested,12 direct test references fake_quantize,6,0,1,6,0,,directly_tested,6 direct test references fft,38,4,1,42,0,f32 fp32 fp64 int32 int64 int8,directly_tested,42 direct test references -flash_attn,145,79,11,224,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp6_e2m3 fp6_e3m2 fp8_e4m3 fp8_e5m2 int32 int64 int8 nvfp4,directly_tested,224 direct test references +flash_attn,145,82,11,227,0,bf16 f16 f32 f64 fp16 fp32 fp4_e2m1 fp64 fp6_e2m3 fp6_e3m2 fp8_e4m3 fp8_e5m2 int32 int64 int8 nvfp4,directly_tested,227 direct test references flatten,7,5,0,12,0,,directly_tested,12 direct test references flip,7,0,0,7,0,,directly_tested,7 direct test references floor,6,0,1,6,0,bf16 f32 fp16 fp32,directly_tested,6 direct test references @@ -232,7 +232,7 @@ info_nce_loss,7,0,0,7,0,,directly_tested,7 direct test references instance_norm,9,0,0,9,0,,directly_tested,9 direct test references interpolate,12,0,1,12,0,fp16,directly_tested,12 direct test references inv,3,0,1,3,0,fp32,directly_tested,3 direct test references -inverse_sqrt_lr,5,0,1,5,0,,directly_tested,5 direct test references +inverse_sqrt_lr,5,0,1,5,0,fp32,directly_tested,5 direct test references irfft,17,2,1,19,0,fp32 fp64 int32 int64 int8,directly_tested,19 direct test references is_concyclic,3,0,0,3,0,,directly_tested,3 direct test references isfinite,7,0,0,7,0,f32 fp16 fp32,directly_tested,7 direct test references @@ -262,7 +262,7 @@ lightning_attention,14,2,0,16,0,,directly_tested,16 direct test references linear_attn,44,8,2,52,0,,directly_tested,52 direct test references linear_attn_state,7,0,1,7,0,fp16,directly_tested,7 direct test references linear_general,15,0,0,15,0,,directly_tested,15 direct test references -linear_warmup_lr,1,0,0,1,1,,structural_only,category default for 'schedule' +linear_warmup_lr,1,0,0,1,1,fp32,structural_only,category default for 'schedule' lion,25,0,2,25,0,fp32 fp64,directly_tested,25 direct test references load_balance_loss,6,0,0,6,0,,directly_tested,6 direct test references load_sharded,3,0,1,3,0,,directly_tested,3 direct test references @@ -312,7 +312,7 @@ moe,32,0,0,32,0,fp16 fp4_e2m1 fp6_e2m3 fp8_e4m3 nvfp4,directly_tested,32 direct moe_combine,10,0,0,10,0,,directly_tested,10 direct test references moe_dispatch,9,3,0,12,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,12 direct test references moe_swiglu_block,14,19,1,33,0,bf16 fp8_e4m3 nvfp4,directly_tested,33 direct test references -momentum,12,1,0,13,0,fp64,directly_tested,13 direct test references +momentum,12,1,0,13,0,fp32 fp64,directly_tested,13 direct test references mor_partition,9,3,1,12,0,fp16 fp32,directly_tested,12 direct test references mor_router,10,4,1,14,0,fp16,directly_tested,14 direct test references mor_scatter,10,2,1,12,0,fp16 fp32,directly_tested,12 direct test references @@ -327,7 +327,7 @@ multi_head_attention,9,0,1,9,0,,directly_tested,9 direct test references muon,15,0,0,15,0,fp32,directly_tested,15 direct test references named_sharding,0,0,0,0,1,,structural_only,category default for 'sharding' ne,3,1,0,4,0,bf16 f32 fp32 int32 int64 int8,directly_tested,4 direct test references -nesterov,12,3,0,15,0,fp64,directly_tested,15 direct test references +nesterov,12,3,0,15,0,fp32 fp64,directly_tested,15 direct test references nonzero,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only norm,6,2,1,8,0,fp32,directly_tested,8 direct test references normalize_group_advantages,12,0,0,12,0,bf16 f16 fp16 fp32 fp64 fp8_e4m3,directly_tested,12 direct test references @@ -351,8 +351,8 @@ pmap,0,0,0,0,1,,structural_only,category default for 'transform' pmax,2,0,0,2,0,,directly_tested,2 direct test references pmean,2,0,0,2,0,,directly_tested,2 direct test references pmin,2,0,0,2,0,,directly_tested,2 direct test references -polyak_avg,1,0,0,1,1,,structural_only,category default for 'grad_transform' -polynomial_lr,1,0,0,1,1,,structural_only,category default for 'schedule' +polyak_avg,1,0,0,1,1,fp32,structural_only,category default for 'grad_transform' +polynomial_lr,1,0,0,1,1,fp32,structural_only,category default for 'schedule' popcount,15,0,0,15,0,bf16 f16 f32 fp16 fp32 fp64 fp8_e4m3 int16 int32 int64 int8,directly_tested,15 direct test references pow,15,0,2,15,0,bf16 f32 fp32,directly_tested,15 direct test references power_attn,12,0,1,12,0,fp16,directly_tested,12 direct test references @@ -422,7 +422,7 @@ segment_reduce,7,0,0,7,0,fp32,directly_tested,7 direct test references select,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only selective_ssm,72,10,4,82,0,bf16 f16 fp16 fp32,directly_tested,82 direct test references seq2seq_loss,7,0,0,7,0,,directly_tested,7 direct test references -sgd,22,10,0,32,0,fp64,directly_tested,32 direct test references +sgd,22,10,0,32,0,fp32 fp64,directly_tested,32 direct test references shard_map,0,0,0,0,1,,structural_only,category default for 'sharding' sharded_dataset,2,0,0,2,0,,directly_tested,2 direct test references sigmoid,17,9,1,26,0,bf16 f16 f32 fp32,directly_tested,26 direct test references @@ -490,11 +490,11 @@ trace,3,0,1,3,0,fp32,directly_tested,3 direct test references training.loss_adamw,1,0,0,1,1,f32,directly_tested,parameterized exact x86/ROCm fused-versus-unfused oracle training.loss_sgd,1,0,0,1,1,f32,directly_tested,parameterized exact x86/ROCm fused-versus-unfused oracle transpose,33,41,0,74,0,bf16 f16 f32 fp16 fp32 int32 int64 int8,directly_tested,74 direct test references -tree_flatten,17,0,1,17,0,,directly_tested,17 direct test references +tree_flatten,17,0,1,17,0,fp32,directly_tested,17 direct test references tree_map,5,0,1,5,0,,directly_tested,5 direct test references tree_reduce,4,0,1,4,0,,directly_tested,4 direct test references tree_transpose,2,0,1,2,0,,directly_tested,2 direct test references -tree_unflatten,13,0,1,13,0,,directly_tested,13 direct test references +tree_unflatten,13,0,1,13,0,fp32,directly_tested,13 direct test references tri_solve,18,2,0,20,0,bf16 f16 f32 fp16 fp32 int32 int64,directly_tested,20 direct test references tridiagonal_solve,13,1,1,14,0,f64 fp32,directly_tested,14 direct test references triplet_loss,10,0,1,10,0,,directly_tested,10 direct test references diff --git a/docs/audit/generated/test_coverage.md b/docs/audit/generated/test_coverage.md index 8e7a87af7..b0cf74a4f 100644 --- a/docs/audit/generated/test_coverage.md +++ b/docs/audit/generated/test_coverage.md @@ -7,11 +7,11 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by ## Headline - **515** ops in `primitive_coverage` registry. -- **5815** total Python-test references, **1463** total lit-fixture references. +- **5829** total Python-test references, **1466** total lit-fixture references. - **70** ops have **zero** references in either test surface. - **113** ops have ≤1 reference ("thinly tested"). - **188** ops have ≥10 references ("well tested"). -- **187** ops have at least one associated `pytest.raises` negative test. +- **188** ops have at least one associated `pytest.raises` negative test. ## Top 20 most-tested ops @@ -19,7 +19,7 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by |----|--------:|---------:|------:|----:|--------| | `matmul` | 490 | 308 | 798 | 24 | `bf16`, `f16`, `f32`, `f64` … | | `add` | 160 | 68 | 228 | 11 | `bf16`, `f16`, `f32`, `f64` … | -| `flash_attn` | 145 | 79 | 224 | 11 | `bf16`, `f16`, `f32`, `f64` … | +| `flash_attn` | 145 | 82 | 227 | 11 | `bf16`, `f16`, `f32`, `f64` … | | `softmax` | 152 | 56 | 208 | 31 | `bf16`, `f16`, `f32`, `f64` … | | `mul` | 155 | 41 | 196 | 13 | `bf16`, `f16`, `f32`, `f64` … | | `relu` | 132 | 39 | 171 | 11 | `bf16`, `f16`, `f32`, `f64` … | diff --git a/docs/spec/LOWERING_PIPELINE_SPEC.md b/docs/spec/LOWERING_PIPELINE_SPEC.md index 136fb5c56..96ab2e88f 100644 --- a/docs/spec/LOWERING_PIPELINE_SPEC.md +++ b/docs/spec/LOWERING_PIPELINE_SPEC.md @@ -580,8 +580,13 @@ The copy descriptor retains the unpadded logical K/V source extent. Dynamic slice coordinates are explicit index operands, so the target copy consumer can zero-fill the final ragged KV block without inventing a computed base pointer. `boundary_mask` carries causal and sliding-window policy with absolute Q/KV -offsets. Non-zero dropout uses `block_dropout`, keyed by the absolute KV offset, -and requires an explicit seed. +offsets. Non-zero dropout uses `block_dropout`, keyed by the absolute KV offset +**and a per-instance `stream_offset`**, and requires an explicit seed. The +stream offset is the second counter axis: a rank-4 attention is distributed +into `B*H` rank-2 instances whose KV offsets each restart at 0, so without it +every instance would replay one identical mask (Decision #18). The producer +passes `(b*H + h) * Sq * Sk_padded` — a disjoint counter block per instance — +and an attention that was never distributed passes 0. For `tessera.matmul` inside a mesh region: ```mlir diff --git a/python/tessera/__init__.py b/python/tessera/__init__.py index 2ada351a6..a7699550c 100644 --- a/python/tessera/__init__.py +++ b/python/tessera/__init__.py @@ -1148,14 +1148,36 @@ def adafactor(params, grads, state=None, column_state=None, **kwargs): # operands. A full state returns ``(new_param, new_state)``; factored # row/column state returns ``(new_param, new_row, new_column)``. The # public tree optimizer and its dictionary state remain unchanged. + # + # ``step`` follows the flat ``adam``/``adamw`` ABI above: a 1-based + # kwarg, because the flat form carries no state dictionary to read it + # from. It selects the bias-corrected second-moment decay + # (`optim.adafactor_decay`); without it the first updates are inflated + # by 1/sqrt(1 - beta2**step). + # + # ABSENT is not the same as 1. adafactor_decay(b2, 1) is exactly 0 -- + # correct for a genuine first step, where v_1 = g^2 -- so defaulting a + # missing ``step`` to 1 would make every call of a stateful caller that + # never passes one discard the moments it just supplied, turning a + # stateful optimizer into a stateless one with no diagnostic. Such a + # caller keeps the legacy uncorrected decay instead: no bias + # correction, exactly as before this ABI grew a step, and strictly + # better than silently resetting its state. if isinstance(state, np.ndarray): + from . import optim as _optim + param = np.asarray(params) grad = np.asarray(grads, dtype=np.float32) row_or_full = np.asarray(state, dtype=np.float32) if param.shape != grad.shape: raise ValueError("flat Adafactor parameter and gradient must match") lr = np.float32(float(kwargs.get("lr", 1.0e-3))) - beta2 = np.float32(float(kwargs.get("beta2", 0.999))) + nominal_beta2 = float(kwargs.get("beta2", 0.999)) + beta2 = np.float32( + _optim.adafactor_decay(nominal_beta2, int(kwargs["step"])) + if "step" in kwargs + else nominal_beta2 + ) eps = np.float32(float(kwargs.get("eps", 1.0e-30))) grad2 = grad * grad if column_state is None: @@ -1179,7 +1201,26 @@ def adafactor(params, grads, state=None, column_state=None, **kwargs): update = grad / (np.sqrt(scale) + eps) return param - lr * update, new_row, new_column from . import optim as _optim - return _optim.adafactor(params, grads, state, **kwargs) + tree_kwargs = dict(kwargs) + declared_step = tree_kwargs.pop("step", None) + if declared_step is not None: + # The tree ABI owns the step counter inside its state dict. A + # caller that also declares one must agree with it: silently + # preferring either would make the two ABIs disagree on which + # decay was applied (Decision #21a — semantic keys never default). + if state is not None and not isinstance(state, dict): + raise ValueError( + "tree Adafactor state must be the optimizer state dict " + f"carrying 'step'; got {type(state).__name__}" + ) + carried = 0 if state is None else int(state.get("step", 0)) + if int(declared_step) != carried + 1: + raise ValueError( + "tree Adafactor derives its step from state['step']; " + f"declared step={int(declared_step)} disagrees with the " + f"carried state step {carried} (expected {carried + 1})" + ) + return _optim.adafactor(params, grads, state, **tree_kwargs) def lion(params, grads, state=None, **kwargs): # The compiler-visible flat ABI carries the moment tensor directly and diff --git a/python/tessera/autodiff/vjp.py b/python/tessera/autodiff/vjp.py index c6ff53c33..249c2026b 100644 --- a/python/tessera/autodiff/vjp.py +++ b/python/tessera/autodiff/vjp.py @@ -4687,10 +4687,17 @@ def forward(p, g, s): ) return d_params, d_grads, d_state + from tessera.optim import adafactor_decay + p = np.asarray(params) g = np.asarray(grads, dtype=np.float64) do = np.asarray(dout, dtype=np.float64) - b2 = float(beta2) + # Differentiate the forward that actually runs: the second-moment decay is + # the bias-corrected, step-dependent one (`optim.adafactor_decay`), not the + # nominal `beta2`. The step is 1-based and comes from the carried state. + b2 = adafactor_decay( + float(beta2), 1 if state is None else int(state.get("step", 0)) + 1 + ) epsilon = float(eps) dparam = do if state is None: diff --git a/python/tessera/compiler/emit/nvidia_solver_krylov.py b/python/tessera/compiler/emit/nvidia_solver_krylov.py index dfec670f2..c412b1242 100644 --- a/python/tessera/compiler/emit/nvidia_solver_krylov.py +++ b/python/tessera/compiler/emit/nvidia_solver_krylov.py @@ -68,8 +68,38 @@ def _source() -> str: return *scalar; } +// Two matvecs, chosen per solver from MEASURED device results rather than from +// the access-pattern argument alone (RTX 5070, sm_120, medians of 9 reps). +// +// solver / shape one thread per row one warp per row +// dense_cg 2049 1.011 ms 2.277 ms (0.44x) +// dense_cg 1025 0.534 ms 1.073 ms (0.50x) +// dense_gmres 2049 0.967 ms 0.621 ms (1.56x) +// dense_gmres 1025 0.579 ms 0.400 ms (1.45x) +// +// The warp form coalesces better -- 32 lanes touch one row instead of 32 rows +// -- but a COOPERATIVE launch caps the grid at what stays resident, so it also +// buys 32x fewer rows in flight at the same block count. GMRES absorbs that and +// wins; CG, which grid-syncs far more often per iteration, does not and loses +// roughly half its throughput. Shipping one form for both regressed CG by up to +// 2.3x, which no amount of reasoning about transactions per load would have +// revealed. template -__device__ void tsr_matvec(const T *a, const float *x, float *y, int n, +__device__ void tsr_matvec_scalar(const T *a, const float *x, float *y, int n, + cg::grid_group grid) { + long gid = (long)blockIdx.x * blockDim.x + threadIdx.x; + long stride = (long)gridDim.x * blockDim.x; + for (long row = gid; row < n; row += stride) { + float sum = 0.0f; + const T *arow = a + row * (long)n; + for (int col = 0; col < n; ++col) sum = fmaf(tsr_load(arow, col), x[col], sum); + y[row] = sum; + } + grid.sync(); +} + +template +__device__ void tsr_matvec_warp(const T *a, const float *x, float *y, int n, cg::grid_group grid) { // One WARP per row, lane `l` walking columns `l, l+32, ...`: the 32 lanes of // a single load instruction then touch 32 consecutive elements of one row @@ -126,7 +156,7 @@ def _source() -> str: grid.sync(); for (int iteration = 0; iteration < max_iterations && *status == 0; ++iteration) { - tsr_matvec(a, p, ap, n, grid); + tsr_matvec_scalar(a, p, ap, n, grid); float local_pap = 0.0f; for (long i = gid; i < n; i += stride) local_pap += p[i] * ap[i]; float pap = tsr_grid_sum(local_pap, partials, &scalars[3], grid); @@ -154,7 +184,7 @@ def _source() -> str: if (scalars[6] != 0.0f) { // A convergence decision is always based on the true residual b-Ax. - tsr_matvec(a, x, workspace, n, grid); + tsr_matvec_scalar(a, x, workspace, n, grid); float local_true = 0.0f; for (long i = gid; i < n; i += stride) { r[i] = tsr_load(b, i) - workspace[i]; @@ -186,7 +216,7 @@ def _source() -> str: if (gid == 0 && *status == 0) *status = 3; grid.sync(); // Return A*x as the final matvec state, not a stale search-direction product. - tsr_matvec(a, x, ap, n, grid); + tsr_matvec_scalar(a, x, ap, n, grid); for (long i = gid; i < n; i += stride) r[i] = tsr_load(b, i) - ap[i]; } @@ -214,7 +244,7 @@ def _source() -> str: grid.sync(); while (*iterations < max_iterations && *status == 0) { - tsr_matvec(a, x, w, n, grid); + tsr_matvec_warp(a, x, w, n, grid); float local_r = 0.0f; for (long i = gid; i < n; i += stride) { r[i] = tsr_load(b, i) - w[i]; local_r += r[i] * r[i]; @@ -237,7 +267,7 @@ def _source() -> str: for (int j = 0; j < restart && *iterations < max_iterations; ++j) { const float *vj = basis + (long)j * n; - tsr_matvec(a, vj, w, n, grid); + tsr_matvec_warp(a, vj, w, n, grid); // Twice-modified Gram-Schmidt protects orthogonality on difficult // matrices while retaining a deterministic reduction order. @@ -316,7 +346,7 @@ def _source() -> str: grid.sync(); // Estimated Givens residual never establishes convergence by itself. - tsr_matvec(a, x, w, n, grid); + tsr_matvec_warp(a, x, w, n, grid); float local_true = 0.0f; for (long k = gid; k < n; k += stride) { r[k] = tsr_load(b, k) - w[k]; local_true += r[k] * r[k]; diff --git a/python/tessera/compiler/scheduled_matmul.py b/python/tessera/compiler/scheduled_matmul.py index 3f5ee4a9e..39948565d 100644 --- a/python/tessera/compiler/scheduled_matmul.py +++ b/python/tessera/compiler/scheduled_matmul.py @@ -17,6 +17,13 @@ _HASH_RE = re.compile(r'tessera\.schedule_hash = "([0-9a-f]{64})"') +# Mirrors kScheduledSm120MatmulPrefix in +# src/compiler/codegen/tessera_gpu_backend_NVIDIA/runtime/cuda/tessera_nvidia_ptx_launch.cpp. +# The runtime selects the scheduled-matmul launcher by this prefix, so it is +# ABI, not cosmetics. +_SM120_SCHEDULED_MATMUL_PREFIX = "nvidia_sm120_scheduled_matmul_" + + @dataclass(frozen=True) class ScheduledMatmulArtifact: graph_ir: str @@ -115,6 +122,15 @@ def lower_scheduled_matmul( targeted = copy.deepcopy(module) targeted.module_attrs["tessera.target"] = f'"{contract[0]}"' targeted.module_attrs["tessera.arch"] = f'"{contract[1]}"' + # The Tile kernel symbol is derived by the C++ passes from this function's + # name, and the runtime dispatches scheduled sm_120 matmuls by name prefix, + # so the prefix has to be applied HERE -- renaming only the Python-side + # descriptor entry desynchronises it from the symbol actually present in the + # compiled PTX ("native PTX is missing entry ..."). + if targeted.functions and contract[0] == "nvidia_sm120": + fn0 = targeted.functions[0] + if not fn0.name.startswith(_SM120_SCHEDULED_MATMUL_PREFIX): + fn0.name = f"{_SM120_SCHEDULED_MATMUL_PREFIX}{fn0.name}" graph_ir = targeted.to_mlir(target=target, canonical=True) schedule_ir = run_tessera_opt(tool, graph_ir, "--tessera-graph-to-schedule") tile_ir = run_tessera_opt(tool, schedule_ir, "--tessera-schedule-to-tile") @@ -353,7 +369,19 @@ def extent(value: object) -> int | None: ) if reduced: suffix += "_outf16" - function_name = f"{function.name}{suffix}" + ( + # The runtime dispatches scheduled sm_120 matmuls by NAME PREFIX + # (kScheduledSm120MatmulPrefix in tessera_nvidia_ptx_launch.cpp). Naming + # the kernel after the caller's Graph function made that dispatch depend + # on what the user happened to call their function: every name without + # the prefix fell through the runtime's strcmp chain and the launch + # returned rc=5. Exactly one place in the tree — a benchmark — named a + # function to satisfy it; every other caller silently could not launch. + # The prefix is part of the ABI, so the compiler emits it rather than + # asking the frontend to spell it. + base_name = function.name + if not base_name.startswith(_SM120_SCHEDULED_MATMUL_PREFIX): + base_name = f"{_SM120_SCHEDULED_MATMUL_PREFIX}{base_name}" + function_name = f"{base_name}{suffix}" + ( "_macro_kernel" if _uses_sm120_macro_cta(m, n, k, storage, accum) else "_kernel" diff --git a/python/tessera/compiler/stateful_training.py b/python/tessera/compiler/stateful_training.py index 808a80d23..ecc9b2e95 100644 --- a/python/tessera/compiler/stateful_training.py +++ b/python/tessera/compiler/stateful_training.py @@ -172,12 +172,25 @@ def build_adafactor_vjp_state_contract( raise ValueError("factored Adafactor requires rank-2+ parameter state") if topology == "full" and len(shape) >= 2: raise ValueError("full Adafactor requires rank-0/1 parameter state") - numeric = { + from tessera.optim import adafactor_decay + + # `beta2` is the caller's nominal asymptotic decay and `step` its 1-based + # update index; `beta2_effective` is the bias-corrected decay the forward + # actually applied (`optim.adafactor_decay`) and therefore the one the + # adjoint must differentiate. Both are recorded so the contract round-trips + # through its own `numeric` dict without re-applying the correction. + nominal_beta2 = float(kwargs.get("beta2", 0.999)) + step = int(kwargs.get("step", 1)) + numeric: dict[str, Any] = { "lr": float(kwargs.get("lr", 1.0e-3)), - "beta2": float(kwargs.get("beta2", 0.999)), + "beta2": nominal_beta2, "eps": float(kwargs.get("eps", 1.0e-30)), + "step": step, + "beta2_effective": adafactor_decay(nominal_beta2, step), } - if not all(math.isfinite(value) for value in numeric.values()): + if not all( + math.isfinite(float(value)) for value in numeric.values() + ): raise ValueError("Adafactor VJP lineage requires finite coefficients") state_shapes = ( (shape[:-1], (shape[-1],)) if topology == "factored" else (shape,) @@ -720,7 +733,7 @@ def lower_scheduled_adafactor_vjp( write_order = ", ".join(map(str, range(len(output_shapes)))) schedule_ir = f'''module attributes {{tessera.target = "{compiler_target}", tessera.arch = "{architecture}"}} {{ func.func @tessera_adafactor_vjp({args}) -> ({", ".join(output_types)}) {{ - %grads:{len(output_types)} = schedule.adafactor_vjp {inputs} {{artifact_hash = "{digest}", lineage_payload = {json.dumps(payload)}, arch = "{architecture}", topology = "{topology}", learning_rate = {float(numeric["lr"]):.9e} : f32, beta2 = {float(numeric["beta2"]):.9e} : f32, epsilon = {float(numeric["eps"]):.9e} : f32, mutation_mode = "{mutation["mode"]}", alias_policy = "{mutation["alias_policy"]}", state_transition = "{mutation["state_transition"]}", ordered_writes = array, workgroup_size = {256 if target != "x86" else 1} : i64}} : {", ".join(input_types)} -> {", ".join(output_types)} + %grads:{len(output_types)} = schedule.adafactor_vjp {inputs} {{artifact_hash = "{digest}", lineage_payload = {json.dumps(payload)}, arch = "{architecture}", topology = "{topology}", learning_rate = {float(numeric["lr"]):.9e} : f32, beta2 = {float(numeric["beta2_effective"]):.9e} : f32, epsilon = {float(numeric["eps"]):.9e} : f32, mutation_mode = "{mutation["mode"]}", alias_policy = "{mutation["alias_policy"]}", state_transition = "{mutation["state_transition"]}", ordered_writes = array, workgroup_size = {256 if target != "x86" else 1} : i64}} : {", ".join(input_types)} -> {", ".join(output_types)} schedule.artifact {{hash = "{digest}", arch = "{architecture}", shape_key = "family=adafactor_vjp;topology={topology};shape={'x'.join(map(str, shape))};storage=f32", numeric_policy = "f32;functional_no_alias"}} return {", ".join(f"%grads#{index}" for index in range(len(output_types)))} : {", ".join(output_types)} }} @@ -769,7 +782,9 @@ def validate_scheduled_adafactor_vjp_metadata( raise ValueError("Adafactor VJP scheduled artifact policy is stale") for name, key in ( ("learning_rate", "lr"), - ("beta2", "beta2"), + # The scheduled artifact carries the bias-corrected decay, which is what + # the physical adjoint consumes; `numeric["beta2"]` stays the nominal. + ("beta2", "beta2_effective"), ("epsilon", "eps"), ): match = re.search(rf"{name} = ([+\-0-9.eE]+) : f32", tile_ir) diff --git a/python/tessera/optim.py b/python/tessera/optim.py index 84b298a4e..45bb6adc9 100644 --- a/python/tessera/optim.py +++ b/python/tessera/optim.py @@ -12,6 +12,8 @@ import math from typing import Any, Callable +import warnings + import numpy as np @@ -351,6 +353,112 @@ def adam( ) +_ADAFACTOR_V_REPRESENTATION = "debiased_v1" + + +def migrate_adafactor_state( + state: dict[str, Any], beta2: float, *, state_dtype: str = "fp32" +) -> dict[str, Any]: + """Convert a pre-bias-correction Adafactor state to the debiased form. + + Call this ONCE on a checkpoint written before `adafactor_decay` existed. + Such a state holds the raw EMA ``v_raw = (1 - b2**t) * v_debiased``, so it + is divided by that factor -- exactly inverting the bias it was written + with. A state already carrying the marker is returned unchanged. + """ + if state.get("v_representation") == _ADAFACTOR_V_REPRESENTATION: + return state + step = int(state.get("step", 0)) + migrated = dict(state) + migrated["v_representation"] = _ADAFACTOR_V_REPRESENTATION + bias = 1.0 - float(beta2) ** step + if step <= 0 or bias <= 0.0: + return migrated + migrated["v"] = _adafactor_tree_map_unary( + lambda s: _adafactor_scale_state(s, 1.0 / bias, state_dtype=state_dtype), + state["v"], + ) + return migrated + + +def _warn_if_adafactor_state_unmarked(state: dict[str, Any]) -> None: + """Flag a state whose second-moment representation is ambiguous. + + ``state["v"]`` carries the DEBIASED estimate since the bias correction + landed; before that it carried the raw EMA. A state without the marker + could be either, and the two cannot be told apart from the values. + + This deliberately does NOT rescale. Auto-migrating on a missing marker + would silently rewrite every hand-built state dict -- which is a worse + failure than the one it fixes, since a test or caller that assembles its + own state has no legacy bias to remove. Callers restoring a genuine + pre-correction checkpoint should call `migrate_adafactor_state` once. + """ + if state.get("v_representation") == _ADAFACTOR_V_REPRESENTATION: + return + if int(state.get("step", 0)) <= 0: + return + warnings.warn( + "Adafactor state carries no 'v_representation' marker at step " + f"{int(state.get('step', 0))}. Since the bias correction landed, " + "state['v'] holds the DEBIASED second moment; a checkpoint written " + "before that holds the raw EMA and will be read as if it were already " + "debiased, erasing most of its history. If this state came from an " + "older checkpoint, pass it through optim.migrate_adafactor_state " + "once; if it was built by hand, set " + "state['v_representation'] = optim._ADAFACTOR_V_REPRESENTATION.", + RuntimeWarning, + stacklevel=3, + ) + + +def adafactor_decay(beta2: float, step: int) -> float: + """Bias-corrected second-moment decay for the Adafactor update at ``step``. + + Adafactor (Shazeer & Stern 2018) does **not** run a fixed second-moment + decay. A raw EMA ``v_t = b2*v_{t-1} + (1-b2)*g^2`` started from + ``v_0 = 0`` is biased low by ``1 - b2**t``, so the first updates are + inflated by ``1/sqrt(1 - b2**t)`` — 31.6x at the default ``b2 = 0.999``, + and still >2x after 1000 steps. + + This returns the step-dependent decay + + b2_t = b2 * (1 - b2**(t-1)) / (1 - b2**t) + + for which the recursion carries the *debiased* estimate directly: + + v_t = b2_t*v_{t-1} + (1 - b2_t)*g_t^2 == EMA_t / (1 - b2**t) + + (at ``t = 1``, ``b2_1 = 0`` and ``v_1 = g_1^2`` exactly). So this is + algebraically the explicit ``1 - beta2**step`` correction that ``adamw`` + already applies above, expressed as a decay rate the way the paper's own + ``1 - t**-0.8`` schedule is. Expressing it as a decay rate is what makes + it landable: every physical Adafactor kernel (AVX-512, gfx1151, sm_120) + already takes ``beta2`` as a scalar, so the correction is applied + host-side and needs no kernel ABI change. Unlike ``1 - t**-0.8`` it also + preserves the caller's ``beta2`` as the asymptotic decay instead of + silently discarding it (Decision #21a). + + ``step`` is 1-based: it is the index of the update being computed, i.e. + ``state["step"] + 1``. + """ + b2 = float(beta2) + t = int(step) + if not 0.0 <= b2 < 1.0: + raise ValueError( + f"adafactor beta2 must lie in [0, 1); got {beta2!r}" + ) + if t < 1: + raise ValueError( + f"adafactor step is 1-based and must be >= 1; got {step!r}" + ) + prev = 1.0 - b2 ** (t - 1) + current = 1.0 - b2**t + if current <= 0.0: + return 0.0 + return b2 * prev / current + + def adafactor( params: Tree, grads: Tree, @@ -371,9 +479,26 @@ def adafactor( """ base_params = _master_tree(params, state, master_dtype) if state is None: - state = {"v": tree_map(lambda p: _adafactor_zero_state(_asarray(p), state_dtype=state_dtype), params), "step": 0} + state = { + "v": tree_map(lambda p: _adafactor_zero_state(_asarray(p), state_dtype=state_dtype), params), + "step": 0, + "v_representation": _ADAFACTOR_V_REPRESENTATION, + } + # `v` changed meaning when the bias correction landed: the step-dependent + # decay makes the recursion carry the DEBIASED estimate directly, where it + # previously carried the raw EMA. A checkpoint written before that holds + # the raw form, and reading it as debiased silently erases most of the + # accumulated history -- at step 2 with the default beta2 the effective + # decay is ~0.5, so a stored 0.001*g1^2 is treated as the whole prior + # estimate. The schema had no marker to tell the two apart, so it gets one, + # and a state without it is migrated rather than misread. + _warn_if_adafactor_state_unmarked(state) + # The tracked step is finally used: without it the zero-initialized second + # moment is biased low and the first updates are inflated by + # 1/sqrt(1 - beta2**step). See `adafactor_decay`. + decay = adafactor_decay(beta2, int(state["step"]) + 1) new_v = _adafactor_tree_map( - lambda s, g: _adafactor_update_state(s, _compute_array(g, compute_dtype), beta2, state_dtype=state_dtype), + lambda s, g: _adafactor_update_state(s, _compute_array(g, compute_dtype), decay, state_dtype=state_dtype), state["v"], grads, ) @@ -384,7 +509,15 @@ def adafactor( ) new_master = tree_map2(lambda p, u: _compute_array(p, compute_dtype) - float(lr) * _compute_array(u, compute_dtype), base_params, updates) new_params = tree_map2(lambda p_new, p_orig: _cast_like_param(p_new, p_orig, cast_updates_to_param_dtype), new_master, params) - return new_params, _attach_master_state({"v": new_v, "step": int(state["step"]) + 1}, new_master, master_dtype) + return new_params, _attach_master_state( + { + "v": new_v, + "step": int(state["step"]) + 1, + "v_representation": _ADAFACTOR_V_REPRESENTATION, + }, + new_master, + master_dtype, + ) def _is_adafactor_slot(x: Any) -> bool: @@ -413,6 +546,33 @@ def _adafactor_zero_state(arr: np.ndarray, *, state_dtype: str = "fp32"): return {"v": np.zeros_like(arr, dtype=_np_dtype(state_dtype)), "factored": False} +def _adafactor_tree_map_unary(fn: Callable[[Any], Any], slot_tree: Tree) -> Tree: + if _is_adafactor_slot(slot_tree): + return fn(slot_tree) + if isinstance(slot_tree, dict): + return {k: _adafactor_tree_map_unary(fn, v) for k, v in slot_tree.items()} + if isinstance(slot_tree, tuple): + return tuple(_adafactor_tree_map_unary(fn, s) for s in slot_tree) + if isinstance(slot_tree, list): + return [_adafactor_tree_map_unary(fn, s) for s in slot_tree] + return fn(slot_tree) + + +def _adafactor_scale_state(state, scale: float, *, state_dtype: str = "fp32"): + """Scale a slot's second moments, preserving the factored/full shape.""" + dtype = _np_dtype(state_dtype) + if state["factored"]: + return { + "row": (np.asarray(state["row"], dtype=dtype) * dtype(scale)).astype(dtype, copy=False), + "col": (np.asarray(state["col"], dtype=dtype) * dtype(scale)).astype(dtype, copy=False), + "factored": True, + } + return { + "v": (np.asarray(state["v"], dtype=dtype) * dtype(scale)).astype(dtype, copy=False), + "factored": False, + } + + def _adafactor_update_state(state, grad: np.ndarray, beta2: float, *, state_dtype: str = "fp32"): grad2 = grad.astype(_np_dtype(state_dtype), copy=False) ** 2 if state["factored"]: diff --git a/python/tessera/runtime.py b/python/tessera/runtime.py index 6cdfa2146..cb3022b52 100644 --- a/python/tessera/runtime.py +++ b/python/tessera/runtime.py @@ -10078,8 +10078,9 @@ def _execute_nvidia_adafactor_backward( from .compiler.nvidia_training import package_adafactor_backward package = package_adafactor_backward( + # The bias-corrected decay the forward applied, not the nominal beta2. topology=topology, lr=float(numeric["lr"]), - beta2=float(numeric["beta2"]), epsilon=float(numeric["eps"]), + beta2=float(numeric["beta2_effective"]), epsilon=float(numeric["eps"]), ) count = int(arrays[operands[0]].size) buffers = { @@ -15890,9 +15891,15 @@ def _execute_x86_compiled_adafactor( operand_names = [str(name) for name in ops[0].get("operands", [])] values = _bind_launch_args(args, names) kwargs = ops[0].get("kwargs") or {} + from .optim import adafactor_decay + parameters = { "lr": float(kwargs.get("lr", 1e-3)), - "beta2": float(kwargs.get("beta2", 0.999)), + # The AVX-512 kernel takes a scalar decay, so the step-dependent bias + # correction is applied here — the kernel ABI is unchanged. + "beta2": adafactor_decay( + float(kwargs.get("beta2", 0.999)), int(kwargs.get("step", 1)) + ), "eps": float(kwargs.get("eps", 1e-30)), } if len(operand_names) == 3: @@ -15951,7 +15958,9 @@ def _execute_x86_compiled_adafactor_backward( state_contract=state_contract, ) lr = float(kwargs.get("lr", 1e-3)) - beta2 = float(kwargs.get("beta2", 0.999)) + # The contract records both the nominal decay and the bias-corrected one the + # forward applied; the adjoint must differentiate the latter. + beta2 = float(kwargs["beta2_effective"]) eps = float(kwargs.get("eps", 1e-30)) g = np.ascontiguousarray(_as_numpy(values[operand_names[1]]), np.float32) dy = np.ascontiguousarray(_as_numpy(values[cotangent_name]), np.float32) @@ -28396,9 +28405,15 @@ def _execute_rocm_compiled_adafactor( ) values = _bind_launch_args(args, arg_names) kwargs = ops[0].get("kwargs") or {} + from .optim import adafactor_decay + parameters = dict( lr=float(kwargs.get("lr", 1e-3)), - beta2=float(kwargs.get("beta2", 0.999)), + # gfx1151 adafactor_row/col/update take a scalar decay, so the + # step-dependent bias correction is applied here (kernel ABI unchanged). + beta2=adafactor_decay( + float(kwargs.get("beta2", 0.999)), int(kwargs.get("step", 1)) + ), eps=float(kwargs.get("eps", 1e-30)), ) if len(names) == 3: @@ -28824,7 +28839,9 @@ def _execute_rocm_compiled_adafactor_backward( ) parameters = dict( lr=float(kwargs.get("lr", 1e-3)), - beta2=float(kwargs.get("beta2", 0.999)), + # The contract records both the nominal decay and the bias-corrected one + # the forward applied; the adjoint must differentiate the latter. + beta2=float(kwargs["beta2_effective"]), eps=float(kwargs.get("eps", 1e-30)), ) if len(names) == 3: diff --git a/scripts/_nvidia_env.sh b/scripts/_nvidia_env.sh new file mode 100755 index 000000000..70d36f552 --- /dev/null +++ b/scripts/_nvidia_env.sh @@ -0,0 +1,77 @@ +#!/usr/bin/env bash +# Shared NVIDIA/CUDA resolver for Tessera's device lanes. Source it BEFORE +# launching python/pytest; it puts the driver shim and the CUDA toolkit on PATH +# and LD_LIBRARY_PATH. +# +# source "$(dirname "$0")/_nvidia_env.sh" +# "${TESSERA_PY}" -m pytest tests/device/nvidia/ -q +# +# Why this exists +# --------------- +# The NVIDIA device gate is `tests/_support/nvidia.py::nvidia_cuda_host_ready`, +# which requires `shutil.which("nvidia-smi")`. Under WSL2 the driver shim lives +# in /usr/lib/wsl/lib -- a directory the interactive .bashrc adds and a +# non-interactive shell (an `ssh `, a CI step, a bare `pytest`) +# never sees. +# +# The failure mode is a GREEN RUN, not an error. Measured on The-Super-Bear +# 2026-08-30: `pytest tests/device/nvidia/` reported **454 passed, 395 skipped, +# exit 0** while executing zero GPU work -- every skip read "host WSL CUDA +# device/toolchain unavailable". The GPU, /dev/dxg and CUDA 13.3 were all +# present and healthy the entire time; only PATH was wrong. Reporting that run +# as sm_120 evidence would have asserted a hardware result that never happened, +# which is the exact trap CLAUDE.md's Working Rules name ("a missing device +# *skips* rather than errors"). +# +# Once the PATH was fixed the suite ran for real and surfaced 80 genuine +# failures that had been invisible behind the skip, two of which were compiler +# defects. So this is not hygiene: the silent skip was actively hiding bugs. +# +# This is the NVIDIA twin of _rocm_env.sh, and it follows the same contract: +# detect by capability, respect anything already exported, and be a SILENT +# NO-OP on a host with no NVIDIA GPU so Mac/ROCm boxes still skip honestly +# (Decision #26) instead of having a device fabricated for them. +# +# Every expansion is nounset-safe (${VAR:-}). A release or CI script sourcing +# this under `set -u` would otherwise abort on the first unset CUDA_HOME or +# LD_LIBRARY_PATH instead of detecting the toolkit. + +# Driver shim (WSL2 puts nvidia-smi and libcuda.so here; native Linux does not +# have this directory at all, and does not need it). +if [ -d /usr/lib/wsl/lib ]; then + case ":${PATH:-}:" in + *":/usr/lib/wsl/lib:"*) ;; + *) PATH="/usr/lib/wsl/lib:${PATH}"; export PATH ;; + esac + case ":${LD_LIBRARY_PATH:-}:" in + *":/usr/lib/wsl/lib:"*) ;; + *) LD_LIBRARY_PATH="/usr/lib/wsl/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + export LD_LIBRARY_PATH ;; + esac +fi + +# CUDA toolkit. An already-exported CUDA_HOME wins; otherwise prefer the +# versioned symlink over a bare guess. +if [ -z "${CUDA_HOME:-}" ]; then + for _tessera_cuda_root in /usr/local/cuda /opt/cuda; do + if [ -x "${_tessera_cuda_root}/bin/nvcc" ]; then + CUDA_HOME="${_tessera_cuda_root}" + export CUDA_HOME + break + fi + done + unset _tessera_cuda_root +fi +if [ -n "${CUDA_HOME:-}" ] && [ -x "${CUDA_HOME}/bin/nvcc" ]; then + case ":${PATH:-}:" in + *":${CUDA_HOME}/bin:"*) ;; + *) PATH="${CUDA_HOME}/bin:${PATH}"; export PATH ;; + esac + if [ -d "${CUDA_HOME}/lib64" ]; then + case ":${LD_LIBRARY_PATH:-}:" in + *":${CUDA_HOME}/lib64:"*) ;; + *) LD_LIBRARY_PATH="${CUDA_HOME}/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}" + export LD_LIBRARY_PATH ;; + esac + fi +fi diff --git a/src/compiler/tile_opt_fa4/include/tessera/Dialect/Attn/Attn.td b/src/compiler/tile_opt_fa4/include/tessera/Dialect/Attn/Attn.td index b57b345b6..cb3692e5d 100644 --- a/src/compiler/tile_opt_fa4/include/tessera/Dialect/Attn/Attn.td +++ b/src/compiler/tile_opt_fa4/include/tessera/Dialect/Attn/Attn.td @@ -191,16 +191,37 @@ def BoundaryMaskOp : Op { // Counter-based dropout for one score block. The absolute KV offset is an SSA // operand, preventing every loop iteration from replaying the first mask tile. +// +// `stream_offset` is the second, orthogonal counter axis: the RNG stream base +// of the *instance* this block belongs to (Decision #18). A rank-4 +// FlashAttention is distributed into B*H rank-2 instances by +// `DistributeRank4FlashAttn`, and the batch/head coordinates at that site are +// `scf.for` induction variables — SSA values, so they cannot ride on the +// `seed` attribute. Without this operand every instance replays one mask and +// dropout is fully correlated across batch and head instead of iid. +// +// The contract is that instances occupy DISJOINT counter ranges: the producer +// passes `(b*H + h) * Sq * Sk_padded`, so a consumer keying on +// `(seed, stream_offset + row*Sk + (kv_offset + col))` gets an independent +// stream per (batch, head). That is exactly the counter the generated gfx1151 +// WMMA attention kernels already build from their launch geometry +// (`((bh*Sq)+q)*Sk + k`, GenerateWMMAFlashAttnKernel.cpp / ...BwdKernel.cpp); +// stating it here makes the shared Tile-IR contract say what the physical +// kernels do instead of understating it (Decisions #29/#32). +// +// A rank-2 attention that was never distributed passes index 0 and keeps its +// previous mask exactly. def BlockDropoutOp : Op { let summary = "Apply deterministic counter-based dropout to one score block"; let arguments = (ins AnyType:$scores, Index:$kv_offset, + Index:$stream_offset, F32Attr:$dropout_p, I64Attr:$seed ); let results = (outs AnyType:$masked_scores); - let assemblyFormat = "$scores `kv_off` `=` $kv_offset `p` `=` $dropout_p `seed` `=` $seed attr-dict `:` type($scores) `->` type($masked_scores)"; + let assemblyFormat = "$scores `kv_off` `=` $kv_offset `stream` `=` $stream_offset `p` `=` $dropout_p `seed` `=` $seed attr-dict `:` type($scores) `->` type($masked_scores)"; let hasVerifier = 1; } diff --git a/src/transforms/lib/TileIRLoweringPass.cpp b/src/transforms/lib/TileIRLoweringPass.cpp index b29d7926b..dbae3aa8d 100644 --- a/src/transforms/lib/TileIRLoweringPass.cpp +++ b/src/transforms/lib/TileIRLoweringPass.cpp @@ -118,6 +118,49 @@ static Operation *emitAttnOp(OpBuilder &b, Location loc, return b.create(st); } +// Recover the flattened (batch, query_head) instance index of a rank-2 +// FlashAttention that `DistributeRank4FlashAttn` produced, as `b*H + h`. +// +// The coordinates are `scf.for` induction variables, so they are SSA values and +// cannot ride on an attribute. Rather than thread a new operand through +// `tessera.flash_attn`, they are DERIVED from the enclosing loops this same +// pass annotated with `tessera.attention_distribution` (Decision #30 — derive, +// don't ask; the annotations already exist and already have a consumer in the +// ROCm recognizer). +// +// Returns a null Value when the op is not enclosed by both distribution loops, +// which is the correct answer for a hand-written rank-2 attention: it is its +// own single instance and gets stream base 0. +static Value deriveDistributionInstanceIndex(OpBuilder &b, Location loc, + Operation *op) { + Value batchIndex; + Value headIndex; + int64_t queryHeads = 0; + for (Operation *parent = op->getParentOp(); parent; + parent = parent->getParentOp()) { + auto forOp = dyn_cast(parent); + if (!forOp) + continue; + auto axis = + forOp->getAttrOfType("tessera.attention_distribution"); + if (!axis) + continue; + if (axis.getValue() == "batch" && !batchIndex) { + batchIndex = forOp.getInductionVar(); + } else if (axis.getValue() == "query_head" && !headIndex) { + headIndex = forOp.getInductionVar(); + if (auto count = + forOp->getAttrOfType("tessera.query_head_count")) + queryHeads = count.getInt(); + } + } + if (!batchIndex || !headIndex || queryHeads <= 0) + return Value(); + Value heads = arith::ConstantIndexOp::create(b, loc, queryHeads); + return arith::AddIOp::create( + b, loc, arith::MulIOp::create(b, loc, batchIndex, heads), headIndex); +} + // ───────────────────────────────────────────────────────────────────────────── // FlashAttn lowering pattern // ───────────────────────────────────────────────────────────────────────────── @@ -496,6 +539,34 @@ struct LowerFlashAttnToTileIR : public RewritePattern { Value upper = arith::ConstantIndexOp::create(rewriter, loc, paddedSk); Value step = arith::ConstantIndexOp::create(rewriter, loc, tkv); + // Per-instance dropout RNG stream base (Decision #18). This slice may be + // one of B*H instances that `DistributeRank4FlashAttn` produced; the KV + // `boundary` iter-arg below restarts at 0 inside every one of them, so + // without a distinct stream base every instance would replay one identical + // mask and dropout would be fully correlated across batch and head. + // + // The base is loop-invariant, so it is materialized once here rather than + // inside the KV loop. + Value dropoutStream = zeroIndex; + if (dropout && dropout.getValueAsDouble() > 0.0) { + Value instance = deriveDistributionInstanceIndex(rewriter, loc, op); + if (instance) { + Value stride = arith::ConstantIndexOp::create( + rewriter, loc, qRows * paddedSk); + dropoutStream = + arith::MulIOp::create(rewriter, loc, instance, stride); + } else if (op->hasAttr("tessera.rank4_distributed")) { + // Fail closed: this op IS one of the distributed instances but the + // batch/head loops this pass annotates are not reachable, so a + // per-instance stream cannot be derived. Emitting stream 0 would + // silently restore the correlated-mask defect (Decisions #21/#21a). + return rewriter.notifyMatchFailure( + op, "distributed rank-4 attention carries dropout but its " + "tessera.attention_distribution batch/query_head loops are " + "not reachable, so no per-instance RNG stream can be derived"); + } + } + auto kvLoop = scf::ForOp::create( rewriter, loc, zeroIndex, upper, step, ValueRange{accInit, negInf, zero, producerInit, consumerInit, @@ -603,7 +674,7 @@ struct LowerFlashAttnToTileIR : public RewritePattern { if (dropout && dropout.getValueAsDouble() > 0.0) { Operation *drop = emitAttnOp( rewriter, loc, "tessera_attn.block_dropout", - {scores, boundary}, {scoresType}, + {scores, boundary, dropoutStream}, {scoresType}, {rewriter.getNamedAttr( "dropout_p", rewriter.getF32FloatAttr( diff --git a/tests/_support/environment.py b/tests/_support/environment.py index a15b71663..31f02e1a5 100644 --- a/tests/_support/environment.py +++ b/tests/_support/environment.py @@ -15,6 +15,9 @@ REPO_ROOT = Path(__file__).resolve().parents[2] PYTHON_ROOT = REPO_ROOT / "python" CUDA_BIN_DIRS = (Path("/usr/local/cuda/bin"), Path("/usr/local/cuda-13.3/bin")) +# WSL2 ships the NVIDIA driver shim (nvidia-smi, libcuda.so) here rather than in +# a system bin directory, and only an interactive shell puts it on PATH. +NVIDIA_DRIVER_DIRS = (Path("/usr/lib/wsl/lib"),) def is_wsl() -> bool: @@ -41,6 +44,42 @@ def ensure_cuda_bin_on_path() -> Path | None: return None +def ensure_nvidia_driver_on_path() -> Path | None: + """Make the NVIDIA driver shim visible to this process. + + Under WSL2 ``nvidia-smi`` lives in ``/usr/lib/wsl/lib``, which the + interactive ``.bashrc`` adds and a non-interactive shell -- an + ``ssh ``, a CI step, a bare ``pytest`` -- never sees. Because + the device gate probes for ``nvidia-smi`` by name, its absence from PATH + made every NVIDIA device test SKIP while the run still exited 0: measured + 2026-08-30 on The-Super-Bear as *454 passed, 395 skipped*, with a healthy + RTX 5070, /dev/dxg and CUDA 13.3 present the whole time. Reporting that as + sm_120 evidence would assert a hardware result that never ran, so the gate + repairs its own PATH rather than trusting the caller to remember. + """ + entries = os.environ.get("PATH", "").split(os.pathsep) + for root in NVIDIA_DRIVER_DIRS: + if (root / "nvidia-smi").is_file(): + if str(root) not in entries: + os.environ["PATH"] = os.pathsep.join([str(root), *filter(None, entries)]) + return root + return None + + +def nvidia_gpu_is_plausibly_present() -> bool: + """Whether this host looks like it has an NVIDIA GPU, ignoring PATH. + + Used to tell "no GPU here, skip honestly" apart from "a GPU is sitting + right there and the environment is hiding it", which is a misconfiguration + worth shouting about rather than skipping past. + """ + if any((root / "nvidia-smi").is_file() for root in NVIDIA_DRIVER_DIRS): + return True + return any( + Path(node).exists() for node in ("/dev/nvidiactl", "/dev/nvidia0", "/dev/dxg") + ) + + def _tool_path(env_name: str, *candidates: Path | str) -> Path | None: configured = os.environ.get(env_name) if configured: diff --git a/tests/_support/nvidia.py b/tests/_support/nvidia.py index 0d0fcf272..bea895eb7 100644 --- a/tests/_support/nvidia.py +++ b/tests/_support/nvidia.py @@ -4,10 +4,15 @@ import os import shutil import subprocess +import warnings from typing import Any import pytest -from tests._support.environment import ensure_cuda_bin_on_path +from tests._support.environment import ( + ensure_cuda_bin_on_path, + ensure_nvidia_driver_on_path, + nvidia_gpu_is_plausibly_present, +) def nvidia_cuda_toolchain_available() -> bool: @@ -49,16 +54,37 @@ def require_nvidia_mma_runtime() -> Any: def nvidia_cuda_host_ready() -> bool: - """Whether CUDA tooling and the NVIDIA driver are reachable from this host.""" - if not nvidia_cuda_toolchain_available() or shutil.which("nvidia-smi") is None: - return False - try: - return subprocess.run( - ["nvidia-smi"], stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, timeout=5, check=False, - ).returncode == 0 - except OSError: - return False + """Whether CUDA tooling and the NVIDIA driver are reachable from this host. + + Repairs PATH first, then WARNS if a GPU is evidently present but still + unreachable. A silent skip is the correct answer for a host with no NVIDIA + hardware; it is the WRONG answer for a host that has one and cannot see it, + because the run then reports success having executed nothing. That is not + hypothetical -- it hid 80 real failures on The-Super-Bear until 2026-08-30, + two of which were compiler defects. + """ + ensure_nvidia_driver_on_path() + ready = nvidia_cuda_toolchain_available() and shutil.which("nvidia-smi") is not None + if ready: + try: + ready = subprocess.run( + ["nvidia-smi"], stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, timeout=5, check=False, + ).returncode == 0 + except OSError: + ready = False + if not ready and nvidia_gpu_is_plausibly_present(): + warnings.warn( + "NVIDIA device lanes are being SKIPPED on a host that appears to " + "have an NVIDIA GPU. This is an environment problem, not an absent " + "device, and the skips will look like a clean run. Source " + "scripts/_nvidia_env.sh (or put the driver shim -- /usr/lib/wsl/lib " + "under WSL2 -- and the CUDA toolkit on PATH) before treating any " + "result from this session as device evidence.", + RuntimeWarning, + stacklevel=2, + ) + return ready def assert_native_gpu(result: dict[str, Any]) -> None: diff --git a/tests/device/nvidia/test_e2e_spine_native.py b/tests/device/nvidia/test_e2e_spine_native.py index 0bb3aacf5..8c1740a6d 100644 --- a/tests/device/nvidia/test_e2e_spine_native.py +++ b/tests/device/nvidia/test_e2e_spine_native.py @@ -1,5 +1,7 @@ from __future__ import annotations +import re + import numpy as np import pytest @@ -439,6 +441,35 @@ def _decode_ue8m0(codes: np.ndarray) -> np.ndarray: return np.ldexp(np.ones(codes.shape, np.float32), codes.astype(np.int32) - 127) +def _assert_canonical_k_loop(bundle, *, expected_k: int) -> None: + """Assert the canonical k-loop contract in whichever form was emitted. + + Two producers satisfy this contract and they state it differently, so + pinning one spelling fails on the other for reasons that have nothing to do + with the loop being canonical: + + * ``tile.matmul_kernel`` DECLARES it as ``tessera.canonical_k_loop`` + + ``tessera.tile_k``. Those attributes are read by + ``MatmulKernelOp::verify`` and ``materializeSm120MatmulKernel``, and by + nothing else. + * the typed-MMA producer emits no ``matmul_kernel`` at all and instead + STATES the same facts structurally: a typed ``mma_desc`` carrying k and + an explicit ``scf.for`` stepping by it. Adding the attribute to this + route would be a declaration with no consumer (Decision #29). + + Which producer runs depends on shape as well as storage -- ragged K falls + back to ``matmul_kernel`` even for fp16/bf16 -- so branch on what was + actually emitted. + """ + tile_text = bundle.tile.text + if "tile.matmul_kernel" in tile_text: + assert "tessera.canonical_k_loop = true" in tile_text + assert f"tessera.tile_k = {expected_k} : i64" in tile_text + else: + assert f"k = {expected_k}" in tile_text + assert re.search(rf"scf\.for\b.*\bstep %c{expected_k}", tile_text) + + @pytest.mark.hardware_nvidia @pytest.mark.parametrize("shape", [(16, 8, 16), (37, 29, 23)]) def test_canonical_sm120_request_packages_registers_launches_and_compares(shape) -> None: @@ -585,11 +616,28 @@ def test_canonical_sm120_k_loop_shape_matrix(storage, shape) -> None: assert bundle.native_image is not None assert bundle.native_image.resource_record is not None assert bundle.launch_descriptor is not None - expected_k = 8 if storage == "tf32" else 16 - assert "tessera.canonical_k_loop = true" in bundle.tile.text - assert "a = " in bundle.tile.text and 'acc = "f32"' in bundle.tile.text - assert f"tessera.tile_k = {expected_k} : i64" in bundle.tile.text - assert 'epilogue = #tile.epilogue' in bundle.tile.text + # The canonical k-loop contract is expressed differently by the two + # producers, and asserting one form against both was the defect here. + # + # tf32 (schedule "direct") -> tile.matmul_kernel, which DECLARES the + # loop with tessera.canonical_k_loop / tile_k. Those attributes are + # read by MatmulKernelOp::verify and materializeSm120MatmulKernel, + # and by nothing else. + # fp16/bf16 (schedule "shared") -> the typed-MMA producer, which emits + # no matmul_kernel at all and instead STATES the same facts in the + # IR: a typed mma_desc carrying k, and an explicit scf.for stepping + # by it. Nothing reads the attribute off this route, so adding it + # here would be a declaration with no consumer (Decision #29). + # + # Assert the structural form where it exists, since a materialised loop + # with the right step is a stronger guarantee than a boolean asserting one. + assert 'acc = "f32"' in bundle.tile.text + _assert_canonical_k_loop(bundle, expected_k=8 if storage == "tf32" else 16) + if "tile.matmul_kernel" in bundle.tile.text: + assert ( + 'epilogue = #tile.epilogue' + in bundle.tile.text + ) assert "scf.for" in bundle.target_ir.text metrics = bundle.native_image.resource_record.metrics assert metrics["spill_store_bytes"] == 0 diff --git a/tests/device/nvidia/test_optimizer_reverse.py b/tests/device/nvidia/test_optimizer_reverse.py index 4f263a48a..ba7bdb5b8 100644 --- a/tests/device/nvidia/test_optimizer_reverse.py +++ b/tests/device/nvidia/test_optimizer_reverse.py @@ -49,7 +49,7 @@ def _adamw(p, g, m1, m2): @ts.jit(target="nvidia_sm120", autodiff="reverse", wrt=("p", "g", "state")) def _adafactor_full(p, g, state): updated, _new_state = ts.ops.adafactor( - p, g, state, lr=0.003, beta2=0.91, eps=1.0e-7 + p, g, state, lr=0.003, beta2=0.91, eps=1.0e-7, step=2 ) return updated @@ -59,7 +59,7 @@ def _adafactor_full(p, g, state): ) def _adafactor_factored(p, g, row, col): updated, _new_row, _new_col = ts.ops.adafactor( - p, g, row, col, lr=0.003, beta2=0.91, eps=1.0e-7 + p, g, row, col, lr=0.003, beta2=0.91, eps=1.0e-7, step=2 ) return updated @@ -162,6 +162,8 @@ def test_sm120_adafactor_full_and_factored_exact_certificates(): from tessera.autodiff.vjp import get_vjp rng = np.random.default_rng(1311) + # The jit'd forwards above declare step=2, matching the `"step": 1` state + # the reference VJPs carry (the update being differentiated is step 2). kwargs = {"lr": 0.003, "beta2": 0.91, "eps": 1.0e-7} for compiled, shape, topology in ( (_adafactor_full, (17,), "full"), diff --git a/tests/tessera-ir/phase3/streaming_attention_dropout_stream.mlir b/tests/tessera-ir/phase3/streaming_attention_dropout_stream.mlir new file mode 100644 index 000000000..3a304e2fe --- /dev/null +++ b/tests/tessera-ir/phase3/streaming_attention_dropout_stream.mlir @@ -0,0 +1,72 @@ +// Per-instance dropout RNG streams for distributed rank-4 FlashAttention. +// +// `DistributeRank4FlashAttn` splits a rank-4 attention into B*H rank-2 +// instances inside two annotated `scf.for` loops, and each instance restarts +// its own KV `boundary` iter-arg at 0. The dropout counter therefore has to +// carry a second, per-instance axis or every one of the B*H instances replays +// one identical mask (Decision #18 stream separation). The batch/head +// coordinates are loop induction variables, so `tessera_attn.block_dropout` +// takes the stream base as an SSA operand rather than folding it into `seed`. +// +// RUN: tessera-opt --tessera-tile-ir-lowering='tile-q=8 tile-kv=8 sm=90' \ +// RUN: --allow-unregistered-dialect %s | FileCheck %s + +// B=2, H=4, Sq=8, Sk=8 -> instance index b*4 + h, stride Sq*Sk_padded = 64. +// CHECK-LABEL: func.func @rank4_dropout_gets_per_instance_streams +// CHECK-DAG: %[[STRIDE:.*]] = arith.constant 64 : index +// CHECK-DAG: %[[HEADS:.*]] = arith.constant 4 : index +// CHECK: scf.for %[[B:.*]] = %{{.*}} to %{{.*}} step +// CHECK: scf.for %[[H:.*]] = %{{.*}} to %{{.*}} step +// CHECK: %[[SCALED:.*]] = arith.muli %[[B]], %[[HEADS]] +// CHECK: %[[INSTANCE:.*]] = arith.addi %[[SCALED]], %[[H]] +// CHECK: %[[STREAM:.*]] = arith.muli %[[INSTANCE]], %[[STRIDE]] +// The stream base is loop-invariant, so it is hoisted above the KV loop. +// CHECK: scf.for %{{.*}} iter_args +// CHECK: tessera_attn.block_dropout %{{.*}} kv_off = %{{.*}} stream = %[[STREAM]] p = 2.500000e-01 seed = 37 +// CHECK: tessera.attention_distribution = "query_head" +// CHECK: tessera.attention_distribution = "batch" +func.func @rank4_dropout_gets_per_instance_streams( + %q: tensor<2x4x8x16xf16>, + %k: tensor<2x4x8x16xf16>, + %v: tensor<2x4x8x16xf16>) -> tensor<2x4x8x16xf32> { + %o = "tessera.flash_attn"(%q, %k, %v) + <{operandSegmentSizes = array}> { + causal = true, + dropout_p = 0.25 : f64, + dropout_seed = 37 : i64, + head_dim = 16 : i64, + scale = 0.25 : f32, + tessera.tile_q = 8 : i32, + tessera.tile_kv = 8 : i32 + } : (tensor<2x4x8x16xf16>, tensor<2x4x8x16xf16>, tensor<2x4x8x16xf16>) + -> tensor<2x4x8x16xf32> + return %o : tensor<2x4x8x16xf32> +} + +// An attention that was never distributed is its own single instance: it must +// keep stream base 0 and therefore the exact mask it drew before, so the fix +// cannot perturb a rank-2 lane. The negative half matters — a producer that +// invented a stream for every attention would silently change every existing +// rank-2 dropout mask. +// CHECK-LABEL: func.func @rank2_dropout_keeps_stream_zero +// CHECK-DAG: %[[ZERO:.*]] = arith.constant 0 : index +// CHECK: tessera_attn.block_dropout %{{.*}} kv_off = %{{.*}} stream = %[[ZERO]] p = 2.500000e-01 seed = 37 +// CHECK-NOT: tessera.attention_distribution +// CHECK-NOT: tessera.flash_attn +func.func @rank2_dropout_keeps_stream_zero( + %q: tensor<8x16xf16>, + %k: tensor<8x16xf16>, + %v: tensor<8x16xf16>) -> tensor<8x16xf32> { + %o = "tessera.flash_attn"(%q, %k, %v) + <{operandSegmentSizes = array}> { + causal = true, + dropout_p = 0.25 : f64, + dropout_seed = 37 : i64, + head_dim = 16 : i64, + scale = 0.25 : f32, + tessera.tile_q = 8 : i32, + tessera.tile_kv = 8 : i32 + } : (tensor<8x16xf16>, tensor<8x16xf16>, tensor<8x16xf16>) + -> tensor<8x16xf32> + return %o : tensor<8x16xf32> +} diff --git a/tests/unit/test_autodiff_stateful_plugin_binding.py b/tests/unit/test_autodiff_stateful_plugin_binding.py index 1a3ef8428..67041a823 100644 --- a/tests/unit/test_autodiff_stateful_plugin_binding.py +++ b/tests/unit/test_autodiff_stateful_plugin_binding.py @@ -113,14 +113,17 @@ def test_flat_adafactor_full_and_factored_match_tree_reference() -> None: g = rng.normal(size=(3, 5)).astype(np.float32) row = np.abs(rng.normal(size=(3,))).astype(np.float32) col = np.abs(rng.normal(size=(5,))).astype(np.float32) + # Step 3, not the default 1: at step 1 the bias-corrected decay is exactly + # 0, so the carried row/col are discarded and the comparison would not + # exercise the state at all. kwargs = {"lr": 0.003, "beta2": 0.91, "eps": 1.0e-7} actual_p, actual_row, actual_col = ts.ops.adafactor( - p, g, row, col, **kwargs + p, g, row, col, step=3, **kwargs ) expected_p, expected_state = ts.optim.adafactor( p, g, - {"v": {"row": row, "col": col, "factored": True}, "step": 0}, + {"v": {"row": row, "col": col, "factored": True}, "step": 2}, **kwargs, ) np.testing.assert_allclose(actual_p, expected_p, rtol=2e-6, atol=2e-6) @@ -130,11 +133,13 @@ def test_flat_adafactor_full_and_factored_match_tree_reference() -> None: vector_p = p[0] vector_g = g[0] full = np.abs(rng.normal(size=(5,))).astype(np.float32) - actual_p, actual_full = ts.ops.adafactor(vector_p, vector_g, full, **kwargs) + actual_p, actual_full = ts.ops.adafactor( + vector_p, vector_g, full, step=3, **kwargs + ) expected_p, expected_state = ts.optim.adafactor( vector_p, vector_g, - {"v": {"v": full, "factored": False}, "step": 0}, + {"v": {"v": full, "factored": False}, "step": 2}, **kwargs, ) np.testing.assert_allclose(actual_p, expected_p, rtol=2e-6, atol=2e-6) diff --git a/tests/unit/test_nvidia_solver_krylov_source.py b/tests/unit/test_nvidia_solver_krylov_source.py index d95717f67..875b918db 100644 --- a/tests/unit/test_nvidia_solver_krylov_source.py +++ b/tests/unit/test_nvidia_solver_krylov_source.py @@ -1,7 +1,7 @@ """Host-free contracts on the emitted dense-Krylov CUDA source. There is no CUDA on the machines that run this suite, so these assert on the -GENERATED TEXT. They pin the memory-access shape of `tsr_matvec`, which is the +GENERATED TEXT. They pin the memory-access shape of `tsr_matvec_warp`, which is the O(n^2) term of every CG/GMRES iteration; nothing here is a device measurement, and no wall-clock claim can be made from it (code review 2026-08-29, P3). """ @@ -37,7 +37,7 @@ def test_matvec_columns_are_walked_by_lane_not_by_one_thread(): transaction-count argument is sound, the wall-clock effect is not available on a host without CUDA. """ - body = _device_fn(kv._source(), "tsr_matvec") + body = _device_fn(kv._source(), "tsr_matvec_warp") assert "for (int col = lane; col < n; col += 32)" in body, ( "column loop is not lane-strided — the loads are uncoalesced") assert "for (int col = 0; col < n; ++col)" not in body @@ -47,7 +47,7 @@ def test_matvec_reduction_is_a_fixed_full_mask_shuffle_tree(): """The per-row sum must stay reproducible for a fixed launch geometry, the same contract `tsr_grid_sum` states. A fixed 5-step butterfly over a fixed 32-lane warp is; a size-dependent or atomic reduction is not.""" - body = _device_fn(kv._source(), "tsr_matvec") + body = _device_fn(kv._source(), "tsr_matvec_warp") assert "for (int off = 16; off; off >>= 1)" in body assert "__shfl_down_sync(0xffffffffu, sum, off)" in body assert "atomicAdd" not in body, "an atomic makes the row sum non-reproducible" @@ -58,7 +58,7 @@ def test_matvec_row_index_is_warp_uniform(): """The full-mask shuffles require all 32 lanes of the warp to reach them, so `row` must be derived from a warp id, not a thread id — otherwise lanes of one warp take different trip counts and the shuffle is undefined.""" - body = _device_fn(kv._source(), "tsr_matvec") + body = _device_fn(kv._source(), "tsr_matvec_warp") assert re.search(r"long warp = .*blockDim\.x \+ threadIdx\.x\) >> 5", body) assert re.search(r"long warps = .*gridDim\.x \* blockDim\.x\) >> 5", body) assert "for (long row = warp; row < n; row += warps)" in body @@ -69,5 +69,25 @@ def test_matvec_row_index_is_warp_uniform(): def test_matvec_still_widens_low_precision_storage_before_multiplying(): """The module contract: f16/bf16 operands convert to f32 before the multiply so the Krylov convergence claim is unambiguous.""" - body = _device_fn(kv._source(), "tsr_matvec") + body = _device_fn(kv._source(), "tsr_matvec_warp") assert "fmaf(tsr_load(arow, col), x[col], sum)" in body + + +def test_cg_and_gmres_select_their_measured_matvec(): + """The two solvers must not share one matvec. + + Measured on an RTX 5070 (sm_120), medians of 9 reps: the warp-per-row form + is 1.24-1.56x faster for GMRES on device_event, and 0.44-0.63x — i.e. up to + 2.3x SLOWER — for CG. A cooperative launch caps the grid at what stays + resident, so warp-per-row also buys 32x fewer rows in flight; GMRES absorbs + that and CG, which grid-syncs far more per iteration, does not. Shipping one + form for both shipped that CG regression. + """ + source = kv._source() + assert "tsr_matvec_scalar" in source and "tsr_matvec_warp" in source + cg = _device_fn(source, "tsr_dense_cg") + gmres = _device_fn(source, "tsr_dense_gmres") + assert "tsr_matvec_scalar(a," in cg + assert "tsr_matvec_warp(a," not in cg + assert "tsr_matvec_warp(a," in gmres + assert "tsr_matvec_scalar(a," not in gmres diff --git a/tests/unit/test_rocm_optimizer_compiled.py b/tests/unit/test_rocm_optimizer_compiled.py index 4f4f49bc2..5ecdcbb72 100644 --- a/tests/unit/test_rocm_optimizer_compiled.py +++ b/tests/unit/test_rocm_optimizer_compiled.py @@ -142,7 +142,9 @@ def test_adafactor_factored_multistep(): row = np.zeros(SHAPE[:-1], np.float32) col = np.zeros(SHAPE[-1], np.float32) state = None - for _ in range(3): + # `step` is 1-based and selects Adafactor's bias-corrected second-moment + # decay, exactly as it does for the flat adamw ABI above. + for step in range(1, 4): g = rng.standard_normal(SHAPE).astype(np.float32) artifact = rt.RuntimeArtifact(metadata={ "target": "rocm", @@ -155,7 +157,9 @@ def test_adafactor_factored_multistep(): "op_name": "tessera.adafactor", "result": "o", "operands": ["p", "g", "row", "col"], - "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + "kwargs": { + "lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": step, + }, }], }) result = rt.launch(artifact, (p, g, row, col)) @@ -195,7 +199,9 @@ def test_adafactor_factored_nan_gradient_propagates_like_reference(): "op_name": "tessera.adafactor", "result": "o", "operands": ["p", "g", "row", "col"], - "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + # The tree reference below carries step 1, so the update being + # computed is step 2. + "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": 2}, }], }) result = rt.launch(artifact, (p, g, row, col)) @@ -223,7 +229,7 @@ def test_adafactor_full_moment_vector_multistep(): p = rng.standard_normal(19).astype(np.float32) moment = np.zeros_like(p) state = None - for _ in range(3): + for step in range(1, 4): g = rng.standard_normal(p.shape).astype(np.float32) artifact = rt.RuntimeArtifact(metadata={ "target": "rocm", @@ -236,7 +242,9 @@ def test_adafactor_full_moment_vector_multistep(): "op_name": "tessera.adafactor", "result": "o", "operands": ["p", "g", "moment"], - "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + "kwargs": { + "lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": step, + }, }], }) result = rt.launch(artifact, (p, g, moment)) @@ -325,7 +333,9 @@ def test_adafactor_factored_backward_executes_on_gfx1151(): target="rocm_gfx1151", parameter_shape=SHAPE, topology="factored", - kwargs={"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + # The reference VJP below carries state step 1, so the update being + # differentiated is step 2 (selects the bias-corrected decay). + kwargs={"lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": 2}, ) artifact = rt.RuntimeArtifact(metadata={ "target": "rocm", @@ -372,7 +382,9 @@ def test_adafactor_full_backward_executes_on_gfx1151(): target="rocm_gfx1151", parameter_shape=p.shape, topology="full", - kwargs={"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + # The reference VJP below carries state step 1, so the update being + # differentiated is step 2 (selects the bias-corrected decay). + kwargs={"lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": 2}, ) artifact = rt.RuntimeArtifact(metadata={ "target": "rocm", diff --git a/tests/unit/test_s10_optim.py b/tests/unit/test_s10_optim.py index 11379474b..7a4b39c67 100644 --- a/tests/unit/test_s10_optim.py +++ b/tests/unit/test_s10_optim.py @@ -3,8 +3,10 @@ from __future__ import annotations import numpy as np +import pytest import tessera as ts +from tessera import optim from tessera.state import tree_flatten, tree_unflatten @@ -110,3 +112,266 @@ def test_ema_polyak_and_optimizer_state_tree_round_trip(): restored = tree_unflatten(treedef, leaves) assert restored["step"] == state["step"] np.testing.assert_allclose(restored["m"]["w"], state["m"]["w"]) + + +# ───────────────────────────────────────────────────────────────────────────── +# Adafactor second-moment bias correction (CODE_REVIEW_2026-08-29, +# `python/tessera/optim.py:427`). The tracked `step` used to be incremented and +# never read, so the zero-initialized second moment was biased low and the first +# ~1/(1-beta2) updates were inflated by 1/sqrt(1 - beta2**step) — 31.6x at the +# default beta2=0.999. `optim.adafactor_decay` is the single shared correction; +# these tests pin it against an independently written reference so a regression +# in either the tree form, the flat compiler ABI, or the analytic VJP shows up. + + +def _uncorrected_adafactor_full(g, *, beta2, eps, steps): + """Textbook Adafactor full-moment update with an EXPLICIT 1-beta2**t debias. + + Written directly from the definition (raw EMA, then divide by 1-beta2**t) + rather than through `adafactor_decay`, so it is an independent check of the + step-dependent-decay formulation rather than a restatement of it. + """ + v = np.zeros_like(g[0], dtype=np.float64) + updates = [] + for t in range(1, steps + 1): + gt = np.asarray(g[t - 1], dtype=np.float64) + v = beta2 * v + (1.0 - beta2) * gt * gt + v_hat = v / (1.0 - beta2**t) + updates.append(gt / (np.sqrt(np.maximum(v_hat, eps)) + eps)) + return updates + + +def test_adafactor_step_one_update_is_not_inflated_by_the_ema_bias(): + """A constant gradient must give a ~1.0 normalized update from step 1.""" + grad = np.full((4, 4), 0.1, dtype=np.float32) + params, state = ts.optim.adafactor( + {"w": np.ones((4, 4), dtype=np.float32)}, {"w": grad}, lr=1.0 + ) + magnitude = float(np.abs(1.0 - np.asarray(params["w"])).mean()) + # Pre-fix this was 1/sqrt(1-0.999) = 31.6227... + assert magnitude == pytest.approx(1.0, rel=1e-4), magnitude + assert state["step"] == 1 + + +def test_adafactor_matches_explicit_bias_corrected_ema_across_steps(): + """Full-moment (rank-1) leaves must track the explicit 1-beta2**t debias.""" + rng = np.random.default_rng(20260830) + beta2, eps, lr, steps = 0.999, 1e-30, 1.0, 6 + grads = [rng.normal(scale=0.3, size=(7,)).astype(np.float32) for _ in range(steps)] + expected = _uncorrected_adafactor_full(grads, beta2=beta2, eps=eps, steps=steps) + + params, state = {"w": np.zeros(7, dtype=np.float32)}, None + for index, grad in enumerate(grads): + previous = np.asarray(params["w"], dtype=np.float64) + params, state = ts.optim.adafactor( + params, {"w": grad}, state, lr=lr, beta2=beta2, eps=eps + ) + applied = (previous - np.asarray(params["w"], dtype=np.float64)) / lr + np.testing.assert_allclose(applied, expected[index], rtol=2e-5, atol=2e-6) + assert state["step"] == index + 1 + + +def test_adafactor_decay_is_the_debiasing_recursion_and_fails_closed(): + beta2 = 0.9 + # t=1 must discard the (empty) prior so v_1 == g_1**2 exactly. + assert ts.optim.adafactor_decay(beta2, 1) == 0.0 + # The recursion must reproduce EMA_t / (1 - beta2**t) exactly. + raw, corrected = 0.0, 0.0 + for t in range(1, 40): + g2 = float(t) ** 2 + raw = beta2 * raw + (1.0 - beta2) * g2 + decay = ts.optim.adafactor_decay(beta2, t) + corrected = decay * corrected + (1.0 - decay) * g2 + assert corrected == pytest.approx(raw / (1.0 - beta2**t), rel=1e-12) + # Asymptotically the caller's beta2 is preserved, not replaced. + assert ts.optim.adafactor_decay(beta2, 500) == pytest.approx(beta2, rel=1e-9) + # Semantic keys fail closed rather than defaulting (Decision #21a). + for bad_step in (0, -1): + with pytest.raises(ValueError, match="1-based"): + ts.optim.adafactor_decay(beta2, bad_step) + for bad_beta2 in (1.0, -0.1, 1.5): + with pytest.raises(ValueError, match=r"beta2"): + ts.optim.adafactor_decay(bad_beta2, 3) + + +def test_flat_adafactor_abi_carries_the_step_like_flat_adam(): + """The flat compiler ABI must agree with the tree form at every step, not + only at step 1 where the correction happens to zero the carried state.""" + rng = np.random.default_rng(90210) + shape = (3, 5) + kwargs = {"lr": 0.003, "beta2": 0.91, "eps": 1.0e-7} + p = rng.normal(size=shape).astype(np.float32) + row = np.zeros(shape[:-1], np.float32) + col = np.zeros(shape[-1], np.float32) + tree_params, tree_state = {"w": p}, None + for step in range(1, 5): + g = rng.normal(scale=0.3, size=shape).astype(np.float32) + p, row, col = ts.ops.adafactor(p, g, row, col, step=step, **kwargs) + tree_params, tree_state = ts.optim.adafactor( + tree_params, {"w": g}, tree_state, **kwargs + ) + assert tree_state["step"] == step + np.testing.assert_allclose(p, tree_params["w"], rtol=2e-6, atol=2e-6) + np.testing.assert_allclose(row, tree_state["v"]["w"]["row"], rtol=2e-6) + np.testing.assert_allclose(col, tree_state["v"]["w"]["col"], rtol=2e-6) + # A declared step that contradicts the tree state must not be silently + # resolved in favour of either side. + with pytest.raises(ValueError, match="disagrees with the carried state"): + ts.ops.adafactor(tree_params, {"w": g}, tree_state, step=99, **kwargs) + + +def test_adafactor_vjp_differentiates_the_corrected_forward(): + """The analytic VJP must track the step-dependent decay, not nominal beta2. + + (A previous batch shipped a defect by fixing an eager path while its VJP + kept the old behaviour; this pins the pair together.)""" + from tessera.autodiff.vjp import get_vjp + + rng = np.random.default_rng(4242) + shape = (17,) + kwargs = {"lr": 0.003, "beta2": 0.91, "eps": 1.0e-7} + p = rng.normal(size=shape).astype(np.float32) + g = rng.normal(scale=0.2, size=shape).astype(np.float32) + dy = rng.normal(size=shape).astype(np.float32) + for carried in (0, 1, 5): + state = { + "v": { + "v": rng.uniform(0.1, 0.3, size=shape).astype(np.float32), + "factored": False, + }, + "step": carried, + } + analytic = get_vjp("adafactor")(dy, p, g, state, **kwargs) + + def forward(gradient): + return np.asarray( + ts.optim.adafactor(p, gradient, state, **kwargs)[0], + dtype=np.float64, + ) + + numeric = np.zeros(shape, dtype=np.float64) + h = 1e-3 + for index in range(shape[0]): + bump = np.zeros(shape, dtype=np.float64) + bump[index] = h + plus = forward((g + bump).astype(np.float32)) + minus = forward((g - bump).astype(np.float32)) + numeric[index] = float( + np.sum(np.asarray(dy, dtype=np.float64) * (plus - minus)) / (2 * h) + ) + # The forward computes in fp32, so central differences sit on a ~1e-4 + # noise floor (the same reason `jvp_adafactor` pins h=1e-3). That is + # still two orders tighter than the error an uncorrected decay causes: + # see `test_adafactor_vjp_rejects_the_nominal_decay` below. + np.testing.assert_allclose( + np.asarray(analytic[1], dtype=np.float64), numeric, rtol=2e-2, atol=1e-4 + ) + + +def test_adafactor_vjp_rejects_the_nominal_decay(): + """Discrimination check for the test above: had the VJP kept differentiating + the *nominal* beta2 (the pre-fix behaviour), the gradient would be wrong by + far more than the fp32 finite-difference noise floor.""" + from tessera.autodiff import vjp as vjp_module + from tessera.autodiff.vjp import get_vjp + + rng = np.random.default_rng(4242) + shape = (17,) + kwargs = {"lr": 0.003, "beta2": 0.91, "eps": 1.0e-7} + p = rng.normal(size=shape).astype(np.float32) + g = rng.normal(scale=0.2, size=shape).astype(np.float32) + dy = rng.normal(size=shape).astype(np.float32) + state = { + "v": { + "v": rng.uniform(0.1, 0.3, size=shape).astype(np.float32), + "factored": False, + }, + "step": 1, + } + corrected = np.asarray(get_vjp("adafactor")(dy, p, g, state, **kwargs)[1]) + + original = vjp_module.adafactor_decay if hasattr(vjp_module, "adafactor_decay") else None + del original + import tessera.optim as optim_module + + saved = optim_module.adafactor_decay + try: # pre-fix behaviour: a fixed, uncorrected decay + optim_module.adafactor_decay = lambda beta2, step: float(beta2) + uncorrected = np.asarray(get_vjp("adafactor")(dy, p, g, state, **kwargs)[1]) + finally: + optim_module.adafactor_decay = saved + relative = float( + np.max(np.abs(uncorrected - corrected) / (np.abs(corrected) + 1e-12)) + ) + assert relative > 0.05, relative + + +# --- PR #644 review: the two Adafactor ABI/representation seams --------------- + +def test_flat_adafactor_without_step_preserves_the_moments_it_was_given(): + """`adafactor_decay(b2, 1)` is exactly 0 -- right for a genuine first step, + where v_1 = g^2. Defaulting an ABSENT step to 1 would therefore make every + call of a stateful caller that never passes one discard the moments it just + supplied, turning a stateful optimizer stateless with no diagnostic.""" + import tessera as ts + + assert optim.adafactor_decay(0.999, 1) == 0.0 + param = np.ones((4, 4), np.float32) + grad = np.full((4, 4), 0.5, np.float32) + carried = np.full((4, 4), 9.0, np.float32) + + no_step = np.asarray(ts.ops.adafactor(param, grad, carried)[1]) + first_step = np.asarray(ts.ops.adafactor(param, grad, carried, step=1)[1]) + + # Omitting step keeps the carried EMA (legacy, uncorrected) ... + assert float(no_step.ravel()[0]) > 8.0 + # ... while an explicit step 1 correctly restarts from g^2. + assert float(first_step.ravel()[0]) == pytest.approx(0.25) + + +def test_unmarked_adafactor_state_warns_rather_than_being_misread(): + """`state["v"]` changed meaning: it now carries the debiased estimate, not + the raw EMA. The two cannot be told apart from the values, so an unmarked + state at step > 0 is flagged. It is deliberately NOT rescaled: auto- + migrating would silently rewrite every hand-built state dict, which is a + worse failure than the one it fixes.""" + params = {"w": np.ones((4, 4), np.float32)} + grads = {"w": np.full((4, 4), 0.5, np.float32)} + _, marked = optim.adafactor(params, grads) + assert marked["v_representation"] == optim._ADAFACTOR_V_REPRESENTATION + + _, marked_2 = optim.adafactor(params, grads, marked) + unmarked = {k: v for k, v in marked.items() if k != "v_representation"} + with pytest.warns(RuntimeWarning, match="v_representation"): + _, from_unmarked = optim.adafactor(params, grads, unmarked) + # Warned, but the values are untouched -- identical to the marked run. + np.testing.assert_array_equal( + np.asarray(from_unmarked["v"]["w"]["row"]), + np.asarray(marked_2["v"]["w"]["row"]), + ) + + +def test_explicit_migration_recovers_a_legacy_checkpoint_and_is_idempotent(): + params = {"w": np.ones((4, 4), np.float32)} + grads = {"w": np.full((4, 4), 0.5, np.float32)} + _, native_1 = optim.adafactor(params, grads) + _, native_2 = optim.adafactor(params, grads, native_1) + + # Forge the pre-correction representation: same step, v scaled back by the + # bias factor the old code left in. + bias = 1.0 - 0.999 ** int(native_1["step"]) + legacy = {k: v for k, v in native_1.items() if k != "v_representation"} + legacy["v"] = optim._adafactor_tree_map_unary( + lambda slot: optim._adafactor_scale_state(slot, bias, state_dtype="fp32"), + native_1["v"], + ) + + migrated = optim.migrate_adafactor_state(legacy, 0.999) + _, resumed = optim.adafactor(params, grads, migrated) + np.testing.assert_allclose( + np.asarray(resumed["v"]["w"]["row"]), + np.asarray(native_2["v"]["w"]["row"]), + rtol=1e-5, + ) + # Migrating an already-marked state is a no-op. + assert optim.migrate_adafactor_state(migrated, 0.999) is migrated diff --git a/tests/unit/test_x86_optimizer_compiled.py b/tests/unit/test_x86_optimizer_compiled.py index 0d4a3f5bf..b74e217ea 100644 --- a/tests/unit/test_x86_optimizer_compiled.py +++ b/tests/unit/test_x86_optimizer_compiled.py @@ -165,6 +165,11 @@ def test_lion_backward_stop_sign_vjp(): def _adafactor_artifact(rt, *, backward: bool, factored: bool, shape): + # Every reference below carries state step 1, so the update under test is + # step 2. `step` is 1-based and selects Adafactor's bias-corrected + # second-moment decay (`optim.adafactor_decay`), exactly as it does for the + # flat adam/adamw ABI above. + numeric = {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6, "step": 2} operands = ["p", "g", "row", "col"] if factored else ["p", "g", "moment"] names = operands + (["dy"] if backward else []) metadata = { @@ -177,7 +182,7 @@ def _adafactor_artifact(rt, *, backward: bool, factored: bool, shape): "arg_names": names, "out_cotangent": "dy", "ops": [{ "op_name": "tessera.adafactor", "operands": operands, - "kwargs": {"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + "kwargs": dict(numeric), }], } if backward: @@ -187,7 +192,7 @@ def _adafactor_artifact(rt, *, backward: bool, factored: bool, shape): target="x86", parameter_shape=shape, topology="factored" if factored else "full", - kwargs={"lr": 1e-2, "beta2": 0.9, "eps": 1e-6}, + kwargs=dict(numeric), ) metadata["state_contract"] = dict(scheduled.state_contract) metadata["scheduled_training"] = scheduled.metadata()