libmir-cuda 0.2.0

CUDA inference backend for libmir
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

use models::semantic::{
    ActivationSpec, AttentionOutputSpec, AttentionSpec, DecoderLayerSpec, DecoderSpec,
    FeedForwardSpec, KeyValueRelation, MixerSpec, NormalizationKind, NormalizationSpec,
    PositionEncodingSpec, QkNormalization, RopeScalingSpec, RotaryLayoutSpec, RotarySpec,
    RoutedExpertsSpec, RouterNormalization, SemanticModelSpec,
};

use super::ClampedRoutedConfig;

#[test]
fn admits_dynamic_geometry_instead_of_one_checkpoint_size() -> crate::Result<()> {
    let spec = spec(64, 96, 4, 2, 16);
    let config = ClampedRoutedConfig::from_semantic(&spec)?;

    assert_eq!(config.hidden, 64);
    assert_eq!(config.intermediate, 96);
    assert_eq!(config.top_k, 2);
    Ok(())
}

#[test]
fn semantic_config_does_not_apply_cuda_kernel_limits() -> crate::Result<()> {
    let config = ClampedRoutedConfig::from_semantic(&spec(64, 96, 4, 2, 258))?;
    assert_eq!(config.head_dim, 258);
    Ok(())
}

fn spec(
    hidden: usize,
    intermediate: usize,
    query_heads: usize,
    key_value_heads: usize,
    head_dim: usize,
) -> SemanticModelSpec {
    let norm = NormalizationSpec {
        kind: NormalizationKind::Rms,
        epsilon: 1.0e-5,
    };
    SemanticModelSpec {
        schema_version: 1,
        decoder: DecoderSpec {
            hidden_size: hidden,
            vocab_size: 128,
            tie_word_embeddings: false,
            final_norm: norm,
            layers: vec![DecoderLayerSpec {
                index: 0,
                input_norm: norm,
                post_attention_norm: norm,
                mixer: MixerSpec::SoftmaxAttention(AttentionSpec {
                    query_heads,
                    key_value_heads,
                    head_dim,
                    key_value_relation: KeyValueRelation::Separate,
                    qk_normalization: QkNormalization::None,
                    projection_bias: true,
                    output: AttentionOutputSpec::Direct,
                    sinks: true,
                    scale: 0.25,
                    window: Some(32),
                    position: PositionEncodingSpec::Rotary(RotarySpec {
                        theta: 150_000.0,
                        partial_factor: 1.0,
                        layout: RotaryLayoutSpec::Standard,
                        algorithm: None,
                        scaling: Some(RopeScalingSpec::Yarn {
                            factor: 32.0,
                            beta_fast: 32.0,
                            beta_slow: 1.0,
                            original_context_len: 4096,
                        }),
                    }),
                }),
                feed_forward: FeedForwardSpec::Routed {
                    routed: RoutedExpertsSpec {
                        expert_count: 8,
                        top_k: 2,
                        intermediate_size: intermediate,
                        activation: ActivationSpec::SwiGlu {
                            alpha: 1.702,
                            clamp: Some(7.0),
                            up_shift: 1.0,
                        },
                        router_normalization: RouterNormalization::UnitTopK,
                    },
                    shared: None,
                },
            }],
        },
    }
}