gllm 0.10.6

Pure Rust library for local embeddings, reranking, and text generation with MoE-optimized inference and aggressive performance tuning
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

use burn::tensor::backend::Backend;
use burn::tensor::Tensor;

pub struct KVCache<B: Backend> {
    keys: Vec<Tensor<B, 4>>,
    values: Vec<Tensor<B, 4>>,
    num_layers: usize,
    max_len: usize,
    current_len: usize,
}

impl<B: Backend> KVCache<B> {
    pub fn preallocate(
        num_layers: usize,
        max_len: usize,
        batch_size: usize,
        num_heads: usize,
        head_dim: usize,
        device: &B::Device,
    ) -> Self {
        let mut keys = Vec::with_capacity(num_layers);
        let mut values = Vec::with_capacity(num_layers);
        for _ in 0..num_layers {
            keys.push(Tensor::zeros([batch_size, num_heads, max_len, head_dim], device));
            values.push(Tensor::zeros([batch_size, num_heads, max_len, head_dim], device));
        }

        Self {
            keys,
            values,
            num_layers,
            max_len,
            current_len: 0,
        }
    }

    pub fn update(
        &mut self,
        layer: usize,
        new_k: Tensor<B, 4>,
        new_v: Tensor<B, 4>,
    ) -> (Tensor<B, 4>, Tensor<B, 4>) {
        if layer >= self.num_layers {
            return (new_k, new_v);
        }

        let [batch_size, num_heads, seq_len, head_dim] = new_k.dims();
        if self.max_len == 0 {
            return (new_k, new_v);
        }

        let start = self.current_len;
        let end = start + seq_len;
        assert!(
            end <= self.max_len,
            "KV cache exceeded max_len ({} > {})",
            end,
            self.max_len
        );

        let updated_k = self.keys[layer].clone().slice_assign(
            [0..batch_size, 0..num_heads, start..end, 0..head_dim],
            new_k,
        );
        let updated_v = self.values[layer].clone().slice_assign(
            [0..batch_size, 0..num_heads, start..end, 0..head_dim],
            new_v,
        );

        self.keys[layer] = updated_k;
        self.values[layer] = updated_v;

        if layer + 1 == self.num_layers {
            self.current_len = end;
        }

        let cached_k = self.keys[layer]
            .clone()
            .slice([0..batch_size, 0..num_heads, 0..end, 0..head_dim]);
        let cached_v = self.values[layer]
            .clone()
            .slice([0..batch_size, 0..num_heads, 0..end, 0..head_dim]);

        (cached_k, cached_v)
    }

    pub fn seq_len(&self) -> usize {
        self.current_len
    }
}