diff --git a/docs/audit/generated/runtime_abi.csv b/docs/audit/generated/runtime_abi.csv index 16e66779a..0e34e35fd 100644 --- a/docs/audit/generated/runtime_abi.csv +++ b/docs/audit/generated/runtime_abi.csv @@ -150,6 +150,10 @@ apple,tessera_apple_gpu_flash_attn_gqa_f16,flash_attn_gqa,f16,src/compiler/codeg apple,tessera_apple_gpu_flash_attn_gqa_f16,flash_attn_gqa,f16,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp apple,tessera_apple_gpu_flash_attn_gqa_f32,flash_attn_gqa,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm apple,tessera_apple_gpu_flash_attn_gqa_f32,flash_attn_gqa,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +apple,tessera_apple_gpu_gated_delta_rule_chunked_f32,gated_delta_rule_chunked,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +apple,tessera_apple_gpu_gated_delta_rule_chunked_f32,gated_delta_rule_chunked,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +apple,tessera_apple_gpu_gated_delta_rule_f32,gated_delta_rule,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +apple,tessera_apple_gpu_gated_delta_rule_f32,gated_delta_rule,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp apple,tessera_apple_gpu_gather_blocks_dev_f32,gather_blocks_dev,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm apple,tessera_apple_gpu_gather_blocks_dev_f32,gather_blocks_dev,f32,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp apple,tessera_apple_gpu_gather_blocks_dev_f32_enc,gather_blocks_dev_f32_enc,,src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm diff --git a/docs/audit/generated/runtime_abi.md b/docs/audit/generated/runtime_abi.md index 2f71a5751..f23812bbb 100644 --- a/docs/audit/generated/runtime_abi.md +++ b/docs/audit/generated/runtime_abi.md @@ -4,9 +4,9 @@ Human-readable view. The canonical machine-readable artifact is `runtime_abi.csv ## Headline -- **276** unique `extern "C" tessera_*` C ABI symbols across all backends. +- **278** unique `extern "C" tessera_*` C ABI symbols across all backends. - **6 / 6** core runtime headers present. -- **116** Apple GPU kernel families with per-dtype variants. +- **118** Apple GPU kernel families with per-dtype variants. ## Core runtime headers @@ -23,7 +23,7 @@ Human-readable view. The canonical machine-readable artifact is `runtime_abi.csv | Backend | Unique tessera_* symbols | |---------|-------------------------:| -| `apple` | 265 | +| `apple` | 267 | | `nvidia` | 3 | | `x86` | 8 | @@ -85,6 +85,8 @@ Human-readable view. The canonical machine-readable artifact is `runtime_abi.csv | `flash_attn` | `bf16`, `f16`, `f32` | | `flash_attn_bias` | `bf16`, `f16`, `f32` | | `flash_attn_gqa` | `bf16`, `f16`, `f32` | +| `gated_delta_rule` | `f32` | +| `gated_delta_rule_chunked` | `f32` | | `gather_blocks_dev` | `f32` | | `gelu` | `bf16`, `f16`, `f32` | | `grouped_gemm` | `f32` | diff --git a/docs/audit/generated/test_coverage.csv b/docs/audit/generated/test_coverage.csv index edb32d744..ba95ad8ad 100644 --- a/docs/audit/generated/test_coverage.csv +++ b/docs/audit/generated/test_coverage.csv @@ -183,7 +183,7 @@ focal_loss,4,0,0,4,0,,covered_by_family,category default for 'loss' fori_loop,19,0,1,19,0,,structural_only,category default for 'control_flow' fused_epilogue,2,6,0,8,0,,structural_only,unclassified — defaults to structural_only gated_attention,6,0,0,6,0,,needs_direct_test,category default for 'attention' -gated_deltanet,7,2,0,9,0,,needs_direct_test,category default for 'attention' +gated_deltanet,13,2,0,15,0,fp64,needs_direct_test,category default for 'attention' gather,9,3,0,12,0,fp32 fp4_e2m1 fp6_e3m2 fp8_e4m3 fp8_e5m2 nvfp4,needs_direct_test,category default for 'layout_transform' ge,1,0,0,1,1,,covered_by_family,category default for 'comparison' gelu,32,17,1,49,0,bf16 f16 f32 fp16 fp32,covered_by_family,category default for 'elementwise' @@ -377,7 +377,7 @@ score_matching_loss,1,0,0,1,1,,covered_by_family,category default for 'loss' sddmm,2,0,0,2,0,,needs_direct_test,category default for 'sparse' segment_reduce,4,0,0,4,0,fp32,structural_only,unclassified — defaults to structural_only select,1,0,0,1,1,,structural_only,unclassified — defaults to structural_only -selective_ssm,36,0,2,36,0,,structural_only,unclassified — defaults to structural_only +selective_ssm,38,5,2,43,0,,structural_only,unclassified — defaults to structural_only seq2seq_loss,3,0,0,3,0,,covered_by_family,category default for 'loss' sgd,6,0,0,6,0,,needs_direct_test,category default for 'functional_optimizer_step' shard_map,0,0,0,0,1,,structural_only,category default for 'sharding' diff --git a/docs/audit/generated/test_coverage.md b/docs/audit/generated/test_coverage.md index c850109ba..d70a61f24 100644 --- a/docs/audit/generated/test_coverage.md +++ b/docs/audit/generated/test_coverage.md @@ -7,10 +7,10 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by ## Headline - **457** ops in `primitive_coverage` registry. -- **2445** total Python-test references, **685** total lit-fixture references. +- **2453** total Python-test references, **690** total lit-fixture references. - **104** ops have **zero** references in either test surface. - **202** ops have ≤1 reference ("thinly tested"). -- **54** ops have ≥10 references ("well tested"). +- **55** ops have ≥10 references ("well tested"). - **61** ops have at least one associated `pytest.raises` negative test. ## Top 20 most-tested ops @@ -32,8 +32,8 @@ Generated from `python/tessera/compiler/test_coverage_audit.py`. Don't edit by | `gelu` | 32 | 17 | 49 | 1 | `bf16`, `f16`, `f32`, `fp16` … | | `cholesky` | 15 | 30 | 45 | 0 | `bf16`, `f16`, `f32`, `fp16` … | | `msa_sparse_attention` | 37 | 8 | 45 | 1 | | +| `selective_ssm` | 38 | 5 | 43 | 2 | | | `linear_attn` | 31 | 8 | 39 | 1 | | -| `selective_ssm` | 36 | 0 | 36 | 2 | | | `cast` | 7 | 27 | 34 | 0 | `fp32` | | `transpose` | 12 | 18 | 30 | 0 | `fp32` | | `moe_swiglu_block` | 10 | 17 | 27 | 1 | `fp8_e4m3`, `nvfp4` | diff --git a/docs/audit/generated/verifier_coverage.csv b/docs/audit/generated/verifier_coverage.csv index 092f6926d..054408f3c 100644 --- a/docs/audit/generated/verifier_coverage.csv +++ b/docs/audit/generated/verifier_coverage.csv @@ -133,6 +133,7 @@ src/compiler/ir/TesseraOps.td,RopeSplitOp,0,0,no_verifier src/compiler/ir/TesseraOps.td,SVDOp,1,1,real src/compiler/ir/TesseraOps.td,ScoreCombineOp,1,1,real src/compiler/ir/TesseraOps.td,SelectOp,0,0,no_verifier +src/compiler/ir/TesseraOps.td,SelectiveSsmOp,1,1,real src/compiler/ir/TesseraOps.td,SigmoidOp,0,0,no_verifier src/compiler/ir/TesseraOps.td,SiluMulOp,0,0,no_verifier src/compiler/ir/TesseraOps.td,SiluOp,0,0,no_verifier diff --git a/docs/audit/generated/verifier_coverage.md b/docs/audit/generated/verifier_coverage.md index 4e786d06a..7f8570313 100644 --- a/docs/audit/generated/verifier_coverage.md +++ b/docs/audit/generated/verifier_coverage.md @@ -6,11 +6,11 @@ Human-readable view. The canonical machine-readable artifact is `verifier_covera | Status | Count | Meaning | |--------|-------|---------| -| `real` | 112 | `hasVerifier = 1;` + substantive `verify()` body. | +| `real` | 113 | `hasVerifier = 1;` + substantive `verify()` body. | | `trivial_stub` | 0 | `hasVerifier = 1;` + trivial `return success();` stub. | | `absent` | 0 | `hasVerifier = 1;` but no `verify()` body (build error risk). | | `no_verifier` | 47 | No verifier declared. TD constraints suffice — fine for many ops. | -| **Total** | 159 | | +| **Total** | 160 | | ## Per-dialect details @@ -152,6 +152,7 @@ Human-readable view. The canonical machine-readable artifact is `verifier_covera | `SVDOp` | `real` | | `ScoreCombineOp` | `real` | | `SelectOp` | `no_verifier` | +| `SelectiveSsmOp` | `real` | | `SigmoidOp` | `no_verifier` | | `SiluMulOp` | `no_verifier` | | `SiluOp` | `no_verifier` | diff --git a/docs/audit/roadmap/FRONTIER_MODEL_SURVEY_2026.md b/docs/audit/roadmap/FRONTIER_MODEL_SURVEY_2026.md index 8daa66692..57265a89a 100644 --- a/docs/audit/roadmap/FRONTIER_MODEL_SURVEY_2026.md +++ b/docs/audit/roadmap/FRONTIER_MODEL_SURVEY_2026.md @@ -25,16 +25,31 @@ already shipped the right validation tool for promotion: the evaluator's DESIL cross-path + metamorphic oracles (`python/tessera/compiler/evaluator.py`) — a fused kernel is auto-provable against its reference. -### Doc-drift correction (fix before building) - -The audit pass surfaced a CLAUDE.md claim that a "dedicated Mamba2 Graph IR op -landed (2026-05-18)". **It did not** — there is no `selective_ssm` op in -`src/compiler/ir/TesseraOps.td`; the registry only flips `graph_ir_lowering = -registered` (intent). Conversely, an earlier internal audit claimed -`gated_deltanet` has *no* ODS op — **also wrong**: `Tessera_GatedDeltaNetOp` -exists at `TesseraOps.td:1071` with a rich operand set. Both errors are the same -hazard (Decision #25/#26): registry/prose intent ≠ materialized compiler surface. -Every status in this doc is grounded to a `file:line`. +### Grounding corrections (verified at source) + +Two claims that floated through the audit, both checked against the tree: + +1. **The false "Mamba2 op landed" claim was real — in `primitive_coverage.py`, + not CLAUDE.md.** CLAUDE.md proper says the op is "*pending*"; but the coverage + registry's comment asserted *"dedicated Mamba2 Graph IR op landed (2026-05-18) + as `tessera.selective_ssm`"* and set `graph_ir_lowering="registered"` while **no + ODS op existed** in `TesseraOps.td` — registry intent that outran the compiler + surface (Decision #25/#26). **L4 closed this for real** (see below): the op is + now materialized + lit-proven, so `registered` is finally honest. (`gated_deltanet` + always did have its ODS op at `TesseraOps.td:1109`.) + +2. **The real correctness finding (now test-proven):** the shipped + `gated_deltanet` / `kimi_delta_attention` / `modified_delta_attention` + reference (`__init__.py::_delta_attention_impl`, lines 1215-1225) computes + `Ŝ_t = α_t·Ŝ_{t-1} + β_t·k_t v_tᵀ` — **gated linear attention, missing the + DeltaNet `(I − β_t k_t k_tᵀ)` erase term**. The ODS summary ("Gated DeltaNet + recurrent attention") and the runtime comment ("the delta recurrence is + algebraically the quadratic form `(QKᵀ⊙mask)@V`", `runtime.py:6428`) describe + the delta rule, but the math is linear attention — and every parity test passes + because the GPU path faithfully matches the mislabeled reference. + `tessera.stdlib.delta_rule` (Track L L1/L2, landed below) adds the genuine rule + and the oracle that locks the distinction + (`tests/unit/test_stdlib_delta_rule.py`). ## Per-model contract → Tessera status (grounded) @@ -145,13 +160,108 @@ scaled execution on Apple GPU gated vs numpy. | L | Title | Definition of done | Oracle | |---|---|---|---| -| **L0** | Doc-drift fix + contract lock | Remove the false "Mamba2 op landed" claim; lock the `gated_deltanet` operand contract; add a failing scaled gate | n/a | -| **L1** | **Wire β/decay/state through the recurrent kernel** | `linear_attn_f32` becomes a true gated-delta-rule step: `S_t = α_t·S_{t-1}(I − β_t k kᵀ) + β_t v kᵀ`, fp32 state accum; decode path correct | metamorphic: decay=1,β=1 ≡ existing linear_attn; vs numpy delta-rule ref | -| **L2** | **Chunked UT-transform prefill (the keystone)** | New lowering: chunk C=64, per-chunk `A=tril(−diag(β)KKᵀ,−1)`, `T=(I−A)⁻¹` as a dedicated C×C triangular-solve tile primitive, WY factors via GEMM, cross-chunk state carry; α-decay folded in | **chunk ≡ recurrent** (DESIL cross-path) — the make-or-break proof | -| **L3** | **Hybrid layer schedule** as a first-class attribute | Lower `layer_types` literally (`(i+1)%period==0 → full_attention`) into the layer schedule; dual KV/recurrent-state cache coexist per layer | full-config artifact lit (Qwen3.6 dims) | -| **L4** | `selective_ssm` (Mamba2) ODS op + chunk-scan lowering | Materialize the op the docs already claim; reuse the L2 chunk machinery | chunk-scan ≡ sequential-scan | +| **L0** ✅ | Grounding correction + contract lock | Correct the propagated misquote; document the *real* finding (delta family = linear attn, no erase); lock it with an oracle | `test_existing_gated_deltanet_is_linear_attention_not_delta` | +| **L1** ✅ | **Genuine gated delta recurrence** (decode form) | `gated_delta_rule_recurrent` adds the `(v_t − α_t v̂_t)` erase, fp32 state, return_state; **not** "wire β/decay" (those were already wired) — the existing reference was missing the erase entirely | vs independent brute-force `(I−βkkᵀ)` recurrence; `erase=False` ≡ existing ref; state-carry | +| **L2** ✅ | **Chunked UT-transform prefill (the keystone)** | `gated_delta_rule_chunked`: chunk C, `Ã=tril(β·γ-ratio·KKᵀ,−1)`, `(I+Ã)⁻¹` via explicit forward substitution (`_forward_substitution`), WY/output as GEMM, γ-decay folding, cross-chunk state carry | **chunk ≡ recurrent** across ungated/β/fully-gated/output-gated + chunk-size-invariant (the make-or-break proof) | +| **L1.1** ✅ | Genuine delta rule on Metal (decode form) | `tessera_apple_gpu_gated_delta_rule_f32` — per-(b,h) sequential MSL scan with the erase; `backend="apple_gpu"` on the recurrent reference | **Metal ≡ numpy** (DESIL) + Metal ≡ L2 chunked (independent routes) | +| **L2.1** ✅ | Chunked UT-transform on Metal (prefill form) | `tessera_apple_gpu_gated_delta_rule_chunked_f32` — one threadgroup per (b,h), the within-chunk `(I+Ã)⁻¹` solve on-device; `backend="apple_gpu"` on the chunked reference | **Metal chunked ≡ numpy** (all chunk sizes incl. partial) + **Metal chunked ≡ Metal recurrent** | +| **L2.2** | Cooperative-parallel chunk kernel | Parallelize the within-chunk solve + state carry across threads (currently lane-0); tensor-core-style GEMM tiles | perf ratchet; correctness already gated | +| **L3** ✅ | **Hybrid layer schedule** as a first-class attribute | `HybridSchedule` lowers `layer_types` literally; reference stack threads the **dual cache** (recurrent Ŝ for linear layers, KV for full layers) | **streaming dual-cache decode ≡ full recompute** + Qwen3.6 full-config schedule check | +| **L3.1** | Promote `gated_deltanet` ODS → true rule | Decide + execute the shipped-numerics change (today's op is linear attn); wire the ODS op to the L1.1/L2.1 kernels | DESIL + existing-test migration | +| **L4** ✅ | `selective_ssm` (Mamba2) ODS op | Materialize the op the registry falsely claimed: `Tessera_SelectiveSsmOp` + verifier; close the drift. Chunk-scan (`_mamba_ssd.py`) + chunk≡sequential oracle already existed | lit roundtrip/verifier + chunk-scan ≡ sequential-scan | +| **L4.1** ✅ | Hybrid SSM mixer (Nemotron) | `linear_mixer="ssm"` adds a Mamba SSM mixer to the L3 hybrid stack; SSM state `h[D,N]` carried in the dual cache alongside attention-anchor KV | dual-cache decode ≡ recompute with SSM layers + `_ssm_scan` ≡ shipped `selective_ssm` | | **L5** | LFM2.5 LIV mixer variant | `Linear→(B⊙x̃)→depthwise-causal-conv(k=3)→(C⊙z)→Linear_out` as a fused mixer over the existing `depthwise_conv1d` | vs numpy LIV ref; scaled exec | +### Landed 2026-06-15 — L0–L2 (reference tier, host-free) + +`python/tessera/stdlib/delta_rule.py` + `tests/unit/test_stdlib_delta_rule.py` +(18 oracles, all green). Mirrors the M-series house pattern: numpy reference + +oracle first, fused MSL kernel (L1.1/L2.1) as the hardware-gated follow-up. What +is proven host-free: the genuine gated delta rule (recurrent ≡ independent +brute-force in the paper's `(I−βkkᵀ)` layout), the chunk-parallel UT-transform +(`chunk ≡ recurrent` across all gating modes, chunk-size-invariant), the +triangular-solve primitive, cross-chunk state carry, and the L0 lock that the +shipped `gated_deltanet` is linear attention (`erase=False`), materially distinct +from the true rule when keys correlate. + +### Landed 2026-06-15 — L1.1 (genuine delta rule on Metal) + +`tessera_apple_gpu_gated_delta_rule_f32` (`apple_gpu_runtime.mm` + non-Darwin +stub parity) — a per-(b,h) sequential MSL scan carrying the `(v_t − α·v̂_t)` +erase, registered in `_apple_gpu_backend` and reachable via +`gated_delta_rule_recurrent(..., backend="apple_gpu")`. +`tests/unit/test_apple_gpu_gated_delta_rule.py` (7 oracles, all run on Metal): +**Metal ≡ numpy** for true delta / β+decay / output-gate / non-square head dims, +`erase=False` ≡ the shipped linear reference, and **Metal recurrent ≡ the L2 +chunked UT-transform** (the genuine rule reached by two fully independent routes). + +**Numerics finding (realism, → `numeric_policy`):** the delta rule is only +well-conditioned with **L2-normalized keys** — then `β·‖k‖²=β<1` makes +`(I−βkkᵀ)` a contraction and f32≡f64. With unnormalized keys (`‖k‖²≫1`) the +recurrence *expands* (eigenvalue `1−β‖k‖²<0`) and f32 legitimately diverges from +f64 (~10% here) — genuine ill-conditioning, not a kernel defect. A production +`gated_delta_rule` op should carry key-normalization in its contract (real models +do), with fp32 state accumulation. + +### Landed 2026-06-15 — L2.1 (chunked UT-transform on Metal) + +`tessera_apple_gpu_gated_delta_rule_chunked_f32` — one threadgroup per (b,h), C +threads (one per token-in-chunk), state Ŝ in threadgroup memory across the chunk +loop. The GEMM-shaped rows (A, W̃, output) parallelize across threads; the +`(I+Ã)⁻¹` forward substitution + rank-1 state carry run on lane 0 (cooperative +parallelization of those is L2.2, a perf follow-up). Reachable via +`gated_delta_rule_chunked(..., backend="apple_gpu")`. Oracles (in the same test +file): **Metal chunked ≡ numpy** across chunk sizes 1/4/8/16/32 (S=20 exercises a +partial last chunk), **Metal chunked ≡ Metal recurrent** (two independent +on-device kernels), output-gate, and `erase=False` ≡ the shipped linear ref. + +### Landed 2026-06-15 — L3 (hybrid layer schedule + dual cache) + +`tessera.stdlib.hybrid` — `HybridSchedule` makes `layer_types` first-class +(`qwen3_6_schedule` = `[lin,lin,lin,full]·N`; `nemotron_schedule` = sparse +anchors) + a reference stack that threads the **dual cache**: constant-size +recurrent Ŝ for linear (genuine gated-delta) layers, growing KV for full-attention +layers. Linear layers L2-normalize keys (the L1.1 conditioning finding). +`tests/unit/test_stdlib_hybrid.py` (9 oracles): the headline **streaming +dual-cache decode ≡ full recompute** across prefill points and schedules +(all-linear, every-other-anchor), schedule validation, and the Qwen3.6 full-config +check (30 linear / 10 full at 40 layers). + +**Not yet done:** L2.2 (cooperative-parallel chunk kernel — perf), L3.1 (promote +the `gated_deltanet` ODS op to the true rule — a shipped-numerics decision), and +MoE/MTP composition into the hybrid stack (`stdlib.moe` exists; wiring is additive). + +### Landed 2026-06-15 — L4 (`selective_ssm` Mamba2 ODS op) + +`Tessera_SelectiveSsmOp` (`src/compiler/ir/TesseraOps.td`) + `SelectiveSsmOp::verify` +(`TesseraOps.cpp`) — `tessera.selective_ssm` is now a genuine Graph IR op +(rank-checked: rank-3 x / shape-equal delta / matching b,c / A rank-1|2 / optional +`gate` shape-equal x / `init` state rank-3). `tessera-opt` rebuilds clean (MLIR +22.1.6); lit fixture `tests/tessera-ir/model_class/selective_ssm.mlir` passes +(model_class sweep 5/5). The **drift is closed**: the coverage registry's +`graph_ir_lowering="registered"` for `selective_ssm` is now backed by a real op +(comment corrected in `primitive_coverage.py`). The chunked-parallel SSD lowering +(`_mamba_ssd.py::selective_ssm_parallel`) and its **chunk ≡ sequential** oracle +(`test_mamba_ssd_gpu.py`, 12 tests) already existed and stay green. + +### Landed 2026-06-15 — L4.1 (hybrid SSM mixer — Nemotron expressible) + +`tessera.stdlib.hybrid` now takes `linear_mixer = "delta" | "ssm"`. The stack was +refactored to per-mixer **span functions** so `hybrid_forward` (one span) and +`hybrid_decode` (streamed spans) run identical per-layer code — the dual-cache +oracle is meaningful for all three mixer types (delta Ŝ, SSM h[D,N], attention +KV). `_ssm_scan` reproduces the shipped `tessera.ops.selective_ssm` (the L4 op's +reference) **and returns the carried state** (which the public reference does +not), so streaming SSM decode is exact. Nemotron is now the second flagship +(after Qwen3.6) expressible end-to-end: `nemotron_schedule` + `linear_mixer="ssm"` += Mamba layers + sparse attention anchors. `tests/unit/test_stdlib_hybrid.py` +(+5 L4.1 oracles): `_ssm_scan ≡ selective_ssm`, Nemotron-shaped dual-cache decode +≡ recompute (across prefill points), and an all-SSM stack. + +**Still open:** MoE/MTP composition into the hybrid stack (`stdlib.moe` exists; +additive), the optional Mamba short-conv (shares the L5 LIV machinery), L2.2 +(perf), L3.1 (`gated_deltanet` ODS→true-rule), L5 (LFM2.5 LIV mixer). + Sequencing: **L1 unblocks L2** (decode state is the chunk carry); **L2 is the keystone** (only the chunked GEMM form is tensor-core-viable for prefill — the papers are unanimous); L3 is independent of L1/L2; L4 reuses L2; L5 is parallel. diff --git a/python/tessera/__init__.py b/python/tessera/__init__.py index 0c9db5423..fa9d7854d 100644 --- a/python/tessera/__init__.py +++ b/python/tessera/__init__.py @@ -1173,6 +1173,14 @@ def _delta_attention_impl(Q, K, V, *, gate=None, beta=None, decay=None, return_state: bool = False, state_dtype: str = "fp32", modified: bool = False): + # NOTE (Track L): this recurrence is `Ŝ_t = α_t·Ŝ_{t-1} + β_t·k_t v_tᵀ`, + # i.e. *gated linear attention* — it omits the DeltaNet + # `(I − β_t k_t k_tᵀ)` erase term, so despite the name it is not the true + # delta rule. The genuine gated delta rule (recurrent + chunk-parallel + # UT-transform) lives in `tessera.stdlib.delta_rule`, with the reduction + # oracle (`erase=False` ≡ this impl) in + # `tests/unit/test_stdlib_delta_rule.py`. Promoting this op to the true + # rule is tracked as Track L L3+ (would change shipped numerics). if hasattr(Q, "_data"): Q = Q._data if hasattr(K, "_data"): K = K._data if hasattr(V, "_data"): V = V._data diff --git a/python/tessera/_apple_gpu_backend.py b/python/tessera/_apple_gpu_backend.py index 97e260779..b88604b92 100644 --- a/python/tessera/_apple_gpu_backend.py +++ b/python/tessera/_apple_gpu_backend.py @@ -118,6 +118,13 @@ def _load(): ("tessera_apple_gpu_ebm_half_mse_f32", [fp, fp, fp, i32]), ("tessera_apple_gpu_ebm_ism_f32", [fp, fp, fp, i32, i32]), ("tessera_apple_gpu_ebm_dsm_f32", [fp, fp, fp, fp, i32, i32, flt]), + # Track L (L1.1) — genuine gated delta rule (recurrent, true erase). + # (Q, K, V, beta, decay, O, B, H, S, D_qk, D_v, erase). + ("tessera_apple_gpu_gated_delta_rule_f32", + [fp, fp, fp, fp, fp, fp, i32, i32, i32, i32, i32, i32]), + # Track L (L2.1) — chunked UT-transform prefill (+chunk arg). + ("tessera_apple_gpu_gated_delta_rule_chunked_f32", + [fp, fp, fp, fp, fp, fp, i32, i32, i32, i32, i32, i32, i32]), ): try: sym = getattr(lib, name) @@ -172,6 +179,52 @@ def _ptr(a: np.ndarray): return a.ctypes.data_as(ctypes.POINTER(ctypes.c_float)) +def gpu_gated_delta_rule(Q: np.ndarray, K: np.ndarray, V: np.ndarray, + beta: np.ndarray, decay: np.ndarray, + erase: bool = True) -> np.ndarray: + """Genuine gated delta rule (recurrent) on the Apple GPU (Track L L1.1). + + Q, K, V : [B, H, S, D] f32 (D_qk = K/Q last dim, D_v = V last dim). + beta, decay : [B, H, S] f32 per-token scalars (pass ones for "absent"). + The output gate is applied by the caller. Matches + ``tessera.stdlib.delta_rule.gated_delta_rule_recurrent``. + """ + lib = _load() + Qf, Kf, Vf = _f32(Q, "Q"), _f32(K, "K"), _f32(V, "V") + bf, df = _f32(beta, "beta"), _f32(decay, "decay") + B, H, S, D_qk = Qf.shape + D_v = Vf.shape[-1] + O = np.empty((B, H, S, D_v), dtype=np.float32) + lib.tessera_apple_gpu_gated_delta_rule_f32( + _ptr(Qf), _ptr(Kf), _ptr(Vf), _ptr(bf), _ptr(df), _ptr(O), + ctypes.c_int32(B), ctypes.c_int32(H), ctypes.c_int32(S), + ctypes.c_int32(D_qk), ctypes.c_int32(D_v), + ctypes.c_int32(1 if erase else 0)) + return O + + +def gpu_gated_delta_rule_chunked(Q: np.ndarray, K: np.ndarray, V: np.ndarray, + beta: np.ndarray, decay: np.ndarray, + chunk: int = 32, erase: bool = True) -> np.ndarray: + """Chunk-parallel UT-transform gated delta rule on the Apple GPU (L2.1). + + Same signature/contract as :func:`gpu_gated_delta_rule` plus ``chunk`` (the + within-chunk tile, ≤ 32). chunk ≡ recurrent. + """ + lib = _load() + Qf, Kf, Vf = _f32(Q, "Q"), _f32(K, "K"), _f32(V, "V") + bf, df = _f32(beta, "beta"), _f32(decay, "decay") + B, H, S, D_qk = Qf.shape + D_v = Vf.shape[-1] + O = np.empty((B, H, S, D_v), dtype=np.float32) + lib.tessera_apple_gpu_gated_delta_rule_chunked_f32( + _ptr(Qf), _ptr(Kf), _ptr(Vf), _ptr(bf), _ptr(df), _ptr(O), + ctypes.c_int32(B), ctypes.c_int32(H), ctypes.c_int32(S), + ctypes.c_int32(D_qk), ctypes.c_int32(D_v), ctypes.c_int32(int(chunk)), + ctypes.c_int32(1 if erase else 0)) + return O + + # ── Bespoke Metal kernels (the back-half) ──────────────────────────────────── diff --git a/python/tessera/compiler/primitive_coverage.py b/python/tessera/compiler/primitive_coverage.py index bf5720fd5..b7bda5796 100644 --- a/python/tessera/compiler/primitive_coverage.py +++ b/python/tessera/compiler/primitive_coverage.py @@ -1667,11 +1667,16 @@ def _graph_ir_lowering_for_category(category: str | None, current: str) -> str: "depthwise_conv1d": "registered", "online_softmax": "registered", "online_softmax_state": "registered", - # selective_ssm — dedicated Mamba2 Graph IR op landed (2026-05-18) as - # `tessera.selective_ssm` (state-space lowering kind, stateful effect). - # The closed-form JVP through the recurrence was already shipped; the - # `registered` flip below completes the Graph IR lowering brick that - # was the last remaining `missing` entry across the registry. + # selective_ssm — dedicated Mamba2 Graph IR op `tessera.selective_ssm` + # (state-space lowering kind, stateful effect). NOTE (Track L L4, + # 2026-06-15): this entry previously asserted the op "landed (2026-05-18)" + # while NO ODS op existed in TesseraOps.td — registry/prose intent that + # outran the compiler surface (Decision #25/#26). The op is now genuinely + # materialized: `Tessera_SelectiveSsmOp` in TesseraOps.td + + # `SelectiveSsmOp::verify` in TesseraOps.cpp, proven by + # `tests/tessera-ir/model_class/selective_ssm.mlir`. The closed-form JVP + # and the chunked-parallel SSD lowering (`_mamba_ssd.py`, chunk ≡ sequential) + # were already shipped; this makes `registered` honest. "selective_ssm": "registered", } diff --git a/python/tessera/stdlib/__init__.py b/python/tessera/stdlib/__init__.py index 53549edad..92e57af3a 100644 --- a/python/tessera/stdlib/__init__.py +++ b/python/tessera/stdlib/__init__.py @@ -11,6 +11,6 @@ from __future__ import annotations -from . import attention, moe, quant +from . import attention, delta_rule, hybrid, moe, quant -__all__ = ["quant", "moe", "attention"] +__all__ = ["quant", "moe", "attention", "delta_rule", "hybrid"] diff --git a/python/tessera/stdlib/delta_rule.py b/python/tessera/stdlib/delta_rule.py new file mode 100644 index 000000000..efec0f446 --- /dev/null +++ b/python/tessera/stdlib/delta_rule.py @@ -0,0 +1,256 @@ +"""``tessera.stdlib.delta_rule`` — the *true* gated delta rule (Track L, L1+L2). + +Background / why this module exists +----------------------------------- +Tessera already ships ``tessera.ops.gated_deltanet`` (and the ``kimi_delta`` / +``modified_delta`` siblings), but their reference recurrence +(``tessera/__init__.py::_delta_attention_impl``) is + + Ŝ_t = α_t · Ŝ_{t-1} + β_t · k_t v_tᵀ # additive accumulation + +which is **gated linear attention**, *not* the delta rule. The DeltaNet update +(Yang et al., "Gated Delta Networks", arXiv:2412.06464) carries an **erase** +term that removes the value currently bound to key ``k_t`` before writing the +new one: + + Ŝ_t = α_t · Ŝ_{t-1} + β_t · k_t · (v_t − α_t · v̂_t)ᵀ , v̂_t = Ŝ_{t-1}ᵀ k_t + +The ``(v_t − α_t v̂_t)`` correction is the generalized-Householder +``(I − β_t k_t k_tᵀ)`` erase in the paper's layout. Dropping it recovers the +existing linear-attention reference exactly — which is the reduction oracle in +``tests/unit/test_stdlib_delta_rule.py``. + +This module provides that genuine rule in two algebraically-equivalent forms: + +* ``gated_delta_rule_recurrent`` — the obviously-correct O(S) sequential + recurrence (the decode form; carries a constant-size ``[d_k, d_v]`` state). +* ``gated_delta_rule_chunked`` — the chunk-parallel **UT-transform** form (the + prefill form): everything is GEMM except one within-chunk unit-lower-triangular + solve ``(I + A)⁻¹`` done by explicit forward substitution + (``_forward_substitution`` — the "triangular-solve tile primitive" a real + kernel would specialize). The headline oracle is **chunk ≡ recurrent**. + +Layout convention (matches ``_delta_attention_impl``): + Q, K : [B, H, S, d_k] V : [B, H, S, d_v] state Ŝ : [B, H, d_k, d_v] + read O_t = q_tᵀ Ŝ_t (state read *after* the t-th update) + β, decay(α) : optional [B, H, S] gate : optional, broadcastable to O + +State accumulates in float64 here (the reference tier); the production contract +is fp32 accumulation regardless of bf16 storage — the erase + rank-update are +numerically sensitive. The fused MSL kernel is the L1.1/L2.1 follow-up; this is +the reference + oracle that proves the math, provable host-free. +""" + +from __future__ import annotations + +import numpy as np + + +def _arr(x) -> np.ndarray: + if hasattr(x, "_data"): + x = x._data + return np.asarray(x, dtype=np.float64) + + +def _per_token(x, B: int, H: int, S: int) -> np.ndarray | None: + """Normalize an optional β / decay argument to a [B, H, S] float64 array.""" + if x is None: + return None + a = _arr(x) + return np.broadcast_to(a, (B, H, S)).astype(np.float64, copy=False) + + +def _apply_gate(O: np.ndarray, gate, out_dtype) -> np.ndarray: + if gate is not None: + g = 1.0 / (1.0 + np.exp(-_arr(gate))) + O = O * np.broadcast_to(g, O.shape) + return O.astype(out_dtype, copy=False) + + +def _state_dtype(state_dtype: str, out_dtype): + return np.float32 if state_dtype in ("fp32", "bf16") else out_dtype + + +# ───────────────────────────────────────────────────────────────────────────── +# L1 — the genuine recurrence (decode form, constant-size state) +# ───────────────────────────────────────────────────────────────────────────── +def gated_delta_rule_recurrent(Q, K, V, *, beta=None, decay=None, gate=None, + state=None, causal: bool = True, + return_state: bool = False, + state_dtype: str = "fp32", + erase: bool = True, backend: str = "numpy"): + """True gated delta rule via the sequential recurrence. + + ``erase=True`` is the DeltaNet rule; ``erase=False`` degenerates to the + existing gated-linear-attention reference (``tessera.ops.gated_deltanet``), + which is the reduction oracle. ``backend="apple_gpu"`` (L1.1) runs the + genuine recurrence on Metal (``tessera_apple_gpu_gated_delta_rule_f32``), + falling back to numpy on a Metal miss or out-of-envelope shape — the DESIL + Metal≡numpy oracle. The GPU path covers ``state=None`` / + ``return_state=False`` (decode-from-zero); other cases use numpy. + """ + if not causal: + raise ValueError("gated_delta_rule_recurrent is a causal recurrence") + Q, K, V = _arr(Q), _arr(K), _arr(V) + if Q.ndim != 4 or K.ndim != 4 or V.ndim != 4: + raise ValueError("delta rule expects rank-4 (B, H, S, D) tensors") + out_dtype = np.result_type(np.asarray(Q), np.asarray(K), np.asarray(V)) + B, H, S, d_k = Q.shape + d_v = V.shape[-1] + beta_a = _per_token(beta, B, H, S) + decay_a = _per_token(decay, B, H, S) + + if backend == "apple_gpu" and state is None and not return_state: + try: + from tessera import _apple_gpu_backend as _agb + ones = np.ones((B, H, S), dtype=np.float32) + O = _agb.gpu_gated_delta_rule( + Q.astype(np.float32), K.astype(np.float32), V.astype(np.float32), + ones if beta_a is None else beta_a.astype(np.float32), + ones if decay_a is None else decay_a.astype(np.float32), + erase=erase).astype(np.float64) + return _apply_gate(O, gate, out_dtype) + except Exception: # noqa: BLE001 — any Metal/load miss → numpy oracle + pass + + if state is None: + Sst = np.zeros((B, H, d_k, d_v), dtype=np.float64) + else: + Sst = _arr(state).copy() + + O = np.zeros((B, H, S, d_v), dtype=np.float64) + for t in range(S): + a = decay_a[:, :, t][:, :, None, None] if decay_a is not None else 1.0 + b = beta_a[:, :, t][:, :, None, None] if beta_a is not None else 1.0 + k_t = K[:, :, t, :] # [B,H,d_k] + v_t = V[:, :, t, :] # [B,H,d_v] + # v̂_t = Ŝ_{t-1}ᵀ k_t — the value currently bound to key k_t. + v_hat = np.einsum("bhd,bhde->bhe", k_t, Sst) # [B,H,d_v] + a_s = decay_a[:, :, t][:, :, None] if decay_a is not None else 1.0 + target = v_t - a_s * v_hat if erase else v_t # erase correction + Sst = a * Sst + b * np.einsum("bhd,bhe->bhde", k_t, target) + O[:, :, t, :] = np.einsum("bhd,bhde->bhe", Q[:, :, t, :], Sst) + + O = _apply_gate(O, gate, out_dtype) + Sst = Sst.astype(_state_dtype(state_dtype, out_dtype), copy=False) + return (O, Sst) if return_state else O + + +# ───────────────────────────────────────────────────────────────────────────── +# L2 — the chunk-parallel UT-transform (prefill form) +# ───────────────────────────────────────────────────────────────────────────── +def _forward_substitution(A_strict: np.ndarray, W: np.ndarray) -> np.ndarray: + """Solve (I + A) U = W for U, where A is strictly lower-triangular (so + (I + A) is unit lower-triangular). Explicit forward substitution — the + "triangular-solve tile primitive"; the only non-GEMM step of the chunk form. + + A_strict : [..., C, C] W : [..., C, d] -> U : [..., C, d] + """ + C = A_strict.shape[-1] + U = np.array(W, dtype=np.float64, copy=True) + for t in range(C): + if t > 0: + # U[t] = W[t] - Σ_{j...d", A_strict[..., t, :t], U[..., :t, :]) + U[..., t, :] = W[..., t, :] - corr + return U + + +def gated_delta_rule_chunked(Q, K, V, *, beta=None, decay=None, gate=None, + state=None, chunk_size: int = 64, + causal: bool = True, return_state: bool = False, + state_dtype: str = "fp32", erase: bool = True, + backend: str = "numpy"): + """True gated delta rule via the chunk-parallel UT-transform. + + Per chunk (carried state Ŝ₀, cumulative within-chunk decay γ_t = Π_{i≤t} α_i): + + Ã[t,j] = β_t (γ_t/γ_j)(k_tᵀ k_j) for jbhtj", Kc, Kc) # [B,H,C,C] + ratio = gamma[:, :, :, None] / gamma[:, :, None, :] # γ_t/γ_j + tril_strict = np.tril(np.ones((C, C)), k=-1) + A_strict = (bc[:, :, :, None] * ratio * KKt) * tril_strict if erase \ + else np.zeros((B, H, C, C)) + + # Target W̃ = β (V − γ · Ŝ₀ᵀ K). + KS0 = np.einsum("bhtd,bhde->bhte", Kc, Sst) # k_tᵀ Ŝ₀ -> [B,H,C,d_v] + if erase: + W = bc[:, :, :, None] * (Vc - gamma[:, :, :, None] * KS0) + else: + W = bc[:, :, :, None] * Vc + U = _forward_substitution(A_strict, W) # [B,H,C,d_v] + + # Output: O = diag(γ) Q Ŝ₀ + (tril(QKᵀ,0) ⊙ Γ) U. + QS0 = np.einsum("bhtd,bhde->bhte", Qc, Sst) # [B,H,C,d_v] + QKt = np.einsum("bhtd,bhjd->bhtj", Qc, Kc) # [B,H,C,C] + tril_incl = np.tril(np.ones((C, C)), k=0) + Mintra = QKt * ratio * tril_incl # γ_t/γ_j on j≤t + O[:, :, c0:c1, :] = (gamma[:, :, :, None] * QS0 + + np.einsum("bhtj,bhje->bhte", Mintra, U)) + + # State carry: Ŝ_C = γ_C Ŝ₀ + Kᵀ diag(γ_C/γ_t) U. + scale = (gamma_C[:, :, None] / gamma) # [B,H,C] + Sst = (gamma_C[:, :, None, None] * Sst + + np.einsum("bhtd,bhte->bhde", Kc, scale[:, :, :, None] * U)) + + O = _apply_gate(O, gate, out_dtype) + Sst = Sst.astype(_state_dtype(state_dtype, out_dtype), copy=False) + return (O, Sst) if return_state else O diff --git a/python/tessera/stdlib/hybrid.py b/python/tessera/stdlib/hybrid.py new file mode 100644 index 000000000..b1740890a --- /dev/null +++ b/python/tessera/stdlib/hybrid.py @@ -0,0 +1,275 @@ +"""``tessera.stdlib.hybrid`` — hybrid linear/attention/SSM schedules (Track L L3+L4.1). + +Frontier hybrids alternate a constant-state mixer with periodic full-attention +"anchor" layers: + * Qwen3.6 — Gated DeltaNet (linear) ×3 : Gated Attention ×1 + * Nemotron — Mamba-2 SSM (linear) : sparse attention anchors + * Mellum2 — sliding-window : full + +The defining *systems* contract is a **dual cache**: linear/SSM layers carry a +fixed-size recurrent state (delta Ŝ[d_k,d_v] / SSM h[D,N]); full-attention layers +carry a growing KV cache. Step-by-step decode equals a full recompute only if +all of them are threaded correctly. + +This module makes the schedule first-class (`HybridSchedule`) and the linear-slot +mixer pluggable (`linear_mixer = "delta" | "ssm"`). The stack is built from +per-mixer **span functions** that handle a token-span of any length carrying +their own cache, so `hybrid_forward` (one span) and `hybrid_decode` (streamed +spans) run identical per-layer code — the oracle (`tests/unit/test_stdlib_hybrid.py`) +is **streaming dual-cache decode ≡ full recompute**. Delta layers L2-normalize +keys (the L1.1 conditioning finding); the SSM step reproduces +`tessera.ops.selective_ssm` (the L4 op) and returns the carried state. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from . import delta_rule as _dr + +LINEAR = "linear" +FULL = "full" +DELTA = "delta" +SSM = "ssm" + + +# ───────────────────────────────────────────────────────────────────────────── +# The schedule — first-class layer_types +# ───────────────────────────────────────────────────────────────────────────── +@dataclass(frozen=True) +class HybridSchedule: + """Per-layer mixer assignment. ``full`` layer iff ``(i + full_offset) % + period == 0``. Qwen3.6 = period 4, full_offset 1 → ``[lin,lin,lin,full]·N``.""" + num_layers: int + period: int = 4 + full_offset: int = 1 + + def __post_init__(self): + if self.num_layers < 1: + raise ValueError("num_layers must be >= 1") + if self.period < 1: + raise ValueError("period must be >= 1") + + def layer_types(self) -> list[str]: + return [FULL if ((i + self.full_offset) % self.period == 0) else LINEAR + for i in range(self.num_layers)] + + def is_full(self, i: int) -> bool: + return (i + self.full_offset) % self.period == 0 + + def full_indices(self) -> list[int]: + return [i for i in range(self.num_layers) if self.is_full(i)] + + def linear_indices(self) -> list[int]: + return [i for i in range(self.num_layers) if not self.is_full(i)] + + def counts(self) -> dict[str, int]: + t = self.layer_types() + return {LINEAR: t.count(LINEAR), FULL: t.count(FULL)} + + +def qwen3_6_schedule(num_layers: int = 40) -> HybridSchedule: + """Qwen3.6-35B-A3B: ``[Gated DeltaNet ×3, Gated Attention] × (N/4)``.""" + return HybridSchedule(num_layers=num_layers, period=4, full_offset=1) + + +def nemotron_schedule(num_layers: int, attn_period: int = 8) -> HybridSchedule: + """Nemotron-3-style: predominantly linear (Mamba), a sparse attention anchor + every ``attn_period`` layers.""" + return HybridSchedule(num_layers=num_layers, period=attn_period, full_offset=1) + + +# ───────────────────────────────────────────────────────────────────────────── +# Reference hybrid stack (the dual-cache contract) +# ───────────────────────────────────────────────────────────────────────────── +@dataclass +class HybridConfig: + d_model: int + num_heads: int + head_dim: int + schedule: HybridSchedule + ffn_mult: int = 2 + linear_mixer: str = DELTA # "delta" (Qwen3.6) | "ssm" (Nemotron/Mamba) + ssm_state: int = 8 # SSM state dim N (scalar-state config) + + @property + def inner(self) -> int: + return self.num_heads * self.head_dim + + def mixer_for(self, i: int) -> str: + """One of ``full`` / ``delta`` / ``ssm`` for layer ``i``.""" + if self.schedule.is_full(i): + return FULL + return self.linear_mixer + + +def synth_weights(cfg: HybridConfig, rng) -> list[dict]: + """Small random per-layer weights (reference tier), keyed by mixer type.""" + Dm, inner, N = cfg.d_model, cfg.inner, cfg.ssm_state + ff = cfg.ffn_mult * Dm + s = 1.0 / np.sqrt(Dm) + layers = [] + for i in range(cfg.schedule.num_layers): + w = dict( + n1=rng.standard_normal(Dm) * 0.1 + 1.0, + n2=rng.standard_normal(Dm) * 0.1 + 1.0, + wo=rng.standard_normal((inner, Dm)) / np.sqrt(inner), + wg=rng.standard_normal((Dm, ff)) * s, + wu=rng.standard_normal((Dm, ff)) * s, + wd=rng.standard_normal((ff, Dm)) / np.sqrt(ff), + ) + mixer = cfg.mixer_for(i) + if mixer == FULL: + w.update(wq=rng.standard_normal((Dm, inner)) * s, + wk=rng.standard_normal((Dm, inner)) * s, + wv=rng.standard_normal((Dm, inner)) * s) + elif mixer == DELTA: + w.update(wq=rng.standard_normal((Dm, inner)) * s, + wk=rng.standard_normal((Dm, inner)) * s, + wv=rng.standard_normal((Dm, inner)) * s, + wbeta=rng.standard_normal((Dm, cfg.num_heads)) * s, + wdecay=rng.standard_normal((Dm, cfg.num_heads)) * s) + else: # SSM + w.update(w_x=rng.standard_normal((Dm, inner)) * s, + w_b=rng.standard_normal((Dm, N)) * s, + w_c=rng.standard_normal((Dm, N)) * s, + w_dt=rng.standard_normal((Dm, inner)) * s, + dt_bias=rng.standard_normal(inner) * 0.1 - 1.0, + a_log=rng.standard_normal(inner) * 0.1) # A = -exp(a_log) < 0 + layers.append(w) + return layers + + +def _rmsnorm(x, g, eps=1e-5): + return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + eps) * g + + +def _swiglu(x, w): + g = x @ w["wg"] + u = x @ w["wu"] + return (g / (1.0 + np.exp(-g)) * u) @ w["wd"] + + +def _softplus(x): + return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0.0) + + +def _heads(x, W, H, Dh): + B, S, _ = x.shape + return np.transpose((x @ W).reshape(B, S, H, Dh), (0, 2, 1, 3)) + + +def _merge(o, Wo): + B, H, S, Dh = o.shape + return np.transpose(o, (0, 2, 1, 3)).reshape(B, S, H * Dh) @ Wo + + +def _norm_last(x): + return x / (np.linalg.norm(x, axis=-1, keepdims=True) + 1e-12) + + +def _per_head_scalar(x, W): + return np.transpose(1.0 / (1.0 + np.exp(-(x @ W))), (0, 2, 1)) # [B,H,S] + + +def _causal_attention(Q, K, V): + Dh = Q.shape[-1] + scores = np.einsum("bhsd,bhtd->bhst", Q, K) / np.sqrt(Dh) + Sq, Sk = Q.shape[2], K.shape[2] + mask = np.tril(np.ones((Sq, Sk), bool), k=Sk - Sq) # aligns q-offset to cached keys + scores = np.where(mask, scores, -1e30) + p = np.exp(scores - scores.max(-1, keepdims=True)) + p /= p.sum(-1, keepdims=True) + return np.einsum("bhst,bhtd->bhsd", p, V) + + +def _ssm_scan(x, A, B, C, dt, h0=None): + """Mamba-2 SSD recurrence, returns (y, h_final). Exactly reproduces + ``tessera.ops.selective_ssm`` (scalar-state A), plus the carried state. + + x, dt : [B, S, D] A : [D] B, C : [B, S, N] h : [B, D, N] + """ + Bsz, Sl, D = x.shape + N = B.shape[-1] + A2d = A[None, :, None] # [1, D, 1] (scalar-state) + h = np.zeros((Bsz, D, N)) if h0 is None else np.array(h0, copy=True) + y = np.zeros((Bsz, Sl, D)) + for t in range(Sl): + A_bar = np.exp(dt[:, t, :, None] * A2d) # [B, D, 1] → bcast N + B_bar = dt[:, t, :, None] * B[:, t, None, :] # [B, D, N] + h = A_bar * h + B_bar * x[:, t, :, None] + y[:, t, :] = np.einsum("bdn,bn->bd", h, C[:, t, :]) + return y, h + + +# ── per-mixer span functions (handle a token-span of any length + cache) ───── +def _delta_span(x, w, cfg, cache): + H, Dh = cfg.num_heads, cfg.head_dim + h = _rmsnorm(x, w["n1"]) + Q = _heads(h, w["wq"], H, Dh) + K = _norm_last(_heads(h, w["wk"], H, Dh)) # L2-normalized keys (L1.1 finding) + V = _heads(h, w["wv"], H, Dh) + beta = _per_head_scalar(h, w["wbeta"]) + decay = _per_head_scalar(h, w["wdecay"]) + o, st = _dr.gated_delta_rule_recurrent( + Q, K, V, beta=beta, decay=decay, state=cache.get("S"), + return_state=True, state_dtype="fp64") + cache["S"] = st + x = x + _merge(np.asarray(o), w["wo"]) + return x + _swiglu(_rmsnorm(x, w["n2"]), w) + + +def _ssm_span(x, w, cfg, cache): + H, Dh = cfg.num_heads, cfg.head_dim + h = _rmsnorm(x, w["n1"]) + x_ssm = h @ w["w_x"] + Bp = h @ w["w_b"] + Cp = h @ w["w_c"] + dt = _softplus(h @ w["w_dt"] + w["dt_bias"]) + A = -np.exp(w["a_log"]) + y, hnew = _ssm_scan(x_ssm, A, Bp, Cp, dt, h0=cache.get("H")) + cache["H"] = hnew + x = x + y @ w["wo"] # y [B,S,inner] -> [B,S,Dm] + return x + _swiglu(_rmsnorm(x, w["n2"]), w) + + +def _full_span(x, w, cfg, cache): + H, Dh = cfg.num_heads, cfg.head_dim + h = _rmsnorm(x, w["n1"]) + q = _heads(h, w["wq"], H, Dh) + k = _heads(h, w["wk"], H, Dh) + v = _heads(h, w["wv"], H, Dh) + cache["K"] = k if cache.get("K") is None else np.concatenate([cache["K"], k], axis=2) + cache["V"] = v if cache.get("V") is None else np.concatenate([cache["V"], v], axis=2) + o = _causal_attention(q, cache["K"], cache["V"]) + x = x + _merge(o, w["wo"]) + return x + _swiglu(_rmsnorm(x, w["n2"]), w) + + +_SPAN = {DELTA: _delta_span, SSM: _ssm_span, FULL: _full_span} + + +def _run_layers(x_span, weights, cfg, caches): + for i, w in enumerate(weights): + x_span = _SPAN[cfg.mixer_for(i)](x_span, w, cfg, caches[i]) + return x_span + + +def hybrid_forward(x, weights, cfg: HybridConfig): + """Full (parallel) forward over the whole sequence; returns hidden [B,S,Dm].""" + return _run_layers(x, weights, cfg, [{} for _ in weights]) + + +def hybrid_decode(x, weights, cfg: HybridConfig, prefill: int = 1): + """Stream the sequence token-by-token (after a `prefill` chunk) carrying the + dual cache — recurrent Ŝ (delta) / h (SSM) for linear layers, KV for full + layers. Equals `hybrid_forward` for the same input (the oracle).""" + B, S, Dm = x.shape + caches = [{} for _ in weights] + out = np.zeros_like(x) + out[:, :prefill] = _run_layers(x[:, :prefill], weights, cfg, caches) + for t in range(prefill, S): + out[:, t:t + 1] = _run_layers(x[:, t:t + 1], weights, cfg, caches) + return out diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm index b939afd16..fcc355307 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime.mm @@ -8397,6 +8397,391 @@ inline void reference_linear_attn_f32(const float* Q, const float* K, reference_linear_attn_f32(Q, K, V, O, B, H, S, D_qk, D_v, feature_map, causal); } +//===---------------------------------------------------------------------===// +// Track L (L1.1) — the *genuine* gated delta rule on Metal. +// +// Unlike linear attention (and the composed (QKᵀ⊙mask)@V path), the DeltaNet +// recurrence carries an erase term and cannot be expressed as a masked QKᵀ: +// +// v̂_t = Ŝ_{t-1}ᵀ k_t (value bound to key k_t) +// Ŝ_t = α_t·Ŝ_{t-1} + β_t·k_t·(v_t − α_t·v̂_t)ᵀ (erase, then write) +// O_t = q_tᵀ Ŝ_t (read after the t-th update) +// +// One thread per (b, h); per-thread state Ŝ[d_k, d_v]; sequential scan (decode +// form). β / decay are per-token scalars supplied as (B,H,S) arrays (caller +// fills 1.0 when absent); ``erase=0`` degenerates to gated linear attention. +// The output gate is applied host-side. Matches +// ``tessera.stdlib.delta_rule.gated_delta_rule_recurrent`` exactly. +//===---------------------------------------------------------------------===// +static bool dispatch_gated_delta_rule_msl( + MetalDeviceContext &ctx, const float* Q, const float* K, const float* V, + const float* beta, const float* decay, float* O, int32_t B, int32_t H, + int32_t S, int32_t D_qk, int32_t D_v, int32_t erase) { + static NSString *const kKernelSource = @R"MSL( +#include +using namespace metal; + +kernel void gated_delta_rule_f32( + device const float* Q [[buffer(0)]], + device const float* K [[buffer(1)]], + device const float* V [[buffer(2)]], + device const float* beta [[buffer(3)]], + device const float* decay [[buffer(4)]], + device float* O [[buffer(5)]], + constant int& B [[buffer(6)]], + constant int& H [[buffer(7)]], + constant int& S [[buffer(8)]], + constant int& D_qk [[buffer(9)]], + constant int& D_v [[buffer(10)]], + constant int& erase [[buffer(11)]], + uint gid [[thread_position_in_grid]]) +{ + if (gid >= (uint)(B * H)) return; + int bh = (int)gid; + if (D_qk > 16 || D_v > 64 || D_qk * D_v > 256) return; + + float state[256]; // Ŝ[d_qk, d_v], row-major + for (int i = 0; i < D_qk * D_v; ++i) state[i] = 0.0f; + float vhat[64]; + + int rqk = S * D_qk, rv = S * D_v; + int qb = bh * rqk, kb = qb, vb = bh * rv, ob = vb, sb = bh * S; + + for (int t = 0; t < S; ++t) { + float a = decay[sb + t]; + float bt = beta[sb + t]; + // v̂_t = Ŝ_{t-1}ᵀ k_t (read from the OLD state, all d_v). + for (int e = 0; e < D_v; ++e) { + float acc = 0.0f; + for (int d = 0; d < D_qk; ++d) acc += K[kb + t * D_qk + d] * state[d * D_v + e]; + vhat[e] = acc; + } + // Ŝ_t = α·Ŝ_{t-1} + β·k_t·(v_t − α·v̂_t)ᵀ. + for (int d = 0; d < D_qk; ++d) { + float k_d = K[kb + t * D_qk + d]; + int ro = d * D_v; + for (int e = 0; e < D_v; ++e) { + float v_e = V[vb + t * D_v + e]; + float tgt = (erase != 0) ? (v_e - a * vhat[e]) : v_e; + state[ro + e] = a * state[ro + e] + bt * k_d * tgt; + } + } + // O_t = q_tᵀ Ŝ_t (after the update). + for (int e = 0; e < D_v; ++e) { + float acc = 0.0f; + for (int d = 0; d < D_qk; ++d) acc += Q[qb + t * D_qk + d] * state[d * D_v + e]; + O[ob + t * D_v + e] = acc; + } + } +} +)MSL"; + + @autoreleasepool { + id pso = + compile_msl_kernel(ctx, kKernelSource, @"gated_delta_rule_f32"); + if (!pso) return false; + + NSUInteger qBytes = sizeof(float) * (NSUInteger)B * H * S * D_qk; + NSUInteger vBytes = sizeof(float) * (NSUInteger)B * H * S * D_v; + NSUInteger sBytes = sizeof(float) * (NSUInteger)B * H * S; + + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufQ, ctx, Q, qBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufK, ctx, K, qBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufV, ctx, V, vBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufBeta, ctx, beta, sBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufDecay, ctx, decay, sBytes); + TS_METAL_BUF_ACQUIRE(bufO, ctx, vBytes); + if (!bufQ || !bufK || !bufV || !bufBeta || !bufDecay || !bufO) return false; + + id cb = [ctx.queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pso]; + [enc setBuffer:bufQ offset:0 atIndex:0]; + [enc setBuffer:bufK offset:0 atIndex:1]; + [enc setBuffer:bufV offset:0 atIndex:2]; + [enc setBuffer:bufBeta offset:0 atIndex:3]; + [enc setBuffer:bufDecay offset:0 atIndex:4]; + [enc setBuffer:bufO offset:0 atIndex:5]; + [enc setBytes:&B length:sizeof(int32_t) atIndex:6]; + [enc setBytes:&H length:sizeof(int32_t) atIndex:7]; + [enc setBytes:&S length:sizeof(int32_t) atIndex:8]; + [enc setBytes:&D_qk length:sizeof(int32_t) atIndex:9]; + [enc setBytes:&D_v length:sizeof(int32_t) atIndex:10]; + [enc setBytes:&erase length:sizeof(int32_t) atIndex:11]; + + NSUInteger total = (NSUInteger)B * H; + MTLSize grid = MTLSizeMake(total, 1, 1); + NSUInteger tg_x = std::min(total, pso.maxTotalThreadsPerThreadgroup); + if (tg_x == 0) tg_x = 1; + MTLSize tg = MTLSizeMake(tg_x, 1, 1); + [enc dispatchThreads:grid threadsPerThreadgroup:tg]; + [enc endEncoding]; + if (!commit_and_wait_with_timeout(ctx, cb, 60000, "gated_delta_rule_msl")) + return false; + std::memcpy(O, [bufO contents], vBytes); + return true; + } +} + +static void reference_gated_delta_rule_f32( + const float* Q, const float* K, const float* V, const float* beta, + const float* decay, float* O, int32_t B, int32_t H, int32_t S, + int32_t D_qk, int32_t D_v, int32_t erase) { + std::vector state((std::size_t)D_qk * D_v, 0.0f); + std::vector vhat((std::size_t)D_v, 0.0f); + for (int32_t b = 0; b < B; ++b) { + for (int32_t h = 0; h < H; ++h) { + std::fill(state.begin(), state.end(), 0.0f); + int qb = (b * H + h) * S * D_qk, kb = qb; + int vb = (b * H + h) * S * D_v, ob = vb, sb = (b * H + h) * S; + for (int32_t t = 0; t < S; ++t) { + float a = decay[sb + t], bt = beta[sb + t]; + for (int32_t e = 0; e < D_v; ++e) { + float acc = 0.0f; + for (int32_t d = 0; d < D_qk; ++d) + acc += K[kb + t * D_qk + d] * state[(std::size_t)d * D_v + e]; + vhat[e] = acc; + } + for (int32_t d = 0; d < D_qk; ++d) { + float k_d = K[kb + t * D_qk + d]; + for (int32_t e = 0; e < D_v; ++e) { + float v_e = V[vb + t * D_v + e]; + float tgt = erase ? (v_e - a * vhat[e]) : v_e; + state[(std::size_t)d * D_v + e] = + a * state[(std::size_t)d * D_v + e] + bt * k_d * tgt; + } + } + for (int32_t e = 0; e < D_v; ++e) { + float acc = 0.0f; + for (int32_t d = 0; d < D_qk; ++d) + acc += Q[qb + t * D_qk + d] * state[(std::size_t)d * D_v + e]; + O[ob + t * D_v + e] = acc; + } + } + } + } +} + +extern "C" void tessera_apple_gpu_gated_delta_rule_f32( + const float* Q, const float* K, const float* V, const float* beta, + const float* decay, float* O, int32_t B, int32_t H, int32_t S, + int32_t D_qk, int32_t D_v, int32_t erase) { + if (D_qk > 16 || D_v > 64 || D_qk * D_v > 256) { + reference_gated_delta_rule_f32(Q, K, V, beta, decay, O, B, H, S, D_qk, D_v, erase); + return; + } + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_gated_delta_rule_msl(ctx, Q, K, V, beta, decay, O, B, H, + S, D_qk, D_v, erase)) + return; + reference_gated_delta_rule_f32(Q, K, V, beta, decay, O, B, H, S, D_qk, D_v, erase); +} + +//===---------------------------------------------------------------------===// +// Track L (L2.1) — the chunk-parallel UT-transform on Metal (prefill form). +// +// One threadgroup per (b, h), C threads (one per token-in-chunk), state Ŝ in +// threadgroup memory across the chunk loop. Per chunk: +// A[t,j] = β_t(γ_t/γ_j)(k_tᵀk_j) (j +using namespace metal; + +#define TS_DR_MAXC 32 +#define TS_DR_MAXD 16 + +kernel void gated_delta_rule_chunked_f32( + device const float* Q [[buffer(0)]], + device const float* K [[buffer(1)]], + device const float* V [[buffer(2)]], + device const float* beta [[buffer(3)]], + device const float* decay [[buffer(4)]], + device float* O [[buffer(5)]], + constant int& B [[buffer(6)]], + constant int& H [[buffer(7)]], + constant int& S [[buffer(8)]], + constant int& D_qk [[buffer(9)]], + constant int& D_v [[buffer(10)]], + constant int& Cc [[buffer(11)]], + constant int& erase [[buffer(12)]], + uint tgid [[threadgroup_position_in_grid]], + uint tid [[thread_position_in_threadgroup]]) +{ + int bh = (int)tgid; + if (bh >= B * H) return; + if (D_qk > TS_DR_MAXD || D_v > TS_DR_MAXD || Cc > TS_DR_MAXC || D_qk * D_v > 256) + return; + + threadgroup float state[256]; + threadgroup float Ks[TS_DR_MAXC * TS_DR_MAXD]; + threadgroup float Qs[TS_DR_MAXC * TS_DR_MAXD]; + threadgroup float Vs[TS_DR_MAXC * TS_DR_MAXD]; + threadgroup float Us[TS_DR_MAXC * TS_DR_MAXD]; + threadgroup float As[TS_DR_MAXC * TS_DR_MAXC]; + threadgroup float gam[TS_DR_MAXC]; + threadgroup float bet[TS_DR_MAXC]; + threadgroup float dec[TS_DR_MAXC]; + + int t = (int)tid; + for (int i = (int)tid; i < D_qk * D_v; i += Cc) state[i] = 0.0f; + threadgroup_barrier(mem_flags::mem_threadgroup); + + int rqk = S * D_qk, rv = S * D_v; + int qb = bh * rqk, kb = qb, vb = bh * rv, ob = vb, sb = bh * S; + + for (int c0 = 0; c0 < S; c0 += Cc) { + int Cn = min(Cc, S - c0); + if (t < Cn) { + for (int d = 0; d < D_qk; ++d) { + Ks[t * D_qk + d] = K[kb + (c0 + t) * D_qk + d]; + Qs[t * D_qk + d] = Q[qb + (c0 + t) * D_qk + d]; + } + for (int e = 0; e < D_v; ++e) Vs[t * D_v + e] = V[vb + (c0 + t) * D_v + e]; + bet[t] = beta[sb + c0 + t]; + dec[t] = decay[sb + c0 + t]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (t == 0) { float g = 1.0f; for (int i = 0; i < Cn; ++i) { g *= dec[i]; gam[i] = g; } } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if (t < Cn) { + // A row t (strictly lower, with γ-ratio decay folding). + for (int j = 0; j < Cn; ++j) { + float a = 0.0f; + if (erase != 0 && j < t) { + float kk = 0.0f; + for (int d = 0; d < D_qk; ++d) kk += Ks[t * D_qk + d] * Ks[j * D_qk + d]; + a = bet[t] * (gam[t] / gam[j]) * kk; + } + As[t * Cc + j] = a; + } + // W̃[t] = β_t(v_t − γ_t Ŝ₀ᵀk_t) (stored into Us, solved in place). + for (int e = 0; e < D_v; ++e) { + float kss = 0.0f; + if (erase != 0) for (int d = 0; d < D_qk; ++d) kss += Ks[t * D_qk + d] * state[d * D_v + e]; + Us[t * D_v + e] = bet[t] * (Vs[t * D_v + e] - (erase != 0 ? gam[t] * kss : 0.0f)); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Forward substitution: U[t] = W̃[t] − Σ_{j pso = + compile_msl_kernel(ctx, kKernelSource, @"gated_delta_rule_chunked_f32"); + if (!pso) return false; + + NSUInteger qBytes = sizeof(float) * (NSUInteger)B * H * S * D_qk; + NSUInteger vBytes = sizeof(float) * (NSUInteger)B * H * S * D_v; + NSUInteger sBytes = sizeof(float) * (NSUInteger)B * H * S; + + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufQ, ctx, Q, qBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufK, ctx, K, qBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufV, ctx, V, vBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufBeta, ctx, beta, sBytes); + TS_METAL_BUF_ACQUIRE_WITH_BYTES(bufDecay, ctx, decay, sBytes); + TS_METAL_BUF_ACQUIRE(bufO, ctx, vBytes); + if (!bufQ || !bufK || !bufV || !bufBeta || !bufDecay || !bufO) return false; + + id cb = [ctx.queue commandBuffer]; + id enc = [cb computeCommandEncoder]; + [enc setComputePipelineState:pso]; + [enc setBuffer:bufQ offset:0 atIndex:0]; + [enc setBuffer:bufK offset:0 atIndex:1]; + [enc setBuffer:bufV offset:0 atIndex:2]; + [enc setBuffer:bufBeta offset:0 atIndex:3]; + [enc setBuffer:bufDecay offset:0 atIndex:4]; + [enc setBuffer:bufO offset:0 atIndex:5]; + [enc setBytes:&B length:sizeof(int32_t) atIndex:6]; + [enc setBytes:&H length:sizeof(int32_t) atIndex:7]; + [enc setBytes:&S length:sizeof(int32_t) atIndex:8]; + [enc setBytes:&D_qk length:sizeof(int32_t) atIndex:9]; + [enc setBytes:&D_v length:sizeof(int32_t) atIndex:10]; + [enc setBytes:&Cc length:sizeof(int32_t) atIndex:11]; + [enc setBytes:&erase length:sizeof(int32_t) atIndex:12]; + + MTLSize grid = MTLSizeMake((NSUInteger)B * H, 1, 1); + MTLSize tg = MTLSizeMake((NSUInteger)Cc, 1, 1); + [enc dispatchThreadgroups:grid threadsPerThreadgroup:tg]; + [enc endEncoding]; + if (!commit_and_wait_with_timeout(ctx, cb, 60000, "gated_delta_rule_chunked_msl")) + return false; + std::memcpy(O, [bufO contents], vBytes); + return true; + } +} + +extern "C" void tessera_apple_gpu_gated_delta_rule_chunked_f32( + const float* Q, const float* K, const float* V, const float* beta, + const float* decay, float* O, int32_t B, int32_t H, int32_t S, + int32_t D_qk, int32_t D_v, int32_t chunk, int32_t erase) { + // chunk ≡ recurrent, so the host fallback is the proven recurrent reference. + if (D_qk > 16 || D_v > 16 || chunk > 32 || chunk < 1 || D_qk * D_v > 256) { + reference_gated_delta_rule_f32(Q, K, V, beta, decay, O, B, H, S, D_qk, D_v, erase); + return; + } + MetalDeviceContext &ctx = deviceContext(); + if (ctx.ok && dispatch_gated_delta_rule_chunked_msl(ctx, Q, K, V, beta, decay, O, + B, H, S, D_qk, D_v, chunk, erase)) + return; + reference_gated_delta_rule_f32(Q, K, V, beta, decay, O, B, H, S, D_qk, D_v, erase); +} + //===---------------------------------------------------------------------===// // attention_variants_plan, MLA-2 — DeepSeek MLA decode runtime entry. // diff --git a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp index d1b9c494c..03558c672 100644 --- a/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp +++ b/src/compiler/codegen/Tessera_Apple_Backend/runtime/apple_gpu_runtime_stub.cpp @@ -1198,6 +1198,58 @@ extern "C" void tessera_apple_gpu_linear_attn_f32(const float* Q, const float* K causal); } +// Track L (L1.1) — genuine gated delta rule, non-Apple twin of the .mm +// host reference. Same math as the MSL kernel + .mm fallback (with the +// (v_t − α·v̂_t) erase term). See apple_gpu_runtime.mm for the contract. +extern "C" void tessera_apple_gpu_gated_delta_rule_f32( + const float* Q, const float* K, const float* V, const float* beta, + const float* decay, float* O, int32_t B, int32_t H, int32_t S, + int32_t D_qk, int32_t D_v, int32_t erase) { + std::vector state(static_cast(D_qk) * D_v, 0.0f); + std::vector vhat(static_cast(D_v), 0.0f); + for (int32_t b = 0; b < B; ++b) { + for (int32_t h = 0; h < H; ++h) { + std::fill(state.begin(), state.end(), 0.0f); + int qb = (b * H + h) * S * D_qk, kb = qb; + int vb = (b * H + h) * S * D_v, ob = vb, sb = (b * H + h) * S; + for (int32_t t = 0; t < S; ++t) { + float a = decay[sb + t], bt = beta[sb + t]; + for (int32_t e = 0; e < D_v; ++e) { + float acc = 0.0f; + for (int32_t d = 0; d < D_qk; ++d) + acc += K[kb + t * D_qk + d] * state[static_cast(d) * D_v + e]; + vhat[e] = acc; + } + for (int32_t d = 0; d < D_qk; ++d) { + float k_d = K[kb + t * D_qk + d]; + for (int32_t e = 0; e < D_v; ++e) { + float v_e = V[vb + t * D_v + e]; + float tgt = erase ? (v_e - a * vhat[e]) : v_e; + state[static_cast(d) * D_v + e] = + a * state[static_cast(d) * D_v + e] + bt * k_d * tgt; + } + } + for (int32_t e = 0; e < D_v; ++e) { + float acc = 0.0f; + for (int32_t d = 0; d < D_qk; ++d) + acc += Q[qb + t * D_qk + d] * state[static_cast(d) * D_v + e]; + O[ob + t * D_v + e] = acc; + } + } + } + } +} + +// Track L (L2.1) — chunked UT-transform, non-Apple twin. chunk ≡ recurrent, +// so the stub forwards to the (proven) recurrent reference above. +extern "C" void tessera_apple_gpu_gated_delta_rule_chunked_f32( + const float* Q, const float* K, const float* V, const float* beta, + const float* decay, float* O, int32_t B, int32_t H, int32_t S, + int32_t D_qk, int32_t D_v, int32_t /*chunk*/, int32_t erase) { + tessera_apple_gpu_gated_delta_rule_f32(Q, K, V, beta, decay, O, B, H, S, + D_qk, D_v, erase); +} + // attention_variants_plan, MLA-2 — non-Apple stub. Numpy-reference path // for the host-only fallback. The Apple .mm file's `reference_mla_decode_f32` // already implements the same math; this is its non-Apple twin. diff --git a/src/compiler/ir/TesseraOps.cpp b/src/compiler/ir/TesseraOps.cpp index 5dd04cf77..ea268a7af 100644 --- a/src/compiler/ir/TesseraOps.cpp +++ b/src/compiler/ir/TesseraOps.cpp @@ -1832,6 +1832,44 @@ LogicalResult ModifiedDeltaAttentionOp::verify() { "modified_delta_attention"); } +LogicalResult SelectiveSsmOp::verify() { + auto xTy = dyn_cast(getX().getType()); + auto bTy = dyn_cast(getB().getType()); + auto cTy = dyn_cast(getC().getType()); + auto dTy = dyn_cast(getDelta().getType()); + auto aTy = dyn_cast(getA().getType()); + // Unranked / dynamic operands — defer (consistent with the attention ops). + if (!xTy || !bTy || !cTy || !dTy || !aTy) + return success(); + if (xTy.getRank() != 3) + return emitOpError("x must be rank-3 (B, S, D)"); + if (dTy.getRank() != 3 || dTy.getShape() != xTy.getShape()) + return emitOpError("delta must be rank-3 and shape-equal to x (B, S, D)"); + if (bTy.getRank() != 3 || cTy.getRank() != 3) + return emitOpError("b and c must be rank-3 (B, S, N)"); + if (bTy.getShape() != cTy.getShape()) + return emitOpError("b and c must have matching shape (B, S, N)"); + if (xTy.getShape()[0] != bTy.getShape()[0] || + xTy.getShape()[1] != bTy.getShape()[1]) + return emitOpError("x and b must share batch and sequence dims"); + if (aTy.getRank() != 1 && aTy.getRank() != 2) + return emitOpError("a must be rank-1 (D) or rank-2 (D, N)"); + int64_t D = xTy.getShape()[2]; + if (!aTy.isDynamicDim(0) && D != ShapedType::kDynamic && aTy.getShape()[0] != D) + return emitOpError("a leading dim must equal x channel dim D"); + if (auto g = getGate()) { + auto gTy = dyn_cast(g.getType()); + if (gTy && gTy.getShape() != xTy.getShape()) + return emitOpError("gate must be shape-equal to x (B, S, D)"); + } + if (auto s = getState()) { + auto sTy = dyn_cast(s.getType()); + if (sTy && sTy.getRank() != 3) + return emitOpError("state must be rank-3 (B, D, N)"); + } + return success(); +} + LogicalResult HybridAttentionOp::verify() { // ``pattern`` on hybrid_attention is a *free-form* model-specific hybrid // variant (e.g. "kimi_kda_mla", "ling_1_7_mla_lightning"), not a closed diff --git a/src/compiler/ir/TesseraOps.td b/src/compiler/ir/TesseraOps.td index f87b34005..b750ce670 100644 --- a/src/compiler/ir/TesseraOps.td +++ b/src/compiler/ir/TesseraOps.td @@ -1119,6 +1119,44 @@ def Tessera_ModifiedDeltaAttentionOp let hasVerifier = 1; } +// Mamba-2 selective state-space model (SSD). Track L L4 — this is the op the +// coverage registry previously *claimed* had landed; it now genuinely exists. +def Tessera_SelectiveSsmOp : Op< + Tessera_Dialect, "selective_ssm", + [Pure, AttrSizedOperandSegments] +> { + let summary = "Mamba-2 selective state-space model (SSD) recurrence"; + let description = [{ + Selective state-space recurrence (Mamba-2 / SSD). Per token t, channel d, + state dim n: + A_bar = exp(delta[b,t,d] * a[d,n]); B_bar = delta[b,t,d] * b[b,t,n] + h[t,d,n] = A_bar * h[t-1,d,n] + B_bar * x[b,t,d] + y[t,d] = sum_n c[b,t,n] * h[t,d,n] + `a` is the (negative) state-matrix diagonal — rank-1 `(D)` (scalar-state, the + common per-head config that lowers to the chunked-parallel SSD matmul form) + or rank-2 `(D,N)`. `b`/`c` are the input/output projections `(B,S,N)`, + `delta` the selective time-step `(B,S,D)`. Optional `gate` (output gate, + shape of `x`) and `state` (initial carry `(B,D,N)`). Reference: + `tessera.ops.selective_ssm`; chunked-parallel lowering: + `python/tessera/_mamba_ssd.py`. + }]; + let arguments = (ins TensorType:$x, + TensorType:$a, + TensorType:$b, + TensorType:$c, + TensorType:$delta, + Optional:$gate, + Optional:$state, + DefaultValuedAttr:$chunk_size); + let results = (outs TensorType:$y); + let assemblyFormat = [{ + $x `,` $a `,` $b `,` $c `,` $delta + (`gate` `(` $gate^ `)`)? (`init` `(` $state^ `)`)? + attr-dict `:` functional-type(operands, results) + }]; + let hasVerifier = 1; +} + def Tessera_HybridAttentionOp : Op< Tessera_Dialect, "hybrid_attention", [Pure] diff --git a/tests/tessera-ir/model_class/selective_ssm.mlir b/tests/tessera-ir/model_class/selective_ssm.mlir new file mode 100644 index 000000000..e64a84e6d --- /dev/null +++ b/tests/tessera-ir/model_class/selective_ssm.mlir @@ -0,0 +1,33 @@ +// RUN: tessera-opt %s | FileCheck %s +// +// Track L (L4) — `tessera.selective_ssm` is now a genuine Graph IR op (Mamba-2 +// SSD recurrence). The coverage registry previously *claimed* this op had +// landed (graph_ir_lowering = "registered") while no ODS op existed; this +// fixture is the proof that the claim is now true. Verifier: +// `SelectiveSsmOp::verify` (rank-3 x / shape-equal delta / matching b,c / A +// rank-1|2 / optional gate shape-equal x / state rank-3). + +// CHECK-LABEL: func.func @ssm_scalar_state +func.func @ssm_scalar_state(%x: tensor<2x16x8xf32>, %a: tensor<8xf32>, + %b: tensor<2x16x4xf32>, %c: tensor<2x16x4xf32>, + %delta: tensor<2x16x8xf32>) -> tensor<2x16x8xf32> { + // CHECK: tessera.selective_ssm %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} {chunk_size = 64 : i64} + %y = tessera.selective_ssm %x, %a, %b, %c, %delta {chunk_size = 64 : i64} + : (tensor<2x16x8xf32>, tensor<8xf32>, tensor<2x16x4xf32>, + tensor<2x16x4xf32>, tensor<2x16x8xf32>) -> tensor<2x16x8xf32> + return %y : tensor<2x16x8xf32> +} + +// Per-state-dim A (D, N) + optional output gate and initial-state carry. +// CHECK-LABEL: func.func @ssm_full_state_gated +func.func @ssm_full_state_gated(%x: tensor<1x32x16xf32>, %a: tensor<16x8xf32>, + %b: tensor<1x32x8xf32>, %c: tensor<1x32x8xf32>, + %delta: tensor<1x32x16xf32>, %g: tensor<1x32x16xf32>, + %s: tensor<1x16x8xf32>) -> tensor<1x32x16xf32> { + // CHECK: tessera.selective_ssm %{{.*}} gate(%{{.*}}) init(%{{.*}}) + %y = tessera.selective_ssm %x, %a, %b, %c, %delta gate(%g) init(%s) + : (tensor<1x32x16xf32>, tensor<16x8xf32>, tensor<1x32x8xf32>, + tensor<1x32x8xf32>, tensor<1x32x16xf32>, tensor<1x32x16xf32>, + tensor<1x16x8xf32>) -> tensor<1x32x16xf32> + return %y : tensor<1x32x16xf32> +} diff --git a/tests/unit/test_apple_gpu_gated_delta_rule.py b/tests/unit/test_apple_gpu_gated_delta_rule.py new file mode 100644 index 000000000..0370d49df --- /dev/null +++ b/tests/unit/test_apple_gpu_gated_delta_rule.py @@ -0,0 +1,164 @@ +"""Track L (L1.1) — the genuine gated delta rule on Metal, DESIL-gated. + +`tessera_apple_gpu_gated_delta_rule_f32` (recurrent, with the `(v_t − α·v̂_t)` +erase) must equal the numpy reference `gated_delta_rule_recurrent`. This is the +*true* DeltaNet recurrence — unlike linear attention it carries an erase term and +cannot be written as a masked `(QKᵀ⊙mask)@V`, so it runs as a per-(b,h) +sequential MSL scan, not a composed bmm. + +Conditioning note: the delta rule is only well-behaved when keys are +L2-normalized, so `β·k·kᵀ` has eigenvalue `β < 1` and `(I − β k kᵀ)` is a +contraction. With unnormalized keys (`‖k‖²≫1`) the recurrence expands and f32 +diverges from f64 — that is genuine ill-conditioning, not a kernel defect, so +these oracles use normalized keys (the regime real models use). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tessera import _apple_gpu_backend as agb +from tessera.stdlib import delta_rule as dr + +_GPU = agb.is_available() +gpu = pytest.mark.skipif(not _GPU, reason="apple_gpu runtime unavailable") + +_B, _H, _S, _D = 2, 3, 16, 16 + + +def _normalize(x): + return x / np.linalg.norm(x, axis=-1, keepdims=True) + + +def _qkv(seed=0, dv=_D): + rng = np.random.default_rng(seed) + # L2-normalized Q/K (the contraction regime the rule is used in). + Q = _normalize(rng.standard_normal((_B, _H, _S, _D))).astype(np.float32) + K = _normalize(rng.standard_normal((_B, _H, _S, _D))).astype(np.float32) + V = rng.standard_normal((_B, _H, _S, dv)).astype(np.float32) + return Q, K, V + + +def _sig(x): + return (1.0 / (1.0 + np.exp(-x))).astype(np.float32) + + +def test_symbol_is_registered(): + """L1.1 lock: the genuine-delta kernel is in the runtime ABI.""" + assert hasattr(agb, "gpu_gated_delta_rule") + + +@gpu +def test_metal_equals_numpy_true_delta(): + """Headline DESIL oracle: Metal genuine delta rule ≡ numpy reference.""" + Q, K, V = _qkv(1) + o_gpu = dr.gated_delta_rule_recurrent(Q, K, V, backend="apple_gpu") + o_ref = dr.gated_delta_rule_recurrent(Q, K, V) + np.testing.assert_allclose(np.asarray(o_gpu), o_ref, rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_equals_numpy_with_beta_and_decay(): + Q, K, V = _qkv(2) + beta = _sig(np.random.default_rng(3).standard_normal((_B, _H, _S))) + decay = _sig(np.random.default_rng(4).standard_normal((_B, _H, _S)) + 2.0) + o_gpu = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay, + backend="apple_gpu") + o_ref = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay) + np.testing.assert_allclose(np.asarray(o_gpu), o_ref, rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_equals_numpy_with_output_gate(): + Q, K, V = _qkv(5) + gate = np.random.default_rng(6).standard_normal((_B, _H, _S, _D)).astype(np.float32) + o_gpu = dr.gated_delta_rule_recurrent(Q, K, V, gate=gate, backend="apple_gpu") + o_ref = dr.gated_delta_rule_recurrent(Q, K, V, gate=gate) + np.testing.assert_allclose(np.asarray(o_gpu), o_ref, rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_erase_off_equals_existing_linear_reference(): + """erase=False on Metal ≡ the shipped (linear-attention) gated_deltanet.""" + from tessera import ops + Q, K, V = _qkv(7) + beta = _sig(np.random.default_rng(8).standard_normal((_B, _H, _S))) + o_gpu = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, erase=False, + backend="apple_gpu") + o_existing = np.asarray(ops.gated_deltanet(Q, K, V, beta=beta), np.float64) + np.testing.assert_allclose(np.asarray(o_gpu), o_existing, rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_equals_chunked_prefill(): + """End-to-end Track L closure: Metal recurrent ≡ the L2 chunked UT-transform + (both are the genuine rule, reached by independent routes).""" + Q, K, V = _qkv(9) + beta = _sig(np.random.default_rng(10).standard_normal((_B, _H, _S))) + decay = _sig(np.random.default_rng(11).standard_normal((_B, _H, _S)) + 2.0) + o_gpu = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay, + backend="apple_gpu") + o_chunk = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, decay=decay, chunk_size=8) + np.testing.assert_allclose(np.asarray(o_gpu), o_chunk, rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_non_square_head_dims(): + """D_qk ≠ D_v (state is rectangular).""" + Q, K, V = _qkv(12, dv=8) + o_gpu = dr.gated_delta_rule_recurrent(Q, K, V, backend="apple_gpu") + o_ref = dr.gated_delta_rule_recurrent(Q, K, V) + np.testing.assert_allclose(np.asarray(o_gpu), o_ref, rtol=1e-4, atol=1e-4) + + +# ── L2.1 — chunk-parallel UT-transform on Metal ────────────────────────────── +def test_chunked_symbol_is_registered(): + assert hasattr(agb, "gpu_gated_delta_rule_chunked") + + +@gpu +@pytest.mark.parametrize("chunk", [1, 4, 8, 16, 32]) +def test_metal_chunked_equals_numpy(chunk): + """L2.1 headline: the on-device chunk UT-transform ≡ numpy recurrent, across + chunk sizes (S=20 exercises a partial last chunk).""" + rng = np.random.default_rng(20) + B, H, S, D = 2, 3, 20, 16 + Q = _normalize(rng.standard_normal((B, H, S, D))).astype(np.float32) + K = _normalize(rng.standard_normal((B, H, S, D))).astype(np.float32) + V = rng.standard_normal((B, H, S, D)).astype(np.float32) + beta = _sig(rng.standard_normal((B, H, S))) + decay = _sig(rng.standard_normal((B, H, S)) + 2.0) + o_gpu = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, decay=decay, + chunk_size=chunk, backend="apple_gpu") + o_ref = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay) + np.testing.assert_allclose(np.asarray(o_gpu), o_ref, rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_chunked_equals_metal_recurrent(): + """The two Metal kernels (sequential L1.1 vs chunked L2.1) agree — same rule, + fully independent on-device routes.""" + Q, K, V = _qkv(21) + beta = _sig(np.random.default_rng(22).standard_normal((_B, _H, _S))) + decay = _sig(np.random.default_rng(23).standard_normal((_B, _H, _S)) + 2.0) + o_rec = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay, backend="apple_gpu") + o_chk = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, decay=decay, + chunk_size=8, backend="apple_gpu") + np.testing.assert_allclose(np.asarray(o_rec), np.asarray(o_chk), rtol=1e-4, atol=1e-4) + + +@gpu +def test_metal_chunked_with_output_gate_and_erase_off(): + Q, K, V = _qkv(24) + gate = np.random.default_rng(25).standard_normal((_B, _H, _S, _D)).astype(np.float32) + o_gpu = dr.gated_delta_rule_chunked(Q, K, V, gate=gate, chunk_size=8, backend="apple_gpu") + o_ref = dr.gated_delta_rule_recurrent(Q, K, V, gate=gate) + np.testing.assert_allclose(np.asarray(o_gpu), o_ref, rtol=1e-4, atol=1e-4) + # erase=False chunked ≡ shipped linear reference. + from tessera import ops + beta = _sig(np.random.default_rng(26).standard_normal((_B, _H, _S))) + o_lin = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, erase=False, + chunk_size=8, backend="apple_gpu") + o_exist = np.asarray(ops.gated_deltanet(Q, K, V, beta=beta), np.float64) + np.testing.assert_allclose(np.asarray(o_lin), o_exist, rtol=1e-4, atol=1e-4) diff --git a/tests/unit/test_stdlib_delta_rule.py b/tests/unit/test_stdlib_delta_rule.py new file mode 100644 index 000000000..7043fcce2 --- /dev/null +++ b/tests/unit/test_stdlib_delta_rule.py @@ -0,0 +1,214 @@ +"""Track L (L0/L1/L2) — the *true* gated delta rule, oracle-gated. + +L0 (discovery lock): the shipped ``tessera.ops.gated_deltanet`` is gated *linear* +attention (no erase) — it equals our ``erase=False`` path and DIFFERS from the +genuine delta rule when keys correlate. + +L1 (recurrence): ``gated_delta_rule_recurrent`` matches an independent +brute-force delta recurrence written in the paper's ``(I − β k kᵀ)`` layout; +``erase=False`` reduces to the existing reference; state carry is consistent. + +L2 (chunk UT-transform, the keystone): ``gated_delta_rule_chunked ≡ +gated_delta_rule_recurrent`` across ungated / β / fully-gated / output-gated +cases and across chunk sizes — the chunk≡recurrent DESIL proof. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tessera import ops +from tessera.stdlib import delta_rule as dr + + +# ── helpers ────────────────────────────────────────────────────────────────── +def _rng(seed): + return np.random.default_rng(seed) + + +def _qkv(rng, B=2, H=3, S=12, d_k=5, d_v=4, scale=1.0): + Q = rng.standard_normal((B, H, S, d_k)) * scale + K = rng.standard_normal((B, H, S, d_k)) * scale + V = rng.standard_normal((B, H, S, d_v)) * scale + return Q, K, V + + +def _sig(x): + return 1.0 / (1.0 + np.exp(-x)) + + +def _bruteforce_delta(Q, K, V, *, beta=None, decay=None, gate=None): + """Independent transcription of the gated delta rule in the paper's layout: + state S ∈ [d_v, d_k], S_t = α_t S_{t-1}(I − β_t k_t k_tᵀ) + β_t v_t k_tᵀ, + O_t = S_t q_t. Deliberately a different layout + explicit Householder matrix + so it is NOT the same code path as the module under test. + """ + Q = np.asarray(Q, np.float64); K = np.asarray(K, np.float64); V = np.asarray(V, np.float64) + B, H, S, d_k = Q.shape + d_v = V.shape[-1] + O = np.zeros((B, H, S, d_v)) + I = np.eye(d_k) + for b in range(B): + for h in range(H): + St = np.zeros((d_v, d_k)) + for t in range(S): + k = K[b, h, t]; v = V[b, h, t]; q = Q[b, h, t] + a = float(decay[b, h, t]) if decay is not None else 1.0 + bt = float(beta[b, h, t]) if beta is not None else 1.0 + St = a * St @ (I - bt * np.outer(k, k)) + bt * np.outer(v, k) + O[b, h, t] = St @ q + if gate is not None: + O = O * _sig(np.asarray(gate, np.float64)) + return O + + +# ── L0 — discovery lock ────────────────────────────────────────────────────── +def test_existing_gated_deltanet_is_linear_attention_not_delta(): + """The shipped op == erase-off (linear) path, and DIFFERS from the true + delta rule once keys carry overlapping directions.""" + rng = _rng(0) + Q, K, V = _qkv(rng) + beta = _sig(rng.standard_normal(Q.shape[:3])) + existing = np.asarray(ops.gated_deltanet(Q, K, V, beta=beta), np.float64) + linear = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, erase=False) + delta = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, erase=True) + # Existing reference IS the linear (no-erase) path. + np.testing.assert_allclose(existing, linear, rtol=1e-10, atol=1e-10) + # ...and the genuine delta rule is materially different. + assert not np.allclose(delta, linear, rtol=1e-3, atol=1e-3) + + +# ── L1 — the genuine recurrence ────────────────────────────────────────────── +def test_recurrent_matches_independent_bruteforce_ungated(): + rng = _rng(1) + Q, K, V = _qkv(rng) + ours = dr.gated_delta_rule_recurrent(Q, K, V) + ref = _bruteforce_delta(Q, K, V) + np.testing.assert_allclose(ours, ref, rtol=1e-9, atol=1e-9) + + +def test_recurrent_matches_bruteforce_fully_gated(): + rng = _rng(2) + Q, K, V = _qkv(rng) + beta = _sig(rng.standard_normal(Q.shape[:3])) + decay = _sig(rng.standard_normal(Q.shape[:3]) + 2.0) # decay near 1 + ours = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay) + ref = _bruteforce_delta(Q, K, V, beta=beta, decay=decay) + np.testing.assert_allclose(ours, ref, rtol=1e-8, atol=1e-8) + + +def test_recurrent_erase_off_equals_existing_reference_with_decay(): + rng = _rng(3) + Q, K, V = _qkv(rng) + beta = _sig(rng.standard_normal(Q.shape[:3])) + decay = _sig(rng.standard_normal(Q.shape[:3]) + 1.5) + existing = np.asarray(ops.gated_deltanet(Q, K, V, beta=beta, decay=decay), np.float64) + ours = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay, erase=False) + np.testing.assert_allclose(existing, ours, rtol=1e-9, atol=1e-9) + + +def test_output_gate_matches_bruteforce(): + rng = _rng(4) + Q, K, V = _qkv(rng, d_v=4) + gate = rng.standard_normal((2, 3, 12, 4)) + ours = dr.gated_delta_rule_recurrent(Q, K, V, gate=gate) + ref = _bruteforce_delta(Q, K, V, gate=gate) + np.testing.assert_allclose(ours, ref, rtol=1e-9, atol=1e-9) + + +def test_return_state_shape_and_cross_call_carry(): + rng = _rng(5) + Q, K, V = _qkv(rng, S=10, d_k=5, d_v=4) + beta = _sig(rng.standard_normal(Q.shape[:3])) + # One pass over the whole sequence... + O_full, S_full = dr.gated_delta_rule_recurrent( + Q, K, V, beta=beta, return_state=True, state_dtype="fp64") + assert S_full.shape == (2, 3, 5, 4) + # ...equals two passes carrying state across the split. + O1, S1 = dr.gated_delta_rule_recurrent( + Q[:, :, :6], K[:, :, :6], V[:, :, :6], beta=beta[:, :, :6], + return_state=True, state_dtype="fp64") + O2 = dr.gated_delta_rule_recurrent( + Q[:, :, 6:], K[:, :, 6:], V[:, :, 6:], beta=beta[:, :, 6:], state=S1) + np.testing.assert_allclose(O_full[:, :, :6], O1, rtol=1e-9, atol=1e-9) + np.testing.assert_allclose(O_full[:, :, 6:], O2, rtol=1e-9, atol=1e-9) + + +# ── L2 — the chunk UT-transform keystone ───────────────────────────────────── +def test_forward_substitution_solves_unit_lower_triangular(): + rng = _rng(6) + C, d = 7, 3 + A = np.tril(rng.standard_normal((C, C)), k=-1) # strictly lower + W = rng.standard_normal((C, d)) + U = dr._forward_substitution(A, W) + np.testing.assert_allclose((np.eye(C) + A) @ U, W, rtol=1e-10, atol=1e-10) + + +def test_chunk_equals_recurrent_ungated(): + rng = _rng(7) + Q, K, V = _qkv(rng, S=20) + rec = dr.gated_delta_rule_recurrent(Q, K, V) + ch = dr.gated_delta_rule_chunked(Q, K, V, chunk_size=8) + np.testing.assert_allclose(ch, rec, rtol=1e-9, atol=1e-9) + + +def test_chunk_equals_recurrent_with_beta(): + rng = _rng(8) + Q, K, V = _qkv(rng, S=20) + beta = _sig(rng.standard_normal(Q.shape[:3])) + rec = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta) + ch = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, chunk_size=8) + np.testing.assert_allclose(ch, rec, rtol=1e-9, atol=1e-9) + + +def test_chunk_equals_recurrent_fully_gated(): + """The decay-folding proof: γ_t/γ_j ratios in the chunk form must reproduce + the per-token decay of the recurrence.""" + rng = _rng(9) + Q, K, V = _qkv(rng, S=24) + beta = _sig(rng.standard_normal(Q.shape[:3])) + decay = _sig(rng.standard_normal(Q.shape[:3]) + 2.0) + rec = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay) + ch = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, decay=decay, chunk_size=8) + np.testing.assert_allclose(ch, rec, rtol=1e-8, atol=1e-8) + + +def test_chunk_equals_recurrent_with_output_gate(): + rng = _rng(10) + Q, K, V = _qkv(rng, S=16, d_v=4) + gate = rng.standard_normal((2, 3, 16, 4)) + rec = dr.gated_delta_rule_recurrent(Q, K, V, gate=gate) + ch = dr.gated_delta_rule_chunked(Q, K, V, gate=gate, chunk_size=8) + np.testing.assert_allclose(ch, rec, rtol=1e-9, atol=1e-9) + + +@pytest.mark.parametrize("chunk_size", [1, 4, 8, 16, 64]) +def test_chunk_size_invariance(chunk_size): + rng = _rng(11) + Q, K, V = _qkv(rng, S=16) + beta = _sig(rng.standard_normal(Q.shape[:3])) + decay = _sig(rng.standard_normal(Q.shape[:3]) + 2.0) + rec = dr.gated_delta_rule_recurrent(Q, K, V, beta=beta, decay=decay) + ch = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, decay=decay, chunk_size=chunk_size) + np.testing.assert_allclose(ch, rec, rtol=1e-8, atol=1e-8) + + +def test_chunk_state_carry_matches_recurrent_state(): + rng = _rng(12) + Q, K, V = _qkv(rng, S=18, d_k=5, d_v=4) + beta = _sig(rng.standard_normal(Q.shape[:3])) + _, S_rec = dr.gated_delta_rule_recurrent( + Q, K, V, beta=beta, return_state=True, state_dtype="fp64") + _, S_ch = dr.gated_delta_rule_chunked( + Q, K, V, beta=beta, chunk_size=7, return_state=True, state_dtype="fp64") + np.testing.assert_allclose(S_ch, S_rec, rtol=1e-8, atol=1e-8) + + +def test_chunk_erase_off_equals_linear_reference(): + rng = _rng(13) + Q, K, V = _qkv(rng, S=16) + beta = _sig(rng.standard_normal(Q.shape[:3])) + existing = np.asarray(ops.gated_deltanet(Q, K, V, beta=beta), np.float64) + ch = dr.gated_delta_rule_chunked(Q, K, V, beta=beta, chunk_size=8, erase=False) + np.testing.assert_allclose(ch, existing, rtol=1e-9, atol=1e-9) diff --git a/tests/unit/test_stdlib_hybrid.py b/tests/unit/test_stdlib_hybrid.py new file mode 100644 index 000000000..a5bf9c7a1 --- /dev/null +++ b/tests/unit/test_stdlib_hybrid.py @@ -0,0 +1,142 @@ +"""Track L (L3) — hybrid linear/attention schedule + dual-cache contract. + +Schedule: `layer_types` is a first-class object (Qwen3.6 = `[lin,lin,lin,full]·N`). +Dual cache: linear layers carry a constant-size recurrent state Ŝ; full-attention +layers carry a growing KV cache. The headline oracle is **streaming dual-cache +decode ≡ full recompute** — it only holds if both caches are threaded correctly. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from tessera.stdlib import hybrid as hy + + +# ── schedule (first-class layer_types) ─────────────────────────────────────── +def test_qwen36_schedule_pattern(): + s = hy.qwen3_6_schedule(num_layers=40) + types = s.layer_types() + assert len(types) == 40 + # [linear, linear, linear, full] repeated. + assert types[:4] == [hy.LINEAR, hy.LINEAR, hy.LINEAR, hy.FULL] + assert s.counts() == {hy.LINEAR: 30, hy.FULL: 10} + assert s.full_indices() == [3, 7, 11, 15, 19, 23, 27, 31, 35, 39] + + +def test_nemotron_schedule_is_mostly_linear(): + s = hy.nemotron_schedule(num_layers=32, attn_period=8) + c = s.counts() + assert c[hy.LINEAR] == 28 and c[hy.FULL] == 4 # sparse anchors + + +def test_schedule_validates(): + with pytest.raises(ValueError): + hy.HybridSchedule(num_layers=0) + with pytest.raises(ValueError): + hy.HybridSchedule(num_layers=4, period=0) + + +def test_full_config_qwen36_dims_no_execution(): + """Full-config artifact check: the schedule lowers at production depth.""" + s = hy.qwen3_6_schedule(num_layers=40) + assert s.counts() == {hy.LINEAR: 30, hy.FULL: 10} + # 75% of token-mixers are linear (Gated DeltaNet), every 4th is full attention. + assert len(s.linear_indices()) == 30 + + +# ── the dual-cache oracle ──────────────────────────────────────────────────── +def _cfg(num_layers=8, H=2, Dh=8, Dm=16): + return hy.HybridConfig(d_model=Dm, num_heads=H, head_dim=Dh, + schedule=hy.HybridSchedule(num_layers=num_layers, + period=4, full_offset=1)) + + +@pytest.mark.parametrize("prefill", [1, 3, 7]) +def test_streaming_dualcache_decode_equals_full_recompute(prefill): + """Headline L3 oracle: token-by-token decode carrying Ŝ (linear layers) + + KV (full layers) reproduces the full parallel forward.""" + rng = np.random.default_rng(0) + cfg = _cfg(num_layers=8) + w = hy.synth_weights(cfg, rng) + x = rng.standard_normal((2, 12, cfg.d_model)) + full = hy.hybrid_forward(x, w, cfg) + stream = hy.hybrid_decode(x, w, cfg, prefill=prefill) + np.testing.assert_allclose(stream, full, rtol=1e-9, atol=1e-9) + + +def test_decode_equals_forward_pure_linear_stack(): + """All-linear stack (period huge → no full layers): pure recurrent-state + carry ≡ full forward.""" + rng = np.random.default_rng(1) + cfg = hy.HybridConfig(d_model=16, num_heads=2, head_dim=8, + schedule=hy.HybridSchedule(num_layers=6, period=999)) + assert cfg.schedule.counts()[hy.FULL] == 0 + w = hy.synth_weights(cfg, rng) + x = rng.standard_normal((1, 10, 16)) + np.testing.assert_allclose(hy.hybrid_decode(x, w, cfg, prefill=2), + hy.hybrid_forward(x, w, cfg), rtol=1e-9, atol=1e-9) + + +def test_decode_equals_forward_with_attention_anchors(): + """Every-other-layer full attention — both caches exercised heavily.""" + rng = np.random.default_rng(2) + cfg = hy.HybridConfig(d_model=24, num_heads=3, head_dim=8, + schedule=hy.HybridSchedule(num_layers=6, period=2, full_offset=1)) + assert cfg.schedule.counts() == {hy.LINEAR: 3, hy.FULL: 3} + w = hy.synth_weights(cfg, rng) + x = rng.standard_normal((2, 9, 24)) + np.testing.assert_allclose(hy.hybrid_decode(x, w, cfg, prefill=1), + hy.hybrid_forward(x, w, cfg), rtol=1e-9, atol=1e-9) + + +# ── L4.1 — SSM (Mamba) mixer (Nemotron) ────────────────────────────────────── +def test_ssm_scan_matches_shipped_selective_ssm_reference(): + """_ssm_scan (which returns the carried state) must reproduce the shipped + tessera.ops.selective_ssm — the L4 op's reference — so the SSM mixer is + anchored to ground truth, not self-consistent.""" + from tessera import ops + rng = np.random.default_rng(30) + B, S, D, N = 2, 14, 6, 4 + x = rng.standard_normal((B, S, D)) + A = -np.exp(rng.standard_normal(D) * 0.1) # negative, scalar-state + Bp = rng.standard_normal((B, S, N)) + Cp = rng.standard_normal((B, S, N)) + dt = np.log1p(np.exp(rng.standard_normal((B, S, D)))) # softplus > 0 + y, h = hy._ssm_scan(x, A, Bp, Cp, dt) + ref = np.asarray(ops.selective_ssm(x, A, Bp, Cp, dt)) + np.testing.assert_allclose(y, ref, rtol=1e-9, atol=1e-9) + assert h.shape == (B, D, N) + + +def _nemotron_cfg(num_layers=8): + return hy.HybridConfig(d_model=16, num_heads=2, head_dim=8, ssm_state=4, + linear_mixer=hy.SSM, + schedule=hy.nemotron_schedule(num_layers, attn_period=4)) + + +@pytest.mark.parametrize("prefill", [1, 2, 5]) +def test_nemotron_ssm_dualcache_decode_equals_recompute(prefill): + """Nemotron-shaped: Mamba SSM linear layers + sparse attention anchors. + Streaming carries SSM state h (linear) + KV (anchors) ≡ full recompute.""" + rng = np.random.default_rng(31) + cfg = _nemotron_cfg(num_layers=8) + assert cfg.mixer_for(0) == hy.SSM and cfg.mixer_for(3) == hy.FULL + w = hy.synth_weights(cfg, rng) + x = rng.standard_normal((2, 11, cfg.d_model)) + np.testing.assert_allclose(hy.hybrid_decode(x, w, cfg, prefill=prefill), + hy.hybrid_forward(x, w, cfg), rtol=1e-9, atol=1e-9) + + +def test_all_ssm_stack_decode_equals_recompute(): + """Pure Mamba stack (no anchors): SSM state carry alone ≡ full forward.""" + rng = np.random.default_rng(32) + cfg = hy.HybridConfig(d_model=16, num_heads=2, head_dim=8, ssm_state=4, + linear_mixer=hy.SSM, + schedule=hy.HybridSchedule(num_layers=5, period=999)) + assert cfg.schedule.counts()[hy.FULL] == 0 + w = hy.synth_weights(cfg, rng) + x = rng.standard_normal((1, 10, 16)) + np.testing.assert_allclose(hy.hybrid_decode(x, w, cfg, prefill=3), + hy.hybrid_forward(x, w, cfg), rtol=1e-9, atol=1e-9)