aprender-serve 0.64.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
63
64
65
66
67
68
69
70


#[test]
#[cfg(feature = "gpu")]
#[serial_test::serial]
fn test_imp_108c_attention_softmax_normalized() {
    // IMP-108c: Verify attention weights sum to 1 for each position
    let config = GGUFConfig {
        architecture: "test".to_string(),
        constraints: crate::gguf::ArchConstraints::from_architecture("test"),
        hidden_dim: 16,
        intermediate_dim: 32,
        num_layers: 1,
        num_heads: 2,
        num_kv_heads: 2,
        vocab_size: 50,
        context_length: 128,
        rope_theta: 10000.0,
        eps: 1e-5,
        rope_type: 0,
            explicit_head_dim: None,
            query_pre_attn_scalar: None,
        bos_token_id: None,
            eos_token_id: None,
    };

    let model = create_test_model_with_config(&config);
    let seq_len = 4;
    let hidden_dim = config.hidden_dim;
    let head_dim = hidden_dim / config.num_heads;

    // Create Q, K with known values to verify softmax
    let q: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 3) as f32) * 0.1)
        .collect();
    let k: Vec<f32> = (0..seq_len * hidden_dim)
        .map(|i| ((i % 5) as f32) * 0.1)
        .collect();

    // Use V = identity-like pattern to extract attention weights
    // V[j] = one-hot at position j within head
    let mut v = vec![0.0f32; seq_len * hidden_dim];
    for pos in 0..seq_len {
        for head in 0..config.num_heads {
            // Set V[pos, head, pos % head_dim] = 1.0
            let idx = pos * hidden_dim + head * head_dim + (pos % head_dim);
            v[idx] = 1.0;
        }
    }

    let output = model
        .batched_causal_attention_gpu(&q, &k, &v, seq_len)
        .expect("test");

    // Output should be valid (finite)
    assert!(
        output.iter().all(|x| x.is_finite()),
        "IMP-108c: All attention outputs should be finite"
    );

    // Output at each position should reflect weighted sum of V
    // Since V entries are 0 or 1, output values should be in [0, 1] range
    // (attention weights are normalized, so weighted sum of [0,1] is in [0,1])
    for &val in &output {
        assert!(
            val >= -0.01 && val <= 1.01,
            "IMP-108c: Attention output {} should be weighted sum of V (in [0,1])",
            val
        );
    }
}