From 69ae983281bd6f4bf8e4498203cf9e2685b7b1ff Mon Sep 17 00:00:00 2001 From: Feathbow Date: Tue, 4 Aug 2026 09:17:55 +0100 Subject: [PATCH] feat(gemma4): local decoder layer and its HF golden oracle Signed-off-by: Feathbow --- Cargo.lock | 1 + docs/models/gemma4/hf-golden.md | 4 +- pegainfer-core/src/ops.rs | 5 + pegainfer-gemma4/Cargo.toml | 3 +- pegainfer-gemma4/src/config.rs | 75 +++++++ pegainfer-gemma4/src/layer.rs | 287 ++++++++++++++++++++++++ pegainfer-gemma4/src/layer_oracle.rs | 176 +++++++++++++++ pegainfer-gemma4/src/lib.rs | 9 +- pegainfer-gemma4/src/manifest/schema.rs | 3 + pegainfer-gemma4/src/testkit.rs | 76 +++++++ pegainfer-gemma4/src/weights.rs | 38 ++-- pegainfer-gemma4/src/weights/load.rs | 89 +++++--- 12 files changed, 714 insertions(+), 52 deletions(-) create mode 100644 pegainfer-gemma4/src/layer.rs create mode 100644 pegainfer-gemma4/src/layer_oracle.rs create mode 100644 pegainfer-gemma4/src/testkit.rs diff --git a/Cargo.lock b/Cargo.lock index 83fb0a688..58f1f1d9e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3543,6 +3543,7 @@ version = "0.1.0" dependencies = [ "anyhow", "cudarc", + "half", "log", "pegainfer-core", "pegainfer-frontend", diff --git a/docs/models/gemma4/hf-golden.md b/docs/models/gemma4/hf-golden.md index 78d603469..74744695b 100644 --- a/docs/models/gemma4/hf-golden.md +++ b/docs/models/gemma4/hf-golden.md @@ -2,8 +2,8 @@ **TL;DR:** `test_data/gemma4-12b-hf-golden.safetensors` is the Hugging Face reference for Gemma 4 12B — 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. Nothing consumes it -yet; it exists so the layer and forward comparisons have something to compare against. +single-token, a nine-token and a 1024-token (exactly the sliding window) case. It is the reference +the in-crate golden gates replay their probes against. Last touched: 2026-08. diff --git a/pegainfer-core/src/ops.rs b/pegainfer-core/src/ops.rs index 3ca524e08..fd24e1618 100644 --- a/pegainfer-core/src/ops.rs +++ b/pegainfer-core/src/ops.rs @@ -30,6 +30,7 @@ pub use pegainfer_kernels::ops::eagle3_rope_into; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::embedding_batch; pub use pegainfer_kernels::ops::embedding_decode_into; +pub use pegainfer_kernels::ops::extract_hidden_rows_raw_into; pub use pegainfer_kernels::ops::extract_vec; pub use pegainfer_kernels::ops::extract_vec_into; pub use pegainfer_kernels::ops::extract_vec_ref; @@ -39,6 +40,7 @@ pub use pegainfer_kernels::ops::f32_to_bf16_hidden_into; pub use pegainfer_kernels::ops::fused_add_rms_norm_batch_into; pub use pegainfer_kernels::ops::fused_add_rms_norm_into; pub use pegainfer_kernels::ops::gather_hidden_tokens_into; +pub use pegainfer_kernels::ops::gelu_tanh_mul_batch_into; pub use pegainfer_kernels::ops::gemm; pub use pegainfer_kernels::ops::gemm_graphsafe_into_checked; pub use pegainfer_kernels::ops::gemm_graphsafe_ref_into_checked; @@ -62,6 +64,7 @@ pub use pegainfer_kernels::ops::pack_lora_b_rows_into; pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd256_into; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::qk_norm_rope_batch_decode_into; +pub use pegainfer_kernels::ops::qk_norm_rope_prefill_hd256_plain_into; pub use pegainfer_kernels::ops::rms_norm; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::rms_norm_batch_into; @@ -69,6 +72,7 @@ pub use pegainfer_kernels::ops::rms_norm_batch_offset_into; pub use pegainfer_kernels::ops::rms_norm_gated_batch_into; pub use pegainfer_kernels::ops::rms_norm_into; pub use pegainfer_kernels::ops::rms_norm_offset_into; +pub use pegainfer_kernels::ops::scale_bf16_in_place; pub use pegainfer_kernels::ops::scale_f32_in_place; pub use pegainfer_kernels::ops::scaled_add_batch_into; pub use pegainfer_kernels::ops::scaled_add_rows_indexed_into; @@ -79,6 +83,7 @@ pub use pegainfer_kernels::ops::silu_mul_batch_into; #[cfg(not(feature = "kernel-call-trace"))] pub use pegainfer_kernels::ops::silu_mul_fused_batch_into; pub use pegainfer_kernels::ops::single_decode_nhd_into; +pub use pegainfer_kernels::ops::single_prefill_hd256_into; pub use pegainfer_kernels::ops::single_prefill_nhd_causal_into; pub use pegainfer_kernels::ops::single_prefill_nhd_noncausal_into; pub use pegainfer_kernels::ops::write_vec_into; diff --git a/pegainfer-gemma4/Cargo.toml b/pegainfer-gemma4/Cargo.toml index a1bafc6a1..25df48fbe 100644 --- a/pegainfer-gemma4/Cargo.toml +++ b/pegainfer-gemma4/Cargo.toml @@ -7,12 +7,13 @@ name = "pegainfer-gemma4" version = "0.1.0" [features] -gemma4 = ["dep:cudarc", "dep:log", "dep:pegainfer-core"] +gemma4 = ["dep:cudarc", "dep:half", "dep:log", "dep:pegainfer-core"] default = [] [dependencies] anyhow = { workspace = true } cudarc = { workspace = true, optional = true } +half = { workspace = true, optional = true } log = { workspace = true, optional = true } pegainfer-core = { workspace = true, optional = true } pegainfer-frontend = { workspace = true } diff --git a/pegainfer-gemma4/src/config.rs b/pegainfer-gemma4/src/config.rs index 83374c691..0925559de 100644 --- a/pegainfer-gemma4/src/config.rs +++ b/pegainfer-gemma4/src/config.rs @@ -17,6 +17,19 @@ pub(crate) enum LayerKind { Global, } +/// First and last sliding layers from the layer map; `None` when the map has +/// no sliding entry. +#[cfg(test)] +pub(crate) fn first_last_sliding(layer_types: &[LayerKind]) -> Option<(usize, usize)> { + let mut sliding = layer_types + .iter() + .enumerate() + .filter(|(_, kind)| matches!(kind, LayerKind::Sliding)) + .map(|(index, _)| index); + let first = sliding.next()?; + Some((first, sliding.next_back().unwrap_or(first))) +} + /// What the manifest is derived from. Only [`Gemma4Config::from_file`] is /// probe-backed; a value built directly is not, so consumers check what they /// depend on. @@ -35,6 +48,15 @@ pub(crate) struct Gemma4Config { pub(crate) tie_word_embeddings: bool, /// The MoE size keeps its dense MLP and adds experts alongside it. pub(crate) moe_enabled: bool, + // Not manifest inputs: until a serving path lands, only the oracle reads + // these three. + #[allow(dead_code)] + pub(crate) rms_norm_eps: f32, + /// The sliding-attention rope theta; the global family reads its own. + #[allow(dead_code)] + pub(crate) sliding_rope_theta: f32, + #[allow(dead_code)] + pub(crate) sliding_window: usize, } #[cfg(feature = "gemma4")] @@ -70,6 +92,17 @@ impl Gemma4Config { "Gemma 4: layer_types has {} entries but num_hidden_layers is {num_hidden_layers}", layer_types.len() ); + let rope = tc + .get("rope_parameters") + .ok_or_else(|| anyhow::anyhow!("Gemma 4: missing text_config.rope_parameters"))?; + let sliding_rope = rope + .get("sliding_attention") + .ok_or_else(|| anyhow::anyhow!("Gemma 4: missing rope_parameters.sliding_attention"))?; + let sliding_window = usize_field(tc, "sliding_window")?; + anyhow::ensure!( + sliding_window > 0, + "Gemma 4: sliding_window must be positive" + ); Ok(Self { hidden_size: usize_field(tc, "hidden_size")?, intermediate_size: usize_field(tc, "intermediate_size")?, @@ -82,10 +115,30 @@ impl Gemma4Config { layer_types, tie_word_embeddings: bool_field(tc, "tie_word_embeddings")?, moe_enabled: bool_field(tc, "enable_moe_block")?, + rms_norm_eps: f32_field(tc, "text_config", "rms_norm_eps")?, + sliding_rope_theta: f32_field(sliding_rope, "sliding_attention", "rope_theta")?, + sliding_window, }) } } +/// Numeric config values land in f32 compute; the checked cast rejects +/// anything the narrowing would turn infinite rather than rounding it in +/// silently. +#[cfg(feature = "gemma4")] +fn f32_field(obj: &serde_json::Value, ctx: &str, field: &str) -> Result { + let value = obj + .get(field) + .and_then(serde_json::Value::as_f64) + .ok_or_else(|| anyhow::anyhow!("Gemma 4: {ctx}.{field} missing or not a number"))?; + let narrowed = value as f32; + anyhow::ensure!( + narrowed.is_finite(), + "Gemma 4: {ctx}.{field} = {value} overflows f32" + ); + Ok(narrowed) +} + #[cfg(feature = "gemma4")] fn usize_field(text_config: &serde_json::Value, field: &str) -> Result { let value = text_config @@ -105,3 +158,25 @@ fn bool_field(text_config: &serde_json::Value, field: &str) -> Result { .and_then(serde_json::Value::as_bool) .ok_or_else(|| anyhow::anyhow!("Gemma 4: text_config.{field} missing or not a boolean")) } + +#[cfg(test)] +mod layer_map_tests { + use super::*; + + #[test] + fn first_last_sliding_handles_edges() { + // The real 12B map is reconciled against the fixture's own parse in + // the oracle; here only the iterator logic is under test. + assert_eq!( + first_last_sliding(&[LayerKind::Sliding, LayerKind::Global, LayerKind::Sliding]), + Some((0, 2)) + ); + assert_eq!(first_last_sliding(&[LayerKind::Sliding]), Some((0, 0))); + assert_eq!( + first_last_sliding(&[LayerKind::Global, LayerKind::Sliding]), + Some((1, 1)) + ); + assert_eq!(first_last_sliding(&[LayerKind::Global]), None); + assert_eq!(first_last_sliding(&[]), None); + } +} diff --git a/pegainfer-gemma4/src/layer.rs b/pegainfer-gemma4/src/layer.rs new file mode 100644 index 000000000..b31857d43 --- /dev/null +++ b/pegainfer-gemma4/src/layer.rs @@ -0,0 +1,287 @@ +//! One Gemma 4 local (sliding-attention) decoder layer, prefill form. +//! +//! The graph is the HF reference's, with the constants a from-the-paper +//! implementation gets wrong: the four norm sites are all norm-then-add +//! (sandwich, not the fused add-then-norm shape), attention is unscaled +//! (`scaling = 1.0` in the reference — not `head_dim**-0.5`), V takes a +//! weightless RMS norm and no RoPE, RoPE rotates the full 256-wide head, +//! and `layer_scalar` multiplies the layer output after both residual adds. +//! +//! There is no KV cache: K and V stay contiguous and feed `single_prefill`, +//! which computes exact sliding attention for any prompt short enough that +//! the window never truncates (`seq_len <= sliding_window`; the window +//! first evicts at `sliding_window + 1` tokens). The window boundary +//! belongs to the SWA kernels and the prefill ladder, not to this layer. + +use anyhow::Context as _; +use anyhow::Result; +use half::bf16; +use pegainfer_core::ops; +use pegainfer_core::tensor::DeviceContext; +use pegainfer_core::tensor::DeviceVec; +use pegainfer_core::tensor::HiddenStates; + +use crate::weights::Gemma4Layer; + +/// The geometry a local layer runs at, read off the validated config. +pub(crate) struct LocalLayerGeometry { + pub(crate) hidden_size: usize, + pub(crate) intermediate_size: usize, + pub(crate) num_q_heads: usize, + pub(crate) num_kv_heads: usize, + pub(crate) head_dim: usize, + pub(crate) rms_norm_eps: f32, +} + +/// Token-major `[num_heads * head_dim, seq_len]` rows into a contiguous HND +/// cache of the same total size: `[head][pos][head_dim]`, `seq_len` rows per +/// head. Each head's column window is extracted to a `[head_dim, seq_len]` +/// block, then placed at its head-major offset. +fn nhd_to_hnd( + ctx: &DeviceContext, + src: &HiddenStates, + num_heads: usize, + head_dim: usize, +) -> Result { + let seq_len = src.seq_len; + anyhow::ensure!( + src.hidden_dim == num_heads * head_dim, + "nhd_to_hnd src.hidden_dim {} != num_heads {} * head_dim {}", + src.hidden_dim, + num_heads, + head_dim + ); + let mut hnd = HiddenStates::zeros(ctx, src.hidden_dim, seq_len)?; + let mut head_block = HiddenStates::zeros(ctx, head_dim, seq_len)?; + let block_len = head_dim * seq_len; + for head in 0..num_heads { + ops::extract_hidden_rows_raw_into( + ctx, + &src.data, + src.hidden_dim, + &mut head_block.data, + head_dim, + head * head_dim, + seq_len, + )?; + let src_view = head_block.data.slice(0..block_len); + let mut dst_view = hnd.data.slice_mut(head * block_len..(head + 1) * block_len); + ctx.stream + .memcpy_dtod(&src_view, &mut dst_view) + .map_err(|e| anyhow::anyhow!("nhd_to_hnd head {head} copy failed: {e}"))?; + } + Ok(hnd) +} + +/// Runs one local layer on `x` (`[hidden_size, seq_len]`), tokens at +/// positions `start_pos..start_pos + seq_len`. Buffers are allocated per +/// call: this is the correctness building block, and the executor that would +/// own persistent buffers does not exist yet. +#[allow(clippy::too_many_arguments)] +pub(crate) fn local_layer_forward( + ctx: &DeviceContext, + layer: &Gemma4Layer, + geom: &LocalLayerGeometry, + x: &HiddenStates, + start_pos: usize, + sliding_window: usize, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + cos_max_pos: usize, +) -> Result { + let seq_len = x.seq_len; + // Probe form: no KV cache exists, so history before start_pos would be + // silently missing, and past the window the full causal mask no longer + // equals sliding attention. Reject both instead of mis-computing. + anyhow::ensure!( + start_pos == 0, + "local_layer_forward is prefill-from-zero only; start_pos {start_pos} needs a KV cache" + ); + anyhow::ensure!( + seq_len <= sliding_window, + "local_layer_forward seq_len {seq_len} exceeds sliding_window {sliding_window}; the window \ + would truncate" + ); + let q_dim = geom.num_q_heads * geom.head_dim; + let kv_dim = geom.num_kv_heads * geom.head_dim; + anyhow::ensure!( + x.hidden_dim == geom.hidden_size, + "local layer x.hidden_dim {} != hidden_size {}", + x.hidden_dim, + geom.hidden_size + ); + anyhow::ensure!( + geom.head_dim == 256, + "local layer head_dim {} != 256, which the prep kernel is instantiated at", + geom.head_dim + ); + let v_proj = layer + .attention + .v_proj + .as_ref() + .context("local layer requires v_proj; only global layers ship without one")?; + + let mut normed_x = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::rms_norm_batch_into( + ctx, + x, + &layer.input_layernorm, + geom.rms_norm_eps, + &mut normed_x, + ); + + let mut q_states = HiddenStates::zeros(ctx, q_dim, seq_len)?; + let mut k_states = HiddenStates::zeros(ctx, kv_dim, seq_len)?; + let mut v_states = HiddenStates::zeros(ctx, kv_dim, seq_len)?; + ops::gemm_rows_into_checked( + ctx, + &layer.attention.q_proj, + 0, + q_dim, + &normed_x, + &mut q_states, + )?; + ops::gemm_rows_into_checked( + ctx, + &layer.attention.k_proj, + 0, + kv_dim, + &normed_x, + &mut k_states, + )?; + ops::gemm_rows_into_checked(ctx, v_proj, 0, kv_dim, &normed_x, &mut v_states)?; + + // Full-head rotation: Gemma 4 local layers have no partial factor, so + // rotary_dim is the head width. + let mut q_prep = HiddenStates::zeros(ctx, q_dim, seq_len)?; + let mut k_prep = HiddenStates::zeros(ctx, kv_dim, seq_len)?; + ops::qk_norm_rope_prefill_hd256_plain_into( + ctx, + &q_states, + &k_states, + &mut q_prep, + &mut k_prep, + &layer.attention.q_norm, + &layer.attention.k_norm, + cos_cache, + sin_cache, + start_pos, + cos_max_pos, + geom.num_q_heads, + geom.num_kv_heads, + geom.head_dim, + geom.rms_norm_eps, + )?; + + // v_norm is weightless (`with_scale=False`): a plain-w RMS norm with a + // ones weight is the same arithmetic. The `[kv_dim, seq_len]` buffer is + // reinterpreted as `[head_dim, seq_len * num_kv_heads]` — heads are + // contiguous within each token row, so the narrower row width makes the + // reduction per (token, head) — then the shape is restored. + let ones = DeviceVec::from_host(ctx, &vec![bf16::from_f32(1.0); geom.head_dim])?; + v_states.hidden_dim = geom.head_dim; + v_states.seq_len = seq_len * geom.num_kv_heads; + let mut v_normed = HiddenStates::zeros(ctx, geom.head_dim, seq_len * geom.num_kv_heads)?; + ops::rms_norm_batch_into(ctx, &v_states, &ones, geom.rms_norm_eps, &mut v_normed); + v_normed.hidden_dim = kv_dim; + v_normed.seq_len = seq_len; + + // single_prefill's contiguous cache is HND — k[head, pos, dim] — while + // the prep and the GEMMs emit token-major rows, so K and V are + // reassembled per head. Per-call copies; the executor owns avoiding + // this once it owns layouts. + let k_hnd = nhd_to_hnd(ctx, &k_prep, geom.num_kv_heads, geom.head_dim)?; + let v_hnd = nhd_to_hnd(ctx, &v_normed, geom.num_kv_heads, geom.head_dim)?; + + // Unscaled attention: the reference sets scaling = 1.0 for both layer + // kinds; sm_scale = rsqrt(head_dim) here would shrink logits to 1/16. + let mut attn = HiddenStates::zeros(ctx, q_dim, seq_len)?; + ops::single_prefill_hd256_into( + ctx, + &q_prep, + 0, + seq_len, + &k_hnd, + &v_hnd, + &mut attn, + geom.num_q_heads, + geom.num_kv_heads, + seq_len, + 1.0, + )?; + + let mut attn_proj = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::gemm_rows_into_checked( + ctx, + &layer.attention.o_proj, + 0, + geom.hidden_size, + &attn, + &mut attn_proj, + )?; + let mut o_normed = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::rms_norm_batch_into( + ctx, + &attn_proj, + &layer.post_attention_layernorm, + geom.rms_norm_eps, + &mut o_normed, + ); + let mut h2 = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::add_batch_into(ctx, x, &o_normed, &mut h2)?; + + let mut mlp_in = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::rms_norm_batch_into( + ctx, + &h2, + &layer.pre_feedforward_layernorm, + geom.rms_norm_eps, + &mut mlp_in, + ); + let mut gate = HiddenStates::zeros(ctx, geom.intermediate_size, seq_len)?; + let mut up = HiddenStates::zeros(ctx, geom.intermediate_size, seq_len)?; + ops::gemm_rows_into_checked( + ctx, + &layer.mlp.gate, + 0, + geom.intermediate_size, + &mlp_in, + &mut gate, + )?; + ops::gemm_rows_into_checked( + ctx, + &layer.mlp.up, + 0, + geom.intermediate_size, + &mlp_in, + &mut up, + )?; + let mut act = HiddenStates::zeros(ctx, geom.intermediate_size, seq_len)?; + ops::gelu_tanh_mul_batch_into(ctx, &gate, &up, &mut act)?; + let mut down = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::gemm_rows_into_checked(ctx, &layer.mlp.down, 0, geom.hidden_size, &act, &mut down)?; + let mut down_normed = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::rms_norm_batch_into( + ctx, + &down, + &layer.post_feedforward_layernorm, + geom.rms_norm_eps, + &mut down_normed, + ); + + let mut out = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; + ops::add_batch_into(ctx, &h2, &down_normed, &mut out)?; + + // layer_scalar multiplies the layer output after both residual adds — + // not either branch. + ops::scale_bf16_in_place(ctx, &mut out, layer.layer_scalar)?; + + Ok(out) +} + +/// The oracle: replay the HF golden fixture's local-layer probes through +/// this layer implementation on the real checkpoint. See +/// `docs/models/gemma4/hf-golden.md` for what the fixture pins. +#[cfg(test)] +#[path = "layer_oracle.rs"] +mod oracle; diff --git a/pegainfer-gemma4/src/layer_oracle.rs b/pegainfer-gemma4/src/layer_oracle.rs new file mode 100644 index 000000000..fc7d4e893 --- /dev/null +++ b/pegainfer-gemma4/src/layer_oracle.rs @@ -0,0 +1,176 @@ +use super::*; +use crate::config::Gemma4Config; +use crate::config::first_last_sliding; +use crate::testkit::GOLDEN_PATH; +use crate::testkit::METADATA_KEY; +use crate::testkit::assert_checkpoint_matches; +use crate::testkit::bf16_tensor; +use crate::testkit::model_path; +use crate::weights::Gemma4Weights; + +/// Declared against the measured error structure on the pinned 12B +/// checkpoint (sm_89): the one-token case is bitwise exact for both +/// probe layers, and the nine-token case shows scattered rounding noise +/// only — worst element 0.25 absolute, token 0 exact, errors appearing +/// once softmax runs over multiple keys. That is the unscaled-attention +/// signature (logits carry no rsqrt damping, so one-ulp bf16 GEMM +/// differences shift softmax weights), not a layout defect, which would +/// scramble whole head blocks by O(1). ABS_TOL leaves 1.6x headroom over +/// the measured worst; wiring-class bugs (a 16x scale error, swapped +/// weights) blow past it by orders of magnitude. +const ABS_TOL: f32 = 0.4; +const REL_TOL: f32 = 2e-2; + +/// Reports the full error structure before asserting, so a failure shows +/// whether it is scattered rounding noise or a structural pattern (whole +/// tokens or channel blocks off — the signature of a layout bug). +fn compare(got: &[f32], expected: &[bf16], hidden_size: usize, what: &str) -> usize { + assert_eq!(got.len(), expected.len(), "{what} length"); + let mut max_abs = 0.0f32; + let mut max_rel = 0.0f32; + let mut worst_idx = 0usize; + let mut violations = 0usize; + let mut per_token_max: Vec = vec![0.0; got.len().div_ceil(hidden_size)]; + for (i, (&g, &e)) in got.iter().zip(expected).enumerate() { + let e = e.to_f32(); + if !g.is_finite() || !e.is_finite() { + violations += 1; + per_token_max[i / hidden_size] = f32::INFINITY; + continue; + } + let abs = (g - e).abs(); + let rel = if e.abs() > 1e-3 { abs / e.abs() } else { 0.0 }; + if abs > ABS_TOL + REL_TOL * e.abs() { + violations += 1; + } + if abs > max_abs { + max_abs = abs; + worst_idx = i; + } + max_rel = max_rel.max(rel); + per_token_max[i / hidden_size] = per_token_max[i / hidden_size].max(abs); + } + eprintln!( + "{what}: max_abs {max_abs} at (token {}, ch {}), max_rel {max_rel}, \ + {violations}/{} over tolerance, per-token max_abs {per_token_max:?}", + worst_idx / hidden_size, + worst_idx % hidden_size, + got.len() + ); + violations +} + +#[test] +#[ignore = "requires the pinned 12B checkpoint via PEGAINFER_TEST_MODEL_PATH and a GPU"] +fn local_layer_matches_hf_probes() { + let dir = model_path(); + let fixture_bytes = std::fs::read(GOLDEN_PATH).expect("read fixture"); + let (_, meta) = + safetensors::SafeTensors::read_metadata(&fixture_bytes).expect("fixture metadata"); + let manifest: serde_json::Value = serde_json::from_str( + meta.metadata() + .as_ref() + .expect("fixture metadata map") + .get(METADATA_KEY) + .expect("gemma4_golden metadata key"), + ) + .expect("parse fixture manifest"); + assert_checkpoint_matches(&manifest, &dir); + let fixture = safetensors::SafeTensors::deserialize(&fixture_bytes).expect("parse fixture"); + + // The layer indices come from parsing the layer map; the fixture + // metadata records the dumper's own parse, so the two independent + // derivations must agree before anything numeric is asserted. + let config = Gemma4Config::from_file(&dir).expect("config"); + let (first, last) = + first_last_sliding(&config.layer_types).expect("12B carries sliding layers"); + let probe_layers = &manifest["probe_layers"]; + assert_eq!( + first as u64, + probe_layers["sliding_first"] + .as_u64() + .expect("sliding_first"), + "first sliding layer disagrees with the fixture's parse" + ); + assert_eq!( + last as u64, + probe_layers["sliding_last"].as_u64().expect("sliding_last"), + "last sliding layer disagrees with the fixture's parse" + ); + + let cut_labels: Vec = manifest["cut_labels"] + .as_array() + .expect("cut_labels") + .iter() + .map(|v| v.as_str().expect("cut label").to_string()) + .collect(); + let cut_index = |label: &str| { + cut_labels + .iter() + .position(|l| l == label) + .unwrap_or_else(|| panic!("cut {label} missing from fixture")) + }; + + let (weights, _) = Gemma4Weights::from_safetensors(&dir, 0).expect("load 12B weights"); + let ctx = DeviceContext::new_with_device(0).expect("device context"); + + let geom = LocalLayerGeometry { + hidden_size: config.hidden_size, + intermediate_size: config.intermediate_size, + num_q_heads: config.num_attention_heads, + num_kv_heads: config.num_key_value_heads, + head_dim: config.head_dim, + rms_norm_eps: config.rms_norm_eps, + }; + let cos_max_pos = 16; + let (cos_cache, sin_cache) = pegainfer_core::rope::precompute_rope( + &ctx, + &pegainfer_core::rope::RopeTableSpec { + rotary_dim: geom.head_dim, + frequency_dim: geom.head_dim, + max_seq_len: cos_max_pos, + theta: config.sliding_rope_theta, + }, + ) + .expect("rope tables"); + + let mut over_tolerance: Vec = Vec::new(); + for case in ["single", "short"] { + let (shape, hidden) = bf16_tensor(&fixture, &format!("{case}_hidden")); + assert_eq!(shape.len(), 3, "{case}_hidden rank"); + assert_eq!(shape[0], cut_labels.len(), "{case}_hidden cut count"); + let (seq_len, hidden_size) = (shape[1], shape[2]); + assert_eq!(hidden_size, geom.hidden_size, "{case}_hidden width"); + let cut = |label: &str| { + let i = cut_index(label); + &hidden[i * seq_len * hidden_size..(i + 1) * seq_len * hidden_size] + }; + + for (name, index) in [("sliding_first", first), ("sliding_last", last)] { + let x = HiddenStates::from_host(&ctx, cut(&format!("{name}_in")), hidden_size, seq_len) + .expect("x H2D"); + let out = local_layer_forward( + &ctx, + &weights.layers[index], + &geom, + &x, + 0, + config.sliding_window, + &cos_cache, + &sin_cache, + cos_max_pos, + ) + .expect("layer forward"); + let got = out.to_host(&ctx).expect("out D2H"); + let expected = cut(&format!("{name}_out")); + let violations = compare(&got, expected, hidden_size, &format!("{case}/{name}")); + if violations > 0 { + over_tolerance.push(format!("{case}/{name} (layer {index})")); + } + } + } + assert!( + over_tolerance.is_empty(), + "comparisons over tolerance: {over_tolerance:?}" + ); +} diff --git a/pegainfer-gemma4/src/lib.rs b/pegainfer-gemma4/src/lib.rs index 8d6014e4b..3e657fc66 100644 --- a/pegainfer-gemma4/src/lib.rs +++ b/pegainfer-gemma4/src/lib.rs @@ -6,8 +6,15 @@ mod config; mod manifest; pub mod model_line; mod probe; +// Dead until an executor calls in; test targets do use them, so an +// `expect(dead_code)` cannot hold in every build. #[cfg(feature = "gemma4")] -#[expect(dead_code, reason = "no consumer until the executor lands")] +#[allow(dead_code)] +mod layer; +#[cfg(all(feature = "gemma4", test))] +mod testkit; +#[cfg(feature = "gemma4")] +#[allow(dead_code)] mod weights; use std::path::Path; diff --git a/pegainfer-gemma4/src/manifest/schema.rs b/pegainfer-gemma4/src/manifest/schema.rs index dfb046913..1ad88c463 100644 --- a/pegainfer-gemma4/src/manifest/schema.rs +++ b/pegainfer-gemma4/src/manifest/schema.rs @@ -271,6 +271,9 @@ pub(super) fn sample_config() -> Gemma4Config { layer_types: vec![LayerKind::Sliding, LayerKind::Sliding, LayerKind::Global], tie_word_embeddings: true, moe_enabled: false, + rms_norm_eps: 1e-6, + sliding_rope_theta: 10_000.0, + sliding_window: 1024, } } diff --git a/pegainfer-gemma4/src/testkit.rs b/pegainfer-gemma4/src/testkit.rs new file mode 100644 index 000000000..b2579caff --- /dev/null +++ b/pegainfer-gemma4/src/testkit.rs @@ -0,0 +1,76 @@ +//! Shared plumbing for the in-crate checkpoint oracles: the golden fixture, +//! its provenance checks, and typed tensor readers. + +use half::bf16; +use sha2::Digest; +use sha2::Sha256; + +pub(crate) const GOLDEN_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../test_data/gemma4-12b-hf-golden.safetensors" +); +pub(crate) const METADATA_KEY: &str = "gemma4_golden"; + +pub(crate) fn model_path() -> String { + std::env::var("PEGAINFER_TEST_MODEL_PATH").expect( + "PEGAINFER_TEST_MODEL_PATH must point at the pinned 12B Gemma 4 \ + checkpoint the fixture was dumped from", + ) +} + +pub(crate) fn sha256_hex(bytes: &[u8]) -> String { + use std::fmt::Write; + Sha256::digest(bytes) + .iter() + .fold(String::with_capacity(64), |mut hex, byte| { + let _ = write!(hex, "{byte:02x}"); + hex + }) +} + +/// Mirrors the dumper: plain sha256 for the two config files, and the +/// sha256 of the safetensors header (8-byte LE length prefix) for the +/// weights file, which pins the tensor layout without reading 22 GiB. +pub(crate) fn assert_checkpoint_matches(manifest: &serde_json::Value, dir: &str) { + let expected = manifest["file_sha256"] + .as_object() + .expect("manifest file_sha256"); + for (name, digest) in expected { + let expected_hex = digest.as_str().expect("sha256 must be a string"); + let actual = if let Some(file) = name.strip_suffix("#header") { + use std::io::Read as _; + let mut handle = + std::fs::File::open(std::path::Path::new(dir).join(file)).expect("open weights"); + let mut len_bytes = [0u8; 8]; + handle.read_exact(&mut len_bytes).expect("header length"); + let mut header = vec![0u8; u64::from_le_bytes(len_bytes) as usize]; + handle.read_exact(&mut header).expect("header bytes"); + sha256_hex(&header) + } else { + let bytes = + std::fs::read(std::path::Path::new(dir).join(name)).expect("read config file"); + sha256_hex(&bytes) + }; + assert_eq!( + &actual, expected_hex, + "{name} does not match the fixture's pinned checkpoint; this \ + oracle runs against that checkpoint only" + ); + } +} + +pub(crate) fn bf16_tensor( + fixture: &safetensors::SafeTensors<'_>, + name: &str, +) -> (Vec, Vec) { + let view = fixture.tensor(name).expect("fixture tensor"); + assert_eq!(view.dtype(), safetensors::Dtype::BF16, "{name} dtype"); + let host = view + .data() + .as_chunks::<2>() + .0 + .iter() + .map(|b| bf16::from_bits(u16::from_le_bytes(*b))) + .collect(); + (view.shape().to_vec(), host) +} diff --git a/pegainfer-gemma4/src/weights.rs b/pegainfer-gemma4/src/weights.rs index f7ddda6ce..3423983fe 100644 --- a/pegainfer-gemma4/src/weights.rs +++ b/pegainfer-gemma4/src/weights.rs @@ -11,31 +11,35 @@ pub(crate) struct Gemma4Weights { config: Gemma4Config, embed_tokens: DeviceMatrix, norm: DeviceVec, - layers: Vec, + pub(crate) layers: Vec, } pub(crate) struct Gemma4Layer { - input_layernorm: DeviceVec, - post_attention_layernorm: DeviceVec, - pre_feedforward_layernorm: DeviceVec, - post_feedforward_layernorm: DeviceVec, - layer_scalar: DeviceVec, - attention: Gemma4Attention, - mlp: Gemma4Mlp, + pub(crate) input_layernorm: DeviceVec, + pub(crate) post_attention_layernorm: DeviceVec, + pub(crate) pre_feedforward_layernorm: DeviceVec, + pub(crate) post_feedforward_layernorm: DeviceVec, + /// The per-layer output multiplier, read to the host at load: it is a + /// one-element constant consumed as a kernel scalar every step, and a + /// device read here would put a synchronous D2H in every layer of every + /// decode token. + pub(crate) layer_scalar: f32, + pub(crate) attention: Gemma4Attention, + pub(crate) mlp: Gemma4Mlp, } pub(crate) struct Gemma4Attention { - q_proj: DeviceMatrix, - k_proj: DeviceMatrix, + pub(crate) q_proj: DeviceMatrix, + pub(crate) k_proj: DeviceMatrix, /// Absent on global layers, which the checkpoint ships without one. - v_proj: Option, - o_proj: DeviceMatrix, - q_norm: DeviceVec, - k_norm: DeviceVec, + pub(crate) v_proj: Option, + pub(crate) o_proj: DeviceMatrix, + pub(crate) q_norm: DeviceVec, + pub(crate) k_norm: DeviceVec, } pub(crate) struct Gemma4Mlp { - gate: DeviceMatrix, - up: DeviceMatrix, - down: DeviceMatrix, + pub(crate) gate: DeviceMatrix, + pub(crate) up: DeviceMatrix, + pub(crate) down: DeviceMatrix, } diff --git a/pegainfer-gemma4/src/weights/load.rs b/pegainfer-gemma4/src/weights/load.rs index 041afb509..8e54f4e50 100644 --- a/pegainfer-gemma4/src/weights/load.rs +++ b/pegainfer-gemma4/src/weights/load.rs @@ -29,7 +29,7 @@ use crate::manifest::validate::ObservedTensor; /// redemption, prefetch join and unmap fall between them, and the allocations /// submitted under `record_api_wall_ms` execute under /// `execute_and_drain_wall_ms`. Only `elapsed_ms` is a total. -struct LoadStats { +pub(crate) struct LoadStats { /// Every required tensor at its dtype. manifest_bytes: usize, /// Free-before minus free-after. Signed: this measures the device, not the @@ -55,7 +55,7 @@ struct LayerSlots { post_attention_layernorm: VecSlotId, pre_feedforward_layernorm: VecSlotId, post_feedforward_layernorm: VecSlotId, - layer_scalar: VecSlotId, + layer_scalar: f32, q_proj: SlotId, k_proj: SlotId, v_proj: Option, @@ -114,7 +114,30 @@ fn classify_checkpoint(manifest: &Manifest, shards: &[SafeTensors]) -> Result Result { +fn read_scalar_bf16(shards: &[SafeTensors], name: &str) -> Result { + for shard in shards { + if let Ok(view) = shard.tensor(name) { + anyhow::ensure!( + view.dtype() == Dtype::BF16 && view.data().len() == 2, + "Gemma 4: '{name}' must be a single bf16 scalar" + ); + let bits = u16::from_le_bytes([view.data()[0], view.data()[1]]); + let value = half::bf16::from_bits(bits).to_f32(); + anyhow::ensure!( + value.is_finite(), + "Gemma 4: '{name}' = {value} is not finite" + ); + return Ok(value); + } + } + anyhow::bail!("Gemma 4: tensor '{name}' missing from every shard") +} + +fn record_plan( + loader: &mut StagedWeightLoader, + shards: &[SafeTensors], + manifest: &Manifest, +) -> Result { let embed_tokens = record_matrix(loader, &manifest.embed_tokens)?; let norm = record_vector(loader, &manifest.norm)?; let mut layers = Vec::with_capacity(manifest.layers.len()); @@ -125,7 +148,7 @@ fn record_plan(loader: &mut StagedWeightLoader, manifest: &Manifest) -> Result Gemma4Weights { - Gemma4Weights { +) -> Result { + Ok(Gemma4Weights { embed_tokens: loader.take(plan.embed_tokens), norm: loader.take_vec(plan.norm), layers: plan .layers .into_iter() - .map(|slots| Gemma4Layer { - input_layernorm: loader.take_vec(slots.input_layernorm), - post_attention_layernorm: loader.take_vec(slots.post_attention_layernorm), - pre_feedforward_layernorm: loader.take_vec(slots.pre_feedforward_layernorm), - post_feedforward_layernorm: loader.take_vec(slots.post_feedforward_layernorm), - layer_scalar: loader.take_vec(slots.layer_scalar), - attention: Gemma4Attention { - q_proj: loader.take(slots.q_proj), - k_proj: loader.take(slots.k_proj), - v_proj: slots.v_proj.map(|slot| loader.take(slot)), - o_proj: loader.take(slots.o_proj), - q_norm: loader.take_vec(slots.q_norm), - k_norm: loader.take_vec(slots.k_norm), - }, - mlp: Gemma4Mlp { - gate: loader.take(slots.gate), - up: loader.take(slots.up), - down: loader.take(slots.down), - }, + .map(|slots| -> Result { + Ok(Gemma4Layer { + input_layernorm: loader.take_vec(slots.input_layernorm), + post_attention_layernorm: loader.take_vec(slots.post_attention_layernorm), + pre_feedforward_layernorm: loader.take_vec(slots.pre_feedforward_layernorm), + post_feedforward_layernorm: loader.take_vec(slots.post_feedforward_layernorm), + layer_scalar: slots.layer_scalar, + attention: Gemma4Attention { + q_proj: loader.take(slots.q_proj), + k_proj: loader.take(slots.k_proj), + v_proj: slots.v_proj.map(|slot| loader.take(slot)), + o_proj: loader.take(slots.o_proj), + q_norm: loader.take_vec(slots.q_norm), + k_norm: loader.take_vec(slots.k_norm), + }, + mlp: Gemma4Mlp { + gate: loader.take(slots.gate), + up: loader.take(slots.up), + down: loader.take(slots.down), + }, + }) }) - .collect(), + .collect::>>()?, config, - } + }) } fn free_device_bytes() -> Result { @@ -203,7 +228,10 @@ impl Gemma4Weights { /// Loads the text tower onto one device. Config, manifest and headers are /// all checked before a device context exists, so a checkpoint that does /// not match its config costs no GPU. - fn from_safetensors(model_path: &str, device_ordinal: usize) -> Result<(Self, LoadStats)> { + pub(crate) fn from_safetensors( + model_path: &str, + device_ordinal: usize, + ) -> Result<(Self, LoadStats)> { let started = Instant::now(); let config = Gemma4Config::from_file(model_path)?; let manifest = Manifest::from_config(&config)?; @@ -221,14 +249,14 @@ impl Gemma4Weights { let mut loader = StagedWeightLoader::new(&ctx, &shards, &weight_map)?; let recording = Instant::now(); - let plan = record_plan(&mut loader, &manifest)?; + let plan = record_plan(&mut loader, &shards, &manifest)?; let record_api_wall_ms = elapsed_ms(recording); let uploading = Instant::now(); loader.finish()?; let execute_and_drain_wall_ms = elapsed_ms(uploading); - let weights = materialize(&mut loader, plan, config); + let weights = materialize(&mut loader, plan, config)?; drop(loader); drop(prefetch); let device_free_bytes = free_device_bytes()?; @@ -309,7 +337,6 @@ mod tests { } assert_eq!(attention.q_proj.cols, config.hidden_size); assert_eq!(attention.o_proj.rows, config.hidden_size); - assert_eq!(layer.layer_scalar.len, 1); } assert_eq!(weights.embed_tokens.rows, config.vocab_size); assert_eq!(weights.embed_tokens.cols, config.hidden_size);