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

use std::time::Instant;

use runtime::{
    Result as RuntimeResult,
    backend::{DecodeBatchOutput, DecodeBatchRequest, DecodeOutput, TokenEvent},
};

use super::{MetalBackend, execution::execution_sampling};
use crate::native::{
    error::Result,
    model::{DecodeInput, NativeOutput},
    output,
};

impl MetalBackend {
    pub fn decode_batch_tokens(
        &self,
        request: &DecodeBatchRequest,
    ) -> RuntimeResult<DecodeBatchOutput> {
        DecodeBatchOutput::new(self.decode_batch_inner(request)?)
    }

    fn decode_batch_inner(&self, request: &DecodeBatchRequest) -> Result<Vec<DecodeOutput>> {
        let started = Instant::now();
        let lookup = request.model().id.clone();
        let sequences = request.sequences().to_vec();
        let device_pipeline = self.config.fusion.device_token_pipeline.enabled();
        self.with_model(&lookup, move |loaded| {
            let inputs = sequences
                .iter()
                .map(|sequence| DecodeInput {
                    session: sequence.session_id,
                    token: sequence.token_id,
                    sampling: execution_sampling(sequence.sampling_logits, device_pipeline),
                })
                .collect::<Vec<_>>();
            let batched = loaded.can_decode_batch(&inputs);
            let native = if batched {
                loaded.decode_batch(&inputs)?
            } else {
                inputs
                    .iter()
                    .map(|input| loaded.decode(input.session, input.token, input.sampling))
                    .collect::<Result<Vec<NativeOutput>>>()?
            };
            let outputs = native
                .into_iter()
                .zip(&sequences)
                .map(|(native, sequence)| {
                    let output = output::materialize(loaded, native, sequence.sampling_logits)?;
                    Ok(DecodeOutput {
                        event: TokenEvent {
                            token_id: output.next_token,
                            text: if batched {
                                "metal.decode=packed-device-token-pipeline".into()
                            } else {
                                "metal.decode=scalar-device-token-replay".into()
                            },
                            finished: false,
                        },
                        logits: output.logits,
                        candidates: output.candidates,
                    })
                })
                .collect::<Result<Vec<_>>>()?;
            tracing::debug!(
                backend = "metal",
                rows = inputs.len(),
                batched,
                execution_ms = started.elapsed().as_secs_f64() * 1_000.0,
                "completed Metal decode batch"
            );
            Ok(outputs)
        })
    }
}