llama-gguf 0.14.0

A high-performance Rust implementation of llama.cpp - LLM inference engine with full GGUF support
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205

//! Debug KV cache to see if attention across positions is working

use llama_gguf::backend::Backend;
use llama_gguf::backend::cpu::CpuBackend;
use llama_gguf::gguf::GgufFile;
use llama_gguf::tensor::{DType, Tensor};
use std::path::Path;

fn load_tensor(gguf: &GgufFile, name: &str) -> Tensor {
    let info = gguf.data.get_tensor(name).unwrap();
    let data = gguf.tensor_data(name).unwrap();
    let shape: Vec<usize> = info.dims.iter().map(|&d| d as usize).collect();
    Tensor::new(data.to_vec(), shape, DType::from(info.dtype)).unwrap()
}

fn try_load_tensor(gguf: &GgufFile, name: &str) -> Option<Tensor> {
    let info = gguf.data.get_tensor(name)?;
    let data = gguf.tensor_data(name)?;
    let shape: Vec<usize> = info.dims.iter().map(|&d| d as usize).collect();
    Tensor::new(data.to_vec(), shape, DType::from(info.dtype)).ok()
}

fn dequant(backend: &CpuBackend, t: &Tensor) -> Vec<f32> {
    if t.dtype() == DType::F32 {
        t.as_f32().unwrap().to_vec()
    } else {
        let mut out = Tensor::zeros(vec![t.numel()], DType::F32);
        backend.dequantize(t, &mut out).unwrap();
        out.as_f32().unwrap().to_vec()
    }
}

fn rms_norm(x: &[f32], w: &[f32], eps: f32) -> Vec<f32> {
    let sum_sq: f32 = x.iter().map(|v| v * v).sum();
    let rms = (sum_sq / x.len() as f32 + eps).sqrt();
    let inv_rms = 1.0 / rms;
    x.iter()
        .zip(w.iter())
        .map(|(v, wt)| v * inv_rms * wt)
        .collect()
}

fn vec_mat(x: &[f32], w: &[f32], k: usize, n: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; n];
    for j in 0..n {
        for i in 0..k {
            out[j] += x[i] * w[i + j * k];
        }
    }
    out
}

fn apply_rope_neox(data: &mut [f32], pos: usize, head_dim: usize, freq_base: f32) {
    let half_dim = head_dim / 2;
    for i in 0..half_dim {
        let freq = 1.0 / freq_base.powf((2 * i) as f32 / head_dim as f32);
        let theta = pos as f32 * freq;
        let (sin_t, cos_t) = theta.sin_cos();
        let x0 = data[i];
        let x1 = data[i + half_dim];
        data[i] = x0 * cos_t - x1 * sin_t;
        data[i + half_dim] = x0 * sin_t + x1 * cos_t;
    }
}

fn softmax(scores: &mut [f32]) {
    let max_score = scores.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
    let mut sum = 0.0f32;
    for s in scores.iter_mut() {
        *s = (*s - max_score).exp();
        sum += *s;
    }
    for s in scores.iter_mut() {
        *s /= sum;
    }
}

