diff --git a/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.json b/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.json index 89e67402e..f8a1d05a8 100644 --- a/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.json +++ b/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.json @@ -1,8 +1,10 @@ { "slice_id": "LARQL-INFERENCE-TRUST-001A-ST4", + "closeout_slice": "LARQL-INFERENCE-TRUST-001A-ST4A", "work_start_sha": "15864226bd3ffb70a2cdc35ca3e52dafe76692c7", - "pr_base_sha": "15864226bd3ffb70a2cdc35ca3e52dafe76692c7", - "head_sha": "pending PR head", + "head_sha": "2889157f29ea85d690db5d08290bad41154a346f", + "merge_commit": "b2584c0acf7691d9963dff6a557ec7bac58755d9", + "ci": "all 11 triggered workflows completed successfully", "intrinsic_window": 512, "layer_pattern": "SSSSG x7, final layer global (34)", "shared_kv_map": { @@ -27,16 +29,70 @@ "observed_attention_ranges": "match expected (t1-t9 PASS)", "range_mismatch_count": 0, "prefill_decode_comparison_count": 4, - "maximum_absolute_error": 1e-5, + "prefill_decode_method": "canonical kv_prefill_run vs kv_prefill_run+kv_decode_step_run (window-512 E2B-like fixture)", + "prefill_decode_abs_tolerance": 1e-5, + "prefill_decode_per_position": { + "511": { + "hidden_max_abs": 1.073e-6, + "hidden_max_rel": 1.023e-6, + "local_kv_tail_max_abs": 0.0, + "local_cache_len_route_a": 512, + "local_cache_len_route_b": 512, + "global_cache_len_route_a": 512, + "global_cache_len_route_b": 512, + "local_absolute_range": "0..=511", + "next_position_after_decode": 512 + }, + "512": { + "hidden_max_abs": 0.0, + "hidden_max_rel": 0.0, + "local_kv_tail_max_abs": 0.0, + "local_cache_len_route_a": 512, + "local_cache_len_route_b": 512, + "global_cache_len_route_a": 513, + "global_cache_len_route_b": 513, + "local_absolute_range": "1..=512", + "next_position_after_decode": 513 + }, + "513": { + "hidden_max_abs": 0.0, + "hidden_max_rel": 0.0, + "local_kv_tail_max_abs": 0.0, + "local_cache_len_route_a": 512, + "local_cache_len_route_b": 512, + "global_cache_len_route_a": 514, + "global_cache_len_route_b": 514, + "local_absolute_range": "2..=513", + "next_position_after_decode": 514 + }, + "1024": { + "hidden_max_abs": 0.0, + "hidden_max_rel": 0.0, + "local_kv_tail_max_abs": 0.0, + "local_cache_len_route_a": 512, + "local_cache_len_route_b": 512, + "global_cache_len_route_a": 1025, + "global_cache_len_route_b": 1025, + "local_absolute_range": "513..=1024", + "next_position_after_decode": 1025 + } + }, + "prefill_decode_maximum_absolute_error": 1.073e-6, "cache_length_mismatch_count": 0, "absolute_position_mismatch_count": 0, "shared_consumer_count": 2, "independent_shared_cache_allocation_count": 0, "shared_consumer_kv_append_count": 0, + "source_append_deltas_per_decode_token": { + "local_source": "+1 (clipped to intrinsic window)", + "global_source": "+1 (unbounded)", + "local_consumer": 0, + "global_consumer": 0 + }, "poison_weight_test_result": "PASS", "incompatible_geometry_result": "fails loudly (panic)", - "qwen_regression_result": "N/A (no Qwen CPU F32 fixture; conventional full-attention retained)", - "q4k_regression_result": "PASS (regression-only, not claiming ST4 coverage)", + "qwen_regression_result": "PASS (Qwen2 full attention: intrinsic_window=None, 600-token prefill retains full prefix, decode grows cache +1, no shared-KV routing, old-key sensitivity abs=3.679e-4 rel=1.334e-2 at position 600 vs position 0 outside a 512-window)", + "q4k_regression_result": "PASS (regression-only, NOT part of the ST4/ST4A GREEN claim)", "semantic_official_model_run": false, "decision": "GREEN", "excluded_q4k_paths": [ @@ -47,9 +103,19 @@ "coarse_prefill", "coarse_decode_step" ], + "st4a": { + "work_start_sha": "eb40807ef0c392cd34456406b936692c13741716", + "tested_implementation_sha": "97ff07bb28e36420e07a015af0e4b32ccc736ef0", + "new_fixtures": [ + "make_synthetic_e2b_like_weights_random_window512", + "make_qwen2_test_weights" + ], + "boundary_positions_tested": [511, 512, 513, 1024], + "decision": "GREEN" + }, "synthetic_tests": { - "st4_attention_semantics": 18, - "st4_shared_kv": 17, + "st4_attention_semantics": 26, + "st4_shared_kv": 28, "st4_e2b_mapping": 5 } } diff --git a/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.md b/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.md index 8b88e29d0..f2c208fff 100644 --- a/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.md +++ b/bench/baselines/gemma4-e2b-f32-attention-semantics-2026-07-12.md @@ -1,6 +1,6 @@ -# Gemma 4 E2B F32 CPU Attention Semantics (ST4) +# Gemma 4 E2B F32 CPU Attention Semantics (ST4 + ST4A closeout) -Slice: `LARQL-INFERENCE-TRUST-001A-ST4` +Slice: `LARQL-INFERENCE-TRUST-001A-ST4` (closed by `…-ST4A`) ## Decision @@ -12,11 +12,22 @@ A GREEN result means only that the F32 CPU path now implements the required local/global attention ranges, absolute positions, and shared-KV routing needed for the next external semantic comparison. -## Revisions +ST4A closed the remaining evidence gaps without reopening or expanding the +ST4 implementation scope: it added the missing canonical-loop prefill/decode +equivalence at the 512-window boundaries (511/512/513/1024), an explicit +Qwen2 full-attention regression, and a source append-count proof, and +corrected this report's ST4 metadata. -- Work-start SHA: `15864226bd3ffb70a2cdc35ca3e52dafe76692c7` -- PR base SHA: `15864226bd3ffb70a2cdc35ca3e52dafe76692c7` -- Head SHA: _filled at PR head_ +## Revisions (corrected by ST4A) + +- ST4 work-start SHA: `15864226bd3ffb70a2cdc35ca3e52dafe76692c7` +- ST4 head SHA: `2889157f29ea85d690db5d08290bad41154a346f` +- ST4 merge commit (squash into main): `b2584c0acf7691d9963dff6a557ec7bac58755d9` +- ST4 CI: all 11 triggered workflows completed successfully. +- ST4A work-start SHA: `eb40807ef0c392cd34456406b936692c13741716` +- ST4A tested implementation commit SHA: `97ff07bb28e36420e07a015af0e4b32ccc736ef0` + (the ST4A tests commit immediately before this evidence commit; checks out + to the ST4 implementation plus the ST4A tests that verify it). ## What changed (F32-only) @@ -45,7 +56,9 @@ routing needed for the next external semantic comparison. These pass `effective_window = None` (full causal) and do not route shared-KV. The intrinsic Gemma 4 local/global and shared-KV parity must be -verified separately before they become semantic-oracle candidates. +verified separately before they become semantic-oracle candidates. **Q4_K +remains excluded from the ST4/ST4A GREEN claim**; it may run as a +regression only. ## E2B local/global pattern + shared-KV map @@ -57,7 +70,7 @@ verified separately before they become semantic-oracle candidates. - every shared layer points to a source of the same attention type - complete shared-KV map result: PASS (5 mapping tests) -## Synthetic test results (ST4 §19) +## Synthetic test results (ST4 §19, refreshed by ST4A) | Group | Count | Result | |-------|-------|--------| @@ -69,39 +82,141 @@ verified separately before they become semantic-oracle candidates. | Regressions (t43, t44) | 2 | PASS | | E2B 35-layer mapping | 5 | PASS | +Totals (current): `st4_attention_semantics` 26, `st4_shared_kv` 28 +(includes 11 ST4A closeout tests), `test_architectures` st4_* 5. + - range mismatch count: 0 - cache-length mismatch count: 0 - absolute-position mismatch count: 0 -- prefill/decode comparison count: 4 boundary positions -- maximum absolute error: ≤ 1e-5 (GQA primitive parity; bit-exact where - operation ordering is identical) +- prefill/decode comparison: real canonical-loop numerical comparison at 4 + boundary positions (see ST4A closeout below for measured values). - shared consumers tested: 2 (sliding L2, global L3 on the 4-layer fixture) - independent shared-cache allocation count: 0 - shared consumer K/V append count: 0 - poison-weight test result: PASS (consumer K/V poison has no effect; source K/V mutation changes output) - incompatible-geometry result: fails loudly (panic) -- Qwen regression result: N/A (no Qwen fixture in the CPU F32 unit suite; - conventional full-attention retained via `intrinsic_window = None`) -- Q4_K regression result: PASS (regression-only, not claiming ST4 coverage) +- Qwen2 regression result: **PASS** (full-attention regression; see ST4A §2) +- Q4_K regression result: PASS (regression-only, NOT part of the GREEN claim) - semantic official-model run = false ## Verification commands and totals ``` cargo fmt --all -- --check # clean -cargo test -p larql-models --lib # 425 passed -cargo test -p larql-compute --lib # 756 passed -cargo test -p larql-inference --lib # 1290 passed -cargo test -p larql-kv --lib # 765 passed -cargo test -p larql-vindex --lib # 1154 passed -cargo test -p larql-cli --bins # 243 passed -cargo test -p larql-compute --test st4_attention_semantics # 18 passed -cargo test -p larql-kv --test st4_shared_kv # 17 passed -cargo test -p larql-models --test test_architectures (st4_*) # 5 passed -cargo clippy -p larql-{models,compute,inference,kv,vindex,cli} --all-targets -- -D warnings # clean -cargo test -p larql-compute q4k / larql-inference q4k / larql-vindex q4k # green (regression) -cargo build -p larql-cli --release # passed +cargo test -p larql-compute --test st4_attention_semantics # 26 passed +cargo test -p larql-kv --test st4_shared_kv # 28 passed +cargo test -p larql-models --test test_architectures # 90 passed (st4_* = 5) +cargo test -p larql-models # green +cargo test -p larql-compute # green +cargo test -p larql-inference # green +cargo test -p larql-kv # green +cargo clippy -p larql-{models,compute,inference,kv} --all-targets -- -D warnings # clean +``` + +CI status (ST4 PR #60): all 11 triggered workflows completed successfully. + +--- + +## ST4A closeout + +- work-start SHA: `eb40807ef0c392cd34456406b936692c13741716` +- tested implementation commit SHA: `97ff07bb28e36420e07a015af0e4b32ccc736ef0` + (the ST4A tests commit immediately before this evidence commit; ST4A is + tests + evidence only — no attention implementation change). +- new fixtures: `make_synthetic_e2b_like_weights_random_window512` + (4-layer E2B-like Gemma 4 with a 512-token intrinsic window) and + `make_qwen2_test_weights` (Qwen2 full-attention with Q/K/V biases). + +### Boundary positions tested + +511, 512, 513, 1024 — the 512-window boundary and the first clipped +position, plus a far-clipped position. + +### Per-position canonical prefill/decode equivalence + +Route A = one-shot `kv_prefill_run` of tokens `0..=target`; Route B = +`kv_prefill_run` of `0..target` then one `kv_decode_step_run` at the true +absolute position. Absolute tolerance `1e-5`. + +| target | hidden max abs | hidden max rel | local K/V tail max abs | local len (A=B) | global len (A=B) | abs position (B) | +|--------|---------------:|---------------:|-----------------------:|----------------:|-----------------:|-----------------:| +| 511 | 1.073e-6 | 1.023e-6 | 0 | 512 | 512 | 512 | +| 512 | 0 | 0 | 0 | 512 | 513 | 513 | +| 513 | 0 | 0 | 0 | 512 | 514 | 514 | +| 1024 | 0 | 0 | 0 | 512 | 1025 | 1025 | + +The largest hidden-state difference across all four positions is `1.073e-6` +(target 511); positions 512/513/1024 are bit-exact between routes. The +local source K/V tail is bit-exact (`0`) at every position, confirming +both routes hold the same absolute RoPE positions after clipping. + +### Local cache evidence + +Every target yields a 512-row local source cache in both routes, matching +the intrinsic window: + +``` +target 511 → 512 rows, absolute range 0..=511 +target 512 → 512 rows, absolute range 1..=512 +target 513 → 512 rows, absolute range 2..=513 +target 1024 → 512 rows, absolute range 513..=1024 ``` -CI status: pending PR creation (all local checks green). +### Global cache evidence + +The global (full-attention) source retains the entire prefix — `target+1` +rows in both routes (512 / 513 / 514 / 1025). + +### Absolute-position evidence + +After the decode, `cache.next_position == target+1` (512 / 513 / 514 / +1025), not the clipped 512 cache length. The bit-exact local K/V tail +between Route A (natural positions) and Route B (`next_position`-derived +RoPE) is the direct proof that the absolute RoPE position is the true +target position. + +### Source append deltas (per decode token, 4-layer E2B-style fixture) + +``` +local source → +1 row (clipped to its intrinsic window) +global source → +1 row (unbounded) +local consumer → 0 rows +global consumer → 0 rows +``` + +Invariant holds on every one of 6 consecutive decode steps. + +### Shared consumer append count + +0 — shared consumer layers hold no independent cache before, during, or +after any decode step (verified by `cache.layers[c].is_none()`). + +### Qwen2 regression + +PASS. On a real Qwen2 architecture (`family == "qwen2"`, Q/K/V biases): + +- `intrinsic_attention_window == None` on every layer; +- prefill of 600 tokens retains the full 600-row causal prefix (no + clipping) on every layer; +- decode grows the cache by exactly one row per token (5 steps); +- no layer activates Gemma 4 shared-KV routing (`kv_shared_source_layer == None`); +- old-key sensitivity: zeroing the K/V of position 0 changes the + position-600 decode output by `abs = 3.679e-4`, `rel = 1.334e-2`. Position + 0 lies outside a 512-token local window for position 600 (window = + 89..600), so this change proves Qwen2 attends to the early key under full + attention. + +### Q4_K scope statement + +Excluded from the ST4/ST4A GREEN claim. Q4_K paths run `effective_window += None` and do not route shared-KV; they pass as regressions only and +require separate semantic verification before becoming oracle candidates. + +### Decision + +**GREEN.** All 15 ST4A decision-gate criteria pass; the canonical +prefill/decode outputs are numerically equivalent (≤ `1.073e-6` abs) at +511/512/513/1024, local/global cache lengths and absolute positions are +correct, Qwen2 full attention is regressed, source append counts are +exact, and Q4_K remains outside the claim. diff --git a/crates/larql-inference/src/test_utils.rs b/crates/larql-inference/src/test_utils.rs index 16d7540cd..b483133bc 100644 --- a/crates/larql-inference/src/test_utils.rs +++ b/crates/larql-inference/src/test_utils.rs @@ -908,8 +908,9 @@ pub fn make_test_gemma4_moe_weights() -> ModelWeights { // so existing `crate::test_utils::*` callers (forward/ple.rs tests) and // downstream test crates keep working. pub use larql_models::test_fixtures::{ - make_synthetic_e2b_like_weights, make_synthetic_e2b_like_weights_random, - synthetic_e2b_like_arch_json, + make_qwen2_test_weights, make_synthetic_e2b_like_weights, + make_synthetic_e2b_like_weights_random, make_synthetic_e2b_like_weights_random_window512, + synthetic_e2b_like_arch_json, synthetic_e2b_like_arch_json_window512, }; /// Bundled fixture for Q4_K decode-path tests. Mirrors `TestFixtures`. pub struct Q4KTestFixtures { diff --git a/crates/larql-kv/tests/st4_shared_kv.rs b/crates/larql-kv/tests/st4_shared_kv.rs index a0a783930..df779a12e 100644 --- a/crates/larql-kv/tests/st4_shared_kv.rs +++ b/crates/larql-kv/tests/st4_shared_kv.rs @@ -6,16 +6,34 @@ //! //! These tests prove F32 CPU attention and cache semantics using synthetic //! fixtures. They do NOT prove official Gemma 4 logits or generated text. +//! +//! ST4A closeout (tests tagged `st4a_`) adds the missing acceptance +//! evidence: real canonical-loop prefill/decode numerical equivalence at +//! the 512-window boundaries (511/512/513/1024), a Qwen2 full-attention +//! regression, and a source append-count proof. +use larql_compute::attention::intrinsic_attention_window; use larql_inference::attention::run_attention_block_decode_step_shared_backend; use larql_inference::ffn::WeightFfn; use larql_inference::forward::hooks::NoopHook; use larql_inference::larql_models::WeightArray; -use larql_inference::test_utils::{make_synthetic_e2b_like_weights_random, make_test_weights}; +use larql_inference::test_utils::{ + make_qwen2_test_weights, make_synthetic_e2b_like_weights_random, + make_synthetic_e2b_like_weights_random_window512, make_test_weights, +}; use larql_inference::{ModelWeights, WeightsView}; use larql_kv::generation::{kv_decode_step_run, kv_prefill_run}; use ndarray::Array2; +/// Absolute tolerance for F32 CPU prefill/decode numerical equivalence. +/// The prefill and decode paths compute each query's windowed attention +/// over the same K/V rows in the same order, so the difference is at most +/// a few ULPs of float round-off; 1e-5 is the existing ST4 tolerance and +/// passes with a wide margin on the synthetic fixtures. +const EQUIV_ABS_TOL: f32 = 1e-5; +/// Relative-difference floor used to avoid division by ~0 hidden units. +const EQUIV_REL_FLOOR: f32 = 1e-6; + /// Embed a deterministic non-zero hidden batch. fn embed_seq(weights: &ModelWeights, n: usize) -> Array2 { Array2::from_shape_fn((n, weights.hidden_size), |(r, c)| { @@ -23,6 +41,27 @@ fn embed_seq(weights: &ModelWeights, n: usize) -> Array2 { }) } +/// Maximum absolute and relative element-wise difference between two +/// equal-shaped row vectors. Relative difference uses `max(|a|, |b|, floor)` +/// as the denominator so near-zero units do not blow up the metric. +fn max_abs_and_rel_diff(a: &Array2, b: &Array2) -> (f32, f32) { + assert_eq!(a.shape(), b.shape(), "shape mismatch in diff"); + let mut max_abs = 0.0f32; + let mut max_rel = 0.0f32; + for (&x, &y) in a.iter().zip(b.iter()) { + let d = (x - y).abs(); + if d > max_abs { + max_abs = d; + } + let denom = x.abs().max(y.abs()).max(EQUIV_REL_FLOOR); + let rel = d / denom; + if rel > max_rel { + max_rel = rel; + } + } + (max_abs, max_rel) +} + // ── §19 Tests 18–24: cache behavior (canonical loops) ────────────────── #[test] @@ -603,3 +642,577 @@ fn t44_caller_specified_bounded_window_remains_green() { } } } + +// ── ST4A §1: canonical prefill/decode equivalence at the 512 boundary ── +// +// Route A: one-shot `kv_prefill_run` of tokens 0..=target → hidden of the +// last (target) position. +// Route B: `kv_prefill_run` of tokens 0..target, then a single +// `kv_decode_step_run` of token `target` at its true absolute +// position → hidden of position target. +// +// The two routes must agree numerically (max abs diff ≤ EQUIV_ABS_TOL) and +// leave identical local/global source cache tails and lengths. The fixture +// is the 4-layer E2B-like Gemma 4 shape (sliding/global/sliding/global, +// KV-shared) but with a 512-token intrinsic sliding window so positions +// 511/512/513/1024 are reachable. + +/// Evidence captured for one target position. Fields are asserted by the +/// per-position tests and recorded into the ST4A closeout report. +struct EquivalenceEvidence { + target: usize, + hidden_max_abs: f32, + hidden_max_rel: f32, + local_kv_tail_max_abs: f32, + local_len_a: usize, + local_len_b: usize, + global_len_a: usize, + global_len_b: usize, + next_position_b: usize, +} + +/// Window-512 E2B-like fixture shared across all boundary positions. +fn window512_weights() -> ModelWeights { + make_synthetic_e2b_like_weights_random_window512() +} + +/// Token id for absolute position `i` (wrap around the 32-token vocab). +fn token_at(weights: &ModelWeights, i: usize) -> u32 { + (i % weights.vocab_size) as u32 +} + +fn run_prefill_decode_equivalence(target: usize) -> EquivalenceEvidence { + let weights = window512_weights(); + let arch = &*weights.arch; + let ffn = WeightFfn { weights: &weights }; + assert_eq!(arch.sliding_window_size(), Some(512), "fixture window"); + + // Local source = first sliding layer; global source = first global layer. + let local_source = (0..weights.num_layers) + .find(|&l| arch.is_sliding_window_layer(l) && arch.kv_shared_source_layer(l).is_none()) + .expect("local source exists"); + let global_source = (0..weights.num_layers) + .find(|&l| !arch.is_sliding_window_layer(l) && arch.kv_shared_source_layer(l).is_none()) + .expect("global source exists"); + + // ── Route A: one-shot prefill of tokens 0..=target ── + let prompt_a: Vec = (0..=target).map(|i| token_at(&weights, i)).collect(); + let (h_a, cache_a) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt_a, + None, + None, + &mut NoopHook, + ) + .expect("Route A prefill"); + + // ── Route B: prefill 0..target, then decode token `target` ── + let prompt_b: Vec = (0..target).map(|i| token_at(&weights, i)).collect(); + let (_h_pre, mut cache_b) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt_b, + None, + None, + &mut NoopHook, + ) + .expect("Route B prefill"); + let h_b = kv_decode_step_run( + &weights, + &ffn, + &mut cache_b, + token_at(&weights, target), + None, + &mut NoopHook, + ) + .expect("Route B decode"); + + // ── Numerical comparison of the post-layer hidden state ── + let (hidden_max_abs, hidden_max_rel) = max_abs_and_rel_diff(&h_a, &h_b); + + // ── Local source K/V tail comparison (same absolute positions) ── + let local_a = cache_a.get_layer(local_source).expect("A local K/V"); + let local_b = cache_b.get_layer(local_source).expect("B local K/V"); + assert_eq!(local_a.0.shape(), local_b.0.shape(), "local K tail shape"); + let (k_abs, _) = max_abs_and_rel_diff(&local_a.0, &local_b.0); + let (v_abs, _) = max_abs_and_rel_diff(&local_a.1, &local_b.1); + let local_kv_tail_max_abs = k_abs.max(v_abs); + + // ── Shared-consumer absence ── + for layer in 0..weights.num_layers { + if arch.kv_shared_source_layer(layer).is_some() { + assert!( + cache_a.layers[layer].is_none(), + "Route A shared consumer L{layer} must have no cache" + ); + assert!( + cache_b.layers[layer].is_none(), + "Route B shared consumer L{layer} must have no cache" + ); + } + } + + // ── Global source retains the full prefix (length = target+1) ── + let global_a = cache_a.get_layer(global_source).expect("A global K/V"); + let global_b = cache_b.get_layer(global_source).expect("B global K/V"); + + EquivalenceEvidence { + target, + hidden_max_abs, + hidden_max_rel, + local_kv_tail_max_abs, + local_len_a: local_a.0.shape()[0], + local_len_b: local_b.0.shape()[0], + global_len_a: global_a.0.shape()[0], + global_len_b: global_b.0.shape()[0], + next_position_b: cache_b.next_position, + } +} + +/// Assert the invariants that hold for every boundary target on the +/// window-512 fixture, then print the measured metrics for the report. +fn assert_equivalence(ev: &EquivalenceEvidence) { + let target = ev.target; + eprintln!( + "st4a target={target}: hidden abs={:.3e} rel={:.3e} | local_kv_tail abs={:.3e} | \ + local_len A={} B={} | global_len A={} B={} | next_pos B={}", + ev.hidden_max_abs, + ev.hidden_max_rel, + ev.local_kv_tail_max_abs, + ev.local_len_a, + ev.local_len_b, + ev.global_len_a, + ev.global_len_b, + ev.next_position_b, + ); + + // Numerical equivalence of the returned hidden state. + assert!( + ev.hidden_max_abs <= EQUIV_ABS_TOL, + "target {target}: hidden max abs diff {:.3e} exceeds {EQUIV_ABS_TOL}", + ev.hidden_max_abs + ); + // The local source K/V tail must agree (same positions, same RoPE). + assert!( + ev.local_kv_tail_max_abs <= EQUIV_ABS_TOL, + "target {target}: local K/V tail max abs diff {:.3e} exceeds {EQUIV_ABS_TOL}", + ev.local_kv_tail_max_abs + ); + + // Local source cache is clipped to the 512-window tail in both routes. + assert_eq!(ev.local_len_a, 512, "target {target}: local A length"); + assert_eq!(ev.local_len_b, 512, "target {target}: local B length"); + + // Global source retains the full prefix: target+1 rows in both routes. + assert_eq!( + ev.global_len_a, + target + 1, + "target {target}: global A length" + ); + assert_eq!( + ev.global_len_b, + target + 1, + "target {target}: global B length" + ); + + // Absolute position pointer = target+1 after the decode (NOT the clipped + // 512 cache length). The hidden equivalence above is the real proof that + // RoPE used the true absolute position; this pins the pointer explicitly. + assert_eq!( + ev.next_position_b, + target + 1, + "target {target}: absolute position pointer after decode" + ); +} + +#[test] +fn st4a_prefill_decode_equivalence_at_511() { + let ev = run_prefill_decode_equivalence(511); + assert_equivalence(&ev); +} + +#[test] +fn st4a_prefill_decode_equivalence_at_512() { + let ev = run_prefill_decode_equivalence(512); + assert_equivalence(&ev); +} + +#[test] +fn st4a_prefill_decode_equivalence_at_513() { + let ev = run_prefill_decode_equivalence(513); + assert_equivalence(&ev); +} + +#[test] +fn st4a_prefill_decode_equivalence_at_1024() { + let ev = run_prefill_decode_equivalence(1024); + assert_equivalence(&ev); +} + +// ── ST4A §1 (extra): absolute RoPE position ≠ clipped cache length ── +// +// At target 1024 the local cache holds 512 rows but the decode token's +// absolute RoPE position must be 1024. Pin both numbers side by side so a +// regression that derived the position from the cache length is caught +// directly (the equivalence test above catches it indirectly via the K/V +// mismatch). +#[test] +fn st4a_absolute_position_independent_of_clipped_cache_length() { + let weights = window512_weights(); + let ffn = WeightFfn { weights: &weights }; + let target = 1024usize; + let prompt: Vec = (0..target).map(|i| token_at(&weights, i)).collect(); + let (_h, mut cache) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("prefill"); + // Before the decode: position pointer is the true 1024, but every local + // source cache is clipped to the 512-window tail. + assert_eq!(cache.next_position, target, "true absolute position"); + let arch = &*weights.arch; + for layer in 0..weights.num_layers { + if arch.kv_shared_source_layer(layer).is_none() && arch.is_sliding_window_layer(layer) { + assert_eq!( + cache.cached_len(layer), + 512, + "local source L{layer} clipped" + ); + } + } + kv_decode_step_run( + &weights, + &ffn, + &mut cache, + token_at(&weights, target), + None, + &mut NoopHook, + ) + .expect("decode"); + assert_eq!( + cache.next_position, + target + 1, + "position advances past clip" + ); +} + +// ── ST4A §2: Qwen2 full-attention regression ── +// +// Qwen2 is a conventional full-attention arch: no intrinsic sliding window, +// no KV sharing. Prove (1) intrinsic_attention_window = None, (2) prefill +// retains the full causal prefix (cache length == prompt length, even past +// 512), (3) decode grows the cache by exactly one row, (4) Gemma 4 +// shared-KV routing is inactive, and (5) a late query still depends on an +// early key (position 0) that a 512-token Gemma local window would have +// evicted. + +fn qwen_weights() -> ModelWeights { + let w = make_qwen2_test_weights(); + let arch = &*w.arch; + assert_eq!(arch.family(), "qwen2", "fixture must be a real Qwen2 arch"); + w +} + +#[test] +fn st4a_qwen_full_attention_arch_properties() { + let weights = qwen_weights(); + let arch = &*weights.arch; + // No intrinsic window on any layer. + for layer in 0..weights.num_layers { + assert_eq!( + intrinsic_attention_window(arch, layer), + None, + "Qwen2 L{layer} must have no intrinsic window" + ); + assert!( + !arch.is_sliding_window_layer(layer), + "Qwen2 L{layer} must not be a sliding-window layer" + ); + assert_eq!( + arch.kv_shared_source_layer(layer), + None, + "Qwen2 L{layer} must not activate Gemma 4 shared-KV routing" + ); + } + assert_eq!( + arch.sliding_window_size(), + None, + "Qwen2 has no sliding window" + ); +} + +#[test] +fn st4a_qwen_prefill_retains_full_causal_prefix() { + let weights = qwen_weights(); + let ffn = WeightFfn { weights: &weights }; + // 600 tokens — past the 512 Gemma local window. + let prompt: Vec = (0..600).map(|i| token_at(&weights, i)).collect(); + let (_h, cache) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("prefill"); + // Full causal prefix retained: cache length == prompt length (600), not + // clipped to any window. + for layer in 0..weights.num_layers { + assert_eq!( + cache.cached_len(layer), + 600, + "Qwen2 L{layer} prefill cache must equal prompt length (full prefix)" + ); + } + assert_eq!( + cache.next_position, 600, + "position pointer == prompt length" + ); +} + +#[test] +fn st4a_qwen_decode_grows_cache_by_one_row_per_token() { + let weights = qwen_weights(); + let ffn = WeightFfn { weights: &weights }; + let prompt: Vec = (0..520).map(|i| token_at(&weights, i)).collect(); + let (_h, mut cache) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("prefill"); + for step in 0..5usize { + let before = cache.cached_len(0); + kv_decode_step_run( + &weights, + &ffn, + &mut cache, + token_at(&weights, 520 + step), + None, + &mut NoopHook, + ) + .unwrap_or_else(|| panic!("decode step {step} returned None")); + for layer in 0..weights.num_layers { + assert_eq!( + cache.cached_len(layer), + before + 1, + "Qwen2 L{layer} decode must grow cache by exactly 1 at step {step}" + ); + } + assert_eq!(cache.next_position, 521 + step, "position advances by 1"); + } +} + +#[test] +fn st4a_qwen_late_query_depends_on_early_key_outside_window() { + // Old-key sensitivity: a late Qwen query (position 600) must still + // attend to an early key (position 0) that a 512-token Gemma local + // window would have evicted (window for position 600 = 89..600). We + // prefill 600 tokens, then decode position 600 twice from cloned + // caches — once normally, once with position-0 K/V zeroed across all + // layers. A measured difference proves the late query read position 0. + let weights = qwen_weights(); + let ffn = WeightFfn { weights: &weights }; + let prompt: Vec = (0..600).map(|i| token_at(&weights, i)).collect(); + let (_h, cache) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("prefill"); + // Sanity: full prefix retained (no window clipped position 0 away). + for layer in 0..weights.num_layers { + assert_eq!(cache.cached_len(layer), 600); + } + + let mut cache_normal = cache.clone(); + let mut cache_poison = cache.clone(); + // Zero the K AND V of the earliest key (position 0) in every layer. + for layer in 0..weights.num_layers { + if let Some((k, v)) = cache_poison.layers.get_mut(layer).and_then(|o| o.as_mut()) { + k.slice_mut(ndarray::s![0, ..]).fill(0.0); + v.slice_mut(ndarray::s![0, ..]).fill(0.0); + } + } + + let h_normal = kv_decode_step_run( + &weights, + &ffn, + &mut cache_normal, + token_at(&weights, 600), + None, + &mut NoopHook, + ) + .expect("decode normal"); + let h_poison = kv_decode_step_run( + &weights, + &ffn, + &mut cache_poison, + token_at(&weights, 600), + None, + &mut NoopHook, + ) + .expect("decode poison"); + + let (max_abs, max_rel) = max_abs_and_rel_diff(&h_normal, &h_poison); + eprintln!( + "st4a qwen old-key sensitivity: position-0 K/V zeroing changes pos-600 decode by abs={:.3e} rel={:.3e}", + max_abs, max_rel + ); + // The output MUST change: position 600 attends to position 0 under full + // attention. (Under a 512-window, position 0 is outside 89..600 → zeroing + // it would be a no-op.) Require a material, non-round-off change. + let mag = h_normal + .iter() + .map(|v| v.abs()) + .fold(0.0f32, f32::max) + .max(1e-6); + assert!( + max_abs > 1e-4 && max_rel > 1e-3, + "Qwen late query must depend on early key 0 (abs={:.3e}, rel={:.3e}, mag={:.3e}); \ + a 512-window would have made this a no-op", + max_abs, + max_rel, + mag + ); +} + +// ── ST4A §3: source append-count proof ── +// +// For one decode token on the 4-layer E2B-style fixture: the local source +// cache grows by exactly 1 and the global source cache grows by exactly 1, +// while the shared consumer layers append zero rows (they hold no +// independent cache). Proven by direct, unambiguous cache-length deltas +// across several decode steps. +// +// The window-512 E2B-style fixture is used with a SHORT prompt (3 tokens) +// so neither source is at its window cap yet — every decode step therefore +// shows a clean +1 delta on both sources (no clipping masks the append). +// The shared consumers are then checked for zero appends on every step. + +#[test] +fn st4a_source_append_count_per_decode_token() { + let weights = window512_weights(); + let arch = &*weights.arch; + let ffn = WeightFfn { weights: &weights }; + assert_eq!(arch.sliding_window_size(), Some(512), "fixture window"); + + let local_source = (0..weights.num_layers) + .find(|&l| arch.is_sliding_window_layer(l) && arch.kv_shared_source_layer(l).is_none()) + .expect("local source"); + let global_source = (0..weights.num_layers) + .find(|&l| !arch.is_sliding_window_layer(l) && arch.kv_shared_source_layer(l).is_none()) + .expect("global source"); + let consumers: Vec = (0..weights.num_layers) + .filter(|&l| arch.kv_shared_source_layer(l).is_some()) + .collect(); + assert!(!consumers.is_empty(), "fixture must have shared consumers"); + + // Short prompt so both sources sit well below the 512 cap and every + // decode step produces an unambiguous +1 delta (no clipping). + let prompt: Vec = vec![0, 1, 2]; + let (_h, mut cache) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("prefill"); + + for &c in &consumers { + assert!(cache.layers[c].is_none(), "consumer L{c} starts cache-less"); + } + let local_start = cache.cached_len(local_source); + let global_start = cache.cached_len(global_source); + assert_eq!(local_start, 3, "local source starts at prompt length"); + assert_eq!(global_start, 3, "global source starts at prompt length"); + + for step in 0..6usize { + kv_decode_step_run( + &weights, + &ffn, + &mut cache, + (step % 4) as u32, + None, + &mut NoopHook, + ) + .unwrap_or_else(|| panic!("decode step {step} None")); + + // Local source: exactly +1 row per token (below the window cap, so + // the delta is an unambiguous append, not a clip artifact). + assert_eq!( + cache.cached_len(local_source), + local_start + step + 1, + "step {step}: local source L{local_source} must append exactly 1" + ); + // Global source: exactly +1 row per token (unbounded). + assert_eq!( + cache.cached_len(global_source), + global_start + step + 1, + "step {step}: global source L{global_source} must append exactly 1" + ); + // Shared consumers: still no independent cache entry (zero appends). + for &c in &consumers { + assert!( + cache.layers[c].is_none(), + "step {step}: shared consumer L{c} must append 0 rows" + ); + } + } +} + +// Companion: prove the local source still appends exactly one row even once +// it reaches its intrinsic window cap (window-4 E2B fixture), then clips +// straight back to the cap. The cache length stays pinned at the window +// size — never 3 (no append) and never above the cap (no double append). +#[test] +fn st4a_local_source_appends_one_then_clips_at_window_cap() { + let weights = e2b_weights(); + let arch = &*weights.arch; + let ffn = WeightFfn { weights: &weights }; + let window = arch.sliding_window_size().expect("E2B window"); + let local_source = (0..weights.num_layers) + .find(|&l| arch.is_sliding_window_layer(l) && arch.kv_shared_source_layer(l).is_none()) + .expect("local source"); + + // Prefill past the window so the local source is pinned at the cap. + let prompt: Vec = (0..(window as u32 * 2)).collect(); + let (_h, mut cache) = kv_prefill_run( + WeightsView::dense(&weights), + &ffn, + &prompt, + None, + None, + &mut NoopHook, + ) + .expect("prefill"); + assert_eq!(cache.cached_len(local_source), window, "pinned at cap"); + + for step in 0..4usize { + kv_decode_step_run(&weights, &ffn, &mut cache, 0u32, None, &mut NoopHook) + .unwrap_or_else(|| panic!("decode step {step} None")); + // Appended one row then clipped back: length stays exactly at the + // window cap (a no-op append would leave it unchanged too, but the + // unclipped test above proves the append primitive fires; here we + // pin that the post-clip length is exactly the cap, not above it). + assert_eq!( + cache.cached_len(local_source), + window, + "step {step}: local source clips to exactly the window cap" + ); + } +} diff --git a/crates/larql-models/src/test_fixtures.rs b/crates/larql-models/src/test_fixtures.rs index c3ab114c2..62cf1a4ee 100644 --- a/crates/larql-models/src/test_fixtures.rs +++ b/crates/larql-models/src/test_fixtures.rs @@ -434,6 +434,107 @@ pub fn make_starcoder2_test_weights() -> ModelWeights { } } +// ── Qwen2 full-attention synthetic fixture (ST4A regression) ── + +/// Build a small synthetic `ModelWeights` configured as a Qwen2 dense arch. +/// +/// Qwen2 is a conventional full-attention architecture (no intrinsic +/// sliding window, no KV sharing). This fixture populates the Qwen2 +/// attention biases (Q/K/V) so the regression exercises the real Qwen2 +/// attention path, not a generic proxy. Dimensions mirror +/// [`make_test_weights`] (hidden=16, 2 layers) so a 600+ token prefill +/// stays cheap on CPU. +pub fn make_qwen2_test_weights() -> ModelWeights { + const VOCAB: usize = 32; + const HIDDEN: usize = 16; + const INTER: usize = 32; + const NUM_Q: usize = 2; + const NUM_KV: usize = 1; + const HEAD_DIM: usize = 8; + const NUM_LAYERS: usize = 2; + + let arch_json = serde_json::json!({ + "model_type": "qwen2", + "hidden_size": HIDDEN, + "num_hidden_layers": NUM_LAYERS, + "intermediate_size": INTER, + "head_dim": HEAD_DIM, + "num_attention_heads": NUM_Q, + "num_key_value_heads": NUM_KV, + "vocab_size": VOCAB, + }); + let arch = detect_from_json(&arch_json); + + let mut tensors: HashMap = HashMap::new(); + let mut vectors: HashMap> = HashMap::new(); + let mut rng_state = 0x51e2_b000_u64; + let mut rand_mat = |rows: usize, cols: usize, scale: f32| -> WeightArray { + let data: Vec = (0..rows * cols) + .map(|_| { + rng_state = rng_state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + (rng_state as u32) as f32 / u32::MAX as f32 * 2.0 * scale - scale + }) + .collect(); + Array2::from_shape_vec((rows, cols), data) + .unwrap() + .into_shared() + }; + + let embed = rand_mat(VOCAB, HIDDEN, 0.1); + let lm_head = rand_mat(VOCAB, HIDDEN, 0.1); + tensors.insert(arch.embed_key().to_string(), embed.clone()); + vectors.insert(arch.final_norm_key().to_string(), vec![1.0; HIDDEN]); + + let q_dim = NUM_Q * HEAD_DIM; + let kv_dim = NUM_KV * HEAD_DIM; + + for layer in 0..NUM_LAYERS { + tensors.insert(arch.attn_q_key(layer), rand_mat(q_dim, HIDDEN, 0.1)); + tensors.insert(arch.attn_k_key(layer), rand_mat(kv_dim, HIDDEN, 0.1)); + tensors.insert(arch.attn_v_key(layer), rand_mat(kv_dim, HIDDEN, 0.1)); + tensors.insert(arch.attn_o_key(layer), rand_mat(HIDDEN, q_dim, 0.1)); + // Qwen2 attention biases (Q/K/V) — the distinguishing Qwen2 path. + if let Some(k) = arch.attn_q_bias_key(layer) { + vectors.insert(k, vec![0.01; q_dim]); + } + if let Some(k) = arch.attn_k_bias_key(layer) { + vectors.insert(k, vec![0.01; kv_dim]); + } + if let Some(k) = arch.attn_v_bias_key(layer) { + vectors.insert(k, vec![0.01; kv_dim]); + } + tensors.insert(arch.ffn_gate_key(layer), rand_mat(INTER, HIDDEN, 0.1)); + tensors.insert(arch.ffn_up_key(layer), rand_mat(INTER, HIDDEN, 0.1)); + tensors.insert(arch.ffn_down_key(layer), rand_mat(HIDDEN, INTER, 0.1)); + vectors.insert(arch.input_layernorm_key(layer), vec![1.0; HIDDEN]); + vectors.insert(arch.post_attention_layernorm_key(layer), vec![1.0; HIDDEN]); + } + + ModelWeights { + tensors, + vectors, + raw_bytes: HashMap::new(), + packed_mmaps: HashMap::new(), + skipped_tensors: Vec::new(), + packed_byte_ranges: HashMap::new(), + embed, + lm_head, + position_embed: None, + arch, + num_layers: NUM_LAYERS, + hidden_size: HIDDEN, + intermediate_size: INTER, + vocab_size: VOCAB, + logical_vocab_size: None, + head_dim: HEAD_DIM, + num_q_heads: NUM_Q, + num_kv_heads: NUM_KV, + rope_base: 10_000.0, + } +} + // ── Gemma 4 E2B-like synthetic fixture (PLE-aware) ── /// Tiny synthetic Gemma-4-E2B-shaped arch with PLE + KV sharing. @@ -476,6 +577,17 @@ pub fn synthetic_e2b_like_arch_json() -> serde_json::Value { }) } +/// ST4A canonical-loop fixture: identical E2B-like shape (4 layers +/// S/G/S/G, PLE, KV-shared, QK-norm, V-norm) but with a **512-token +/// intrinsic sliding window** instead of 4. Used by the canonical +/// prefill/decode equivalence tests at boundary positions 511, 512, +/// 513, 1024 — a 4-token window cannot exercise those positions. +pub fn synthetic_e2b_like_arch_json_window512() -> serde_json::Value { + let mut json = synthetic_e2b_like_arch_json(); + json["text_config"]["sliding_window"] = serde_json::json!(512); + json +} + /// Build minimal `ModelWeights` matching the synthetic E2B-like arch. /// Tensors zero-filled — fixture's job is to satisfy presence checks /// (PLE keys, KV-shared sources) so per-layer-embedding code paths fire. @@ -572,7 +684,20 @@ pub fn make_synthetic_e2b_like_weights() -> ModelWeights { /// Same 4-layer E2B-like shape: sliding_window=4, num_kv_shared_layers=2 /// (layers 2/3 are shared consumers of 0/1), PLE + QK-norm + V-norm. pub fn make_synthetic_e2b_like_weights_random() -> ModelWeights { - let arch = detect_from_json(&synthetic_e2b_like_arch_json()); + make_synthetic_e2b_like_weights_random_from_arch(synthetic_e2b_like_arch_json()) +} + +/// ST4A sibling of [`make_synthetic_e2b_like_weights_random`] using the +/// 512-token intrinsic window ([`synthetic_e2b_like_arch_json_window512`]). +/// Same 4-layer S/G/S/G shape with KV-shared consumers (2/3 ← 0/1), PLE, +/// QK-norm, V-norm, and distinct random per-layer weights. Sized so a +/// full 1025-token prefill through all four layers stays cheap on CPU. +pub fn make_synthetic_e2b_like_weights_random_window512() -> ModelWeights { + make_synthetic_e2b_like_weights_random_from_arch(synthetic_e2b_like_arch_json_window512()) +} + +fn make_synthetic_e2b_like_weights_random_from_arch(arch_json: serde_json::Value) -> ModelWeights { + let arch = detect_from_json(&arch_json); let num_layers = 4; let hidden = 8; let intermediate = 16;