spark_model/layers/
ple.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! PLE — hashed n-gram injection into the hyper-connection highway.
4//!
5//! Qwen3.8-Flash-Next runs this on ONE layer (`ple_layer_ids` is 1-indexed,
6//! so `[2]` means model layer 1). From the reference's own docstring:
7//!
8//! > PLE projects each token's concatenated n-gram embedding to a shared
9//! > value and one key per residual stream. The normalized stream activations
10//! > gate those values, then a dilated depthwise convolution adds local
11//! > lexical context.
12//!
13//! That is **cross-attention into an n-gram table**, not an additive
14//! embedding — the reading that cost real time on LongCat. The forward, from
15//! `Qwen4ExpTextPLELayer.forward`:
16//!
17//! ```text
18//! embeddings   = ple_embedding(input_ids)                  # [T, 2560]
19//! key_normed   = norm_key(key_proj(emb)) -> [T, hc, H]
20//! value        = value_proj(emb)                           # [T, 2560]
21//! query_normed = norm_query(hidden)      -> [T, hc, H]     # hidden is [T, 10240]
22//! gate  = (key_normed * query_normed).sum(-1) / sqrt(H)    # [T, hc]
23//! gate  = sign(gate) * sqrt(max(|gate|, 1e-6))             # SIGNED SQRT
24//! gated = sigmoid(gate) * value                            # [T, hc, H]
25//! out   = gated.flatten() + silu(conv1d(norm_conv(gated.flatten())))
26//! ```
27//!
28//! and the decoder layer adds it to the highway BEFORE that layer's
29//! attention hyper-connection: `hidden_states = hidden_states + ple(...)`.
30//!
31//! Three things here bite, all quietly:
32//!
33//! 1. **The signed square root** on the gate. Nobody would guess it; omit it
34//!    and the gate distribution is wrong but perfectly finite.
35//! 2. **`conv1d` is depthwise AND dilated** — `groups = 10240`,
36//!    `kernel_size = 4`, `dilation = ngram_size = 3`, so the state is
37//!    `(4-1)*3 = 9` steps, not 3.
38//! 3. **All three norms are the offset-from-1 form** (`normed * (1 + w)`,
39//!    `Qwen4ExpTextRMSNorm`) and **grouped** with `group_size = hidden_size`
40//!    — four independent 2560-wide norms inside the 10240 vector, same as
41//!    `hc_norm`. See `bench/qwen4_exp/ARCHITECTURE.md` §6.
42//!
43//! The n-gram table is ~320M rows x 160 dims. It is NOT resident: the row
44//! cache, pinned arena and deferred-load path from #746 serve it off NVMe.
45//! What does NOT transfer from #746 is the id computation — see `ids.rs`.
46
47#[path = "ple/ids.rs"]
48pub mod ids;
49
50#[cfg(test)]
51#[path = "ple/tests.rs"]
52mod tests;
53
54#[path = "ple/dump.rs"]
55pub mod dump;
56
57#[path = "ple/layer.rs"]
58mod layer;
59
60pub use ids::{PleIdDims, ple_ngram_ids};
61pub use layer::{PleLayer, PleSeqState, PleWeights};
spark_model/layers/ple.rs

spark_model/layers/
ple.rs
