/home/noah/src/realizar/src/gguf/inference/forward/core.rs
Line | Count | Source |
1 | | //! Forward pass implementations for OwnedQuantizedModel |
2 | | //! |
3 | | //! Contains forward, forward_cached methods. |
4 | | //! These are the core inference entry points. |
5 | | |
6 | | use crate::error::Result; |
7 | | use crate::gguf::OwnedQKVWeights; |
8 | | use crate::gguf::ops; |
9 | | use crate::gguf::OwnedQuantizedModel; |
10 | | |
11 | | impl OwnedQuantizedModel { |
12 | | /// Forward pass with fused Q4_K operations (IMP-100) |
13 | | /// |
14 | | /// This is 1.37x faster than dequantized f32 due to reduced memory bandwidth. |
15 | | /// |
16 | | /// # Arguments |
17 | | /// |
18 | | /// * `token_ids` - Input token IDs |
19 | | /// |
20 | | /// # Returns |
21 | | /// |
22 | | /// Logits for next token prediction [vocab_size] |
23 | | /// |
24 | | /// # Errors |
25 | | /// |
26 | | /// Returns error if tensor operations fail |
27 | 4 | pub fn forward(&self, token_ids: &[u32]) -> Result<Vec<f32>> { |
28 | 4 | let hidden_dim = self.config.hidden_dim; |
29 | | // Note: intermediate_dim is encoded in layer weight tensors (in_dim/out_dim) |
30 | 4 | let _ = self.config.intermediate_dim; |
31 | | |
32 | | // 1. Token embedding lookup (f32, fast) |
33 | 4 | let mut hidden = self.embed(token_ids); |
34 | | |
35 | | // Detect if model uses RMSNorm (LLaMA-style) or LayerNorm (phi-2 style) |
36 | | // LLaMA models have ffn_gate_weight (SwiGLU) and no bias in norms |
37 | 4 | let use_rmsnorm = self |
38 | 4 | .layers |
39 | 4 | .first() |
40 | 4 | .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias0 .is_none0 ()); |
41 | | |
42 | | // 2. Process through transformer layers with FUSED Q4_K ops |
43 | 4 | let cpu_debug_layers = std::env::var("CPU_DEBUG_LAYERS").is_ok(); |
44 | 4 | for (layer_idx, layer) in self.layers.iter().enumerate() { |
45 | | // 2a. Attention layer norm (RMSNorm for LLaMA, LayerNorm for others) |
46 | 4 | let normed = if use_rmsnorm { |
47 | 0 | ops::rms_norm(&hidden, &layer.attn_norm_weight, self.config.eps) |
48 | | } else { |
49 | 4 | ops::layer_norm( |
50 | 4 | &hidden, |
51 | 4 | &layer.attn_norm_weight, |
52 | 4 | layer.attn_norm_bias.as_deref(), |
53 | 4 | self.config.eps, |
54 | | ) |
55 | | }; |
56 | | |
57 | | // CORRECTNESS-011: CPU intermediate debug at L0 |
58 | 4 | if cpu_debug_layers && layer_idx < 20 { |
59 | 0 | eprintln!( |
60 | 0 | "[CPU-L{}] RMSNorm: first 3 = [{:.4}, {:.4}, {:.4}]", |
61 | 0 | layer_idx, normed[0], normed[1], normed[2] |
62 | 0 | ); |
63 | 4 | } |
64 | | |
65 | | // 2b. QKV projection with FUSED dequant+dot (1.37x faster) |
66 | | // Note: qkv_dim may differ from 3*hidden_dim for GQA models |
67 | 4 | let qkv_dim = layer.qkv_weight.out_dim(); |
68 | 4 | let q_dim = layer.qkv_weight.q_dim(); |
69 | | // For GQA, k_dim and v_dim may be smaller than q_dim |
70 | 4 | let k_dim = match &layer.qkv_weight { |
71 | 4 | OwnedQKVWeights::Fused(_) => q_dim, |
72 | 0 | OwnedQKVWeights::Separate { k, .. } => k.out_dim, |
73 | | }; |
74 | 4 | let v_dim = match &layer.qkv_weight { |
75 | 4 | OwnedQKVWeights::Fused(_) => q_dim, |
76 | 0 | OwnedQKVWeights::Separate { v, .. } => v.out_dim, |
77 | | }; |
78 | 4 | let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0 ; |
79 | 4 | if let Some(ref bias0 ) = layer.qkv_bias { |
80 | 0 | ops::add_bias(&mut qkv, bias); |
81 | 4 | } |
82 | | |
83 | | // CORRECTNESS-011: Q, K, V before RoPE (after bias) |
84 | 4 | if cpu_debug_layers && (layer_idx < 20 || layer_idx == 40 || layer_idx == 50 ) { |
85 | 0 | eprintln!( |
86 | 0 | "[CPU-L{}] Q (before RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]", |
87 | 0 | layer_idx, qkv[0], qkv[1], qkv[2], qkv[3], qkv[4] |
88 | 0 | ); |
89 | 0 | // K starts at q_dim offset |
90 | 0 | eprintln!( |
91 | 0 | "[CPU-L{}] K (before RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]", |
92 | 0 | layer_idx, |
93 | 0 | qkv[q_dim], |
94 | 0 | qkv[q_dim + 1], |
95 | 0 | qkv[q_dim + 2], |
96 | 0 | qkv[q_dim + 3], |
97 | 0 | qkv[q_dim + 4] |
98 | 0 | ); |
99 | 0 | // V starts at q_dim + k_dim offset |
100 | 0 | let v_offset = q_dim + k_dim; |
101 | 0 | eprintln!( |
102 | 0 | "[CPU-L{}] V (before RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]", |
103 | 0 | layer_idx, |
104 | 0 | qkv[v_offset], |
105 | 0 | qkv[v_offset + 1], |
106 | 0 | qkv[v_offset + 2], |
107 | 0 | qkv[v_offset + 3], |
108 | 0 | qkv[v_offset + 4] |
109 | 0 | ); |
110 | 4 | } |
111 | | |
112 | | // 2c. Proper attention with RoPE and causal mask (IMP-101) |
113 | 4 | let seq_len = token_ids.len(); |
114 | | |
115 | | // Extract Q, K, V and apply RoPE to Q and K |
116 | 4 | let mut q_all = Vec::with_capacity(seq_len * q_dim); |
117 | 4 | let mut k_all = Vec::with_capacity(seq_len * k_dim); |
118 | 4 | let mut v_all = Vec::with_capacity(seq_len * v_dim); |
119 | | |
120 | 8 | for s in 0..seq_len4 { |
121 | 8 | let qkv_start = s * qkv_dim; |
122 | | |
123 | | // Extract Q, K, V for this position (QKV layout: [Q..., K..., V...]) |
124 | 8 | let mut q = qkv[qkv_start..qkv_start + q_dim].to_vec(); |
125 | 8 | let mut k = qkv[qkv_start + q_dim..qkv_start + q_dim + k_dim].to_vec(); |
126 | 8 | let v = &qkv[qkv_start + q_dim + k_dim..qkv_start + q_dim + k_dim + v_dim]; |
127 | | |
128 | | // Apply RoPE to Q and K (position-dependent rotation) |
129 | | // GQA: Q has num_heads, K has num_kv_heads |
130 | 8 | self.apply_rope(&mut q, s, self.config.num_heads); |
131 | 8 | self.apply_rope(&mut k, s, self.config.num_kv_heads); |
132 | | |
133 | | // CORRECTNESS-011: Q after RoPE at position 0 |
134 | 8 | if cpu_debug_layers && layer_idx < 20 && s == 00 { |
135 | 0 | eprintln!( |
136 | 0 | "[CPU-L{}] Q (after RoPE): first 3 = [{:.4}, {:.4}, {:.4}]", |
137 | 0 | layer_idx, q[0], q[1], q[2] |
138 | 0 | ); |
139 | 0 | eprintln!( |
140 | 0 | "[CPU-L{}] K (after RoPE): first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]", |
141 | 0 | layer_idx, k[0], k[1], k[2], k[3], k[4] |
142 | 0 | ); |
143 | 0 | eprintln!( |
144 | 0 | "[CPU-L{}] V: first 5 = [{:.4}, {:.4}, {:.4}, {:.4}, {:.4}]", |
145 | 0 | layer_idx, v[0], v[1], v[2], v[3], v[4] |
146 | 0 | ); |
147 | 8 | } |
148 | | |
149 | 8 | q_all.extend_from_slice(&q); |
150 | 8 | k_all.extend_from_slice(&k); |
151 | 8 | v_all.extend_from_slice(v); |
152 | | } |
153 | | |
154 | | // Compute scaled dot-product attention with causal mask |
155 | 4 | let attn_out = self.causal_attention(&q_all, &k_all, &v_all, seq_len); |
156 | | |
157 | | // CORRECTNESS-011: Attention output |
158 | 4 | if cpu_debug_layers && layer_idx < 20 { |
159 | 0 | eprintln!( |
160 | 0 | "[CPU-L{}] Attn output: first 3 = [{:.4}, {:.4}, {:.4}]", |
161 | 0 | layer_idx, attn_out[0], attn_out[1], attn_out[2] |
162 | 0 | ); |
163 | 4 | } |
164 | | |
165 | | // 2d. Attention output projection with FUSED ops |
166 | | // Input is q_dim (attention output), projects back to hidden_dim |
167 | 4 | let mut attn_output = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0 ; |
168 | 4 | if let Some(ref bias0 ) = layer.attn_output_bias { |
169 | 0 | ops::add_bias(&mut attn_output, bias); |
170 | 4 | } |
171 | | |
172 | | // 2e. Residual connection |
173 | 480 | for i in 0..hidden4 .len4 () { |
174 | 480 | hidden[i] += attn_output[i]; |
175 | 480 | } |
176 | | |
177 | | // 2f. Pre-FFN layer norm (LLaMA uses separate ffn_norm with RMSNorm) |
178 | 4 | let ffn_input = if let Some(ref ffn_norm0 ) = layer.ffn_norm_weight { |
179 | | // LLaMA-style: separate FFN layer norm (use RMSNorm for LLaMA) |
180 | 0 | if use_rmsnorm { |
181 | 0 | ops::rms_norm(&hidden, ffn_norm, self.config.eps) |
182 | | } else { |
183 | 0 | ops::layer_norm( |
184 | 0 | &hidden, |
185 | 0 | ffn_norm, |
186 | 0 | layer.ffn_norm_bias.as_deref(), |
187 | 0 | self.config.eps, |
188 | | ) |
189 | | } |
190 | | } else { |
191 | | // phi-2 style: no separate FFN norm, use hidden directly |
192 | | // (some models apply attn_norm again, but we've already done residual) |
193 | 4 | hidden.clone() |
194 | | }; |
195 | | |
196 | | // 2g. FFN with SwiGLU or GELU activation |
197 | 4 | let ffn_activated = if let Some(ref gate_weight0 ) = layer.ffn_gate_weight { |
198 | | // SwiGLU path (LLaMA, TinyLlama, Mistral, etc.) |
199 | | // output = down(gate(x) * silu(up(x))) |
200 | 0 | let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?; |
201 | 0 | if let Some(ref bias) = layer.ffn_up_bias { |
202 | 0 | ops::add_bias(&mut ffn_up, bias); |
203 | 0 | } |
204 | | |
205 | 0 | let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?; |
206 | 0 | if let Some(ref bias) = layer.ffn_gate_bias { |
207 | 0 | ops::add_bias(&mut ffn_gate, bias); |
208 | 0 | } |
209 | | |
210 | | // SwiGLU: down(silu(gate(x)) * up(x)) |
211 | | // Apply SiLU to GATE projection, not up |
212 | 0 | ops::silu(&mut ffn_gate); |
213 | | |
214 | | // Element-wise multiply: silu(gate) * up |
215 | 0 | for i in 0..ffn_gate.len() { |
216 | 0 | ffn_gate[i] *= ffn_up[i]; |
217 | 0 | } |
218 | | |
219 | 0 | ffn_gate |
220 | | } else { |
221 | | // GELU path (phi-2, GPT-2, etc.) |
222 | 4 | let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?0 ; |
223 | 4 | if let Some(ref bias0 ) = layer.ffn_up_bias { |
224 | 0 | ops::add_bias(&mut ffn_hidden, bias); |
225 | 4 | } |
226 | 4 | ops::gelu(&mut ffn_hidden); |
227 | 4 | ffn_hidden |
228 | | }; |
229 | | |
230 | | // 2g. FFN down projection with FUSED ops |
231 | 4 | let mut ffn_output = self.fused_matmul(&ffn_activated, &layer.ffn_down_weight)?0 ; |
232 | 4 | if let Some(ref bias0 ) = layer.ffn_down_bias { |
233 | 0 | ops::add_bias(&mut ffn_output, bias); |
234 | 4 | } |
235 | | |
236 | | // Residual connection |
237 | 480 | for i in 0..hidden4 .len4 () { |
238 | 480 | hidden[i] += ffn_output[i]; |
239 | 480 | } |
240 | | |
241 | | // CORRECTNESS-011: Per-layer CPU debug output |
242 | 4 | if cpu_debug_layers { |
243 | 0 | let seq_len = token_ids.len(); |
244 | 0 | let last_hidden_start = (seq_len - 1) * hidden_dim; |
245 | 0 | let last_h = &hidden[last_hidden_start..last_hidden_start + hidden_dim]; |
246 | 0 | let sum: f32 = last_h.iter().sum(); |
247 | 0 | let sq_sum: f32 = last_h.iter().map(|x| x * x).sum(); |
248 | 0 | let rms = (sq_sum / last_h.len() as f32).sqrt(); |
249 | 0 | eprintln!( |
250 | 0 | "[CPU-L{}] After layer: first 3 = [{:.4}, {:.4}, {:.4}], sum = {:.4}, rms = {:.4}", |
251 | 0 | layer_idx, last_h[0], last_h[1], last_h[2], sum, rms |
252 | | ); |
253 | 4 | } |
254 | | } |
255 | | |
256 | | // CORRECTNESS-011: CPU hidden state debug output (compare with GPU) |
257 | 4 | if std::env::var("CPU_DEBUG").is_ok() { |
258 | 0 | let seq_len = token_ids.len(); |
259 | 0 | let last_hidden_start = (seq_len - 1) * hidden_dim; |
260 | 0 | let last_hidden_raw = &hidden[last_hidden_start..last_hidden_start + hidden_dim]; |
261 | | |
262 | 0 | let sum: f32 = last_hidden_raw.iter().sum(); |
263 | 0 | let sq_sum: f32 = last_hidden_raw.iter().map(|x| x * x).sum(); |
264 | 0 | let rms = (sq_sum / last_hidden_raw.len() as f32).sqrt(); |
265 | | |
266 | 0 | eprintln!("[CORRECTNESS-011] CPU Hidden before output_norm:"); |
267 | 0 | eprintln!( |
268 | 0 | " first 5 = {:?}", |
269 | 0 | &last_hidden_raw[..5.min(last_hidden_raw.len())] |
270 | | ); |
271 | 0 | eprintln!(" sum = {:.4}, rms = {:.4}", sum, rms); |
272 | 0 | eprintln!(" (GPU shows: sum=466.2486, rms=39.4793)"); |
273 | 4 | } |
274 | | |
275 | | // 3. Final layer norm (RMSNorm for LLaMA, LayerNorm for others) |
276 | 4 | let normed = if use_rmsnorm { |
277 | 0 | ops::rms_norm(&hidden, &self.output_norm_weight, self.config.eps) |
278 | | } else { |
279 | 4 | ops::layer_norm( |
280 | 4 | &hidden, |
281 | 4 | &self.output_norm_weight, |
282 | 4 | self.output_norm_bias.as_deref(), |
283 | 4 | self.config.eps, |
284 | | ) |
285 | | }; |
286 | | |
287 | | // CORRECTNESS-011: CPU normed hidden state debug output |
288 | 4 | if std::env::var("CPU_DEBUG").is_ok() { |
289 | 0 | let seq_len = token_ids.len(); |
290 | 0 | let last_normed_start = (seq_len - 1) * hidden_dim; |
291 | 0 | let last_normed = &normed[last_normed_start..last_normed_start + hidden_dim]; |
292 | | |
293 | 0 | let sum: f32 = last_normed.iter().sum(); |
294 | 0 | let sq_sum: f32 = last_normed.iter().map(|x| x * x).sum(); |
295 | 0 | let rms = (sq_sum / last_normed.len() as f32).sqrt(); |
296 | | |
297 | 0 | eprintln!("[CORRECTNESS-011] CPU Normed hidden:"); |
298 | 0 | eprintln!(" first 5 = {:?}", &last_normed[..5.min(last_normed.len())]); |
299 | 0 | eprintln!(" sum = {:.4}, rms = {:.4}", sum, rms); |
300 | 0 | eprintln!(" (GPU shows: sum=107.5945, rms=4.6616)"); |
301 | 4 | } |
302 | | |
303 | | // 4. LM head projection with FUSED ops (only last token) |
304 | 4 | let seq_len = token_ids.len(); |
305 | 4 | let last_hidden_start = (seq_len - 1) * hidden_dim; |
306 | 4 | let last_hidden = &normed[last_hidden_start..last_hidden_start + hidden_dim]; |
307 | | |
308 | | // Compute logits using fused op |
309 | 4 | let mut logits = self.fused_matmul(last_hidden, &self.lm_head_weight)?0 ; |
310 | | |
311 | 4 | if let Some(ref bias0 ) = self.lm_head_bias { |
312 | 0 | ops::add_bias(&mut logits, bias); |
313 | 4 | } |
314 | | |
315 | 4 | Ok(logits) |
316 | 4 | } |
317 | | |
318 | | /// Forward pass with KV cache for efficient autoregressive decoding |
319 | | /// |
320 | | /// This method properly handles both architectures: |
321 | | /// - LLaMA-style: RMSNorm, SwiGLU FFN, GQA attention |
322 | | /// - phi-2 style: LayerNorm, GELU FFN, MHA attention |
323 | | /// |
324 | | /// Uses O(n) per-token cost instead of O(n²) by caching K/V. |
325 | | /// |
326 | | /// # Arguments |
327 | | /// * `token_id` - Token to process |
328 | | /// * `cache` - KV cache for all layers |
329 | | /// * `position` - Position in sequence for RoPE |
330 | | /// |
331 | | /// # Returns |
332 | | /// Logits for next token prediction [vocab_size] |
333 | 16 | pub fn forward_cached( |
334 | 16 | &self, |
335 | 16 | token_id: u32, |
336 | 16 | cache: &mut crate::gguf::OwnedQuantizedKVCache, |
337 | 16 | position: usize, |
338 | 16 | ) -> Result<Vec<f32>> { |
339 | 16 | let hidden_dim = self.config.hidden_dim; |
340 | | |
341 | | // Detect architecture: LLaMA uses RMSNorm (no bias) and SwiGLU (has gate weight) |
342 | 16 | let use_rmsnorm = self |
343 | 16 | .layers |
344 | 16 | .first() |
345 | 16 | .is_some_and(|l| l.ffn_gate_weight.is_some() && l.attn_norm_bias0 .is_none0 ()); |
346 | | |
347 | | // 1. Token embedding lookup |
348 | 16 | let mut hidden = self.embed(&[token_id]); |
349 | | |
350 | | // PAR-052: Debug output for OwnedQuantizedModel forward path |
351 | 16 | let debug_forward = std::env::var("REALIZAR_DEBUG_FORWARD").is_ok(); |
352 | 16 | if debug_forward && position == 00 { |
353 | 0 | eprintln!("[PAR-052] OwnedQuantizedModel::forward_cached"); |
354 | 0 | eprintln!("[PAR-052] Token ID: {}, Position: {}", token_id, position); |
355 | 0 | eprintln!("[PAR-052] use_rmsnorm: {}", use_rmsnorm); |
356 | 0 | eprintln!( |
357 | 0 | "[PAR-052] Embedding[0..8]: {:?}", |
358 | 0 | &hidden[..8.min(hidden.len())] |
359 | 0 | ); |
360 | 0 | eprintln!("[PAR-052] Embedding sum: {:.6}", hidden.iter().sum::<f32>()); |
361 | 16 | } |
362 | | |
363 | | // 2. Process through transformer layers |
364 | 16 | for (layer_idx, layer) in self.layers.iter().enumerate() { |
365 | | // 2a. Attention layer norm (RMSNorm for LLaMA, LayerNorm for phi-2) |
366 | 16 | let normed = if use_rmsnorm { |
367 | 0 | ops::rms_norm(&hidden, &layer.attn_norm_weight, self.config.eps) |
368 | | } else { |
369 | 16 | ops::layer_norm( |
370 | 16 | &hidden, |
371 | 16 | &layer.attn_norm_weight, |
372 | 16 | layer.attn_norm_bias.as_deref(), |
373 | 16 | self.config.eps, |
374 | | ) |
375 | | }; |
376 | | |
377 | | // PAR-052: Debug layer 0 normed and QKV values |
378 | 16 | if debug_forward && layer_idx == 00 && position == 00 { |
379 | 0 | eprintln!( |
380 | 0 | "[PAR-052-L0] attn_norm[0..8]: {:?}", |
381 | 0 | &layer.attn_norm_weight[..8.min(layer.attn_norm_weight.len())] |
382 | 0 | ); |
383 | 0 | eprintln!( |
384 | 0 | "[PAR-052-L0] normed[0..8]: {:?}", |
385 | 0 | &normed[..8.min(normed.len())] |
386 | 0 | ); |
387 | 0 | eprintln!("[PAR-052-L0] normed sum: {:.6}", normed.iter().sum::<f32>()); |
388 | 16 | } |
389 | | |
390 | | // 2b. QKV projection |
391 | 16 | let _qkv_dim = layer.qkv_weight.out_dim(); |
392 | 16 | let q_dim = layer.qkv_weight.q_dim(); |
393 | 16 | let k_dim = match &layer.qkv_weight { |
394 | 16 | OwnedQKVWeights::Fused(_) => q_dim, |
395 | 0 | OwnedQKVWeights::Separate { k, .. } => k.out_dim, |
396 | | }; |
397 | 16 | let v_dim = match &layer.qkv_weight { |
398 | 16 | OwnedQKVWeights::Fused(_) => q_dim, |
399 | 0 | OwnedQKVWeights::Separate { v, .. } => v.out_dim, |
400 | | }; |
401 | | |
402 | 16 | let mut qkv = self.qkv_matmul(&normed, &layer.qkv_weight)?0 ; |
403 | 16 | if let Some(ref bias0 ) = layer.qkv_bias { |
404 | 0 | ops::add_bias(&mut qkv, bias); |
405 | 16 | } |
406 | | |
407 | | // PAR-052: Debug QKV after projection |
408 | 16 | if debug_forward && layer_idx == 00 && position == 00 { |
409 | 0 | eprintln!( |
410 | 0 | "[PAR-052-L0] QKV dims: q={}, k={}, v={}, total={}", |
411 | 0 | q_dim, |
412 | 0 | k_dim, |
413 | 0 | v_dim, |
414 | 0 | qkv.len() |
415 | 0 | ); |
416 | 0 | eprintln!("[PAR-052-L0] QKV sum: {:.6}", qkv.iter().sum::<f32>()); |
417 | 0 | eprintln!("[PAR-052-L0] Q[0..8]: {:?}", &qkv[..8.min(q_dim)]); |
418 | 16 | } |
419 | | |
420 | | // 2c. Extract Q, K, V and apply RoPE |
421 | 16 | let mut q = qkv[0..q_dim].to_vec(); |
422 | 16 | let mut k = qkv[q_dim..q_dim + k_dim].to_vec(); |
423 | 16 | let v = qkv[q_dim + k_dim..q_dim + k_dim + v_dim].to_vec(); |
424 | | |
425 | | // Apply RoPE with correct head counts for GQA |
426 | 16 | self.apply_rope(&mut q, position, self.config.num_heads); |
427 | 16 | self.apply_rope(&mut k, position, self.config.num_kv_heads); |
428 | | |
429 | | // PAR-052: Debug Q after RoPE |
430 | 16 | if debug_forward && layer_idx == 00 && position == 00 { |
431 | 0 | eprintln!( |
432 | 0 | "[PAR-052-L0] Q after RoPE[0..8]: {:?}", |
433 | 0 | &q[..8.min(q.len())] |
434 | 0 | ); |
435 | 0 | eprintln!( |
436 | 0 | "[PAR-052-L0] K after RoPE[0..4]: {:?}", |
437 | 0 | &k[..4.min(k.len())] |
438 | 0 | ); |
439 | 16 | } |
440 | | |
441 | | // 2d. Compute attention using cached K/V |
442 | 16 | let k_cache = cache.get_k(layer_idx); |
443 | 16 | let v_cache = cache.get_v(layer_idx); |
444 | | |
445 | 16 | let attn_out = if k_cache.is_empty() { |
446 | | // First token - just use V directly (self-attention with single token) |
447 | | // Expand V if GQA (num_kv_heads < num_heads) |
448 | 3 | if self.config.num_kv_heads < self.config.num_heads { |
449 | 0 | let head_dim = hidden_dim / self.config.num_heads; |
450 | 0 | let group_size = self.config.num_heads / self.config.num_kv_heads; |
451 | 0 | (0..self.config.num_heads) |
452 | 0 | .flat_map(|h| { |
453 | 0 | let kv_head = h / group_size; |
454 | 0 | let start = kv_head * head_dim; |
455 | 0 | v[start..start + head_dim].iter().copied() |
456 | 0 | }) |
457 | 0 | .collect() |
458 | | } else { |
459 | 3 | v.clone() |
460 | | } |
461 | | } else { |
462 | | // Use cached K/V for attention with GQA support |
463 | 13 | self.attention_with_cache_gqa(&q, k_cache, v_cache, &k, &v) |
464 | | }; |
465 | | |
466 | | // 2e. Store K and V in cache (store original size, not expanded) |
467 | 16 | cache.append(layer_idx, &k, &v); |
468 | | |
469 | | // 2f. Attention output projection |
470 | 16 | let mut attn_output = self.fused_matmul(&attn_out, &layer.attn_output_weight)?0 ; |
471 | 16 | if let Some(ref bias0 ) = layer.attn_output_bias { |
472 | 0 | ops::add_bias(&mut attn_output, bias); |
473 | 16 | } |
474 | | |
475 | | // 2g. Residual connection |
476 | 1.02k | for i in 0..hidden_dim16 { |
477 | 1.02k | hidden[i] += attn_output[i]; |
478 | 1.02k | } |
479 | | |
480 | | // 2h. Pre-FFN layer norm (LLaMA has separate ffn_norm) |
481 | 16 | let ffn_input = if let Some(ref ffn_norm0 ) = layer.ffn_norm_weight { |
482 | 0 | if use_rmsnorm { |
483 | 0 | ops::rms_norm(&hidden, ffn_norm, self.config.eps) |
484 | | } else { |
485 | 0 | ops::layer_norm( |
486 | 0 | &hidden, |
487 | 0 | ffn_norm, |
488 | 0 | layer.ffn_norm_bias.as_deref(), |
489 | 0 | self.config.eps, |
490 | | ) |
491 | | } |
492 | | } else { |
493 | 16 | hidden.clone() |
494 | | }; |
495 | | |
496 | | // 2i. FFN with SwiGLU or GELU |
497 | 16 | let ffn_output = if let Some(ref gate_weight0 ) = layer.ffn_gate_weight { |
498 | | // SwiGLU path (LLaMA) |
499 | 0 | let mut ffn_up = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?; |
500 | 0 | if let Some(ref bias) = layer.ffn_up_bias { |
501 | 0 | ops::add_bias(&mut ffn_up, bias); |
502 | 0 | } |
503 | | |
504 | 0 | let mut ffn_gate = self.fused_matmul(&ffn_input, gate_weight)?; |
505 | 0 | if let Some(ref bias) = layer.ffn_gate_bias { |
506 | 0 | ops::add_bias(&mut ffn_gate, bias); |
507 | 0 | } |
508 | | |
509 | | // SiLU on gate, then multiply with up |
510 | 0 | ops::silu(&mut ffn_gate); |
511 | 0 | for i in 0..ffn_gate.len() { |
512 | 0 | ffn_gate[i] *= ffn_up[i]; |
513 | 0 | } |
514 | | |
515 | 0 | let mut output = self.fused_matmul(&ffn_gate, &layer.ffn_down_weight)?; |
516 | 0 | if let Some(ref bias) = layer.ffn_down_bias { |
517 | 0 | ops::add_bias(&mut output, bias); |
518 | 0 | } |
519 | 0 | output |
520 | | } else { |
521 | | // GELU path (phi-2) |
522 | 16 | let mut ffn_hidden = self.fused_matmul(&ffn_input, &layer.ffn_up_weight)?0 ; |
523 | 16 | if let Some(ref bias0 ) = layer.ffn_up_bias { |
524 | 0 | ops::add_bias(&mut ffn_hidden, bias); |
525 | 16 | } |
526 | 16 | ops::gelu(&mut ffn_hidden); |
527 | | |
528 | 16 | let mut output = self.fused_matmul(&ffn_hidden, &layer.ffn_down_weight)?0 ; |
529 | 16 | if let Some(ref bias0 ) = layer.ffn_down_bias { |
530 | 0 | ops::add_bias(&mut output, bias); |
531 | 16 | } |
532 | 16 | output |
533 | | }; |
534 | | |
535 | | // 2j. Residual connection |
536 | 1.02k | for i in 0..hidden_dim16 { |
537 | 1.02k | hidden[i] += ffn_output[i]; |
538 | 1.02k | } |
539 | | } |
540 | | |
541 | | // Advance cache position |
542 | 16 | cache.advance(); |
543 | | |
544 | | // 3. Final layer norm |
545 | 16 | let normed = if use_rmsnorm { |
546 | 0 | ops::rms_norm(&hidden, &self.output_norm_weight, self.config.eps) |
547 | | } else { |
548 | 16 | ops::layer_norm( |
549 | 16 | &hidden, |
550 | 16 | &self.output_norm_weight, |
551 | 16 | self.output_norm_bias.as_deref(), |
552 | 16 | self.config.eps, |
553 | | ) |
554 | | }; |
555 | | |
556 | | // PAR-052: Debug final hidden state |
557 | 16 | if debug_forward && position == 00 { |
558 | 0 | eprintln!( |
559 | 0 | "[PAR-052] Final hidden sum: {:.6}", |
560 | 0 | hidden.iter().sum::<f32>() |
561 | 0 | ); |
562 | 0 | eprintln!( |
563 | 0 | "[PAR-052] Final hidden[0..8]: {:?}", |
564 | 0 | &hidden[..8.min(hidden.len())] |
565 | 0 | ); |
566 | 0 | eprintln!( |
567 | 0 | "[PAR-052] After output_norm sum: {:.6}", |
568 | 0 | normed.iter().sum::<f32>() |
569 | 0 | ); |
570 | 0 | eprintln!( |
571 | 0 | "[PAR-052] output_norm_weight[0..4]: {:?}", |
572 | 0 | &self.output_norm_weight[..4.min(self.output_norm_weight.len())] |
573 | 0 | ); |
574 | 0 | eprintln!( |
575 | 0 | "[PAR-052] LM head weight dims: in={}, out={}", |
576 | 0 | self.lm_head_weight.in_dim, self.lm_head_weight.out_dim |
577 | 0 | ); |
578 | 16 | } |
579 | | |
580 | | // 4. LM head projection |
581 | 16 | let mut logits = self.fused_matmul(&normed, &self.lm_head_weight)?0 ; |
582 | 16 | if let Some(ref bias0 ) = self.lm_head_bias { |
583 | 0 | ops::add_bias(&mut logits, bias); |
584 | 16 | } |
585 | | |
586 | | // PAR-052: Debug final logits |
587 | 16 | if debug_forward && position == 00 { |
588 | | // Find top-5 logits |
589 | 0 | let mut indexed: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect(); |
590 | 0 | indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); |
591 | 0 | eprintln!("[PAR-052] Top-5 logits:"); |
592 | 0 | for (idx, val) in indexed.iter().take(5) { |
593 | 0 | eprintln!(" Token {}: {:.6}", idx, val); |
594 | 0 | } |
595 | 0 | eprintln!("[PAR-052] Logits sum: {:.6}", logits.iter().sum::<f32>()); |
596 | 16 | } |
597 | | |
598 | 16 | Ok(logits) |
599 | 16 | } |
600 | | } |