neura_decoder/
cache.rs

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

neura_decoder/
cache.rs
