From d0d0d7c723bcebb018374fe7471a9d0ccc73c951 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Wed, 19 Aug 2026 12:17:42 +0100 Subject: [PATCH 1/2] test(gemma4): long context waypoints answer to the reference Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 125 ++++++++++++++++ .../gemma4-12b-hf-longctx-golden.safetensors | Bin 0 -> 206992 bytes tools/accuracy/dump_gemma4_longctx_golden.py | 138 ++++++++++++++++++ 3 files changed, 263 insertions(+) create mode 100644 test_data/gemma4-12b-hf-longctx-golden.safetensors create mode 100644 tools/accuracy/dump_gemma4_longctx_golden.py diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index 2f7c01fa..f49b4147 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/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, 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" + ); + 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 0000000000000000000000000000000000000000..3dc8827e2a37108979b11e4cfa836572bff1cdc0 GIT binary patch literal 206992 zcmeI5dvI0Ny~kG!RYQGHp*oE)#?(htBCh~Rrqp7q)_W^jnz4@68v{9DiX=1%Qt?tj ztuH9jir|B{pizQ`_d~+@B}ATq5FX)?Kwdy3ylg>5Tj74cXP;fNTied`{&Q#U*O~9E zwb$NfpY{8$-(F{Z_6p;V+quYpPbOw&4i6P(WoH#;Wo9O(Kbbf*G<EFl*3bq3Gd4mejqilRtLH3BOgsfb2U=*$mCFB$)49O`h=$a7u zX0H?Zp~rIya`JM$+xG3=cS!fFUV~DzdJXQImEAkFPbj;4k5JFPp{yZ;va(aNdiKrE z?mnci_t}EHk@)sZ#F^)Pd}vpcWt!8nFWQR5fSH785R1S zIgFp}lM*@1&&n-8NW(+z|c zgZkSj-;b)qO0QpHrSqWBGMtNQ-g309#_<}ouSFd?v<}*Y`ZiB@)L7|9&>eelycgHi zS?M<#a18yn5!XVYCY*x~HRF2@TIs8Y(bi(6Bac|=r_j7sE4}nM+M%+ORyyvF*6*@Y z)^GMHD?NJJO0W6UO0R`RLi3-2)AUa-hef z8e{1bpNC$BrdaxA)1Z0K+c*_2$7kVGsdI0C^QBn;!*J$Y{paD%dhE71a2T?+owctL~ z$2flkYPIw=(12r>{$J4h&~J|8+T*tP;;rGu7axEQhaYU$*B)&50Q4A?XUR{Gv*Zh& zgMNqlBxo}9Dl`?EhW6Ol5bytmeKk~y`WlpB zr~>EK;dq0mCX`Qa#62o4IdKclLr+(EO5BEPpz>;*-wxH_n%$N>2wJoU*X)JrEO`#J z2-*()3A&;lpM&N=HlW^Q@%7MchtL*P_%N<(LHiN3LDw9$_?Mx}Pgwj1CoTS}(-xlx zg`wL%wfJtITl~p07XQi@7QgvRT>BM{p+7@Up2ab=`>cg8oE{EexHSC8GhOYGXYPTr zA}T3+#!C7>4~@5yvMDI1p`3|wHZ;dd#?6C@(DtsC#LP$A0vwm1z7Sf3cBubiXa%$q zS`C$=J&f}f)moHL)&`v0h;vm|(g?*=TgmM0D0iaVg|-?i`PM#M2mPwnN^aV3B{|S+ zXdzT~z)Gs2jrCUYod#S39Xn_x&mXdqyAR|1$5!(C5u9tqXQ3WYe<-WfX0N&-eBZ;T z!bR_2Z$Ym?MM*6{eJQjI^))ENkfC0V za;;@twI13CRibS>+ICpR)Lp38puQI+bXA>YEZ>jz1D4UZ!7@H?w2V>BXoE%_LfL{g zsPD(P23ijN{0Ppq;T42RrK5NCJpND=2O~G+gqo<<&CNv%Gvrx`OInRn$%|}^+EJ9K3IkNHk8%4 z2Ko|8-;U#`_U%ABl)e)s^d*#DgF3WoH;$u9-)qHhK&|_5zgjDPsm_Ww?YBF>njgON zt5e~(&VI+r;3(+EJ@_ni7xWaAxfjQEHnXxk zJhSpJbRvA;kY0A*kb9v_=+~CCc)TU`nu3y=iu&u8R4^Tyjbo_bZItsZ>D2ouOK`jp z*DOZ46m84VwgShiP_DM5j_c44EnaU)uWZ2aMx3KIL0izi73Zr^-)2caggRDR(lgsp z?(|fE@)c-&jU^4+ZAm+z+C7%^^L>_71zlZhNegSO=#Dy1BkJshX?2z~xZaZ1Ll-sR z-i`Py^d@u)ifyvocJsrz?UqC9!<{kbbuN3=I`^M~l9~;@h2wWTl@;N5KI#jgg{Uv_ z)PFh3(x}2Hsq(1STIZMw>pTuBqta+GU5wXZG}YoBoCZxteJ09xEbfi>EG~Wlj-fY7P%c5a6z5i;eI@E2qFjyRGMuYG z`H{t~T8FmvC?U*sap{|(N{hpM7dLtfJ_D_S9^49TgSO*3=;oa`55@1o_0W4z`x=XD zjSBI{weGPvtbgLRL$}vj+@Ly(yK=w9rPkYvuhiR3R&R0b8Z9mX8rbA%Kr`-j(0(xQ z*Wp+9<)(=@tLR!_gT>rCHVH2H=u+v8gZ@(*EgdMO=fw>Q_5j0 zDuhnA;QYr{lmeZG3XhQRQj7=DEUf zhaZ-QU${2SUbyyW&`@X$^c$<_@RC(bdKsFG`V43$>a$SJMpUp+!g6O+yR4#nja7_*I_$KSAUAY#=Pz}p>qt(Ml|XjI26tvfUca>wxW^T7CIG9 zKl`7Sel`ah2K~mmJoP;EJL_`p_s}0upNjJJsNO(X6jiZx2`@x_k#)IkIkX0KgUZmp z7Uep8cC)8*l{miz^{ptYa1IJ@vo6=QD9=9%e$L-mw*TTqc)`+!FM5j=mS_ZAe z`Bf;P#UG+vjq__zu8nFlN~mB9KED;(hGS|6>buZhV~tqrGxuJgEv=cfX43DYejok0QEVj^)$gNzAH|R2N4=Meb(#w`7iuolT&THFbD{Q{ zoYP#WxlnVV=0eSdnhX8E*Wez4-XZH9vfd$UkJZlsy*JT&6S0C=L98HF5G#lk#0vU5 zcjkcBg<2PCU8r@T)`gl2wJy}UQ0qdi3$-rPzhmX!ooijFb)nXUS{G_vsJT$%!psV*Q ziXZjwbo?9;KL+PV@uT=r{3w1DKZ+m4kK#x1V{m;WeiT3Y?4bQe?LUeigY%>KQT!-= z6hDd|#gF1g@uT=LxIPj;iXX*~+JE$OK>Qe-AH|R2NAaWhQT!-=6hDd|#gD=Dk@!*k z=(B_NAGQA|ehki!;z#kL_)+{QeiT27AH|R2$Kd)%{3w1DKWhKc&jImcaDEg&iXX*~ z;z#kL_)+{QeiT0j*GJ+<@uSZU+JDskqxdm6KZ+m4kK#x1qxez$D1H<_iXVgPBk`m7 zQT(X=M?VL|kHPs-{3w1DKZ+m4kK#x1qxez$7+fETAH|P8J81t=`;X$s;QT0l6hDd| z#gF1g@uT=r{3w15u8+iz;z#kL_8 z`s|?nN9{j~AA|Fw_)+{QeiT27AH|R2NAaWhF}OYwKZ+m4kJ^9qb3ptUoFB!H;z#kL z_)+{QeiT27AH|Qs^^y2d{OGfT_8+zXD1HpikK#x1qxez$D1H<_iXX*~|JV4jX0r zu-}Y_{hEPx`olQK-s}1H#4Rr31wEmkjF#^Gkcdck=3U5t^j(i3`D zmhSmSPuMq(CF7!3elAtv+!Z)xzB5k7LGS1l$ByerJ`c@f;U~tA&kvs!=03e*JoJt{ zPH*^3ejA^mNA%_|knaI=pU+F%*y#o1rU#6L{bt;Z(;rXf06n2cjGKNi4#vvdV9wAB zddIlfe|iw{3C|+qqDSlxJ!EY3ik`CnjF}$X@{#wfM#dR_G4DB+^o{YS;+i;!zA{EW z@3*5L%!8?Drw@#qa|JzS&Yv6&pDu;28OelWgqpT`8|363xO#<5{vhQTZ541Hq1=<|z+ zb0z%YxY8fy8)G?v^Yn*f?av9E2Xi3C$?>N*9Dn{!&|}Uge1_;F=Tbf&9Miwee~!8D z1LODOW{ixT@zD=@z;R?AaJ>C-W~}rkGA8hwW68J}C;ea_85{kfCybYVF;+gq%nSO# zoMIp8A7kfp!MTt=(I0xk=YcWM$H@GMk~x>3>&4Hxj(N|#V_qP?jWx%f zK5%>)Kl6dH(g*g9G1Cji!TezC^nl~R_~{MDo-xr2ei!4WC;T4z!T!=4&d-dWUeG`K z%im}|&m1@U#^;nV#i33=82^e3@oeLma~;R{UiiVd_`GnO>F*T~y-B&^+u9Z`&##a44#x9BIiZ~3nptZmt(o+0`|scT>z%^+`BC!(_O`GWq?}Mr=slC(GilAF zxln5+t(o+99pXo=3-x|g@6yh5f!?orUMMG&6MAo=_a<#wQGeU1zwOlTqkbR#zs(Ro ziXX*~;zzN9SV62HRuC(E{lCG~zWe$4QS*f63C$ClC)(~m>z#t$Dd?So{>_h{1A3>R zcM5u^aDL6K-$(sE`g5b+x$B*~ejoMwsNYBZKI&bX-lgeXn%<@9U7GiAIodu8`hC>z zqkbRt``C6)(BH4>?^pFMP4Ck5E=_w)eh%nen%<>}6~qc+1+juyLGRM^E=}*!m;+iD zYF(&xq1J_37iuolx=`yvtqZj-)VeVE|C`i)n)cJQpQimZe=kLUo36i2*LxGaH_>|& zy*JT&6TLUldlS7k(R&lUH_>|&y*D8nYF(&xq1J_37iwLoxlrputqZj-)VfgX!r=R2 z?Wbu!P5Wv3_i26(Xg@7DKZ+m4kK#x1qxez$D1H<_iXVgPBk`m7(Psz!J*NI1Q~Vg5 zAH|R2NAaWhQT!-=6hDd|#gD=Dk@!*kD1Oxbqn`uf$Kd=ZeiT27AH|R2NAaWhQT!-= z46cvFkK#w49kliXXNA=;wg=F*rYpAH|R2NAaWhQT!-=6hDd| zgX<&lqxjKh2kk#<|55xHoFB!H;z#kL_)+{QeiT27AH|Qs^^y2d{3w3Z{-d7*;>Y0p zD1H<_iXX*~;z#kL_)+{QehjXU#E;@fpB=RSsQpLrV{m>HKZ+m4kK#x1qxez$D1H<_ z2G>X8NAaWhQTva64u~Iv^P~7t{3w1DKZ+m4kK#x1qxdnnJ`z8QAANSv{-gFE#gD=H zQT!-=6hDd|#gF1g@uT=r{1{vxi66y};z#X2`Z*we49<_@NAaWhQT!-=6hDd|#gF30 z;QC1XD1P+WLHm!|e-u9k=ST6Q_)+{QeiT27AH|R2NAY8DeI$MqKZ+l<|LEs{_%S#? ziXX*~;z#kL_)+{QeiT27AA{>7@uT?BX9w*+YX4FE7@Qx)kK#x1qxez$D1H<_iXX*~ z!S#{&QT!-=)c&KN1LDWv{3w1DKZ+m4kK#x1qxez$D1HpCkHnASN1q+E|ET>(@ndj) z6hDd|#gF1g@uT=r{3w1DKL*!F;z#kL_)+_heh!EqgY%>KQT!-=6hDd|#gF1g@uT=L zxIPj;iXVM;(Eg+LAH|Qs`BD5ReiT27AH|R2NAaWhQT!NOABi8ukK#w|Kl(W!ehki! z;z#kL_)+{QeiT27AH|R2$Kd)%{3w3(*+Kh{+J6*32Ioidqxez$D1H<_iXX*~;z#jg zaD60x6hDd|wg2enfcP;uKZ+m4kK#x1qxez$D1H<_iXVgPBk`m7(PszkKWhI`{1}`c z#gF1g@uT=r{3w3>|GJJ-;HI@-UB>jpvxaPACr z0y+o%1lN3p@1NVxxs55#J#r7~xPCY^1ARJy&m6&Zm!KTh&bhJoI`^$>oqM9Ua|_UR z*H!p1u1$h|3ROGju0wnIoz8W`=kH$_>!#v(MpKNt7dXfG89O~-+`b2l zlO8a3_L)BLyXXnyq8E&TKF}NX*Y|?`VvOu7J)#%%fL<^r_H}qi&nw1CZ|EO=poff^ zelUi`_0}l?)tYcaO-ir0pmj- zi!+_8LFtb@W2YC4hcUC?z8AbF`{R4U_!vKZVa$x5zOe7~gmE(VPBXoJ(i{56IO!c@ zUU8vwOD}Pe+@Uv&o&GX+O5gVUh{I=6p=O->Yc62??u92$AdZuxJ#RVg^oemZ=G)N^ zjxl{MOLwj}?$5krUfhUl=oNFpL5!KP`#vysj+Z~?>^sMjInRDFKIR2IVBGY`AGb01 zUalROWAuz;L?7rCpQ~Lc+1F~6jFCRDAM6{)i_a1J#^>NP^!@Ih=Nu>cd=0LxMMWukji1~8~Zch*CP5PxS;uGWN*fBSLgg)%SXXtU$DDPSE<8C|t9B2B#_?QRu zg7Gj0_Jj9h{LDXk!SUwUGFIj;zn|m9++*zQZzM0!Mn6xYjq%e9<^%gqAK6#-jd{e} zVE%9n86)H3e0&Pm^EqJ7(|6_q$DGdrbDPf$bL$YUrJu}o<`Kt`eqV^T7PNf=mEb%7 zHh#v@Z#>{-DKYUg>R`h`Vp;wHTeWz#a7sr== zF2bD1$QW#PKUQNOr z$e7tz`tohWcmb}XKYT_Q7rkKoe6Hvb{i1KpW4&jD@pB&Fb94p77@6brhVgNpV?J@* z=^>vj#!g?D5B~W3^MF6*?4R$29}nYX{Pclwvd@f}v9YgyPS7vL%{b@}pKr#&e*4cO z`^P+@Z|oyIpg;7LW6W{n_&wLdn@3(i-hML~rxzr)L~<&UN&TIb!i%jxaCi2fd|FjDeo;x#>5^%Lo1j^7&x?$6tke vz#Gn~WAL3EYkEo#ImYy61J)9=@x4 \ + --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() From 761c7dc03ef7b8cd7cb130c41dbb69f52c89e96c Mon Sep 17 00:00:00 2001 From: Feathbow Date: Wed, 19 Aug 2026 21:51:41 +0100 Subject: [PATCH 2/2] test(gemma4): the longctx fixture shares the reference pin Signed-off-by: Feathbow --- docs/index.md | 2 +- docs/models/gemma4/hf-golden.md | 37 +++++++++++++++--- pegainfer-gemma4/src/serve_oracle.rs | 4 ++ .../gemma4-12b-hf-longctx-golden.safetensors | Bin 206992 -> 206992 bytes 4 files changed, 37 insertions(+), 6 deletions(-) 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 f49b4147..c9037ae5 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -265,6 +265,10 @@ fn longctx_waypoints_match_hf() { 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) = diff --git a/test_data/gemma4-12b-hf-longctx-golden.safetensors b/test_data/gemma4-12b-hf-longctx-golden.safetensors index 3dc8827e2a37108979b11e4cfa836572bff1cdc0..712c7b7750643136ab917e62b08dd9b96223d8ce 100644 GIT binary patch delta 30 kcmbPmk!Qk1o()$SnGN*}nr|?+-(Un`rtLQvnRloG0KM}Ig8%>k delta 30 kcmbPmk!Qk1o()$SnN9Qzn{P0--(Un`rtLQvnRloG0KP#BhX4Qo