avarok_core/kimi_k3/cpu_forward/
stream.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Intra-block AttnRes stream (completed blocks + running partial).
4
5use super::super::attnres::attnres_mix;
6
7/// Matches HF `KimiDecoderLayer._forward_attn_residual`: mix the incoming
8/// prefix with already-archived blocks, then at `layer_idx % block_size == 0`
9/// archive that incoming prefix and reset the intra-block sum. Layer 0
10/// therefore archives the embedding as its own source, not `embed + mixer`.
11#[derive(Clone, Debug)]
12pub struct AttnResStream {
13    pub completed: Vec<Vec<f32>>,
14    pub partial: Vec<f32>,
15    block_size: usize,
16}
17
18impl AttnResStream {
19    pub fn new(hidden: usize, block_size: usize) -> Self {
20        Self {
21            completed: Vec::new(),
22            partial: vec![0.0; hidden],
23            block_size: block_size.max(1),
24        }
25    }
26
27    fn sources(&self) -> Vec<Vec<f32>> {
28        // sources[0] is the skip (current prefix). Mix=0 must return this,
29        // not the first archived block.
30        let mut s = vec![self.partial.clone()];
31        s.extend(self.completed.iter().cloned());
32        s
33    }
34
35    pub fn mix(&self, query: &[f32], norm_w: &[f32], eps: f32, mix: f32) -> Vec<f32> {
36        attnres_mix(&self.sources(), query, norm_w, eps, mix)
37    }
38
39    pub(super) fn add(&mut self, delta: &[f32]) {
40        for (p, d) in self.partial.iter_mut().zip(delta) {
41            *p += *d;
42        }
43    }
44
45    pub(super) fn archive_incoming_at_block_start(&mut self, layer_idx: usize) {
46        if layer_idx.is_multiple_of(self.block_size) {
47            self.completed.push(self.partial.clone());
48            self.partial.fill(0.0);
49        }
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56
57    #[test]
58    fn attnres_archives_incoming_at_block_start() {
59        let mut s = AttnResStream::new(2, 4);
60        s.partial = vec![1.0, 2.0];
61        s.archive_incoming_at_block_start(0);
62        assert_eq!(s.completed, vec![vec![1.0, 2.0]]);
63        assert_eq!(s.partial, vec![0.0, 0.0]);
64        s.add(&[0.5, 0.25]);
65        assert_eq!(s.partial, vec![0.5, 0.25]);
66        assert_eq!(
67            s.completed[0],
68            vec![1.0, 2.0],
69            "archive is embed, not embed+mixer"
70        );
71        s.archive_incoming_at_block_start(1);
72        assert_eq!(s.completed.len(), 1, "non-boundary layer must not archive");
73        s.archive_incoming_at_block_start(4);
74        assert_eq!(s.completed.len(), 2);
75        assert_eq!(s.completed[1], vec![0.5, 0.25]);
76        assert_eq!(s.partial, vec![0.0, 0.0]);
77    }
78}
avarok_core/kimi_k3/cpu_forward/stream.rs

avarok_core/kimi_k3/cpu_forward/
stream.rs
