libmir-metal 0.3.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

use std::time::{Duration, Instant};

use super::{Array, BatchAttentionExecution, KvContext, Result, Stream, execute_measured};

pub(super) fn measure(
    execution: BatchAttentionExecution,
    queries: &[&Array],
    contexts: &[&KvContext],
    scale: f32,
    causal: bool,
    stream: &Stream,
) -> Result<Duration> {
    let config = &stream.config().tuning;
    for _ in 0..config.warmup_iterations {
        let output = execute_measured(execution, queries, contexts, scale, causal, stream, true)?;
        evaluate(&output, stream)?;
    }
    let iterations = config.measurement_iterations.max(1);
    let started = Instant::now();
    for _ in 0..iterations {
        let output = execute_measured(execution, queries, contexts, scale, causal, stream, true)?;
        evaluate(&output, stream)?;
    }
    Ok(started.elapsed() / iterations)
}

fn evaluate(output: &[Array], stream: &Stream) -> Result<()> {
    let refs = output.iter().collect::<Vec<_>>();
    Array::concatenate(&refs, 0, stream)?.async_eval()?;
    stream.synchronize()
}