1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
//! Kimi K3's cross-layer "block residual" mixing mechanism
//! (`_apply_attn_res` in the real `modeling_kimi_linear.py`), a real
//! architectural feature discovered by reading `KimiDecoderLayer.forward`
//! in full rather than assuming a standard pre-norm residual: every
//! layer blends its running residual stream with a growing set of
//! saved checkpoints from earlier layers in the same block, using a
//! learned RMSNorm-projected softmax score to weight the blend --
//! structurally a tiny self-attention over `{checkpoints..., current}`.
//!
//! Confirmed real and active (not a dead/optional code path) via
//! `config.json`'s `attn_res_block_size`=12: every layer calls this
//! twice (once before self-attention using `self_attention_res_norm`/
//! `self_attention_res_proj`, once before the FFN using
//! `mlp_res_norm`/`mlp_res_proj`), every 12th layer (0-indexed
//! `layer_idx % 12 == 0`) additionally commits the layer's pre-blend
//! input as a new checkpoint into the block's growing residual set, and
//! the whole model applies one final blend
//! (`output_attn_res_norm`/`output_attn_res_proj`) before the final
//! norm. Real per-layer tensor names (`self_attention_res_norm.weight`,
//! `self_attention_res_proj.weight`, `mlp_res_norm.weight`,
//! `mlp_res_proj.weight`) confirmed directly against a real Kimi K3
//! shard header fetched earlier this session -- this isn't a rarely-used
//! feature, every layer has these weights.
/// One `_apply_attn_res` call: blends `prefix_sum` (`[hidden_dim]`, the
/// layer's current running residual) with `block_residual` (`[n_blocks,
/// hidden_dim]` flattened, the checkpoints saved so far in this block --
/// may be empty). `norm_weight`/`proj_weight` are both `[hidden_dim]`
/// (the real `proj` is `Linear(hidden_dim, 1, bias=false)`, so its
/// weight is a single `[hidden_dim]` row, not a matrix).
///
/// Real math: normalize each candidate (`prefix_sum` and every saved
/// checkpoint) by its own RMS scale factor only (no elementwise weight
/// yet), score it by dotting with `norm_weight * proj_weight`
/// (mathematically identical to `proj(rms_norm(candidate))` since `proj`
/// has no bias), softmax the scores across candidates, then return the
/// softmax-weighted sum of the *raw* (non-normalized) candidates.
pub fn apply_attn_res(
prefix_sum: &[f32],
block_residual: &[f32],
norm_weight: &[f32],
proj_weight: &[f32],
eps: f32,
) -> Vec<f32> {
let score_weight: Vec<f32> = norm_weight
.iter()
.zip(proj_weight.iter())
.map(|(n, p)| n * p)
.collect();
apply_attn_res_prescored(prefix_sum, block_residual, &score_weight, eps)
}
/// Same computation as `apply_attn_res`, but takes the already-fused
/// `norm_weight * proj_weight` product directly instead of the two
/// separate factors -- what a real Kimi K3 GGUF checkpoint actually
/// stores (`blk.{bid}.attn_res_score`/`ffn_res_score`/
/// `output_res_score`, real names confirmed against
/// `ggml-org/llama.cpp#26185`'s `conversion/kimi_k3.py`, whose
/// `_try_fuse_res` fuses this exact product at GGUF-conversion time
/// since `_apply_attn_res` in the real `modeling_kimi_linear.py` never
/// uses the two factors separately -- see docs/MODELS.md). The
/// safetensors checkpoint stores the two factors separately instead, so
/// `apply_attn_res` fuses them itself and delegates here.
pub fn apply_attn_res_prescored(
prefix_sum: &[f32],
block_residual: &[f32],
score_weight: &[f32],
eps: f32,
) -> Vec<f32> {
let hidden_dim = prefix_sum.len();
assert_eq!(score_weight.len(), hidden_dim);
assert!(block_residual.len().is_multiple_of(hidden_dim));
let n_blocks = block_residual.len() / hidden_dim;
let n_candidates = n_blocks + 1;
let candidate = |c: usize| -> &[f32] {
if c < n_blocks {
&block_residual[c * hidden_dim..(c + 1) * hidden_dim]
} else {
prefix_sum
}
};
let mut scores = vec![0f32; n_candidates];
for (c, score) in scores.iter_mut().enumerate() {
let v = candidate(c);
let mean_sq: f32 = v.iter().map(|x| x * x).sum::<f32>() / hidden_dim as f32;
let rstd = 1.0 / (mean_sq + eps).sqrt();
*score = v
.iter()
.zip(score_weight.iter())
.map(|(x, w)| (x * rstd) * w)
.sum();
}
let max = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
let mut probs = vec![0f32; n_candidates];
let mut sum = 0f32;
for (c, p) in probs.iter_mut().enumerate() {
*p = (scores[c] - max).exp();
sum += *p;
}
for p in probs.iter_mut() {
*p /= sum;
}
let mut out = vec![0f32; hidden_dim];
for (c, &w) in probs.iter().enumerate() {
let v = candidate(c);
for (o, x) in out.iter_mut().zip(v.iter()) {
*o += w * x;
}
}
out
}
#[cfg(test)]
mod tests {
use super::*;
// Generated by an independent Python reference -- do not hand-edit.
// Golden values computed from the same formula
// transcribed from the real `_apply_attn_res` source.
const PREFIX_SUM: [f32; 4] = [0.00123015, 0.298746, -0.274138, -0.890592];
const BLOCK_RESIDUAL: [f32; 8] = [
-0.454671, -0.991647, 0.0601436, 1.34022, -0.492207, -0.620475, 0.489842, 0.356887,
];
const NORM_WEIGHT: [f32; 4] = [1.01054, 0.906953, 0.997075, 1.06953];
const PROJ_WEIGHT: [f32; 4] = [-1.34421, -0.457616, -1.90122, -1.28954];
const EXPECTED: [f32; 4] = [-0.0107297, 0.271161, -0.26009, -0.847388];
#[test]
fn matches_independent_python_reference() {
let got = apply_attn_res(
&PREFIX_SUM,
&BLOCK_RESIDUAL,
&NORM_WEIGHT,
&PROJ_WEIGHT,
1e-5,
);
for (i, (a, b)) in got.iter().zip(EXPECTED.iter()).enumerate() {
assert!((a - b).abs() < 1e-4, "element {i}: rust={a} python={b}");
}
}
#[test]
fn with_no_saved_checkpoints_softmax_over_one_candidate_is_the_identity() {
// n_blocks=0 -> exactly one candidate (prefix_sum itself) ->
// softmax over a single score is always 1.0 -> output must equal
// prefix_sum unchanged, regardless of norm/proj weights.
let prefix_sum = [1.0f32, -2.0, 3.5];
let norm_weight = [0.7f32, 1.3, -0.2];
let proj_weight = [2.0f32, -1.0, 0.5];
let got = apply_attn_res(&prefix_sum, &[], &norm_weight, &proj_weight, 1e-5);
for (a, b) in got.iter().zip(prefix_sum.iter()) {
assert!((a - b).abs() < 1e-5);
}
}
}