vyre_libs/nn/attention/
kv_cache.rs

1//! Explicit immutable-prior to next-generation KV cache updates.
2
3use thiserror::Error;
4use vyre_foundation::ir::{BufferAccess, BufferDecl, DataType, Expr, Node, Program};
5
6use crate::region::wrap_anonymous;
7
8const OP_ID: &str = "vyre-libs::nn::kv_cache_append";
9
10/// Invalid cache append dimensions.
11#[derive(Debug, Clone, PartialEq, Eq, Error)]
12pub enum KvCacheAppendError {
13    /// A required dimension is zero.
14    #[error(
15        "KV cache append requires nonzero batch, heads, capacity, chunk length, and head dimension"
16    )]
17    EmptyShape,
18    /// The appended chunk exceeds the cache generation.
19    #[error(
20        "KV cache append range offset={offset}, chunk_len={chunk_len} exceeds capacity={capacity}"
21    )]
22    Range {
23        /// First cache token replaced by the chunk.
24        offset: u32,
25        /// Logical chunk token count.
26        chunk_len: u32,
27        /// Cache token capacity.
28        capacity: u32,
29    },
30    /// Flattened element counts exceeded addressable IR indexing.
31    #[error("KV cache append element count overflows u32; shard the cache")]
32    ElementCountOverflow,
33    /// The cache dtype cannot represent floating attention activations.
34    #[error("KV cache append requires F16, BF16, or F32 elements; got {dtype:?}")]
35    UnsupportedDtype {
36        /// Rejected cache element dtype.
37        dtype: DataType,
38    },
39}
40
41/// Copy an immutable prior cache and replace one contiguous token range.
42///
43/// Layout is `[batch, heads, tokens, head_dim]` for both the cache and chunk.
44/// `prior` is never modified. `next` contains the complete successor generation.
45#[allow(clippy::too_many_arguments)]
46pub fn kv_cache_append(
47    prior: &str,
48    chunk: &str,
49    next: &str,
50    batch: u32,
51    heads: u32,
52    capacity: u32,
53    chunk_len: u32,
54    head_dim: u32,
55    offset: u32,
56) -> Result<Program, KvCacheAppendError> {
57    kv_cache_append_typed(
58        prior,
59        chunk,
60        next,
61        batch,
62        heads,
63        capacity,
64        chunk_len,
65        head_dim,
66        offset,
67        DataType::F32,
68    )
69}
70
71/// Typed immutable-prior to next-generation KV cache update.
72#[allow(clippy::too_many_arguments)]
73pub fn kv_cache_append_typed(
74    prior: &str,
75    chunk: &str,
76    next: &str,
77    batch: u32,
78    heads: u32,
79    capacity: u32,
80    chunk_len: u32,
81    head_dim: u32,
82    offset: u32,
83    dtype: DataType,
84) -> Result<Program, KvCacheAppendError> {
85    if batch == 0 || heads == 0 || capacity == 0 || chunk_len == 0 || head_dim == 0 {
86        return Err(KvCacheAppendError::EmptyShape);
87    }
88    if !matches!(dtype, DataType::F16 | DataType::BF16 | DataType::F32) {
89        return Err(KvCacheAppendError::UnsupportedDtype { dtype });
90    }
91    let end = offset
92        .checked_add(chunk_len)
93        .filter(|end| *end <= capacity)
94        .ok_or(KvCacheAppendError::Range {
95            offset,
96            chunk_len,
97            capacity,
98        })?;
99    let checked = |values: &[u32]| {
100        values.iter().try_fold(1_u32, |product, value| {
101            product
102                .checked_mul(*value)
103                .ok_or(KvCacheAppendError::ElementCountOverflow)
104        })
105    };
106    let cache_count = checked(&[batch, heads, capacity, head_dim])?;
107    let chunk_count = checked(&[batch, heads, chunk_len, head_dim])?;
108    let cache_head_span = capacity
109        .checked_mul(head_dim)
110        .ok_or(KvCacheAppendError::ElementCountOverflow)?;
111    let chunk_head_span = chunk_len
112        .checked_mul(head_dim)
113        .ok_or(KvCacheAppendError::ElementCountOverflow)?;
114    let index = Expr::var("index");
115    let dimension = Expr::rem(index.clone(), Expr::u32(head_dim));
116    let head_row = Expr::div(index.clone(), Expr::u32(cache_head_span));
117    let token = Expr::rem(
118        Expr::div(index.clone(), Expr::u32(head_dim)),
119        Expr::u32(capacity),
120    );
121    let chunk_index = Expr::add(
122        Expr::mul(head_row, Expr::u32(chunk_head_span)),
123        Expr::add(
124            Expr::mul(
125                Expr::sub(token.clone(), Expr::u32(offset)),
126                Expr::u32(head_dim),
127            ),
128            dimension,
129        ),
130    );
131    let in_chunk = Expr::and(
132        Expr::ge(token.clone(), Expr::u32(offset)),
133        Expr::lt(token, Expr::u32(end)),
134    );
135    let body = vec![
136        Node::let_bind("index", Expr::InvocationId { axis: 0 }),
137        Node::if_then(
138            Expr::lt(index.clone(), Expr::u32(cache_count)),
139            vec![Node::if_then_else(
140                in_chunk,
141                vec![Node::Store {
142                    buffer: next.into(),
143                    index: index.clone(),
144                    value: Expr::load(chunk, chunk_index),
145                }],
146                vec![Node::Store {
147                    buffer: next.into(),
148                    index: index.clone(),
149                    value: Expr::load(prior, index),
150                }],
151            )],
152        ),
153    ];
154    Ok(Program::wrapped(
155        vec![
156            BufferDecl::storage(prior, 0, BufferAccess::ReadWrite, dtype.clone())
157                .with_count(cache_count),
158            BufferDecl::storage(chunk, 1, BufferAccess::ReadOnly, dtype.clone())
159                .with_count(chunk_count),
160            BufferDecl::output(next, 2, dtype).with_count(cache_count),
161        ],
162        [64, 1, 1],
163        vec![wrap_anonymous(OP_ID, body)],
164    ))
165}
vyre_libs/nn/attention/kv_cache.rs

vyre_libs/nn/attention/
kv_cache.rs
