frink_models/
skip_stream.rs

1//! **THE EMBEDDING SKIP STREAM** -- Talkie's second residual: the normed
2//! embedding, added back into every layer's output through a learned
3//! per-layer scalar.
4//!
5//! # What it is
6//!
7//! `src/models/talkie.cpp:50-52` norm the embeddings before layer 0
8//! (`build_norm(inpL, nullptr, nullptr, LLM_NORM_RMS, -1)`) and keep the
9//! result as `embd_skip`; every layer then computes its attention and
10//! FFN on the ordinary residual `inpL` and, AFTER the FFN residual add,
11//! adds `embd_skip * out_scale` (`:123-126`), `out_scale` being the
12//! layer's `blk.N.layer_output_scale.weight`, a `{1}` tensor (`:32`,
13//! REQUIRED). So the embedding reaches every layer twice: once through
14//! the residual stream as usual, and once more, scaled, from the side.
15//!
16//! # Reach -- MEASURED
17//!
18//! `LLM_TENSOR_LAYER_OUT_SCALE` is created by three graphs (`talkie.cpp`,
19//! `gemma4.cpp`, `gemma4-assistant.cpp`), and only `talkie` feeds it
20//! from the embedding: Gemma-4 (`gemma4.cpp:365-366`) multiplies its
21//! own layer output by it, on its own engine, which has done so since
22//! it existed. The embedding norm at `:50` is `talkie` alone. So
23//! [`SKIP_STREAM_ARCHS`] has one row, and the fact is one `bool` on
24//! `ModelConfig` covering both halves, because neither exists without
25//! the other in any graph.
26//!
27//! # Where it lives
28//!
29//! `Decoder::embed_token` is the ONE embedding site (the batch form
30//! delegates to it), so the norm sits there and the vector it returns
31//! IS the skip source. The three host bodies capture it once, before
32//! layer 0, and hand it to the FFN body as a [`SkipStream`]; the two
33//! FFN bodies add `skip * out_scale` after the residual add and before
34//! the loop norm (`crate::layer_loops`; no graph has both, the order is
35//! a convention). `LayerWeights::out_scale` is the per-layer scalar,
36//! loaded only when the config says so and REQUIRED then. Every fused
37//! Metal launch refuses the model: none norms the embedding or carries
38//! a second residual.
39
40use crate::loader::load_f32_vec;
41use crate::LoadError;
42use frink_gguf::TensorSource;
43
44/// Architectures whose graph adds the normed embedding into every
45/// layer's output, with the lines.
46pub const SKIP_STREAM_ARCHS: &[(&str, &str)] = &[("talkie", "src/models/talkie.cpp:50-52,123-126")];
47
48/// Whether this architecture norms its embeddings and keeps them as a
49/// skip stream.
50pub fn has_skip_stream(arch: &str) -> bool {
51    SKIP_STREAM_ARCHS.iter().any(|(name, _)| *name == arch)
52}
53
54/// Layer `l`'s `layer_output_scale`, REQUIRED for a skip-stream model
55/// (`talkie.cpp:32`) and untouched for every other, so a tensor of that
56/// name on an architecture whose graph has no such op stays UNREAD and
57/// is refused as such.
58pub fn load_out_scale(
59    file: &impl TensorSource,
60    arch: &str,
61    skip_stream: bool,
62    l: usize,
63) -> Result<Option<f32>, LoadError> {
64    if !skip_stream {
65        return Ok(None);
66    }
67    let name = format!("blk.{l}.layer_output_scale.weight");
68    let v = load_f32_vec(file, &name)?;
69    if v.len() != 1 {
70        return Err(LoadError::UnsupportedFeature(
71            arch.to_string(),
72            format!("{name} has {} entries, expected 1", v.len()),
73        ));
74    }
75    Ok(Some(v[0]))
76}
77
78/// The skip source for one forward pass: the normed embedding rows the
79/// bodies captured before layer 0, `[rows, hidden_dim]`. `None` for a
80/// model without the stream -- an `Option` the FFN bodies take as an
81/// argument, so a body cannot be reached with the question unasked.
82#[derive(Debug, Clone, Copy)]
83pub struct SkipStream<'a> {
84    pub rows: &'a [f32],
85}
86
87#[cfg(test)]
88mod tests {
89    use super::*;
90
91    #[test]
92    fn only_talkie_has_the_stream() {
93        assert!(has_skip_stream("talkie"));
94        for arch in ["llama", "gemma3", "gemma4", "olmo", "bitnet"] {
95            assert!(!has_skip_stream(arch), "{arch}");
96        }
97    }
98
99    #[test]
100    fn every_table_row_is_an_audited_generic_row() {
101        for (arch, line) in SKIP_STREAM_ARCHS {
102            let profile = crate::capability::resolve_profile(arch)
103                .unwrap_or_else(|| panic!("`{arch}` ({line}) is not a registered architecture"));
104            assert!(matches!(
105                profile.path,
106                crate::capability::ArchPath::GenericGqa { .. }
107            ));
108            assert!(crate::capability::AUDITED_GENERIC_GQA.contains(arch));
109        }
110    }
111}
frink_models/skip_stream.rs

frink_models/
skip_stream.rs
