Coverage Report

Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/home/noah/src/realizar/src/apr_transformer/config.rs
Line
Count
Source
1
//! APR Transformer Configuration Types (PMAT-802)
2
//!
3
//! Configuration structs for APR transformer:
4
//! - AprKVCache: KV cache for efficient autoregressive generation
5
//! - GenerateConfig: Generation parameters
6
//! - AprTransformerConfig: Model architecture configuration
7
//! - AprTransformerLayer: Per-layer weights
8
//! - Q4KLayerWeights: Q4K quantized layer weights
9
10
use serde::{Deserialize, Serialize};
11
12
// ============================================================================
13
14
/// KV Cache for efficient autoregressive generation (Y4)
15
///
16
/// Pre-allocates storage for keys and values to avoid allocations during decode.
17
/// Each layer has separate K and V caches stored contiguously.
18
///
19
/// # Memory Layout
20
///
21
/// For each layer: `[K_pos0, K_pos1, ..., K_posN, V_pos0, V_pos1, ..., V_posN]`
22
/// where each K/V entry has shape `[num_kv_heads * head_dim]`.
23
#[derive(Debug, Clone)]
24
pub struct AprKVCache {
25
    /// Number of layers
26
    num_layers: usize,
27
    /// Number of KV heads
28
    num_kv_heads: usize,
29
    /// Head dimension
30
    head_dim: usize,
31
    /// Maximum context length (pre-allocated capacity)
32
    capacity: usize,
33
    /// Current sequence length (positions filled)
34
    len: usize,
35
    /// K cache per layer: [num_layers][capacity * num_kv_heads * head_dim]
36
    k_cache: Vec<Vec<f32>>,
37
    /// V cache per layer: [num_layers][capacity * num_kv_heads * head_dim]
38
    v_cache: Vec<Vec<f32>>,
39
}
40
41
impl AprKVCache {
42
    /// Create a new KV cache with pre-allocated capacity
43
    ///
44
    /// # Arguments
45
    ///
46
    /// * `config` - Transformer configuration
47
    ///
48
    /// # Returns
49
    ///
50
    /// Empty KV cache with capacity for full context length
51
    #[must_use]
52
26
    pub fn new(config: &AprTransformerConfig) -> Self {
53
26
        let num_layers = config.num_layers;
54
26
        let num_kv_heads = config.num_kv_heads;
55
26
        let head_dim = config.hidden_dim / config.num_heads;
56
26
        let capacity = config.context_length;
57
58
        // Pre-allocate full capacity for each layer
59
26
        let kv_size = capacity * num_kv_heads * head_dim;
60
38
        let 
k_cache26
=
(0..num_layers)26
.
map26
(|_| vec![0.0f32; kv_size]).
collect26
();
61
38
        let 
v_cache26
=
(0..num_layers)26
.
map26
(|_| vec![0.0f32; kv_size]).
collect26
();
62
63
26
        Self {
64
26
            num_layers,
65
26
            num_kv_heads,
66
26
            head_dim,
67
26
            capacity,
68
26
            len: 0,
69
26
            k_cache,
70
26
            v_cache,
71
26
        }
72
26
    }
73
74
    /// Get current sequence length (number of cached positions)
75
    #[must_use]
76
105
    pub fn len(&self) -> usize {
77
105
        self.len
78
105
    }
79
80
    /// Check if cache is empty
81
    #[must_use]
82
3
    pub fn is_empty(&self) -> bool {
83
3
        self.len == 0
84
3
    }
85
86
    /// Get pre-allocated capacity
87
    #[must_use]
88
2
    pub fn capacity(&self) -> usize {
89
2
        self.capacity
90
2
    }
91
92
    /// Get number of KV heads
93
    #[must_use]
94
1
    pub fn num_kv_heads(&self) -> usize {
95
1
        self.num_kv_heads
96
1
    }
97
98
    /// Get head dimension
99
    #[must_use]
100
1
    pub fn head_dim(&self) -> usize {
101
1
        self.head_dim
102
1
    }
103
104
    /// Append K and V for a single position
105
    ///
106
    /// # Arguments
107
    ///
108
    /// * `layer` - Layer index
109
    /// * `k` - Key tensor `[num_kv_heads * head_dim]`
110
    /// * `v` - Value tensor `[num_kv_heads * head_dim]`
111
    ///
112
    /// # Panics
113
    ///
114
    /// Panics if layer index is out of bounds or cache is full
115
132
    pub fn append(&mut self, layer: usize, k: &[f32], v: &[f32]) {
116
132
        assert!(layer < self.num_layers, 
"Layer index out of bounds"1
);
117
131
        assert!(self.len < self.capacity, 
"KV cache is full"0
);
118
119
131
        let kv_size = self.num_kv_heads * self.head_dim;
120
131
        let offset = self.len * kv_size;
121
122
        // Copy K and V into pre-allocated storage
123
131
        self.k_cache[layer][offset..offset + kv_size].copy_from_slice(k);
124
131
        self.v_cache[layer][offset..offset + kv_size].copy_from_slice(v);
125
126
        // Only increment len on first layer to keep consistent
127
131
        if layer == 0 {
128
110
            self.len += 1;
129
110
        
}21
130
131
    }
131
132
    /// Get cached K and V for a layer
133
    ///
134
    /// # Arguments
135
    ///
136
    /// * `layer` - Layer index
137
    ///
138
    /// # Returns
139
    ///
140
    /// Tuple of (K cache slice, V cache slice) containing all cached positions
141
    #[must_use]
142
49
    pub fn get(&self, layer: usize) -> (&[f32], &[f32]) {
143
49
        let kv_size = self.num_kv_heads * self.head_dim;
144
49
        let used_size = self.len * kv_size;
145
146
49
        (
147
49
            &self.k_cache[layer][..used_size],
148
49
            &self.v_cache[layer][..used_size],
149
49
        )
150
49
    }
151
152
    /// Clear the cache (reset to empty without deallocating)
153
1
    pub fn clear(&mut self) {
154
1
        self.len = 0;
155
        // No need to zero memory - will be overwritten on next append
156
1
    }
157
}
158
159
/// Configuration for text generation
160
#[derive(Debug, Clone)]
161
pub struct GenerateConfig {
162
    /// Maximum number of tokens to generate
163
    pub max_tokens: usize,
164
    /// Temperature for sampling (0.0 = greedy)
165
    pub temperature: f32,
166
    /// Top-p nucleus sampling threshold (optional)
167
    pub top_p: f32,
168
    /// Top-k sampling (0 = disabled)
169
    pub top_k: usize,
170
    /// Repetition penalty (1.0 = no penalty)
171
    pub repetition_penalty: f32,
172
}
173
174
impl Default for GenerateConfig {
175
2
    fn default() -> Self {
176
2
        Self {
177
2
            max_tokens: 32,
178
2
            temperature: 1.0,
179
2
            top_p: 0.9,
180
2
            top_k: 0,
181
2
            repetition_penalty: 1.0,
182
2
        }
183
2
    }
184
}
185
186
/// Configuration for APR Transformer models
187
///
188
/// Mirrors `GGUFConfig` for compatibility but is serializable to APR format.
189
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
190
pub struct AprTransformerConfig {
191
    /// Model architecture name (e.g., "phi2", "llama", "qwen2")
192
    pub architecture: String,
193
    /// Embedding/hidden dimension
194
    pub hidden_dim: usize,
195
    /// Number of transformer layers
196
    pub num_layers: usize,
197
    /// Number of attention heads
198
    pub num_heads: usize,
199
    /// Number of key-value heads (for GQA)
200
    pub num_kv_heads: usize,
201
    /// Vocabulary size
202
    pub vocab_size: usize,
203
    /// FFN intermediate dimension
204
    pub intermediate_dim: usize,
205
    /// Maximum context length
206
    pub context_length: usize,
207
    /// RoPE theta for position encoding
208
    pub rope_theta: f32,
209
    /// Layer norm epsilon
210
    pub eps: f32,
211
}
212
213
impl Default for AprTransformerConfig {
214
1
    fn default() -> Self {
215
1
        Self {
216
1
            architecture: "unknown".to_string(),
217
1
            hidden_dim: 512,
218
1
            num_layers: 6,
219
1
            num_heads: 8,
220
1
            num_kv_heads: 8,
221
1
            vocab_size: 32000,
222
1
            intermediate_dim: 2048,
223
1
            context_length: 2048,
224
1
            rope_theta: 10000.0,
225
1
            eps: 1e-5,
226
1
        }
227
1
    }
228
}
229
230
/// Weights for a single transformer layer (all F32)
231
#[derive(Debug, Clone, Serialize, Deserialize)]
232
pub struct AprTransformerLayer {
233
    /// Attention norm weight [hidden_dim]
234
    pub attn_norm_weight: Vec<f32>,
235
    /// Attention norm bias (optional) [hidden_dim]
236
    pub attn_norm_bias: Option<Vec<f32>>,
237
    /// QKV projection weight [hidden_dim, 3*hidden_dim]
238
    pub qkv_weight: Vec<f32>,
239
    /// QKV projection bias (optional) [3*hidden_dim]
240
    pub qkv_bias: Option<Vec<f32>>,
241
    /// Attention output projection weight [hidden_dim, hidden_dim]
242
    pub attn_output_weight: Vec<f32>,
243
    /// Attention output projection bias (optional) [hidden_dim]
244
    pub attn_output_bias: Option<Vec<f32>>,
245
    /// FFN gate weight for SwiGLU (optional) [hidden_dim, intermediate_dim]
246
    pub ffn_gate_weight: Option<Vec<f32>>,
247
    /// FFN gate bias (optional) [intermediate_dim]
248
    pub ffn_gate_bias: Option<Vec<f32>>,
249
    /// FFN up projection weight [hidden_dim, intermediate_dim]
250
    pub ffn_up_weight: Vec<f32>,
251
    /// FFN up projection bias (optional) [intermediate_dim]
252
    pub ffn_up_bias: Option<Vec<f32>>,
253
    /// FFN down projection weight [intermediate_dim, hidden_dim]
254
    pub ffn_down_weight: Vec<f32>,
255
    /// FFN down projection bias (optional) [hidden_dim]
256
    pub ffn_down_bias: Option<Vec<f32>>,
257
    /// FFN norm weight (optional) [hidden_dim]
258
    pub ffn_norm_weight: Option<Vec<f32>>,
259
    /// FFN norm bias (optional) [hidden_dim]
260
    pub ffn_norm_bias: Option<Vec<f32>>,
261
}
262
263
impl AprTransformerLayer {
264
    /// Create an empty layer with given dimensions (non-GQA: num_kv_heads == num_heads)
265
4
    pub fn empty(hidden_dim: usize, intermediate_dim: usize) -> Self {
266
4
        Self {
267
4
            attn_norm_weight: vec![1.0; hidden_dim],
268
4
            attn_norm_bias: None,
269
4
            qkv_weight: vec![0.0; hidden_dim * 3 * hidden_dim],
270
4
            qkv_bias: None,
271
4
            attn_output_weight: vec![0.0; hidden_dim * hidden_dim],
272
4
            attn_output_bias: None,
273
4
            ffn_gate_weight: None,
274
4
            ffn_gate_bias: None,
275
4
            ffn_up_weight: vec![0.0; hidden_dim * intermediate_dim],
276
4
            ffn_up_bias: None,
277
4
            ffn_down_weight: vec![0.0; intermediate_dim * hidden_dim],
278
4
            ffn_down_bias: None,
279
4
            ffn_norm_weight: None,
280
4
            ffn_norm_bias: None,
281
4
        }
282
4
    }
283
284
    /// Create an empty layer with GQA dimensions (num_kv_heads < num_heads)
285
    ///
286
    /// # Arguments
287
    /// * `hidden_dim` - Hidden dimension (num_heads * head_dim)
288
    /// * `num_heads` - Number of query heads
289
    /// * `num_kv_heads` - Number of key/value heads (< num_heads for GQA)
290
    /// * `intermediate_dim` - FFN intermediate dimension
291
1
    pub fn empty_gqa(
292
1
        hidden_dim: usize,
293
1
        num_heads: usize,
294
1
        num_kv_heads: usize,
295
1
        intermediate_dim: usize,
296
1
    ) -> Self {
297
1
        let head_dim = hidden_dim / num_heads;
298
1
        let kv_dim = num_kv_heads * head_dim;
299
        // QKV weight: [hidden_dim, Q_dim + K_dim + V_dim] = [hidden_dim, hidden_dim + 2*kv_dim]
300
1
        let qkv_out_dim = hidden_dim + 2 * kv_dim;
301
302
1
        Self {
303
1
            attn_norm_weight: vec![1.0; hidden_dim],
304
1
            attn_norm_bias: None,
305
1
            qkv_weight: vec![0.0; hidden_dim * qkv_out_dim],
306
1
            qkv_bias: None,
307
1
            attn_output_weight: vec![0.0; hidden_dim * hidden_dim],
308
1
            attn_output_bias: None,
309
1
            ffn_gate_weight: None,
310
1
            ffn_gate_bias: None,
311
1
            ffn_up_weight: vec![0.0; hidden_dim * intermediate_dim],
312
1
            ffn_up_bias: None,
313
1
            ffn_down_weight: vec![0.0; intermediate_dim * hidden_dim],
314
1
            ffn_down_bias: None,
315
1
            ffn_norm_weight: None,
316
1
            ffn_norm_bias: None,
317
1
        }
318
1
    }
319
320
    /// Get total number of parameters in this layer
321
    #[must_use]
322
24
    pub fn num_parameters(&self) -> usize {
323
24
        let mut count = 0;
324
24
        count += self.attn_norm_weight.len();
325
24
        count += self.attn_norm_bias.as_ref().map_or(0, Vec::len);
326
24
        count += self.qkv_weight.len();
327
24
        count += self.qkv_bias.as_ref().map_or(0, Vec::len);
328
24
        count += self.attn_output_weight.len();
329
24
        count += self.attn_output_bias.as_ref().map_or(0, Vec::len);
330
24
        count += self.ffn_gate_weight.as_ref().map_or(0, Vec::len);
331
24
        count += self.ffn_gate_bias.as_ref().map_or(0, Vec::len);
332
24
        count += self.ffn_up_weight.len();
333
24
        count += self.ffn_up_bias.as_ref().map_or(0, Vec::len);
334
24
        count += self.ffn_down_weight.len();
335
24
        count += self.ffn_down_bias.as_ref().map_or(0, Vec::len);
336
24
        count += self.ffn_norm_weight.as_ref().map_or(0, Vec::len);
337
24
        count += self.ffn_norm_bias.as_ref().map_or(0, Vec::len);
338
24
        count
339
24
    }
340
}
341
342
/// Q4K/Q6K raw weights for fused kernel inference (F-GPU-130)
343
///
344
/// When present, matmul operations use fused kernels (matmul_q4k_f32, matmul_q6k_f32)
345
/// instead of the F32 path, avoiding full dequantization overhead.
346
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
347
pub struct Q4KLayerWeights {
348
    /// QKV projection weight in Q4K format (combined, legacy)
349
    pub qkv_weight: Option<Vec<u8>>,
350
    /// Q projection weight in Q4K format (PMAT-103: separate for fused kernel)
351
    pub attn_q_weight: Option<Vec<u8>>,
352
    /// K projection weight in Q4K format (PMAT-103: separate for fused kernel)
353
    pub attn_k_weight: Option<Vec<u8>>,
354
    /// V projection weight in Q4K/Q6K format (PMAT-103: separate for fused kernel)
355
    pub attn_v_weight: Option<Vec<u8>>,
356
    /// V projection weight in Q6K format (when Q4K not available)
357
    pub attn_v_weight_q6k: Option<Vec<u8>>,
358
    /// Attention output projection in Q4K format
359
    pub attn_output_weight: Option<Vec<u8>>,
360
    /// FFN gate weight in Q4K format (for SwiGLU)
361
    pub ffn_gate_weight: Option<Vec<u8>>,
362
    /// FFN up projection in Q4K format
363
    pub ffn_up_weight: Option<Vec<u8>>,
364
    /// FFN down projection in Q4K format
365
    pub ffn_down_weight: Option<Vec<u8>>,
366
    /// FFN down projection in Q6K format (when Q4K not available)
367
    pub ffn_down_weight_q6k: Option<Vec<u8>>,
368
    /// FFN up projection in Q6K format (when Q4K not available)
369
    pub ffn_up_weight_q6k: Option<Vec<u8>>,
370
}
371