aprender-serve 0.34.0

Pure Rust ML inference engine built from scratch - model serving for GGUF and safetensors
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

//! Fresh CPU vs GPU comparison
#[cfg(feature = "cuda")]
use realizar::gguf::OwnedQuantizedModelCuda;
use realizar::gguf::{MappedGGUFModel, OwnedQuantizedModel, QuantizedGenerateConfig};

fn main() {
    let path = std::env::args()
        .nth(1)
        .unwrap_or("/home/noah/models/qwen2.5-coder-1.5b-instruct-q4_k_m.gguf".to_string());
    let mapped = MappedGGUFModel::from_path(&path).expect("Failed to load");

    // Test token sequence
    let tokens: Vec<u32> = vec![791]; // Single token

    println!("Running forward pass on CPU (batch)...");
    let model = OwnedQuantizedModel::from_mapped(&mapped).expect("Failed to parse");
    let cpu_logits = model.forward(&tokens).expect("CPU forward failed");
    let cpu_argmax = cpu_logits
        .iter()
        .enumerate()
        .max_by(|a, b| a.1.partial_cmp(b.1).unwrap())
        .map(|(i, _)| i)
        .unwrap();
    println!(
        "CPU argmax: {} (logit: {:.4})",
        cpu_argmax, cpu_logits[cpu_argmax]
    );

    println!("\nRunning forward pass on CPU (with cache)...");
    let config = QuantizedGenerateConfig {
        max_tokens: 1,
        temperature: 0.0,
        top_k: 1,
        ..Default::default()
    };
    let output = model
        .generate_with_cache(&tokens, &config)
        .expect("CPU generate failed");
    println!("CPU generated token: {:?}", output.last());

    #[cfg(feature = "cuda")]
    {
        println!("\nRunning forward pass on GPU...");
        let model = OwnedQuantizedModel::from_mapped(&mapped).expect("Failed to parse");
        let mut cuda_model = OwnedQuantizedModelCuda::new(model, 0).expect("CUDA init failed");

        let gpu_output = cuda_model
            .generate_gpu_resident(&tokens, &config)
            .expect("GPU forward failed");
        println!("GPU generated token: {:?}", gpu_output.last());

        let cpu_last = *output.last().unwrap_or(&0);
        let gpu_last = *gpu_output.last().unwrap_or(&0);

        println!("\nComparison:");
        if cpu_last == gpu_last {
            println!("✓ MATCH: both generated {}", cpu_last);
        } else {
            println!("✗ DIVERGE: CPU={}, GPU={}", cpu_last, gpu_last);
        }
    }
}