feat(inference): prove Gemma 4 F32 first-token parity (ST5) - #63
Merged
Conversation
Add the ST5 first-token semantic-parity harness and prove LARQL's canonical F32 CPU forward matches the pinned Transformers CPU float32 eager oracle for all four ST3 prompts (first_divergence: null, top-1 exact, top-10 overlap >= 9/10, logit nrmse ~1e-6). Trace capture + interchange format (larql-inference::parity): - format.rs: deterministic manifest binding each tensor to (prompt, stage, layer, shape, dtype, element_count, sha256) with raw LE f32 files; rejects missing/duplicate/shape/dtype/truncated/hash/non-finite. - compare.rs: committed numerical policy (coarse nrmse<=1e-4, cosine >=0.99999; final logits + top-1 exact + top-10 overlap>=9/10) with first-divergence drill-down in execution order; f64 accumulation. - capture.rs: runs the production forward via forward_raw_logits_traced and writes the last-token row of every coarse boundary. Substrate hooks (larql-compute): LayerHook gains read-only on_post_ffn / on_post_ple boundary taps (zero-cost no-op defaults); forward_raw_logits_traced reuses the canonical embed/layer/ple/norm/ lm-head primitives. Oracle (scripts/gemma4_first_token_oracle.py): forward-only, eager, torch.float32, CPU, use_cache=false, never generate(); non-invasive pre-hook capture. Proportional-RoPE correction (the trace-proven layer-4 first divergence): Gemma 4 global layers ship HF rope_type='proportional' — inv_freq exponents over the full head_dim, zero-padded to head_dim/2, half-split over the full head. LARQL divided by rotary_dim and half-split within it. Add RopeFreqMode and apply it across the CPU block, decode, and kv-prefill (gpu.rs CPU fallback) paths. Sliding layers are unaffected, so ST4/ST4A local/global + shared-KV regressions stay green. Tests: 9 synthetic diagnostic self-tests (ST5 section 7) + proportional regression tests; env-gated gemma4_first_token_parity soft-skips CI.
The ST5 proportional-RoPE change added the rope_freq_mode_for_layer trait default to config.rs; without a caller it dropped config.rs line coverage to 89.47%, tripping the per-file 90% gate. Exercise the default (Standard) via the generic-architecture trait-contract test.
The ST5 traced forward added forward_raw_logits_traced + TracedTail to predict/raw.rs; exercised only by the larql-inference parity tests, it dropped raw.rs line coverage to 77.52%, tripping the per-file 90% gate. Add a lib unit test that runs the traced forward with a RecordHook and asserts every per-layer boundary + tail stage matches forward_raw_logits (same math).
The ST5 test/format additions were committed before a final per-package rustfmt pass; cargo fmt --all and cargo fmt -p <crate> resolved a few assert/call-chain wrappings differently. Reformat so every per-package fmt check (the CI gate) passes alongside the workspace check.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
LARQL-INFERENCE-TRUST-001A-ST5 — F32 CPU first-token semantic parity
Decision: GREEN — LARQL's canonical F32 CPU forward matches the pinned Transformers CPU float32 eager oracle for all four ST3 prompts.
Evidence
All four prompts:
first_divergence: null, top-1 exact, top-10 overlap 10/10, logit nrmse ~1e-6 (float32 reduction-order noise):What this PR adds
larql-inference::parity— trace interchange format (deterministic manifest + raw LE f32 tensors), the committed numerical comparator (coarse nrmse≤1e-4 / cosine≥0.99999; final logits + top-1 exact + top-10 overlap≥9/10), first-divergence drill-down in execution order, and LARQL F32 production-path capture.larql-computehooks — read-onlyon_post_ffn/on_post_pleboundary taps (zero-cost no-op defaults) +forward_raw_logits_tracedreusing the canonical primitives.scripts/gemma4_first_token_oracle.py— forward-only, eager,torch.float32, CPU,use_cache=false, nevergenerate(); non-invasive pre-hook capture.gemma4_first_token_paritytest (soft-skips CI without the 18.65 GB artifact).Correction made (trace-proven layer-4 first divergence)
Gemma 4 global (full_attention) layers ship HF
rope_type='proportional': RoPEinv_freqexponents are divided by the full head_dim (512) and zero-padded tohead_dim/2, then half-split over the full head. LARQL divided byrotary_dim(128) and half-split within it — so the first global-attention layer (layer 4) diverged (cosine ~0.9997) while sliding layers (0–3) matched. Fixed by addingRopeFreqMode::{Standard, Proportional}and applying it consistently across the CPU block / decode / kv-prefill (gpu.rs CPU fallback) paths. ST4/ST4A local/global + shared-KV regressions remain green.Validation
cargo fmt --all -- --check— cleancargo test -p {larql-models,larql-compute,larql-inference,larql-kv,larql-vindex}— all pass, 0 failurescargo clippy -p {larql-models,larql-compute,larql-inference,larql-kv,larql-vindex} --all-targets -- -D warnings— cleancargo build -p larql-cli --release— passesScope exclusions
multi-token generation, sampling, KV-cached decode, Q4_K, CUDA, Vulkan, Metal, performance optimization, production quantization, multimodal inputs, tools/thinking-enabled prompts.
Recommended next slice
LARQL-INFERENCE-TRUST-001A-ST6 — Production Q4_K semantic parity against the proven F32 reference.