test(gemma4): long context waypoints answer to the reference

pegainfer-gemma4/src/serve_oracle.rs

@@ -200,6 +200,131 @@ fn score_rows(
     (max_abs, top1)
 }
 
+const LONGCTX_FIXTURE: &str = concat!(
+    env!("CARGO_MANIFEST_DIR"),
+    "/../test_data/gemma4-12b-hf-longctx-golden.safetensors"
+);
+
+/// A case's sdpa rows with a borrowed tolerance: where eager could not fit
+/// next to the tower, the widest dual-backend case lends its floor — the
+/// reference says what agreement is reachable, not that less is correct.
+fn reference_sdpa_only(
+    fixture: &safetensors::SafeTensors<'_>,
+    case: &str,
+    tolerance: f32,
+    backend_top1_share: f64,
+) -> (Vec<i32>, Vec<f32>, usize, usize, f32, usize) {
+    let (shape, ids) = i32_tensor(fixture, &format!("{case}_sdpa_ids"));
+    let (_, lps) = f32_tensor(fixture, &format!("{case}_sdpa_logprobs"));
+    let (positions, top_k) = (shape[0], shape[1]);
+    #[allow(clippy::cast_sign_loss, clippy::cast_possible_truncation)]
+    let backend_top1 = (backend_top1_share * positions as f64).floor() as usize;
+    (ids, lps, positions, top_k, tolerance, backend_top1)
+}
+
+/// The raised ceiling's numeric waypoints: 16384 and 32768 teacher-forced
+/// against the Hugging Face reference — proportional rope and global
+/// attention far past the window fixture's 4096 — each gated at twice its
+/// own backends' measured gap, the widest dual-backend floor standing in
+/// where eager could not fit. The widest case runs again in 2048-token
+/// chunks, the raised ceiling's production prefill shape.
+#[test]
+#[ignore = "requires the pinned 12B checkpoint and the longctx fixture"]
+fn longctx_waypoints_match_hf() {
+    // A 32776-token prefill holds every local page at once (append then
+    // attend); pools sized for one request plus padding.
+    let (ctx, serve, dir) = stack_with(32900, 2200);
+    let bytes = std::fs::read(LONGCTX_FIXTURE).expect("read longctx fixture");
+    let golden = fixture_manifest(
+        &std::fs::read(GOLDEN_PATH).expect("read golden"),
+        METADATA_KEY,
+    );
+    assert_checkpoint_matches(&golden, &dir);
+    let manifest = fixture_manifest(&bytes, "gemma4_longctx_golden");
+    assert_eq!(
+        manifest["revision"], golden["revision"],
+        "the longctx fixture was dumped from a different revision than the golden one"
+    );
+    let eager_skipped: Vec<String> = manifest["eager_skipped"]
+        .as_array()
+        .expect("eager_skipped list")
+        .iter()
+        .map(|v| v.as_str().expect("case name").to_string())
+        .collect();
+    let fixture = safetensors::SafeTensors::deserialize(&bytes).expect("parse fixture");
+    let page = serve.local_pool.layout().page_size;
+
+    // Neither waypoint fits eager next to the tower on this device, so no
+    // in-fixture dual-backend floor exists; the window fixture's deepest
+    // dual case lends its own — the widest measured agreement bound
+    // available. A depth-grown gap past it fails loud and is widened only
+    // with a written justification.
+    let window_bytes = std::fs::read(WINDOW_FIXTURE).expect("read window fixture");
+    let window_manifest = fixture_manifest(&window_bytes, "gemma4_window_golden");
+    assert_eq!(
+        window_manifest["revision"], golden["revision"],
+        "the window fixture was dumped from a different revision than the golden one"
+    );
+    let window_fixture =
+        safetensors::SafeTensors::deserialize(&window_bytes).expect("parse window fixture");
+    let (_, _, donor_positions, _, donor_tolerance, donor_top1) =
+        reference(&window_fixture, "w4096");
+    #[allow(clippy::cast_precision_loss)]
+    let donor_share = donor_top1 as f64 / donor_positions as f64;
+
+    let mut over: Vec<String> = Vec::new();
+    for (case, chunk) in [("w16384", 0), ("w32768", 0), ("w32768", 2048)] {
+        let label = if chunk == 0 {
+            case.to_string()
+        } else {
+            format!("{case}-chunked")
+        };
+        let (ref_ids, ref_lps, positions, top_k, tolerance, backend_top1) =
+            if eager_skipped.contains(&case.to_string()) {
+                reference_sdpa_only(&fixture, case, donor_tolerance, donor_share)
+            } else {
+                reference(&fixture, case)
+            };
+        let run = run_case(&ctx, &serve, &fixture, case, chunk);
+        assert_eq!(run.rows.len(), positions, "{label}: fixture positions");
+        assert_eq!(
+            chunk > 0,
+            run.shifted_multi_token,
+            "{label}: a shifted multi-token step is exactly what chunking adds"
+        );
+
+        let (max_abs, top1) = score_rows(&run.rows, &ref_ids, &ref_lps, top_k, &label);
+        eprintln!(
+            "{label}: max |dlogprob| {max_abs} (tol {tolerance:.2}), top-1 {top1}/{positions} \
+             (backend bar {backend_top1}/{positions}), local pages {}, global {}",
+            run.local_pages, run.global_pages
+        );
+        assert!(
+            top1 >= backend_top1,
+            "{label}: top-1 {top1}/{positions} below the backends' own \
+             {backend_top1}/{positions}"
+        );
+        let released = run.kv_len.saturating_sub(serve.sliding_window) / page;
+        assert_eq!(
+            run.local_pages,
+            run.kv_len.div_ceil(page) - released,
+            "{label}: resident pages after {released} released"
+        );
+        assert_eq!(
+            run.global_pages,
+            run.kv_len.div_ceil(page),
+            "{label}: the global family must keep every page"
+        );
+        if max_abs > tolerance {
+            over.push(format!("{label} ({max_abs} > {tolerance})"));
+        }
+    }
+    assert!(
+        over.is_empty(),
+        "cases over their calibrated floor: {over:?}"
+    );
+}
+
 /// Gated distribution-level because a greedy chain is not reachable at this
 /// depth: the reference's own backends continue the same prompt in different
 /// directions, so each case is gated at twice its measured sdpa-vs-eager

