/home/noah/src/realizar/src/cli/inference.rs
Line | Count | Source |
1 | | //! Inference runner functions for CLI commands |
2 | | //! |
3 | | //! Contains run_gguf_inference, run_gguf_inference_gpu, run_safetensors_inference, |
4 | | //! and run_apr_inference - extracted from main.rs (PMAT-802). |
5 | | |
6 | | #![allow(missing_docs)] |
7 | | #![allow(clippy::missing_errors_doc)] |
8 | | #![allow(clippy::too_many_arguments)] |
9 | | #![allow(clippy::needless_pass_by_value)] |
10 | | |
11 | | use crate::error::Result; |
12 | | |
13 | 0 | pub fn run_gguf_inference( |
14 | 0 | model_ref: &str, |
15 | 0 | _file_data: &[u8], |
16 | 0 | prompt: &str, |
17 | 0 | max_tokens: usize, |
18 | 0 | temperature: f32, |
19 | 0 | format: &str, |
20 | 0 | force_gpu: bool, |
21 | 0 | verbose: bool, |
22 | 0 | ) -> Result<()> { |
23 | | use crate::gguf::{MappedGGUFModel, OwnedQuantizedKVCache, OwnedQuantizedModel}; |
24 | | use std::time::Instant; |
25 | | |
26 | | // Handle --gpu flag warning when CUDA not available |
27 | | #[cfg(not(feature = "cuda"))] |
28 | 0 | if force_gpu { |
29 | 0 | eprintln!("Warning: --gpu flag requires 'cuda' feature. Falling back to CPU."); |
30 | 0 | eprintln!("Build with: cargo build --features cuda"); |
31 | 0 | eprintln!(); |
32 | 0 | } |
33 | | // Suppress unused warning when cuda feature not enabled |
34 | | #[cfg(not(feature = "cuda"))] |
35 | 0 | let _ = force_gpu; |
36 | | |
37 | 0 | let load_start = Instant::now(); |
38 | | |
39 | | // Load model using memory-mapped file (same path as working examples) |
40 | 0 | let mapped = MappedGGUFModel::from_path(model_ref).map_err(|e| { |
41 | 0 | crate::error::RealizarError::UnsupportedOperation { |
42 | 0 | operation: "mmap_gguf".to_string(), |
43 | 0 | reason: format!("Failed to mmap GGUF: {e}"), |
44 | 0 | } |
45 | 0 | })?; |
46 | | |
47 | | // GPU path: Use OwnedQuantizedModel with CUDA acceleration |
48 | | #[cfg(feature = "cuda")] |
49 | | if force_gpu { |
50 | | return run_gguf_inference_gpu( |
51 | | &mapped, |
52 | | prompt, |
53 | | max_tokens, |
54 | | temperature, |
55 | | format, |
56 | | load_start, |
57 | | verbose, |
58 | | ); |
59 | | } |
60 | | |
61 | | // PAR-126: Five-Whys fix - use OwnedQuantizedModel for fast CPU inference |
62 | | // Root cause analysis: |
63 | | // Why-1: CPU path was 14 tok/s vs Ollama's 200 tok/s |
64 | | // Why-2: Old mmap-based transformers use per-matmul allocations |
65 | | // Why-3: Each of 196 matmuls per token allocates/frees Vec |
66 | | // Why-4: Vec allocation overhead + cache pollution from mmap page faults |
67 | | // Why-5: OwnedQuantizedModel copies weights to RAM but uses _into methods |
68 | | // Solution: Use OwnedQuantizedModel - slower loading but faster inference |
69 | 0 | let model = crate::gguf::OwnedQuantizedModel::from_mapped(&mapped).map_err(|e| { |
70 | 0 | crate::error::RealizarError::UnsupportedOperation { |
71 | 0 | operation: "load_model".to_string(), |
72 | 0 | reason: format!("Failed to load model: {e}"), |
73 | 0 | } |
74 | 0 | })?; |
75 | | |
76 | 0 | let load_time = load_start.elapsed(); |
77 | 0 | if verbose { |
78 | 0 | println!("Backend: CPU (AVX2 + SIMD)"); |
79 | 0 | println!("Model loaded in {:.2}ms", load_time.as_secs_f64() * 1000.0); |
80 | 0 | } |
81 | | |
82 | | // Tokenize prompt using GGUF vocabulary |
83 | 0 | let mut prompt_tokens: Vec<u32> = mapped |
84 | 0 | .model |
85 | 0 | .encode(prompt) |
86 | 0 | .unwrap_or_else(|| prompt.chars().map(|c| c as u32).collect()); |
87 | | |
88 | | // Prepend BOS token if available |
89 | 0 | if let Some(bos) = mapped.model.bos_token_id() { |
90 | 0 | prompt_tokens.insert(0, bos); |
91 | 0 | } |
92 | 0 | let prompt_len = prompt_tokens.len(); |
93 | | |
94 | | // Get EOS token for stopping |
95 | 0 | let eos_token_id = mapped.model.eos_token_id(); |
96 | | |
97 | | // Debug: show model info and encoded tokens |
98 | 0 | let config = model.config(); |
99 | 0 | if verbose { |
100 | 0 | println!( |
101 | 0 | "Architecture: {:?}, Hidden: {}, Layers: {}, Heads: {}/{} (KV)", |
102 | 0 | mapped.model.architecture(), |
103 | 0 | config.hidden_dim, |
104 | 0 | config.num_layers, |
105 | 0 | config.num_heads, |
106 | 0 | config.num_kv_heads |
107 | 0 | ); |
108 | 0 | println!( |
109 | 0 | "Prompt tokens: {} (BOS={:?}, EOS={:?})", |
110 | 0 | prompt_len, |
111 | 0 | mapped.model.bos_token_id(), |
112 | 0 | eos_token_id |
113 | 0 | ); |
114 | 0 | println!("Temperature: {:.1}", temperature); |
115 | 0 | println!(); |
116 | 0 | } |
117 | | |
118 | | // Run inference with KV cache for O(n) per-token cost |
119 | 0 | let gen_start = Instant::now(); |
120 | 0 | let max_seq_len = prompt_tokens.len() + max_tokens; |
121 | 0 | let mut cache = OwnedQuantizedKVCache::from_config(config, max_seq_len); |
122 | 0 | let mut all_tokens = prompt_tokens.clone(); |
123 | | |
124 | | // Prefill: process prompt tokens to populate KV cache |
125 | | // We process all tokens but keep the logits from the last one |
126 | 0 | let mut logits: Vec<f32> = vec![]; |
127 | 0 | for (pos, &token_id) in prompt_tokens.iter().enumerate() { |
128 | 0 | logits = model.forward_cached(token_id, &mut cache, pos)?; |
129 | | } |
130 | | |
131 | | // PAR-051: Diagnostic - show top5 logits after prefill |
132 | | // Re-enable by changing `false` to `true` for debugging |
133 | | #[allow(clippy::never_loop)] |
134 | 0 | if std::env::var("CPU_DEBUG") |
135 | 0 | .map(|v| v == "1") |
136 | 0 | .unwrap_or(false) |
137 | | { |
138 | 0 | let mut top5: Vec<(usize, f32)> = logits.iter().copied().enumerate().collect(); |
139 | 0 | top5.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)); |
140 | 0 | top5.truncate(5); |
141 | 0 | eprintln!("[PAR-051] Prompt tokens: {:?}", prompt_tokens); |
142 | 0 | eprintln!("[PAR-051] Logits top5 after prefill: {:?}", top5); |
143 | 0 | let greedy_token = top5[0].0 as u32; |
144 | 0 | let decoded = mapped.model.decode(&[greedy_token]); |
145 | 0 | eprintln!( |
146 | 0 | "[PAR-051] Greedy next token: {} = {:?}", |
147 | | greedy_token, decoded |
148 | | ); |
149 | | |
150 | | // Check embedding - compare first 5 values |
151 | 0 | let last_token = prompt_tokens[prompt_tokens.len() - 1]; |
152 | 0 | let embed = model.embed(&[last_token]); |
153 | 0 | eprintln!( |
154 | 0 | "[PAR-051] Last token {} embed[0..5]: {:?}", |
155 | | last_token, |
156 | 0 | &embed[..5.min(embed.len())] |
157 | | ); |
158 | 0 | eprintln!("[PAR-051] embed sum: {:.6}", embed.iter().sum::<f32>()); |
159 | | |
160 | | // Check model config |
161 | 0 | let cfg = model.config(); |
162 | 0 | eprintln!( |
163 | 0 | "[PAR-051] Config: hidden={}, heads={}/{}, layers={}, vocab={}, eps={:e}", |
164 | | cfg.hidden_dim, |
165 | | cfg.num_heads, |
166 | | cfg.num_kv_heads, |
167 | | cfg.num_layers, |
168 | | cfg.vocab_size, |
169 | | cfg.eps |
170 | | ); |
171 | | |
172 | | // Check logit at index 29906 (token "2") |
173 | 0 | let logit_2 = logits.get(29906).copied().unwrap_or(f32::NAN); |
174 | 0 | eprintln!("[PAR-051] Logit for token 29906 ('2'): {:.6}", logit_2); |
175 | 0 | } |
176 | | |
177 | | // Decode: generate new tokens one at a time |
178 | | // First iteration uses logits from prefill, subsequent use new logits |
179 | 0 | for i in 0..max_tokens { |
180 | | // For first iteration, use logits from prefill; otherwise compute new ones |
181 | 0 | if i > 0 { |
182 | 0 | let position = prompt_tokens.len() + i - 1; |
183 | 0 | let last_token = *all_tokens |
184 | 0 | .last() |
185 | 0 | .expect("all_tokens should not be empty during generation"); |
186 | 0 | logits = model.forward_cached(last_token, &mut cache, position)?; |
187 | 0 | } |
188 | | |
189 | | // Sample next token |
190 | 0 | let next_token = if temperature <= 0.01 { |
191 | | // Greedy decoding |
192 | 0 | logits |
193 | 0 | .iter() |
194 | 0 | .enumerate() |
195 | 0 | .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)) |
196 | 0 | .map_or(0, |(idx, _)| idx as u32) |
197 | | } else { |
198 | | // Temperature sampling |
199 | 0 | OwnedQuantizedModel::sample_topk(&logits, temperature, 40) |
200 | | }; |
201 | | |
202 | | // PERF-002: Debug code removed (was PAR-058-DEBUG and PAR-060) |
203 | | |
204 | | // Stop on EOS |
205 | 0 | if let Some(eos) = eos_token_id { |
206 | 0 | if next_token == eos { |
207 | | // PERF-002: eprintln removed for performance |
208 | 0 | break; |
209 | 0 | } |
210 | 0 | } |
211 | | |
212 | 0 | all_tokens.push(next_token); |
213 | | } |
214 | | |
215 | 0 | let generated = all_tokens; |
216 | 0 | let gen_time = gen_start.elapsed(); |
217 | | |
218 | 0 | let tokens_generated = generated.len() - prompt_len; |
219 | 0 | let tokens_per_sec = if gen_time.as_secs_f64() > 0.0 { |
220 | 0 | tokens_generated as f64 / gen_time.as_secs_f64() |
221 | | } else { |
222 | 0 | 0.0 |
223 | | }; |
224 | | |
225 | | // Decode output using GGUF vocabulary, replacing SentencePiece markers with spaces |
226 | 0 | let output_text = mapped |
227 | 0 | .model |
228 | 0 | .decode(&generated[prompt_len..]) |
229 | 0 | .replace('▁', " "); |
230 | | |
231 | 0 | match format { |
232 | 0 | "json" => { |
233 | 0 | let json = serde_json::json!({ |
234 | 0 | "model": model_ref, |
235 | 0 | "prompt": prompt, |
236 | 0 | "generated_text": output_text, |
237 | 0 | "tokens_generated": tokens_generated, |
238 | 0 | "generation_time_ms": gen_time.as_secs_f64() * 1000.0, |
239 | 0 | "tokens_per_second": tokens_per_sec, |
240 | 0 | "temperature": temperature, |
241 | 0 | }); |
242 | 0 | println!( |
243 | 0 | "{}", |
244 | 0 | serde_json::to_string_pretty(&json).unwrap_or_default() |
245 | 0 | ); |
246 | 0 | }, |
247 | | _ => { |
248 | 0 | if verbose { |
249 | 0 | println!( |
250 | 0 | "Generated ({tokens_generated} tokens in {:.2}ms):", |
251 | 0 | gen_time.as_secs_f64() * 1000.0 |
252 | 0 | ); |
253 | 0 | println!("{prompt}{output_text}"); |
254 | 0 | println!(); |
255 | 0 | println!("Performance: {:.1} tok/s", tokens_per_sec); |
256 | 0 | } else { |
257 | 0 | // Ollama-style clean output: just the response |
258 | 0 | println!("{output_text}"); |
259 | 0 | } |
260 | | }, |
261 | | } |
262 | | |
263 | 0 | Ok(()) |
264 | 0 | } |
265 | | |
266 | | /// Run GGUF inference with CUDA GPU acceleration |
267 | | /// |
268 | | /// Uses OwnedQuantizedModel with CUDA backend for high-performance inference. |
269 | | /// Called when --gpu flag is specified and CUDA feature is enabled. |
270 | | #[cfg(feature = "cuda")] |
271 | | #[allow(clippy::too_many_arguments)] |
272 | | pub fn run_gguf_inference_gpu( |
273 | | mapped: &crate::gguf::MappedGGUFModel, |
274 | | prompt: &str, |
275 | | max_tokens: usize, |
276 | | temperature: f32, |
277 | | format: &str, |
278 | | load_start: std::time::Instant, |
279 | | verbose: bool, |
280 | | ) -> Result<()> { |
281 | | use crate::gguf::{OwnedQuantizedModel, OwnedQuantizedModelCuda, QuantizedGenerateConfig}; |
282 | | use std::time::Instant; |
283 | | |
284 | | if verbose { |
285 | | println!("Backend: CUDA (GPU)"); |
286 | | println!("Creating quantized model with CUDA acceleration..."); |
287 | | } |
288 | | |
289 | | // Create owned quantized model (required for CUDA - can't use borrowed mmap data) |
290 | | let quantized_model = OwnedQuantizedModel::from_mapped(mapped).map_err(|e| { |
291 | | crate::error::RealizarError::UnsupportedOperation { |
292 | | operation: "create_quantized".to_string(), |
293 | | reason: format!("Failed to create quantized model: {e}"), |
294 | | } |
295 | | })?; |
296 | | |
297 | | // Get config info before wrapping |
298 | | let vocab_size = quantized_model.config.vocab_size; |
299 | | let hidden_dim = quantized_model.config.hidden_dim; |
300 | | let num_layers = quantized_model.layers.len(); |
301 | | |
302 | | // PAR-046: Create OwnedQuantizedModelCuda wrapper for actual GPU acceleration |
303 | | // The previous implementation used OwnedQuantizedModel.enable_cuda() which only |
304 | | // initialized the executor but forward_cached still used CPU code paths. |
305 | | let max_seq_len = 256 + max_tokens; // Allow for prompt + generation |
306 | | let mut cuda_model = OwnedQuantizedModelCuda::with_max_seq_len(quantized_model, 0, max_seq_len) |
307 | | .map_err(|e| crate::error::RealizarError::UnsupportedOperation { |
308 | | operation: "OwnedQuantizedModelCuda::new".to_string(), |
309 | | reason: format!("CUDA initialization failed: {e}"), |
310 | | })?; |
311 | | if verbose { |
312 | | println!(" CUDA enabled on GPU: {}", cuda_model.device_name()); |
313 | | } |
314 | | |
315 | | let load_time = load_start.elapsed(); |
316 | | if verbose { |
317 | | println!("Model loaded in {:.2}ms", load_time.as_secs_f64() * 1000.0); |
318 | | } |
319 | | |
320 | | // Tokenize prompt using GGUF vocabulary |
321 | | let mut prompt_tokens: Vec<u32> = mapped |
322 | | .model |
323 | | .encode(prompt) |
324 | | .unwrap_or_else(|| prompt.chars().map(|c| c as u32).collect()); |
325 | | |
326 | | // Prepend BOS token if available |
327 | | if let Some(bos) = mapped.model.bos_token_id() { |
328 | | prompt_tokens.insert(0, bos); |
329 | | } |
330 | | let prompt_len = prompt_tokens.len(); |
331 | | |
332 | | // Get EOS token for stopping |
333 | | let eos_token_id = mapped.model.eos_token_id(); |
334 | | |
335 | | if verbose { |
336 | | println!( |
337 | | "Vocab size: {}, Hidden dim: {}, Layers: {}", |
338 | | vocab_size, hidden_dim, num_layers |
339 | | ); |
340 | | println!( |
341 | | "Prompt tokens: {} (BOS={:?}, EOS={:?})", |
342 | | prompt_len, |
343 | | mapped.model.bos_token_id(), |
344 | | eos_token_id |
345 | | ); |
346 | | println!("Temperature: {:.1}", temperature); |
347 | | println!(); |
348 | | } |
349 | | |
350 | | // PAR-046: Use CUDA-accelerated generation with GPU-resident KV cache |
351 | | // This calls generate_cuda_with_cache -> forward_single_cuda_with_cache -> GPU kernels |
352 | | let gen_start = Instant::now(); |
353 | | |
354 | | // Build stop tokens list |
355 | | let mut stop_tokens = Vec::new(); |
356 | | if let Some(eos) = eos_token_id { |
357 | | stop_tokens.push(eos); |
358 | | } |
359 | | |
360 | | // Configure CUDA generation |
361 | | let gen_config = QuantizedGenerateConfig { |
362 | | max_tokens, |
363 | | temperature, |
364 | | top_k: if temperature <= 0.01 { 1 } else { 40 }, |
365 | | stop_tokens, |
366 | | }; |
367 | | |
368 | | // PAR-047: Use generate_full_cuda_with_cache for maximum GPU acceleration |
369 | | // This path uses: |
370 | | // - GPU matmul for QKV, output projection, and FFN |
371 | | // - GPU incremental_attention_gpu with GQA support (PAR-021) |
372 | | // - Proper SwiGLU activation (PAR-015) |
373 | | // PAR-057: Use GPU-resident path for maximum performance (pre-uploads weights, minimal syncs) |
374 | | // Falls back to generate_full_cuda_with_cache if architecture not supported |
375 | | // PAR-058: Test GPU-resident vs standard CUDA path |
376 | | // PHASE-13: Skip GPU-resident if SKIP_GPU_RESIDENT=1 (for debugging) |
377 | | let skip_gpu_resident = std::env::var("SKIP_GPU_RESIDENT") |
378 | | .map(|v| v == "1") |
379 | | .unwrap_or(false); |
380 | | let generated = if cuda_model.supports_gpu_resident() && !skip_gpu_resident { |
381 | | if verbose { |
382 | | println!("Using GPU-resident path (pre-uploaded weights, ~2 syncs/token)"); |
383 | | } |
384 | | cuda_model |
385 | | .generate_gpu_resident(&prompt_tokens, &gen_config) |
386 | | .map_err(|e| crate::error::RealizarError::UnsupportedOperation { |
387 | | operation: "generate_gpu_resident".to_string(), |
388 | | reason: format!("GPU-resident generation failed: {e}"), |
389 | | })? |
390 | | } else { |
391 | | if verbose { |
392 | | println!("Using standard CUDA path"); |
393 | | } |
394 | | cuda_model |
395 | | .generate_full_cuda_with_cache(&prompt_tokens, &gen_config) |
396 | | .map_err(|e| crate::error::RealizarError::UnsupportedOperation { |
397 | | operation: "generate_full_cuda_with_cache".to_string(), |
398 | | reason: format!("CUDA generation failed: {e}"), |
399 | | })? |
400 | | }; |
401 | | let gen_time = gen_start.elapsed(); |
402 | | |
403 | | let tokens_generated = generated.len().saturating_sub(prompt_len); |
404 | | let tokens_per_sec = if gen_time.as_secs_f64() > 0.0 { |
405 | | tokens_generated as f64 / gen_time.as_secs_f64() |
406 | | } else { |
407 | | 0.0 |
408 | | }; |
409 | | |
410 | | // Decode output using GGUF vocabulary, replacing SentencePiece markers with spaces |
411 | | let output_text = mapped |
412 | | .model |
413 | | .decode(&generated[prompt_len..]) |
414 | | .replace('▁', " "); |
415 | | |
416 | | match format { |
417 | | "json" => { |
418 | | let json = serde_json::json!({ |
419 | | "model": "GGUF (CUDA)", |
420 | | "backend": "GPU", |
421 | | "prompt": prompt, |
422 | | "generated_text": output_text, |
423 | | "tokens_generated": tokens_generated, |
424 | | "generation_time_ms": gen_time.as_secs_f64() * 1000.0, |
425 | | "tokens_per_second": tokens_per_sec, |
426 | | "temperature": temperature, |
427 | | "cuda_enabled": true, |
428 | | "cuda_device": cuda_model.device_name(), |
429 | | }); |
430 | | println!( |
431 | | "{}", |
432 | | serde_json::to_string_pretty(&json).unwrap_or_default() |
433 | | ); |
434 | | }, |
435 | | _ => { |
436 | | if verbose { |
437 | | println!( |
438 | | "Generated ({tokens_generated} tokens in {:.2}ms):", |
439 | | gen_time.as_secs_f64() * 1000.0 |
440 | | ); |
441 | | println!("{prompt}{output_text}"); |
442 | | println!(); |
443 | | println!("Performance: {:.1} tok/s (GPU)", tokens_per_sec); |
444 | | } else { |
445 | | // Ollama-style clean output: just the response |
446 | | println!("{output_text}"); |
447 | | } |
448 | | }, |
449 | | } |
450 | | |
451 | | Ok(()) |
452 | | } |
453 | | |
454 | | /// Run SafeTensors inference with performance timing |
455 | 0 | pub fn run_safetensors_inference( |
456 | 0 | model_ref: &str, |
457 | 0 | prompt: &str, |
458 | 0 | max_tokens: usize, |
459 | 0 | _temperature: f32, |
460 | 0 | format: &str, |
461 | 0 | ) -> Result<()> { |
462 | | use crate::apr::AprV2Model; |
463 | | use crate::safetensors_infer::SafetensorsToAprConverter; |
464 | | use std::path::Path; |
465 | | use std::time::Instant; |
466 | | |
467 | 0 | let load_start = Instant::now(); |
468 | 0 | let model_path = Path::new(model_ref); |
469 | | |
470 | | // Convert SafeTensors to AprTransformer (F32 weights) |
471 | 0 | let transformer = SafetensorsToAprConverter::convert(model_path).map_err(|e| { |
472 | 0 | crate::error::RealizarError::UnsupportedOperation { |
473 | 0 | operation: "convert_safetensors".to_string(), |
474 | 0 | reason: format!("Failed to convert SafeTensors: {e}"), |
475 | 0 | } |
476 | 0 | })?; |
477 | | |
478 | 0 | let load_time = load_start.elapsed(); |
479 | 0 | println!("Model loaded in {:.2}ms", load_time.as_secs_f64() * 1000.0); |
480 | 0 | println!( |
481 | 0 | "Architecture: {} ({} layers, vocab_size={})", |
482 | | transformer.config.architecture, |
483 | | transformer.config.num_layers, |
484 | | transformer.config.vocab_size |
485 | | ); |
486 | | |
487 | | // Use proper tokenizer from sibling tokenizer.json |
488 | 0 | let prompt_tokens = AprV2Model::encode_text(model_path, prompt).unwrap_or_else(|| { |
489 | | // Fallback: simple char tokenization |
490 | 0 | prompt.chars().map(|c| c as u32).collect() |
491 | 0 | }); |
492 | 0 | let prompt_len = prompt_tokens.len(); |
493 | | |
494 | 0 | println!("Prompt tokens: {}", prompt_len); |
495 | 0 | println!(); |
496 | | |
497 | | // Run inference |
498 | 0 | let gen_start = Instant::now(); |
499 | 0 | let generated = transformer.generate(&prompt_tokens, max_tokens)?; |
500 | 0 | let gen_time = gen_start.elapsed(); |
501 | | |
502 | 0 | let tokens_generated = generated.len().saturating_sub(prompt_len); |
503 | 0 | let tokens_per_sec = if gen_time.as_secs_f64() > 0.0 { |
504 | 0 | tokens_generated as f64 / gen_time.as_secs_f64() |
505 | | } else { |
506 | 0 | 0.0 |
507 | | }; |
508 | | |
509 | | // Decode output using proper tokenizer |
510 | 0 | let output_tokens = &generated[prompt_len..]; |
511 | 0 | let output_text = if let Some(tokenizer) = AprV2Model::load_tokenizer(model_path) { |
512 | 0 | tokenizer.decode(output_tokens) |
513 | | } else { |
514 | 0 | output_tokens |
515 | 0 | .iter() |
516 | 0 | .map(|&t| char::from_u32(t.min(127)).unwrap_or('?')) |
517 | 0 | .collect() |
518 | | }; |
519 | | |
520 | 0 | match format { |
521 | 0 | "json" => { |
522 | 0 | let json = serde_json::json!({ |
523 | 0 | "model": model_ref, |
524 | 0 | "format": "SafeTensors", |
525 | 0 | "prompt": prompt, |
526 | 0 | "generated_text": output_text, |
527 | 0 | "tokens_generated": tokens_generated, |
528 | 0 | "generation_time_ms": gen_time.as_secs_f64() * 1000.0, |
529 | 0 | "tokens_per_second": tokens_per_sec, |
530 | 0 | }); |
531 | 0 | println!( |
532 | 0 | "{}", |
533 | 0 | serde_json::to_string_pretty(&json).unwrap_or_default() |
534 | 0 | ); |
535 | 0 | }, |
536 | 0 | _ => { |
537 | 0 | println!( |
538 | 0 | "Generated ({tokens_generated} tokens in {:.2}ms):", |
539 | 0 | gen_time.as_secs_f64() * 1000.0 |
540 | 0 | ); |
541 | 0 | println!("{output_text}"); |
542 | 0 | println!(); |
543 | 0 | println!("Performance: {:.1} tok/s", tokens_per_sec); |
544 | 0 | }, |
545 | | } |
546 | | |
547 | 0 | Ok(()) |
548 | 0 | } |
549 | | |
550 | | /// Run APR inference with performance timing |
551 | | /// |
552 | | /// Supports both CPU and GPU backends (PMAT-106). |
553 | 0 | pub fn run_apr_inference( |
554 | 0 | model_ref: &str, |
555 | 0 | file_data: &[u8], |
556 | 0 | prompt: &str, |
557 | 0 | max_tokens: usize, |
558 | 0 | temperature: f32, |
559 | 0 | format: &str, |
560 | 0 | force_gpu: bool, |
561 | 0 | verbose: bool, |
562 | 0 | ) -> Result<()> { |
563 | | use crate::apr::AprV2Model; |
564 | | use crate::apr_transformer::AprTransformer; |
565 | | use std::path::Path; |
566 | | use std::time::Instant; |
567 | | |
568 | | // Handle --gpu flag warning when CUDA not available |
569 | | #[cfg(not(feature = "cuda"))] |
570 | 0 | if force_gpu { |
571 | 0 | eprintln!("Warning: --gpu flag requires 'cuda' feature. Falling back to CPU."); |
572 | 0 | eprintln!("Build with: cargo build --features cuda"); |
573 | 0 | eprintln!(); |
574 | 0 | } |
575 | | #[cfg(not(feature = "cuda"))] |
576 | 0 | let _ = (force_gpu, verbose); |
577 | | |
578 | | // PMAT-106: GPU path for APR models |
579 | | #[cfg(feature = "cuda")] |
580 | | if force_gpu { |
581 | | return run_apr_inference_gpu(model_ref, file_data, prompt, max_tokens, temperature, format, verbose); |
582 | | } |
583 | | |
584 | 0 | let load_start = Instant::now(); |
585 | | |
586 | | // Load APR transformer (CPU path) |
587 | 0 | let transformer = AprTransformer::from_apr_bytes(file_data).map_err(|e| { |
588 | 0 | crate::error::RealizarError::UnsupportedOperation { |
589 | 0 | operation: "parse_apr".to_string(), |
590 | 0 | reason: format!("Failed to parse APR: {e}"), |
591 | 0 | } |
592 | 0 | })?; |
593 | | |
594 | 0 | let load_time = load_start.elapsed(); |
595 | 0 | if verbose { |
596 | 0 | println!("Backend: CPU (AVX2 + SIMD)"); |
597 | 0 | println!("Model loaded in {:.2}ms", load_time.as_secs_f64() * 1000.0); |
598 | 0 | } |
599 | | |
600 | | // Use proper tokenizer from sibling tokenizer.json |
601 | 0 | let model_path = Path::new(model_ref); |
602 | 0 | let prompt_tokens = AprV2Model::encode_text(model_path, prompt).unwrap_or_else(|| { |
603 | | // Fallback: simple char tokenization |
604 | 0 | prompt.chars().map(|c| c as u32).collect() |
605 | 0 | }); |
606 | 0 | let prompt_len = prompt_tokens.len(); |
607 | | |
608 | 0 | if verbose { |
609 | 0 | println!("Prompt tokens: {}", prompt_len); |
610 | 0 | println!("Temperature: {:.1} (using greedy decoding)", temperature); |
611 | 0 | println!(); |
612 | 0 | } |
613 | | |
614 | | // Run inference with timing |
615 | | // PMAT-103 FIX: Use generate_with_cache for O(n) instead of O(n²) complexity |
616 | 0 | let gen_config = crate::apr_transformer::GenerateConfig { |
617 | 0 | max_tokens, |
618 | 0 | temperature, |
619 | 0 | ..Default::default() |
620 | 0 | }; |
621 | 0 | let gen_start = Instant::now(); |
622 | 0 | let generated = transformer.generate_with_cache(&prompt_tokens, &gen_config)?; |
623 | 0 | let gen_time = gen_start.elapsed(); |
624 | | |
625 | 0 | let tokens_generated = generated.len().saturating_sub(prompt_len); |
626 | 0 | let tokens_per_sec = if gen_time.as_secs_f64() > 0.0 { |
627 | 0 | tokens_generated as f64 / gen_time.as_secs_f64() |
628 | | } else { |
629 | 0 | 0.0 |
630 | | }; |
631 | | |
632 | | // Decode output using proper tokenizer |
633 | 0 | let output_tokens = &generated[prompt_len..]; |
634 | 0 | let output_text = if let Some(tokenizer) = AprV2Model::load_tokenizer(model_path) { |
635 | 0 | tokenizer.decode(output_tokens) |
636 | | } else { |
637 | | // Fallback: simple ASCII |
638 | 0 | output_tokens |
639 | 0 | .iter() |
640 | 0 | .map(|&t| char::from_u32(t.min(127)).unwrap_or('?')) |
641 | 0 | .collect() |
642 | | }; |
643 | | |
644 | 0 | match format { |
645 | 0 | "json" => { |
646 | 0 | let json = serde_json::json!({ |
647 | 0 | "model": model_ref, |
648 | 0 | "format": "APR", |
649 | 0 | "backend": "CPU", |
650 | 0 | "prompt": prompt, |
651 | 0 | "generated_text": output_text, |
652 | 0 | "tokens_generated": tokens_generated, |
653 | 0 | "generation_time_ms": gen_time.as_secs_f64() * 1000.0, |
654 | 0 | "tokens_per_second": tokens_per_sec, |
655 | 0 | "temperature": temperature, |
656 | 0 | }); |
657 | 0 | println!( |
658 | 0 | "{}", |
659 | 0 | serde_json::to_string_pretty(&json).unwrap_or_default() |
660 | 0 | ); |
661 | 0 | }, |
662 | | _ => { |
663 | 0 | if verbose { |
664 | 0 | println!( |
665 | 0 | "Generated ({tokens_generated} tokens in {:.2}ms):", |
666 | 0 | gen_time.as_secs_f64() * 1000.0 |
667 | 0 | ); |
668 | 0 | println!("{output_text}"); |
669 | 0 | println!(); |
670 | 0 | println!("Performance: {:.1} tok/s", tokens_per_sec); |
671 | 0 | } else { |
672 | 0 | // Clean output: just the response |
673 | 0 | println!("{output_text}"); |
674 | 0 | } |
675 | | }, |
676 | | } |
677 | | |
678 | 0 | Ok(()) |
679 | 0 | } |
680 | | |
681 | | /// Run APR inference with CUDA GPU acceleration (PMAT-106) |
682 | | /// |
683 | | /// Uses the APR F32 GPU adapter to convert weights to GpuModel format |
684 | | /// for high-performance inference. |
685 | | #[cfg(feature = "cuda")] |
686 | | #[allow(clippy::too_many_arguments)] |
687 | | pub fn run_apr_inference_gpu( |
688 | | model_ref: &str, |
689 | | file_data: &[u8], |
690 | | prompt: &str, |
691 | | max_tokens: usize, |
692 | | temperature: f32, |
693 | | format: &str, |
694 | | verbose: bool, |
695 | | ) -> Result<()> { |
696 | | use crate::apr::AprV2Model; |
697 | | use crate::apr_transformer::AprTransformer; |
698 | | use crate::gpu::adapters::AprF32ToGpuAdapter; |
699 | | use crate::gpu::GpuGenerateConfig; |
700 | | use std::path::Path; |
701 | | use std::time::Instant; |
702 | | |
703 | | let load_start = Instant::now(); |
704 | | |
705 | | if verbose { |
706 | | println!("Backend: CUDA (GPU)"); |
707 | | println!("Loading APR model for GPU inference..."); |
708 | | } |
709 | | |
710 | | // Load APR as F32 transformer |
711 | | let transformer = AprTransformer::from_apr_bytes(file_data).map_err(|e| { |
712 | | crate::error::RealizarError::UnsupportedOperation { |
713 | | operation: "parse_apr".to_string(), |
714 | | reason: format!("Failed to parse APR: {e}"), |
715 | | } |
716 | | })?; |
717 | | |
718 | | // Debug: Check if gate weights exist in the loaded transformer |
719 | | if verbose { |
720 | | let has_gate = transformer.layers.first().map_or(false, |l| l.ffn_gate_weight.is_some()); |
721 | | eprintln!("[DEBUG-SwiGLU] APR transformer has gate weight: {}", has_gate); |
722 | | if has_gate { |
723 | | let gate_len = transformer.layers[0].ffn_gate_weight.as_ref().map_or(0, |v| v.len()); |
724 | | eprintln!("[DEBUG-SwiGLU] Gate weight elements: {} (expected: {}x{}={})", |
725 | | gate_len, |
726 | | transformer.config.hidden_dim, |
727 | | transformer.config.intermediate_dim, |
728 | | transformer.config.hidden_dim * transformer.config.intermediate_dim); |
729 | | } |
730 | | } |
731 | | |
732 | | // Convert to GpuModel using F32 adapter |
733 | | let mut gpu_model = AprF32ToGpuAdapter::to_gpu_model(&transformer).map_err(|e| { |
734 | | crate::error::RealizarError::UnsupportedOperation { |
735 | | operation: "apr_to_gpu".to_string(), |
736 | | reason: format!("Failed to convert APR to GPU format: {e}"), |
737 | | } |
738 | | })?; |
739 | | |
740 | | // Debug: Check if gate weights exist in GpuModel |
741 | | if verbose { |
742 | | let has_gpu_gate = gpu_model.block_weights.first().map_or(false, |b| b.ffn_gate_weight.is_some()); |
743 | | eprintln!("[DEBUG-SwiGLU] GpuModel has gate weight: {}", has_gpu_gate); |
744 | | |
745 | | // Compare weights: CPU vs GPU |
746 | | if let Some(layer0) = transformer.layers.first() { |
747 | | eprintln!("[DEBUG-WEIGHT] CPU qkv_weight first 5: {:?}", &layer0.qkv_weight[0..5.min(layer0.qkv_weight.len())]); |
748 | | eprintln!("[DEBUG-WEIGHT] GPU qkv_weight first 5: {:?}", &gpu_model.block_weights[0].qkv_weight[0..5.min(gpu_model.block_weights[0].qkv_weight.len())]); |
749 | | eprintln!("[DEBUG-WEIGHT] CPU fc1 (up) first 5: {:?}", &layer0.ffn_up_weight[0..5.min(layer0.ffn_up_weight.len())]); |
750 | | eprintln!("[DEBUG-WEIGHT] GPU fc1 (up) first 5: {:?}", &gpu_model.block_weights[0].ffn_fc1_weight[0..5.min(gpu_model.block_weights[0].ffn_fc1_weight.len())]); |
751 | | |
752 | | // Debug: Compare matmul output for first token embedding |
753 | | let hidden_dim = transformer.config.hidden_dim; |
754 | | let test_embedding = &transformer.token_embedding[0..hidden_dim]; |
755 | | |
756 | | // CPU matmul: y = x @ W where W is [out_dim, in_dim] (transposed internally) |
757 | | let cpu_qkv_dim = layer0.qkv_weight.len() / hidden_dim; |
758 | | let mut cpu_qkv = vec![0.0f32; cpu_qkv_dim]; |
759 | | for o in 0..cpu_qkv_dim { |
760 | | let w_start = o * hidden_dim; |
761 | | let mut sum = 0.0f32; |
762 | | for i in 0..hidden_dim { |
763 | | sum += test_embedding[i] * layer0.qkv_weight[w_start + i]; |
764 | | } |
765 | | cpu_qkv[o] = sum; |
766 | | } |
767 | | |
768 | | // GPU matmul: y = x @ W_t where W_t is [in_dim, out_dim] (already transposed) |
769 | | let gpu_qkv_weight = &gpu_model.block_weights[0].qkv_weight; |
770 | | let gpu_qkv_dim = gpu_qkv_weight.len() / hidden_dim; |
771 | | let mut gpu_qkv = vec![0.0f32; gpu_qkv_dim]; |
772 | | for j in 0..gpu_qkv_dim { |
773 | | let mut sum = 0.0f32; |
774 | | for i in 0..hidden_dim { |
775 | | // GPU weight is [hidden_dim, qkv_dim] row-major: W_t[i,j] = W_t[i * qkv_dim + j] |
776 | | sum += test_embedding[i] * gpu_qkv_weight[i * gpu_qkv_dim + j]; |
777 | | } |
778 | | gpu_qkv[j] = sum; |
779 | | } |
780 | | |
781 | | eprintln!("[DEBUG-MATMUL] CPU QKV first 5: {:?}", &cpu_qkv[0..5.min(cpu_qkv.len())]); |
782 | | eprintln!("[DEBUG-MATMUL] GPU QKV first 5: {:?}", &gpu_qkv[0..5.min(gpu_qkv.len())]); |
783 | | |
784 | | // Compare |
785 | | let max_diff: f32 = cpu_qkv.iter().zip(gpu_qkv.iter()) |
786 | | .map(|(c, g)| (c - g).abs()) |
787 | | .fold(0.0f32, f32::max); |
788 | | eprintln!("[DEBUG-MATMUL] Max diff: {}", max_diff); |
789 | | if max_diff > 0.01 { |
790 | | eprintln!("[DEBUG-MATMUL] WARNING: CPU vs GPU QKV mismatch!"); |
791 | | } |
792 | | } |
793 | | } |
794 | | |
795 | | let load_time = load_start.elapsed(); |
796 | | if verbose { |
797 | | println!("Model loaded in {:.2}ms", load_time.as_secs_f64() * 1000.0); |
798 | | } |
799 | | |
800 | | // Use proper tokenizer from sibling tokenizer.json |
801 | | let model_path = Path::new(model_ref); |
802 | | let prompt_tokens = AprV2Model::encode_text(model_path, prompt).unwrap_or_else(|| { |
803 | | prompt.chars().map(|c| c as u32).collect() |
804 | | }); |
805 | | let prompt_len = prompt_tokens.len(); |
806 | | |
807 | | if verbose { |
808 | | println!("Prompt tokens: {}", prompt_len); |
809 | | println!("Temperature: {:.1}", temperature); |
810 | | println!(); |
811 | | } |
812 | | |
813 | | // Configure generation |
814 | | let gen_config = GpuGenerateConfig { |
815 | | max_tokens, |
816 | | temperature, |
817 | | top_k: if temperature <= 0.01 { 1 } else { 40 }, |
818 | | stop_tokens: vec![], |
819 | | }; |
820 | | |
821 | | // Convert prompt tokens to usize for GpuModel |
822 | | let prompt_tokens_usize: Vec<usize> = prompt_tokens.iter().map(|&t| t as usize).collect(); |
823 | | |
824 | | // Run inference |
825 | | let gen_start = Instant::now(); |
826 | | let generated = gpu_model.generate(&prompt_tokens_usize, &gen_config).map_err(|e| { |
827 | | crate::error::RealizarError::UnsupportedOperation { |
828 | | operation: "gpu_generate".to_string(), |
829 | | reason: format!("GPU generation failed: {e}"), |
830 | | } |
831 | | })?; |
832 | | let gen_time = gen_start.elapsed(); |
833 | | |
834 | | let tokens_generated = generated.len().saturating_sub(prompt_len); |
835 | | let tokens_per_sec = if gen_time.as_secs_f64() > 0.0 { |
836 | | tokens_generated as f64 / gen_time.as_secs_f64() |
837 | | } else { |
838 | | 0.0 |
839 | | }; |
840 | | |
841 | | // Decode output (convert usize back to u32) |
842 | | let output_tokens: Vec<u32> = generated[prompt_len..].iter().map(|&t| t as u32).collect(); |
843 | | let output_text = if let Some(tokenizer) = AprV2Model::load_tokenizer(model_path) { |
844 | | tokenizer.decode(&output_tokens) |
845 | | } else { |
846 | | output_tokens |
847 | | .iter() |
848 | | .map(|&t| char::from_u32(t.min(127)).unwrap_or('?')) |
849 | | .collect() |
850 | | }; |
851 | | |
852 | | match format { |
853 | | "json" => { |
854 | | let json = serde_json::json!({ |
855 | | "model": model_ref, |
856 | | "format": "APR", |
857 | | "backend": "CUDA", |
858 | | "prompt": prompt, |
859 | | "generated_text": output_text, |
860 | | "tokens_generated": tokens_generated, |
861 | | "generation_time_ms": gen_time.as_secs_f64() * 1000.0, |
862 | | "tokens_per_second": tokens_per_sec, |
863 | | "temperature": temperature, |
864 | | }); |
865 | | println!( |
866 | | "{}", |
867 | | serde_json::to_string_pretty(&json).unwrap_or_default() |
868 | | ); |
869 | | }, |
870 | | _ => { |
871 | | if verbose { |
872 | | println!( |
873 | | "Generated ({tokens_generated} tokens in {:.2}ms):", |
874 | | gen_time.as_secs_f64() * 1000.0 |
875 | | ); |
876 | | println!("{output_text}"); |
877 | | println!(); |
878 | | println!("Performance: {:.1} tok/s (GPU)", tokens_per_sec); |
879 | | } else { |
880 | | println!("{output_text}"); |
881 | | } |
882 | | }, |
883 | | } |
884 | | |
885 | | Ok(()) |
886 | | } |
887 | | |