From d033bd1cb966a0b67c23e67a179b679577c568b7 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Tue, 4 Aug 2026 11:37:21 +0100 Subject: [PATCH] feat(gemma4): global decoder layer and its HF golden oracle probes Signed-off-by: Feathbow --- pegainfer-core/src/ops.rs | 2 + pegainfer-gemma4/src/config.rs | 93 ++++-- pegainfer-gemma4/src/layer.rs | 388 +++++++++++++++++++----- pegainfer-gemma4/src/layer_oracle.rs | 68 ++++- pegainfer-gemma4/src/manifest/schema.rs | 2 + 5 files changed, 446 insertions(+), 107 deletions(-) diff --git a/pegainfer-core/src/ops.rs b/pegainfer-core/src/ops.rs index fd24e161..dd316046 100644 --- a/pegainfer-core/src/ops.rs +++ b/pegainfer-core/src/ops.rs @@ -62,6 +62,7 @@ pub use pegainfer_kernels::ops::lora_decode_fused_delta_group3_into; pub use pegainfer_kernels::ops::lora_decode_fused_delta_into; pub use pegainfer_kernels::ops::pack_lora_b_rows_into; pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd256_into; +pub use pegainfer_kernels::ops::qk_norm_partial_rope_batched_decode_hd512_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; @@ -84,6 +85,7 @@ pub use pegainfer_kernels::ops::silu_mul_batch_into; 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_hd512_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/src/config.rs b/pegainfer-gemma4/src/config.rs index 0925559d..1e8dfdd5 100644 --- a/pegainfer-gemma4/src/config.rs +++ b/pegainfer-gemma4/src/config.rs @@ -17,17 +17,15 @@ 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 +#[cfg(all(feature = "gemma4", test))] +pub(crate) fn first_last(layer_types: &[LayerKind], kind: LayerKind) -> Option<(usize, usize)> { + let mut matching = layer_types .iter() .enumerate() - .filter(|(_, kind)| matches!(kind, LayerKind::Sliding)) + .filter(|(_, k)| **k == kind) .map(|(index, _)| index); - let first = sliding.next()?; - Some((first, sliding.next_back().unwrap_or(first))) + let first = matching.next()?; + Some((first, matching.next_back().unwrap_or(first))) } /// What the manifest is derived from. Only [`Gemma4Config::from_file`] is @@ -49,7 +47,7 @@ pub(crate) struct Gemma4Config { /// 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. + // these. #[allow(dead_code)] pub(crate) rms_norm_eps: f32, /// The sliding-attention rope theta; the global family reads its own. @@ -57,6 +55,13 @@ pub(crate) struct Gemma4Config { pub(crate) sliding_rope_theta: f32, #[allow(dead_code)] pub(crate) sliding_window: usize, + #[allow(dead_code)] + pub(crate) global_rope_theta: f32, + /// `partial_rotary_factor * global_head_dim`, validated to land on a + /// positive even width within the head — the active band of the + /// proportional rope tables. + #[allow(dead_code)] + pub(crate) global_rotary_dim: usize, } #[cfg(feature = "gemma4")] @@ -98,6 +103,36 @@ impl Gemma4Config { let sliding_rope = rope .get("sliding_attention") .ok_or_else(|| anyhow::anyhow!("Gemma 4: missing rope_parameters.sliding_attention"))?; + let global_rope = rope + .get("full_attention") + .ok_or_else(|| anyhow::anyhow!("Gemma 4: missing rope_parameters.full_attention"))?; + rope_type_field(sliding_rope, "sliding_attention", "default")?; + rope_type_field(global_rope, "full_attention", "proportional")?; + let sliding_rope_theta = f32_field(sliding_rope, "sliding_attention", "rope_theta")?; + let global_rope_theta = f32_field(global_rope, "full_attention", "rope_theta")?; + anyhow::ensure!( + sliding_rope_theta > 0.0 && global_rope_theta > 0.0, + "Gemma 4: rope_theta must be positive (sliding {sliding_rope_theta}, global \ + {global_rope_theta})" + ); + let global_head_dim = usize_field(tc, "global_head_dim")?; + let partial = global_rope + .get("partial_rotary_factor") + .and_then(serde_json::Value::as_f64) + .ok_or_else(|| { + anyhow::anyhow!( + "Gemma 4: full_attention.partial_rotary_factor missing or not a number" + ) + })?; + let rotary = partial * global_head_dim as f64; + anyhow::ensure!( + rotary > 0.0 + && rotary.fract() == 0.0 + && rotary as usize <= global_head_dim + && (rotary as usize).is_multiple_of(2), + "Gemma 4: partial_rotary_factor {partial} of global_head_dim {global_head_dim} must \ + land on a positive even width within the head, got {rotary}" + ); let sliding_window = usize_field(tc, "sliding_window")?; anyhow::ensure!( sliding_window > 0, @@ -111,13 +146,15 @@ impl Gemma4Config { num_key_value_heads: usize_field(tc, "num_key_value_heads")?, num_global_key_value_heads: usize_field(tc, "num_global_key_value_heads")?, head_dim: usize_field(tc, "head_dim")?, - global_head_dim: usize_field(tc, "global_head_dim")?, + global_head_dim, 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_rope_theta, sliding_window, + global_rope_theta, + global_rotary_dim: rotary as usize, }) } } @@ -159,24 +196,18 @@ fn bool_field(text_config: &serde_json::Value, field: &str) -> Result { .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); - } +/// `rope_type` selects the table-generation algorithm; a value this engine +/// has not wired for that family must fail here, not silently get the other +/// family's tables. +#[cfg(feature = "gemma4")] +fn rope_type_field(rope_group: &serde_json::Value, ctx: &str, implemented: &str) -> Result<()> { + let value = rope_group + .get("rope_type") + .and_then(serde_json::Value::as_str) + .ok_or_else(|| anyhow::anyhow!("Gemma 4: {ctx}.rope_type missing or not a string"))?; + anyhow::ensure!( + value == implemented, + "Gemma 4: {ctx}.rope_type {value:?} is not the implemented {implemented:?}" + ); + Ok(()) } diff --git a/pegainfer-gemma4/src/layer.rs b/pegainfer-gemma4/src/layer.rs index b31857d4..97d9b98f 100644 --- a/pegainfer-gemma4/src/layer.rs +++ b/pegainfer-gemma4/src/layer.rs @@ -1,17 +1,18 @@ -//! One Gemma 4 local (sliding-attention) decoder layer, prefill form. +//! Gemma 4 decoder layers, prefill form: the local (sliding-attention) and +//! global (full-attention) kinds, sharing one attention epilogue. //! //! 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. +//! weightless RMS norm and no RoPE (on global layers V is `k_proj`'s raw +//! output, forked before `k_norm`), RoPE rotates the full head width (the +//! global family's partiality lives in its tables), 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. +//! exact from position zero — for the local kind only while the prompt fits +//! the sliding window. Both forwards reject anything outside that domain. use anyhow::Context as _; use anyhow::Result; @@ -23,8 +24,9 @@ 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 { +/// The geometry a layer runs at, read off the validated config — the local +/// and global kinds differ only in head width and KV head count. +pub(crate) struct LayerGeometry { pub(crate) hidden_size: usize, pub(crate) intermediate_size: usize, pub(crate) num_q_heads: usize, @@ -33,6 +35,69 @@ pub(crate) struct LocalLayerGeometry { pub(crate) rms_norm_eps: f32, } +/// Cos/sin tables for the global layers' proportional RoPE, in the same +/// `[pos * head_dim + d]` layout. The HF reference +/// (`_compute_proportional_rope_parameters`) is not the qwen35-style +/// leading-block partial rotation: the first `rotary_dim / 2` inverse +/// frequencies use the FULL head_dim as the exponent denominator +/// (`theta^(2i/head_dim)`, not `/rotary_dim`), the remaining band is +/// zero-padded, and `rotate_half` then pairs `(d, d + head_dim/2)` across +/// the whole head — zero frequency makes the un-rotated band an exact +/// identity (cos 1, sin 0). The prep kernel therefore runs at +/// `rotary_dim = head_dim`; the partiality lives in these tables. +pub(crate) fn build_proportional_rope_tables( + ctx: &DeviceContext, + rope_theta: f32, + head_dim: usize, + rotary_dim: usize, + max_pos: usize, +) -> Result<(DeviceVec, DeviceVec)> { + anyhow::ensure!( + rope_theta.is_finite() && rope_theta > 0.0, + "proportional rope theta {rope_theta} must be positive and finite" + ); + anyhow::ensure!( + head_dim > 0 + && head_dim.is_multiple_of(2) + && rotary_dim > 0 + && rotary_dim.is_multiple_of(2), + "proportional rope dims must be positive and even: head_dim {head_dim}, rotary_dim \ + {rotary_dim}" + ); + anyhow::ensure!( + rotary_dim <= head_dim, + "proportional rope rotary_dim {rotary_dim} exceeds head_dim {head_dim}" + ); + anyhow::ensure!(max_pos > 0, "proportional rope max_pos must be positive"); + let table_len = max_pos.checked_mul(head_dim).ok_or_else(|| { + anyhow::anyhow!("proportional rope table size {max_pos} x {head_dim} overflows") + })?; + let rope_angles = rotary_dim / 2; + let half_dim = head_dim / 2; + let mut cos = vec![bf16::from_f32(0.0); table_len]; + let mut sin = vec![bf16::from_f32(0.0); table_len]; + for pos in 0..max_pos { + for i in 0..half_dim { + let inv_freq = if i < rope_angles { + 1.0f32 / rope_theta.powf(2.0 * i as f32 / head_dim as f32) + } else { + 0.0 + }; + let angle = pos as f32 * inv_freq; + let row = pos * head_dim; + let (c, s_) = (bf16::from_f32(angle.cos()), bf16::from_f32(angle.sin())); + cos[row + i] = c; + cos[row + i + half_dim] = c; + sin[row + i] = s_; + sin[row + i + half_dim] = s_; + } + } + Ok(( + DeviceVec::from_host(ctx, &cos)?, + DeviceVec::from_host(ctx, &sin)?, + )) +} + /// 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]` @@ -73,6 +138,113 @@ fn nhd_to_hnd( Ok(hnd) } +/// 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. +fn weightless_value_norm( + ctx: &DeviceContext, + mut v_states: HiddenStates, + num_kv_heads: usize, + head_dim: usize, + rms_norm_eps: f32, +) -> Result { + let seq_len = v_states.seq_len; + let kv_dim = v_states.hidden_dim; + anyhow::ensure!( + kv_dim == num_kv_heads * head_dim, + "weightless_value_norm v.hidden_dim {kv_dim} != num_kv_heads {num_kv_heads} * head_dim \ + {head_dim}" + ); + let ones = DeviceVec::from_host(ctx, &vec![bf16::from_f32(1.0); head_dim])?; + v_states.hidden_dim = head_dim; + v_states.seq_len = seq_len * num_kv_heads; + let mut v_normed = HiddenStates::zeros(ctx, head_dim, seq_len * num_kv_heads)?; + ops::rms_norm_batch_into(ctx, &v_states, &ones, rms_norm_eps, &mut v_normed); + v_normed.hidden_dim = kv_dim; + v_normed.seq_len = seq_len; + Ok(v_normed) +} + +/// Everything downstream of attention — o_proj through the `layer_scalar` +/// multiply (applied after both residual adds, not either branch) — is +/// identical for both layer kinds; one implementation keeps the two +/// forwards' numerics from drifting apart. +pub(crate) fn attention_epilogue( + ctx: &DeviceContext, + layer: &Gemma4Layer, + geom: &LayerGeometry, + x: &HiddenStates, + attn: &HiddenStates, +) -> Result { + let seq_len = x.seq_len; + 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)?; + ops::scale_bf16_in_place(ctx, &mut out, layer.layer_scalar)?; + + Ok(out) +} + /// 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 @@ -81,7 +253,7 @@ fn nhd_to_hnd( pub(crate) fn local_layer_forward( ctx: &DeviceContext, layer: &Gemma4Layer, - geom: &LocalLayerGeometry, + geom: &LayerGeometry, x: &HiddenStates, start_pos: usize, sliding_window: usize, @@ -173,18 +345,13 @@ pub(crate) fn local_layer_forward( 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; + let v_normed = weightless_value_norm( + ctx, + v_states, + geom.num_kv_heads, + geom.head_dim, + geom.rms_norm_eps, + )?; // 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 @@ -210,77 +377,154 @@ pub(crate) fn local_layer_forward( 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, + attention_epilogue(ctx, layer, geom, x, &attn) +} + +/// Runs one global (full-attention) layer on `x`; same per-call-buffer probe +/// form as [`local_layer_forward`]. +#[allow(clippy::too_many_arguments)] +pub(crate) fn global_layer_forward( + ctx: &DeviceContext, + layer: &Gemma4Layer, + geom: &LayerGeometry, + x: &HiddenStates, + start_pos: usize, + cos_cache: &DeviceVec, + sin_cache: &DeviceVec, + cos_max_pos: usize, +) -> Result { + let seq_len = x.seq_len; + anyhow::ensure!( + start_pos == 0, + "global_layer_forward is prefill-from-zero only; start_pos {start_pos} needs a KV cache" + ); + anyhow::ensure!( + start_pos + .checked_add(seq_len) + .is_some_and(|end| end <= cos_max_pos), + "global layer positions {start_pos}..{start_pos}+{seq_len} exceed the rope table's \ + cos_max_pos {cos_max_pos}; the prep kernel traps on out-of-range positions" + ); + 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, + "global layer x.hidden_dim {} != hidden_size {}", + x.hidden_dim, + geom.hidden_size + ); + anyhow::ensure!( + geom.head_dim == 512, + "global layer head_dim {} != 512, which the prep kernel is instantiated at", + geom.head_dim + ); + anyhow::ensure!( + layer.attention.v_proj.is_none(), + "global layer must not carry a v_proj; the checkpoint ships its \ + full_attention layers without one" ); - 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)?; + let mut normed_x = HiddenStates::zeros(ctx, geom.hidden_size, seq_len)?; ops::rms_norm_batch_into( ctx, - &h2, - &layer.pre_feedforward_layernorm, + x, + &layer.input_layernorm, geom.rms_norm_eps, - &mut mlp_in, + &mut normed_x, ); - let mut gate = HiddenStates::zeros(ctx, geom.intermediate_size, seq_len)?; - let mut up = HiddenStates::zeros(ctx, geom.intermediate_size, seq_len)?; + + let mut q_states = HiddenStates::zeros(ctx, q_dim, seq_len)?; + let mut k_states = HiddenStates::zeros(ctx, kv_dim, seq_len)?; ops::gemm_rows_into_checked( ctx, - &layer.mlp.gate, + &layer.attention.q_proj, 0, - geom.intermediate_size, - &mlp_in, - &mut gate, + q_dim, + &normed_x, + &mut q_states, )?; ops::gemm_rows_into_checked( ctx, - &layer.mlp.up, + &layer.attention.k_proj, 0, - geom.intermediate_size, - &mlp_in, - &mut up, + kv_dim, + &normed_x, + &mut k_states, )?; - 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( + + // The K=V fork: value is k_proj's raw output, copied BEFORE k_norm and + // RoPE touch K — a runtime fork of the projection, not shared storage; + // after the norms the two tensors differ bitwise. + let mut v_states = HiddenStates::zeros(ctx, kv_dim, seq_len)?; + ctx.stream + .memcpy_dtod(&k_states.data, &mut v_states.data) + .map_err(|e| anyhow::anyhow!("global layer V fork copy failed: {e}"))?; + + // Norm + rotate Q and K. The decode-shaped prep entry is the contiguous + // one (its prefill sibling writes a paged pool); positions are the + // explicit per-token array it expects. rotary_dim is the FULL head — the + // proportional tables carry the partiality (see the table builder). + let positions: Vec = (0..seq_len) + .map(|t| { + i32::try_from(start_pos + t) + .map_err(|_| anyhow::anyhow!("position {} does not fit i32", start_pos + t)) + }) + .collect::>()?; + let positions_d = ctx + .stream + .clone_htod(&positions) + .map_err(|e| anyhow::anyhow!("positions H2D failed: {e}"))?; + let mut q_prep = HiddenStates::zeros(ctx, q_dim, seq_len)?; + ops::qk_norm_partial_rope_batched_decode_hd512_into( ctx, - &down, - &layer.post_feedforward_layernorm, + &q_states, + &mut q_prep, + &mut k_states, + &layer.attention.q_norm, + &layer.attention.k_norm, + cos_cache, + sin_cache, + &positions_d, + cos_max_pos, + geom.num_q_heads, + geom.num_kv_heads, + geom.head_dim, 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)?; + let v_normed = weightless_value_norm( + ctx, + v_states, + geom.num_kv_heads, + geom.head_dim, + geom.rms_norm_eps, + )?; - // layer_scalar multiplies the layer output after both residual adds — - // not either branch. - ops::scale_bf16_in_place(ctx, &mut out, layer.layer_scalar)?; + // single_prefill_hd512 reads HND; reassemble per head (at one KV head + // this is a plain copy, but the layout contract stays explicit). + let k_hnd = nhd_to_hnd(ctx, &k_states, geom.num_kv_heads, geom.head_dim)?; + let v_hnd = nhd_to_hnd(ctx, &v_normed, geom.num_kv_heads, geom.head_dim)?; - Ok(out) + let mut attn = HiddenStates::zeros(ctx, q_dim, seq_len)?; + ops::single_prefill_hd512_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, + )?; + + attention_epilogue(ctx, layer, geom, x, &attn) } -/// The oracle: replay the HF golden fixture's local-layer probes through -/// this layer implementation on the real checkpoint. See +/// The oracle: replay the HF golden fixture's layer probes through this +/// implementation on the real checkpoint. See /// `docs/models/gemma4/hf-golden.md` for what the fixture pins. #[cfg(test)] #[path = "layer_oracle.rs"] diff --git a/pegainfer-gemma4/src/layer_oracle.rs b/pegainfer-gemma4/src/layer_oracle.rs index fc7d4e89..2d5d7fce 100644 --- a/pegainfer-gemma4/src/layer_oracle.rs +++ b/pegainfer-gemma4/src/layer_oracle.rs @@ -1,6 +1,7 @@ use super::*; use crate::config::Gemma4Config; -use crate::config::first_last_sliding; +use crate::config::LayerKind; +use crate::config::first_last; use crate::testkit::GOLDEN_PATH; use crate::testkit::METADATA_KEY; use crate::testkit::assert_checkpoint_matches; @@ -62,7 +63,7 @@ fn compare(got: &[f32], expected: &[bf16], hidden_size: usize, what: &str) -> us #[test] #[ignore = "requires the pinned 12B checkpoint via PEGAINFER_TEST_MODEL_PATH and a GPU"] -fn local_layer_matches_hf_probes() { +fn layers_match_hf_probes() { let dir = model_path(); let fixture_bytes = std::fs::read(GOLDEN_PATH).expect("read fixture"); let (_, meta) = @@ -83,7 +84,9 @@ fn local_layer_matches_hf_probes() { // 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"); + first_last(&config.layer_types, LayerKind::Sliding).expect("12B carries sliding layers"); + let (gfirst, glast) = + first_last(&config.layer_types, LayerKind::Global).expect("12B carries global layers"); let probe_layers = &manifest["probe_layers"]; assert_eq!( first as u64, @@ -97,6 +100,22 @@ fn local_layer_matches_hf_probes() { probe_layers["sliding_last"].as_u64().expect("sliding_last"), "last sliding layer disagrees with the fixture's parse" ); + assert_eq!( + gfirst as u64, + probe_layers["global_first"].as_u64().expect("global_first"), + "first global layer disagrees with the fixture's parse" + ); + assert_eq!( + glast as u64, + probe_layers["global_last"].as_u64().expect("global_last"), + "last global layer disagrees with the fixture's parse" + ); + assert_eq!( + glast, + last + 1, + "global_last must sit directly after sliding_last for the \ + dedup label below to be its input" + ); let cut_labels: Vec = manifest["cut_labels"] .as_array() @@ -114,7 +133,7 @@ fn local_layer_matches_hf_probes() { let (weights, _) = Gemma4Weights::from_safetensors(&dir, 0).expect("load 12B weights"); let ctx = DeviceContext::new_with_device(0).expect("device context"); - let geom = LocalLayerGeometry { + let geom = LayerGeometry { hidden_size: config.hidden_size, intermediate_size: config.intermediate_size, num_q_heads: config.num_attention_heads, @@ -122,6 +141,14 @@ fn local_layer_matches_hf_probes() { head_dim: config.head_dim, rms_norm_eps: config.rms_norm_eps, }; + let global_geom = LayerGeometry { + hidden_size: config.hidden_size, + intermediate_size: config.intermediate_size, + num_q_heads: config.num_attention_heads, + num_kv_heads: config.num_global_key_value_heads, + head_dim: config.global_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, @@ -133,6 +160,14 @@ fn local_layer_matches_hf_probes() { }, ) .expect("rope tables"); + let (gcos_cache, gsin_cache) = build_proportional_rope_tables( + &ctx, + config.global_rope_theta, + global_geom.head_dim, + config.global_rotary_dim, + cos_max_pos, + ) + .expect("proportional rope tables"); let mut over_tolerance: Vec = Vec::new(); for case in ["single", "short"] { @@ -168,6 +203,31 @@ fn local_layer_matches_hf_probes() { over_tolerance.push(format!("{case}/{name} (layer {index})")); } } + + for (name, index, in_label) in [ + ("global_first", gfirst, "global_first_in"), + ("global_last", glast, "sliding_last_out"), + ] { + let x = + HiddenStates::from_host(&ctx, cut(in_label), hidden_size, seq_len).expect("x H2D"); + let out = global_layer_forward( + &ctx, + &weights.layers[index], + &global_geom, + &x, + 0, + &gcos_cache, + &gsin_cache, + cos_max_pos, + ) + .expect("global 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(), diff --git a/pegainfer-gemma4/src/manifest/schema.rs b/pegainfer-gemma4/src/manifest/schema.rs index 1ad88c46..ec61f34e 100644 --- a/pegainfer-gemma4/src/manifest/schema.rs +++ b/pegainfer-gemma4/src/manifest/schema.rs @@ -274,6 +274,8 @@ pub(super) fn sample_config() -> Gemma4Config { rms_norm_eps: 1e-6, sliding_rope_theta: 10_000.0, sliding_window: 1024, + global_rope_theta: 1_000_000.0, + global_rotary_dim: 128, } }