vyre-libs 0.7.2

vyre Category A library ecosystem - pure-IR compositions over foundation IR and primitive-owned kernels
Documentation 
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

//! Explicit K/V cache successor execution contracts.

#![forbid(unsafe_code)]

use vyre_libs::nn::attention::{kv_cache_append, KvCacheAppendError};
use vyre_reference::value::Value;

fn bytes(values: &[f32]) -> Vec<u8> {
    values
        .iter()
        .flat_map(|value| value.to_le_bytes())
        .collect()
}

fn decode(value: &Value) -> Vec<f32> {
    value
        .to_bytes()
        .chunks_exact(4)
        .map(|word| f32::from_le_bytes(word.try_into().expect("Fix: exact f32 word")))
        .collect()
}

#[allow(clippy::too_many_arguments)]
fn execute(
    prior: &[f32],
    chunk: &[f32],
    batch: u32,
    heads: u32,
    capacity: u32,
    chunk_len: u32,
    dim: u32,
    offset: u32,
) -> Vec<f32> {
    let program = kv_cache_append(
        "prior", "chunk", "next", batch, heads, capacity, chunk_len, dim, offset,
    )
    .expect("Fix: valid cache append fixture must build");
    let outputs = vyre_reference::reference_eval(
        &program,
        &[
            Value::from(bytes(prior)),
            Value::from(bytes(chunk)),
            Value::from(vec![0; prior.len() * 4]),
        ],
    )
    .expect("Fix: cache append must execute");
    assert_eq!(decode(&outputs[0]), prior);
    decode(&outputs[1])
}

/// Proves a middle decode chunk replaces only its exact token interval.
#[test]
fn middle_chunk_preserves_prefix_and_suffix_bytes() {
    assert_eq!(
        execute(
            &[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0],
            &[100.0, 101.0, 102.0, 103.0],
            1,
            1,
            4,
            2,
            2,
            1,
        ),
        vec![0.0, 1.0, 100.0, 101.0, 102.0, 103.0, 6.0, 7.0]
    );
}

/// Locks independent batch/head rows so one sequence cannot corrupt another cache.
#[test]
fn batch_and_head_rows_receive_only_their_corresponding_chunk() {
    let prior = (0..12).map(|value| value as f32).collect::<Vec<_>>();
    assert_eq!(
        execute(&prior, &[100.0, 200.0, 300.0, 400.0], 2, 2, 3, 1, 1, 2),
        vec![0.0, 1.0, 100.0, 3.0, 4.0, 200.0, 6.0, 7.0, 300.0, 9.0, 10.0, 400.0,]
    );
}

/// Proves prefill replaces the complete initialized cache without retaining stale bytes.
#[test]
fn full_prefill_replaces_every_cache_element() {
    assert_eq!(
        execute(
            &[-1.0; 6],
            &[1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
            1,
            1,
            3,
            3,
            2,
            0
        ),
        vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
    );
}

/// Ensures empty, overflowing, and out-of-capacity transitions fail before execution.
#[test]
fn invalid_cache_transitions_fail_closed() {
    assert_eq!(
        kv_cache_append("p", "c", "n", 0, 1, 1, 1, 1, 0),
        Err(KvCacheAppendError::EmptyShape)
    );
    assert_eq!(
        kv_cache_append("p", "c", "n", 1, 1, 4, 2, 1, 3),
        Err(KvCacheAppendError::Range {
            offset: 3,
            chunk_len: 2,
            capacity: 4
        })
    );
    assert_eq!(
        kv_cache_append("p", "c", "n", u32::MAX, 2, 1, 1, 1, 0),
        Err(KvCacheAppendError::ElementCountOverflow)
    );
}