libmir-metal 0.2.0

Metal 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
92
93
94
95

use models::layout::{AttentionLayerType, DecoderConfig, EncoderConfig};
use runtime::trace::TraceDecoder;

pub(super) fn execution(
    decoder: Option<&DecoderConfig>,
    encoder: Option<&EncoderConfig>,
) -> TraceDecoder {
    if let Some(decoder) = decoder {
        let (full_attention_layers, sliding_attention_layers) = attention_counts(decoder);
        return TraceDecoder {
            layers: decoder.num_hidden_layers,
            hidden_size: decoder.hidden_size,
            intermediate_size: decoder.intermediate_size,
            vocab_size: decoder.vocab_size,
            attention_heads: decoder.num_attention_heads,
            kv_heads: decoder.num_key_value_heads,
            head_dim: decoder.head_dim,
            global_head_dim: decoder.global_head_dim,
            global_kv_heads: decoder.num_global_key_value_heads,
            full_attention_layers,
            sliding_attention_layers,
            max_position_embeddings: decoder.max_position_embeddings,
            rope_theta: decoder.rope_theta,
            full_attention_rope_theta: decoder.full_attention_rope_theta,
            sliding_attention_rope_theta: decoder.sliding_attention_rope_theta,
            sliding_window: decoder.sliding_window,
            num_experts: decoder.num_experts,
            top_k_experts: decoder.top_k_experts,
            moe_intermediate_size: decoder.moe_intermediate_size,
            hidden_activation: decoder.hidden_activation.clone(),
            final_logit_softcapping: decoder.final_logit_softcapping,
        };
    }
    encoder.map_or_else(empty_shape, encoder_shape)
}

fn empty_shape() -> TraceDecoder {
    TraceDecoder {
        layers: 0,
        hidden_size: 0,
        intermediate_size: 0,
        vocab_size: 0,
        attention_heads: 0,
        kv_heads: 0,
        head_dim: 0,
        global_head_dim: None,
        global_kv_heads: None,
        full_attention_layers: 0,
        sliding_attention_layers: 0,
        max_position_embeddings: 0,
        rope_theta: None,
        full_attention_rope_theta: None,
        sliding_attention_rope_theta: None,
        sliding_window: None,
        num_experts: None,
        top_k_experts: None,
        moe_intermediate_size: None,
        hidden_activation: None,
        final_logit_softcapping: None,
    }
}

fn encoder_shape(encoder: &EncoderConfig) -> TraceDecoder {
    TraceDecoder {
        layers: encoder.num_hidden_layers,
        hidden_size: encoder.hidden_size,
        intermediate_size: encoder.intermediate_size,
        vocab_size: encoder.vocab_size,
        attention_heads: encoder.num_attention_heads,
        kv_heads: encoder.num_attention_heads,
        head_dim: encoder.head_dim,
        global_head_dim: None,
        global_kv_heads: None,
        full_attention_layers: encoder.num_hidden_layers,
        sliding_attention_layers: 0,
        max_position_embeddings: encoder.max_position_embeddings,
        rope_theta: encoder.rope_theta,
        full_attention_rope_theta: None,
        sliding_attention_rope_theta: None,
        sliding_window: None,
        num_experts: None,
        top_k_experts: None,
        moe_intermediate_size: None,
        hidden_activation: Some(encoder.hidden_activation.clone()),
        final_logit_softcapping: None,
    }
}

fn attention_counts(decoder: &DecoderConfig) -> (usize, usize) {
    decoder.layer_types.iter().fold((0, 0), |(full, sliding), layer| match layer {
        AttentionLayerType::Full => (full + 1, sliding),
        AttentionLayerType::Linear => (full, sliding),
        AttentionLayerType::Sliding => (full, sliding + 1),
    })
}