neura_decoder/
cache.rs

1use neura_abi::Element;
2use neura_graph::{AttentionOptions, Graph, Init, Shape, Value};
3
4pub struct Cache {
5    keys: Value,
6    values: Value,
7    seat: Value,
8    heads: u32,
9    width: u32,
10    capacity: u32,
11    rows: u32,
12}
13
14impl Cache {
15    pub fn new(graph: &Graph, heads: u32, width: u32, capacity: u32, rows: u32) -> Self {
16        assert!(
17            heads > 0 && width > 0,
18            "a cache of {heads} heads of {width} numbers holds no key",
19        );
20        assert!(capacity > 0, "a cache of {capacity} positions holds no key");
21        assert!(
22            rows > 0 && rows <= capacity,
23            "a cache of {capacity} positions takes {rows} rows at a time, and a chunk of more rows than the cache holds walks past the slots it fills",
24        );
25        let held = Shape::of([heads, 1, capacity, width]);
26        Self {
27            keys: graph.state(held, Init::Zero, Element::Single),
28            values: graph.state(held, Init::Zero, Element::Single),
29            seat: crate::layout::seat(graph, heads, capacity, rows),
30            heads,
31            width,
32            capacity,
33            rows,
34        }
35    }
36
37    pub fn states(&self) -> [Value; 2] {
38        [self.keys, self.values]
39    }
40
41    pub fn rows(&self) -> u32 {
42        self.rows
43    }
44
45    pub fn write(&self, graph: &Graph, cursor: Value, keys: Value, values: Value) {
46        for (held, what) in [(keys, "key"), (values, "value")] {
47            assert_eq!(
48                graph.shape(held).dims(),
49                [self.heads, 1, self.rows, self.width],
50                "a {what} of {} heads of {} numbers walks {:?} where a cache of {} positions takes {} rows of a chunk",
51                self.heads,
52                self.width,
53                graph.shape(held).dims(),
54                self.capacity,
55                self.rows,
56            );
57        }
58        let slots = graph.add(self.seat, cursor);
59        for (held, stored) in [(keys, self.keys), (values, self.values)] {
60            graph.write_into(stored, slots, held);
61        }
62    }
63
64    pub fn attend(&self, graph: &Graph, query: Value, cursor: Value, scale: Value) -> Value {
65        let shape = graph.shape(query).dims();
66        assert_eq!(
67            shape[2], self.rows,
68            "a query of {} rows weighs the keys a cache of {} rows writes, and {:?} walks another count",
69            shape[2], self.rows, shape,
70        );
71        assert_eq!(
72            shape[3], self.width,
73            "a query of {} numbers walks {shape:?} where the cache it reads holds keys of {} numbers",
74            shape[3], self.width,
75        );
76        assert!(
77            shape[0].is_multiple_of(self.heads),
78            "a query of {} heads reads a cache of {} key heads, and every key head serves a whole group of queries",
79            shape[0],
80            self.heads,
81        );
82        graph.attention(
83            query,
84            self.keys,
85            self.values,
86            AttentionOptions {
87                scale: Some(scale),
88                causal: true,
89                origin: Some(cursor),
90                segments: None,
91                reach: None,
92                query_segments: None,
93                bias: None,
94            },
95        )
96    }
97}
neura_decoder/cache.rs

neura_decoder/
cache.rs
