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

use super::{Storage, lock};
use crate::engine::{
    Array, Error, Result, Stream,
    kernels::{
        PageWriteOptions, PreparedPageWrite, PreparedQuantizedPageWrite, QuantizedPageWriteOptions,
    },
};

pub(super) fn native(
    keys: &Array,
    values: &Array,
    offset: usize,
    stream: &Stream,
    storage: &Storage,
    prepared: &mut PreparedPageWrite,
) -> Result<()> {
    let mut arena = lock(&storage.arena)?;
    let [page_keys, page_values] = stream.page_write(
        [
            keys.native(),
            values.native(),
            arena.keys.native(),
            arena.values.native(),
            storage.table.native(),
        ],
        PageWriteOptions {
            sequence: keys.native().shape()?.dimensions()[2],
            offset,
            kv_heads: arena.kv_heads,
            page_capacity: arena.capacity,
            page_size: arena.page_size,
            head_dim: arena.head_dim,
        },
        prepared,
    )?;
    arena.keys = Array::from_native(page_keys)?;
    arena.values = Array::from_native(page_values)?;
    drop(arena);
    Ok(())
}

pub(super) fn quantized(
    keys: &Array,
    values: &Array,
    offset: usize,
    stream: &Stream,
    storage: &Storage,
    prepared: &mut PreparedQuantizedPageWrite,
) -> Result<()> {
    let mut arena = lock(&storage.arena)?;
    let key_scales = arena.key_scales.as_ref().ok_or(Error::NullHandle("K/V key scales"))?;
    let value_scales = arena.value_scales.as_ref().ok_or(Error::NullHandle("K/V value scales"))?;
    let [page_keys, page_values, next_key_scales, next_value_scales] = stream
        .quantized_page_write(
            [
                keys.native(),
                values.native(),
                arena.keys.native(),
                arena.values.native(),
                key_scales.native(),
                value_scales.native(),
                storage.table.native(),
            ],
            QuantizedPageWriteOptions {
                sequence: keys.native().shape()?.dimensions()[2],
                offset,
                kv_heads: arena.kv_heads,
                page_capacity: arena.capacity,
                page_size: arena.page_size,
                head_dim: arena.head_dim,
            },
            prepared,
        )?;
    arena.keys = Array::from_native(page_keys)?;
    arena.values = Array::from_native(page_values)?;
    arena.key_scales = Some(Array::from_native(next_key_scales)?);
    arena.value_scales = Some(Array::from_native(next_value_scales)?);
    drop(arena);
    Ok(())
}