libmir-metal 0.1.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
96
97
98
99
100
101
102
103
104

use runtime::backend::{CandidateLogitsTrace, LogitsTrace, SamplingLogits};

use super::{
    error::Result,
    model::{LoadedModel, NativeOutput},
};
use crate::engine::{DeviceSampling, sample_u32};
#[derive(Debug)]
pub(super) struct SamplingOutput {
    pub next_token: Option<u32>,
    pub logits: Option<LogitsTrace>,
    pub candidates: Option<CandidateLogitsTrace>,
}

pub(super) fn materialize(
    model: &LoadedModel,
    output: NativeOutput,
    sampling: SamplingLogits,
) -> Result<SamplingOutput> {
    match output {
        NativeOutput::Greedy(next_token) => Ok(SamplingOutput {
            next_token: Some(next_token),
            logits: None,
            candidates: None,
        }),
        NativeOutput::Logits(logits) if sampling == SamplingLogits::None => Ok(SamplingOutput {
            next_token: Some(logits.argmax_u32(model.stream())?),
            logits: None,
            candidates: None,
        }),
        NativeOutput::Logits(logits) if let SamplingLogits::TopK { k, vocab_size } = sampling => {
            top_k(model, &logits, k, vocab_size)
        },
        NativeOutput::Logits(logits)
            if let SamplingLogits::SampleTopK { k, vocab_size, temperature, draw } = sampling =>
        {
            sampled(model, &logits, vocab_size, k, 1.0, temperature, draw)
        },
        NativeOutput::Logits(logits)
            if let SamplingLogits::Sample {
                vocab_size,
                temperature,
                top_p,
                top_k,
                draw,
            } = sampling =>
        {
            sampled(model, &logits, vocab_size, top_k, top_p, temperature, draw)
        },
        NativeOutput::Logits(logits) => {
            let shape = logits.shape()?;
            let values = logits.to_vec_f32_on_stream(model.stream())?;
            Ok(SamplingOutput {
                next_token: None,
                logits: Some(LogitsTrace { shape, values }),
                candidates: None,
            })
        },
    }
}

fn sampled(
    model: &LoadedModel,
    logits: &crate::engine::Array,
    vocab_size: usize,
    top_k: usize,
    top_p: f32,
    temperature: f32,
    draw: f32,
) -> Result<SamplingOutput> {
    let next_token = sample_u32(
        logits,
        DeviceSampling {
            vocab_size,
            top_k,
            top_p,
            temperature,
            draw,
        },
        model.stream(),
    )?;
    Ok(SamplingOutput {
        next_token: Some(next_token),
        logits: None,
        candidates: None,
    })
}

fn top_k(
    model: &LoadedModel,
    logits: &crate::engine::Array,
    k: usize,
    vocab_size: usize,
) -> Result<SamplingOutput> {
    let candidates = logits.top_k(k, vocab_size, model.stream())?;
    Ok(SamplingOutput {
        next_token: None,
        logits: None,
        candidates: Some(CandidateLogitsTrace {
            token_ids: candidates.token_ids.to_vec_u32_on_stream(model.stream())?,
            scores: candidates.scores.to_vec_f32_on_stream(model.stream())?,
        }),
    })
}