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
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

mod gated_delta;
mod page_write;
mod paged_attention;

pub(super) use page_write::{PageWriteOptions, PreparedPageWrite};

use super::Result;

mirtal::metal_kernel! {
    fn gated_delta_gates {
        name: "mirmir_gated_delta_gates",
        templates: [HV: int = 16],
        inputs: [alpha: float, beta: float, a_log: float, dt_bias: float],
        outputs: [decay: f32, update: f32],
        source: file "kernels/gated_delta_gates.metal",
        header: inline "",
        row_contiguous: true,
        atomic_outputs: false,
    }
}

mirtal::metal_library! {
    fn paged_kv_library {
        name: "mirmir_paged_kv",
        source: file "kernels/paged_kv.metal",
    }
}

mirtal::metal_kernel! {
    fn gated_delta_recurrence {
        name: "mirmir_gated_delta",
        templates: [
            InT: dtype = bf16, StT: dtype = f32, DK: int = 128, DV: int = 128,
            HK: int = 8, HV: int = 16, STEPS: int = 4,
        ],
        inputs: [query: InT, key: InT, value: InT, decay: f32, update: f32, state: StT],
        outputs: [output: InT, next_state: StT],
        source: file "kernels/gated_delta_recurrence.metal",
        header: inline "",
        row_contiguous: true,
        atomic_outputs: false,
    }
}

mirtal::metal_kernel! {
    fn gated_delta_decode {
        name: "mirmir_gated_delta_decode_headwide_gates",
        templates: [
            InT: dtype = bf16, StT: dtype = f32, DK: int = 128, DV: int = 128,
            HK: int = 8, HV: int = 16, NORMALIZE: bool = true,
        ],
        inputs: [
            query: InT, key: InT, value: InT, alpha: float, beta: float,
            a_log: float, dt_bias: float, state: StT,
        ],
        outputs: [output: InT, next_state: StT],
        source: file "kernels/gated_delta_decode.metal",
        header: inline "",
        row_contiguous: true,
        atomic_outputs: false,
    }
}

mirtal::metal_kernel! {
    fn paged_attention {
        name: "mirmir_paged_sdpa",
        templates: [
            T: dtype = bf16, QUERY_HEADS: int = 32, KV_HEADS: int = 16,
            PAGE_CAPACITY: int = 128, HEAD_DIM: int = 128,
            QK_PER_THREAD: int = 4, V_PER_THREAD: int = 4, PAGE_SIZE: int = 64,
        ],
        inputs: [
            queries: T, key_pages: T, value_pages: T, page_table: u32,
            page_dependency: u32, attention_scale: scalar<f32>,
        ],
        outputs: [output: T],
        source: file "kernels/paged_attention.metal",
        header: inline "",
        row_contiguous: true,
        atomic_outputs: false,
    }
}

mirtal::metal_library! {
    fn paged_attention_partial_library {
        name: "mirmir_paged_sdpa_partial",
        source: file "kernels/paged_attention/partial.metal",
    }
}

mirtal::metal_kernel! {
    fn paged_attention_reduce {
        name: "mirmir_paged_sdpa_reduce",
        templates: [T: dtype = bf16, HEAD_DIM: int = 128, BLOCKS: int = 64],
        inputs: [partials: T, sums: f32, maximums: f32],
        outputs: [output: T],
        source: file "kernels/paged_attention/reduce.metal",
        header: inline "",
        row_contiguous: true,
        atomic_outputs: false,
    }
}

#[derive(Debug)]
pub(super) struct Kernels {
    gated_delta_gates: mirtal::MetalKernel<4, 2>,
    gated_delta_recurrence: mirtal::MetalKernel<6, 2>,
    gated_delta_decode: mirtal::MetalKernel<8, 2>,
    paged_attention: mirtal::MetalKernel<6, 1>,
    paged_attention_partial: mirtal::MetalLibrary,
    paged_attention_reduce: mirtal::MetalKernel<3, 1>,
    paged_kv: mirtal::MetalLibrary,
}

impl Kernels {
    pub(super) fn new() -> Result<Self> {
        Ok(Self {
            gated_delta_gates: gated_delta_gates()?,
            gated_delta_recurrence: gated_delta_recurrence()?,
            gated_delta_decode: gated_delta_decode()?,
            paged_attention: paged_attention()?,
            paged_attention_partial: paged_attention_partial_library()?,
            paged_attention_reduce: paged_attention_reduce()?,
            paged_kv: paged_kv_library()?,
        })
    }
}

fn template(name: &'static str, value: usize) -> Result<mirtal::TemplateArg> {
    Ok(mirtal::TemplateArg::int(name, i32::try_from(value)?))
}

pub(super) fn new_gated_delta_decode_kernel() -> Result<mirtal::MetalKernel<8, 2>> {
    Ok(gated_delta_decode()?)
}