fn main() {
    let model_path = "/home/joseph/Models/qwen2.5-0.5b-instruct-q4_k_m.gguf";

    eprintln!("Loading model...");
    let gguf = GgufFile::open(Path::new(model_path)).expect("Failed to open GGUF");
    let backend = CpuBackend::new();

    let emb = dequant(&backend, &load_tensor(&gguf, "token_embd.weight"));

    let hidden_size = 896;
    let num_heads = 14;
    let head_dim = 64;
    let num_kv_heads = 2;
    let queries_per_kv = num_heads / num_kv_heads;
    let eps = 1e-6f32;
    let freq_base = 1000000.0f32;
    let scale = 1.0 / (head_dim as f32).sqrt();

    // Load layer 0 weights
    let attn_norm_w = dequant(&backend, &load_tensor(&gguf, "blk.0.attn_norm.weight"));
    let wq = dequant(&backend, &load_tensor(&gguf, "blk.0.attn_q.weight"));
    let wk = dequant(&backend, &load_tensor(&gguf, "blk.0.attn_k.weight"));
    let wv = dequant(&backend, &load_tensor(&gguf, "blk.0.attn_v.weight"));
    let wo = dequant(&backend, &load_tensor(&gguf, "blk.0.attn_output.weight"));

    let q_bias = try_load_tensor(&gguf, "blk.0.attn_q.bias").map(|t| dequant(&backend, &t));
    let k_bias = try_load_tensor(&gguf, "blk.0.attn_k.bias").map(|t| dequant(&backend, &t));
    let v_bias = try_load_tensor(&gguf, "blk.0.attn_v.bias").map(|t| dequant(&backend, &t));

    // Process "1+1=" = [16, 10, 16, 28]
    let tokens = [16u32, 10, 16, 28];
    println!("=== Processing tokens {:?} ===\n", tokens);

    // Store KV cache for all positions
    let mut k_cache: Vec<Vec<f32>> = Vec::new(); // [num_kv_heads * head_dim] per position
    let mut v_cache: Vec<Vec<f32>> = Vec::new();

    // Process each position
    for (pos, &token) in tokens.iter().enumerate() {
        let hidden = emb[token as usize * hidden_size..(token as usize + 1) * hidden_size].to_vec();
        let normed = rms_norm(&hidden, &attn_norm_w, eps);

        // Q, K, V projections with bias BEFORE RoPE
        let mut q = vec_mat(&normed, &wq, hidden_size, num_heads * head_dim);
        let mut k = vec_mat(&normed, &wk, hidden_size, num_kv_heads * head_dim);
        let mut v = vec_mat(&normed, &wv, hidden_size, num_kv_heads * head_dim);

        if let Some(ref bias) = q_bias {
            for (qi, bi) in q.iter_mut().zip(bias.iter()) {
                *qi += *bi;
            }
        }
        if let Some(ref bias) = k_bias {
            for (ki, bi) in k.iter_mut().zip(bias.iter()) {
                *ki += *bi;
            }
        }
        if let Some(ref bias) = v_bias {
            for (vi, bi) in v.iter_mut().zip(bias.iter()) {
                *vi += *bi;
            }
        }

        // Apply RoPE (NeoX style)
        for head in 0..num_heads {
            apply_rope_neox(
                &mut q[head * head_dim..(head + 1) * head_dim],
                pos,
                head_dim,
                freq_base,
            );
        }
        for kv_head in 0..num_kv_heads {
            apply_rope_neox(
                &mut k[kv_head * head_dim..(kv_head + 1) * head_dim],
                pos,
                head_dim,
                freq_base,
            );
        }

        // Store in cache
        k_cache.push(k.clone());
        v_cache.push(v.clone());

        // Compute attention for head 0
        println!("=== Position {} (token {}) ===", pos, token);

        let kv_head = 0; // head 0 uses kv_head 0 (since 14/7 = 2)
        let q_head = &q[0..head_dim];

        // Compute scores against all cached K
        println!("Attention scores (head 0):");
        let mut scores = Vec::new();
        for kv_pos in 0..=pos {
            let k_head = &k_cache[kv_pos][kv_head * head_dim..(kv_head + 1) * head_dim];
            let dot: f32 = q_head.iter().zip(k_head).map(|(qi, ki)| qi * ki).sum();
            let score = dot * scale;
            scores.push(score);
            println!("  pos {} score: {:.4}", kv_pos, score);
        }

        // Softmax
        softmax(&mut scores);
        println!(
            "Softmax weights: {:?}",
            scores
                .iter()
                .map(|s| format!("{:.3}", s))
                .collect::<Vec<_>>()
        );

        // Weighted V
        let mut attn_out = vec![0.0f32; head_dim];
        for kv_pos in 0..=pos {
            let v_head = &v_cache[kv_pos][kv_head * head_dim..(kv_head + 1) * head_dim];
            for d in 0..head_dim {
                attn_out[d] += scores[kv_pos] * v_head[d];
            }
        }
        println!("Attention output sum: {:.4}", attn_out.iter().sum::<f32>());
        println!();
    }

    println!("=== Key observations ===");
    println!("At position 3 (after '='), the attention should weight the '1+1' context");
    println!("to predict '2'. If attention weights are wrong, the prediction will fail.");
}