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
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

use super::BoundLinear;
use crate::engine::{
    FusedAttention, FusedExpertGateUp, FusedGateUp, FusedKeyValue, Result, Stream,
};

impl BoundLinear {
    pub(in crate::engine) fn fuse_gate_up(
        &self,
        up: &Self,
        stream: &Stream,
    ) -> Result<Option<FusedGateUp>> {
        match (self, up) {
            (Self::Affine(gate), Self::Affine(up)) => gate.fuse_gate_up(up, stream),
            _ => Ok(None),
        }
    }

    pub(in crate::engine) fn fused_gate_up_bytes(&self, up: &Self) -> Result<Option<usize>> {
        match (self, up) {
            (Self::Affine(gate), Self::Affine(up)) => gate.fused_gate_up_bytes(up),
            _ => Ok(None),
        }
    }

    pub(in crate::engine) fn fuse_expert_gate_up(
        &self,
        up: &Self,
        stream: &Stream,
    ) -> Result<Option<FusedExpertGateUp>> {
        match (self, up) {
            (Self::Affine(gate), Self::Affine(up)) => gate.fuse_expert_gate_up(up, stream),
            _ => Ok(None),
        }
    }

    pub(in crate::engine) fn fused_expert_gate_up_bytes(&self, up: &Self) -> Result<Option<usize>> {
        match (self, up) {
            (Self::Affine(gate), Self::Affine(up)) => gate.fused_expert_gate_up_bytes(up),
            _ => Ok(None),
        }
    }

    pub(in crate::engine) fn fuse_attention(
        &self,
        key: &Self,
        value: Option<&Self>,
        stream: &Stream,
    ) -> Result<Option<FusedAttention>> {
        match (self, key, value) {
            (Self::Affine(query), Self::Affine(key), Some(Self::Affine(value))) => {
                query.fuse_attention(key, Some(value), stream)
            },
            (Self::Affine(query), Self::Affine(key), None) => {
                query.fuse_attention(key, None, stream)
            },
            _ => Ok(None),
        }
    }

    pub(in crate::engine) fn fuse_key_value(
        &self,
        value: Option<&Self>,
        stream: &Stream,
    ) -> Result<Option<FusedKeyValue>> {
        match (self, value) {
            (Self::Affine(key), Some(Self::Affine(value))) => {
                key.fuse_key_value(Some(value), stream)
            },
            _ => Ok(None),
        }
    }
}