tools/accuracy/dump_gemma4_longctx_golden.py

@@ -0,0 +1,138 @@
+"""Dump the Gemma 4 long-context logprob reference.
+
+    python tools/accuracy/dump_gemma4_longctx_golden.py <model-dir> <out.safetensors> \
+        --source-repo google/gemma-4-12B-it \
+        --revision <same revision as the window fixture>
+
+Two waypoints far past the sliding window (16384, 32768), each followed by 8
+teacher-forced continuation tokens from the corpus. Every position records
+the top-64 ids and logprobs under both sdpa and eager; eager materialises the
+full attention matrix per layer, so a waypoint where it cannot fit is
+recorded sdpa-only and named in the manifest — the consumer borrows the
+widest dual-backend floor instead.
+
+Refuses to write rather than record something unusable: non-finite scores at
+any position abort the dump.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+
+import torch
+import torch.nn.functional as F
+from safetensors.torch import save_file
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+TOP_K = 64
+TEACHER_STEPS = 8
+METADATA_KEY = "gemma4_longctx_golden"
+CASES = {"w16384": 16384, "w32768": 32768}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model_dir")
+    parser.add_argument("out")
+    parser.add_argument("--source-repo", required=True)
+    parser.add_argument("--revision", required=True)
+    parser.add_argument("--device", default="cuda:0")
+    return parser.parse_args()
+
+
+def corpus_ids(tokenizer, min_len: int) -> list[int]:
+    rows = "  <tr>\n    <td>item</td>\n  </tr>\n" * 7000
+    ids = tokenizer(f"<table>\n{rows}", return_tensors="pt").input_ids[0].tolist()
+    if len(ids) < min_len:
+        raise SystemExit(f"corpus tokenizes to {len(ids)} < {min_len} tokens")
+    return ids
+
+
+def main() -> None:
+    args = parse_args()
+    tokenizer = AutoTokenizer.from_pretrained(args.model_dir)
+    ids = corpus_ids(tokenizer, max(CASES.values()) + TEACHER_STEPS)
+
+    tensors: dict[str, torch.Tensor] = {}
+    eager_skipped: list[str] = []
+    for impl in ("sdpa", "eager"):
+        model = AutoModelForCausalLM.from_pretrained(
+            args.model_dir, dtype=torch.bfloat16, attn_implementation=impl
+        ).to(args.device)
+        model.eval()
+        for case, length in CASES.items():
+            seq = torch.tensor([ids[: length + TEACHER_STEPS]], device=args.device)
+
+            def rows_of() -> torch.Tensor:
+                # Keep only the recorded positions' logits: the full
+                # 262k-vocab head over 32k positions is ~17 GiB in bf16 and
+                # does not fit next to the 22 GiB tower.
+                with torch.no_grad():
+                    return model(seq, logits_to_keep=TEACHER_STEPS + 1).logits[0].float()
+
+            try:
+                rows = rows_of()
+            except torch.cuda.OutOfMemoryError:
+                if impl == "sdpa":
+                    raise SystemExit(f"case {case}: sdpa itself cannot fit — no reference")
+                eager_skipped.append(case)
+                torch.cuda.empty_cache()
+                print(f"eager {case}: does not fit, recorded sdpa-only")
+                continue
+            if not torch.equal(rows, rows_of()):
+                raise SystemExit(f"case {case} ({impl}): forward is not reproducible")
+            if not torch.isfinite(rows).all():
+                raise SystemExit(f"case {case} ({impl}): non-finite logits")
+            logprobs = F.log_softmax(rows, dim=-1)
+            top = logprobs.topk(TOP_K, dim=-1)
+            tensors[f"{case}_{impl}_ids"] = top.indices.to(torch.int32).cpu()
+            tensors[f"{case}_{impl}_logprobs"] = top.values.cpu()
+            print(f"{impl} {case}: recorded {rows.shape[0]} positions x top-{TOP_K}")
+        del model
+        torch.cuda.empty_cache()
+
+    for case, length in CASES.items():
+        tensors[f"{case}_prompt"] = torch.tensor(ids[:length], dtype=torch.int32)
+        tensors[f"{case}_teacher"] = torch.tensor(
+            ids[length : length + TEACHER_STEPS], dtype=torch.int32
+        )
+        if case in eager_skipped:
+            continue
+        # The consumer's floor: how far the two backends sit apart on the
+        # recorded positions, evaluated at the sdpa top-64 ids.
+        sdpa_ids = tensors[f"{case}_sdpa_ids"]
+        sdpa_lp = tensors[f"{case}_sdpa_logprobs"]
+        eager_ids = tensors[f"{case}_eager_ids"]
+        eager_lp = tensors[f"{case}_eager_logprobs"]
+        floor = 0.0
+        agree = 0
+        for pos in range(sdpa_ids.shape[0]):
+            eager_map = {
+                int(t): float(v) for t, v in zip(eager_ids[pos].tolist(), eager_lp[pos].tolist())
+            }
+            for t, v in zip(sdpa_ids[pos].tolist(), sdpa_lp[pos].tolist()):
+                if int(t) in eager_map:
+                    floor = max(floor, abs(float(v) - eager_map[int(t)]))
+            agree += int(sdpa_ids[pos][0] == eager_ids[pos][0])
+        print(
+            f"case {case}: sdpa-vs-eager floor max|dlogprob| {floor:.2f}, "
+            f"top-1 agree {agree}/{sdpa_ids.shape[0]}"
+        )
+
+    manifest = {
+        "source_repo": args.source_repo,
+        "revision": args.revision,
+        "transformers": __import__("transformers").__version__,
+        "cases": CASES,
+        "teacher_steps": TEACHER_STEPS,
+        "top_k": TOP_K,
+        "eager_skipped": sorted(eager_skipped),
+        "corpus": "an unnumbered <table> row repeated, trimmed per case",
+        "reference": "teacher-forced top-64 logprobs under sdpa and, where it fits, eager",
+    }
+    save_file(tensors, args.out, metadata={METADATA_KEY: json.dumps(manifest, sort_keys=True)})
+
+
+if __name__ == "__main__":
+    main()