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
96

use std::time::Instant;

use models::vision::{PooledPreprocessedImage, PooledPromptTokens};
use runtime::{
    Result as RuntimeResult,
    backend::{ModelHandle, PrefillOutput, SamplingLogits},
    kv::BlockTable,
};
use uuid::Uuid;

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

impl MetalBackend {
    #[allow(clippy::too_many_arguments)]
    pub fn prefill_pooled_vision_with_progress(
        &self,
        model: &ModelHandle,
        session_id: Uuid,
        prompt: &PooledPromptTokens,
        image: &PooledPreprocessedImage,
        block_table: &BlockTable,
        sampling_logits: SamplingLogits,
        progress: &mut dyn FnMut(MetalProgressEvent),
    ) -> RuntimeResult<PrefillOutput> {
        let started = Instant::now();
        let lookup = model.id.clone();
        let model_id = lookup.clone();
        let prompt = prompt.clone();
        let image = image.clone();
        let blocks = block_table.blocks().len();
        let execution_sampling =
            execution_sampling(sampling_logits, self.config.fusion.device_token_pipeline.enabled());
        Ok(self.with_model_progress(
            &lookup,
            move |loaded, worker_progress| {
                execute(
                    loaded,
                    &model_id,
                    session_id,
                    &prompt,
                    &image,
                    blocks,
                    sampling_logits,
                    execution_sampling,
                    started,
                    worker_progress,
                )
            },
            progress,
        )?)
    }
}

#[allow(clippy::too_many_arguments)]
fn execute(
    loaded: &mut LoadedModel,
    model_id: &str,
    session_id: Uuid,
    prompt: &PooledPromptTokens,
    image: &PooledPreprocessedImage,
    blocks: usize,
    sampling_logits: SamplingLogits,
    execution_sampling: SamplingLogits,
    started: Instant,
    progress: &mut dyn FnMut(MetalProgressEvent),
) -> Result<PrefillOutput> {
    let native =
        loaded.prefill_pooled_vision(session_id, prompt, image, execution_sampling, progress)?;
    let output = output::materialize(loaded, native.output, sampling_logits)?;
    let cached_tokens = loaded.session_cached_tokens(session_id)?;
    let trace = format!(
        "native pooled vision prefill: {} tokens, prefix cache disabled, {} runtime KV blocks, {} cached tokens, {:.3}ms",
        prompt.token_ids.len(),
        blocks,
        cached_tokens,
        started.elapsed().as_secs_f64() * 1_000.0
    );
    tracing::debug!(
        model_id,
        session_id = %session_id,
        prompt_tokens = prompt.token_ids.len(),
        cached_tokens,
        "native pooled vision Metal prefill completed"
    );
    Ok(PrefillOutput {
        accepted_tokens: prompt.token_ids.len(),
        next_token: output.next_token,
        trace: Some(trace),
        logits: output.logits,
        candidates: output.candidates,
    })
}