diff --git a/docs/index.md b/docs/index.md index 0a6e0a12..784497f1 100644 --- a/docs/index.md +++ b/docs/index.md @@ -65,7 +65,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | --- | --- | | `models/gemma4/tokenizer.md` | Gemma 4 tokenizer/chat-template contracts, gated against a Hugging Face reference (token ids plus all five chat renders, content flattened to strings): BOS comes only from the standalone `chat_template.jinja` (which opens a thought channel and accepts a native system role), EOS is declared three times with three values, the published defaults are sampled rather than greedy, image/audio tokens encode straight from user text so text-only serving must reject them at admission, and one divergence stays open — the server's default content format adds a trailing space to system turns. | | `models/gemma4/serving.md` | What the engine promises under load: iteration level scheduling with a ceiling on prompt plus output (8192 by default), prompts prefilled whole at a step boundary by default, requests beyond the configured decode slots (16 by default) queued rather than refused, and the two KV families budgeted separately (7.27 GiB sliding + 2.00 GiB global at 12B with the chunk knob off; the sliding budget shrinks to window plus segment under it). The default configuration needs a 48 GiB card: 32.2 GiB resident before the first request. Clients must send `` themselves or the model degenerates. A row is bit-identical when its companions' content and lengths change under a fixed batch width trajectory, and moves when arrivals or retirements change that trajectory, so greedy output is reproducible for a workload rather than across workloads. An opt-in conversation prefix cache (`PEGAINFER_PREFIX_CACHE=K`) resumes multi-turn prompts from captured prompt state at a pre-allocated page cost. An opt-in chunked walk (`PEGAINFER_MIX_CHUNK_TOKENS=N`) walks admissions through shared segment steps with round-by-round page reservation, and `PEGAINFER_MAX_CONTEXT` raises the ceiling to the checkpoint's 262144 while `PEGAINFER_DECODE_SLOTS` trades concurrency for the memory that buys. Open: no cross-request prefix sharing, single GPU. | -| `models/gemma4/hf-golden.md` | Two Hugging Face references for 12B. The base fixture carries layer-boundary activations at both ends of both layer types plus top-64 logprobs, over a single-token, a nine-token and a 1024-token (exactly the sliding window) case, and pins three facts the forward path has to match — the embedding scale is bf16 62.0 rather than `sqrt(3840)`, text attention is causal, and `layer_scalar` applies to the layer output after both residual adds. The window fixture goes past the window (1023/1024/1025/4096, teacher-forced) under both sdpa and eager. Regeneration is byte-identical and checked with sha256. | +| `models/gemma4/hf-golden.md` | Three Hugging Face references for 12B. The base fixture carries layer-boundary activations at both ends of both layer types plus top-64 logprobs, over a single-token, a nine-token and a 1024-token (exactly the sliding window) case, and pins three facts the forward path has to match — the embedding scale is bf16 62.0 rather than `sqrt(3840)`, text attention is causal, and `layer_scalar` applies to the layer output after both residual adds. The window fixture goes past the window (1023/1024/1025/4096, teacher-forced) under both sdpa and eager. The long-context fixture takes the same comparison to 16384/32768 for the raised ceiling, sdpa-only, with the window fixture's dual-backend floor on loan. Regeneration is byte-identical and checked with sha256. | ## models / glm52 diff --git a/docs/models/gemma4/hf-golden.md b/docs/models/gemma4/hf-golden.md index 6c5446be..c2cf41da 100644 --- a/docs/models/gemma4/hf-golden.md +++ b/docs/models/gemma4/hf-golden.md @@ -1,10 +1,11 @@ # Gemma 4 HF golden fixtures -**TL;DR:** two Hugging Face references for Gemma 4 12B. `test_data/gemma4-12b-hf-golden.safetensors` +**TL;DR:** three Hugging Face references for Gemma 4 12B. `test_data/gemma4-12b-hf-golden.safetensors` covers the window and everything below it — layer-boundary activations at both ends of both layer types, plus top-64 logprobs, over a single-token, a nine-token and a 1024-token case. `test_data/gemma4-12b-hf-window-golden.safetensors` goes past it, recorded under both attention -backends. +backends. `test_data/gemma4-12b-hf-longctx-golden.safetensors` takes the same teacher-forced +comparison to 16384 and 32768 tokens for the raised serving ceiling. Last touched: 2026-08. @@ -87,6 +88,25 @@ Per case: `{case}_prompt` and `{case}_teacher` (int32), plus `{case}_sdpa_ids` / `{case}_sdpa_logprobs` and `{case}_eager_ids` / `{case}_eager_logprobs` (`[9, 64]`, int32 and fp32). The name follows `qwen35-*-hf-long-golden`: one model line, a second context régime. +## The long-context fixture + +`gemma4-12b-hf-longctx-golden.safetensors` extends the window fixture's question to the raised +serving ceiling: 16384- and 32768-token prompts from the same corpus, each followed by eight +teacher-forced continuation tokens, with the top-64 ids and logprobs recorded at the last prompt +position and after each forced token — proportional RoPE and global attention far past the window +fixture's 4096. + +At these depths eager does not fit next to the reference tower on the dump device, so both cases +record sdpa alone and the manifest names them in `eager_skipped`. A case without its own +dual-backend pair borrows the window fixture's deepest dual case (`w4096`) for its tolerance and +top-1 floor — the widest measured agreement bound available. That loan is only meaningful if both +fixtures were dumped under the same reference release, so the gate requires the two manifests to +name the same Transformers version, on top of the same checkpoint revision. The gate also runs the +widest case a second time in 2048-token chunks, the raised ceiling's production prefill shape. + +Per case: `{case}_prompt` and `{case}_teacher` (int32), plus `{case}_sdpa_ids` / +`{case}_sdpa_logprobs` (`[9, 64]`, int32 and fp32). + ## Regenerating ```bash @@ -97,6 +117,10 @@ python tools/accuracy/dump_gemma4_hf_golden.py \ python tools/accuracy/dump_gemma4_window_golden.py \ test_data/gemma4-12b-hf-window-golden.safetensors \ --source-repo google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 + +python tools/accuracy/dump_gemma4_longctx_golden.py \ + test_data/gemma4-12b-hf-longctx-golden.safetensors \ + --source-repo google/gemma-4-12B-it --revision 707f0a3b8a3c7ad586ed01e27eafbad8a27dd0f7 ``` Two runs against the same checkpoint produce the same bytes, so regeneration is checked with @@ -106,6 +130,7 @@ Two runs against the same checkpoint produce the same bytes, so regeneration is | --- | --- | | `gemma4-12b-hf-golden.safetensors` | `c30a338d499512e6f0505bd12b184ebb5af9d7536f0b7fc9ea2bdfdb18b1a46d` | | `gemma4-12b-hf-window-golden.safetensors` | `b72edd51a5977592f3d4b637152aaf794a33e356c9599ab14035dacbb9574c0e` | +| `gemma4-12b-hf-longctx-golden.safetensors` | `1c8442a51913f858af6bc7205bd85c5b67db91373d7bdbdc034bf1ce2e86889d` | That only holds because the metadata is a **single sorted-JSON key**. safetensors serializes its metadata map in randomized order, so a multi-key block makes two runs differ byte for byte while @@ -115,9 +140,11 @@ is not. Provenance is passed in, not inferred: a checkpoint directory carries no record of where it came from. The base fixture's metadata records sha256 of `config.json`, `generation_config.json` and the safetensors *header* — the header pins the tensor layout without reading 22 GiB of payload, the -revision pins the payload. The window fixture records only the source repo and revision: its gate -validates the running checkpoint against the base fixture's hashes, then requires the two fixtures -to name the same revision, so one set of hashes covers both. **Transformers 5.11.0** is verified to load `gemma4_unified`; the +revision pins the payload. The window and long-context fixtures record the source repo, the +revision and the Transformers release: their gates validate the running checkpoint against the base +fixture's hashes, then require every fixture to name the same revision, so one set of hashes covers +all three. The long-context gate additionally requires its Transformers release to match the window +fixture's, because it borrows that fixture's floor. **Transformers 5.11.0** is verified to load `gemma4_unified`; the checkpoint declares `5.10.0.dev0`, a development build that was never released, so the pin is the release that was tested rather than a guess at what that build became. diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index 2f7c01fa..c9037ae5 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -200,6 +200,135 @@ 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, Vec, 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 = 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" + ); + assert_eq!( + manifest["transformers"], window_manifest["transformers"], + "a borrowed floor is only meaningful under the donor's own reference release" + ); + 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 = 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 diff --git a/test_data/gemma4-12b-hf-longctx-golden.safetensors b/test_data/gemma4-12b-hf-longctx-golden.safetensors new file mode 100644 index 00000000..712c7b77 Binary files /dev/null and b/test_data/gemma4-12b-hf-longctx-golden.safetensors differ diff --git a/tools/accuracy/dump_gemma4_longctx_golden.py b/tools/accuracy/dump_gemma4_longctx_golden.py new file mode 100644 index 00000000..aff039fb --- /dev/null +++ b/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 \ + --source-repo google/gemma-4-12B-it \ + --revision + +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 = " \n item\n \n" * 7000 + ids = tokenizer(f"\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
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()