libmir-cuda 0.2.0

CUDA 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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

use mircuda::{DeviceBuffer, LaunchConfig, Stream, bf16, cuda_export};

use super::{ClampedRoutedKernels, narrow};
use crate::Result;

cuda_export!(pub(super) GateUpKernel = "libmir_cuda_clamped_routed_mxfp4_gate_up_bf16"(
    input: &DeviceBuffer<bf16>, blocks: &DeviceBuffer<u8>, scales: &DeviceBuffer<u8>,
    bias: &DeviceBuffer<bf16>, selected: &DeviceBuffer<u32>,
    output: &mut DeviceBuffer<bf16>, tokens: u32, top_k: u32, hidden: u32,
    intermediate: u32, limit: f32,
));
cuda_export!(pub(super) DownKernel = "libmir_cuda_clamped_routed_mxfp4_down_bf16"(
    input: &DeviceBuffer<bf16>, blocks: &DeviceBuffer<u8>, scales: &DeviceBuffer<u8>,
    bias: &DeviceBuffer<bf16>, selected: &DeviceBuffer<u32>,
    routing: &DeviceBuffer<bf16>, output: &mut DeviceBuffer<bf16>, tokens: u32,
    top_k: u32, hidden: u32, intermediate: u32,
));
cuda_export!(pub(super) GateUpMlxKernel = "libmir_cuda_clamped_routed_mlx_mxfp4_gate_up_bf16"(
    input: &DeviceBuffer<bf16>, gate_blocks: &DeviceBuffer<u32>,
    gate_scales: &DeviceBuffer<u8>, gate_bias: &DeviceBuffer<bf16>,
    up_blocks: &DeviceBuffer<u32>, up_scales: &DeviceBuffer<u8>,
    up_bias: &DeviceBuffer<bf16>, selected: &DeviceBuffer<u32>,
    output: &mut DeviceBuffer<bf16>, tokens: u32, top_k: u32, hidden: u32,
    intermediate: u32, limit: f32,
));
cuda_export!(pub(super) DownMlxKernel = "libmir_cuda_clamped_routed_mlx_mxfp4_down_bf16"(
    input: &DeviceBuffer<bf16>, blocks: &DeviceBuffer<u32>, scales: &DeviceBuffer<u8>,
    bias: &DeviceBuffer<bf16>, selected: &DeviceBuffer<u32>,
    routing: &DeviceBuffer<bf16>, output: &mut DeviceBuffer<bf16>, tokens: u32,
    top_k: u32, hidden: u32, intermediate: u32,
));

impl ClampedRoutedKernels {
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn gate_up_native(
        &self,
        stream: &Stream,
        input: &DeviceBuffer<bf16>,
        blocks: &DeviceBuffer<u8>,
        scales: &DeviceBuffer<u8>,
        bias: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        Ok(self.gate_up.launch(
            stream,
            self.gate_launch()?,
            (
                input,
                blocks,
                scales,
                bias,
                selected,
                output,
                narrow(self.spec.tokens)?,
                narrow(self.spec.top_k)?,
                narrow(self.spec.hidden)?,
                narrow(self.spec.intermediate)?,
                self.spec.swiglu_limit,
            ),
        )?)
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn down_native(
        &self,
        stream: &Stream,
        input: &DeviceBuffer<bf16>,
        blocks: &DeviceBuffer<u8>,
        scales: &DeviceBuffer<u8>,
        bias: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        routing: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        Ok(self.down.launch(
            stream,
            self.down_launch()?,
            (
                input,
                blocks,
                scales,
                bias,
                selected,
                routing,
                output,
                narrow(self.spec.tokens)?,
                narrow(self.spec.top_k)?,
                narrow(self.spec.hidden)?,
                narrow(self.spec.intermediate)?,
            ),
        )?)
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn gate_up_mlx(
        &self,
        stream: &Stream,
        input: &DeviceBuffer<bf16>,
        gate_blocks: &DeviceBuffer<u32>,
        gate_scales: &DeviceBuffer<u8>,
        gate_bias: &DeviceBuffer<bf16>,
        up_blocks: &DeviceBuffer<u32>,
        up_scales: &DeviceBuffer<u8>,
        up_bias: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        Ok(self.gate_up_mlx.launch(
            stream,
            self.gate_launch()?,
            (
                input,
                gate_blocks,
                gate_scales,
                gate_bias,
                up_blocks,
                up_scales,
                up_bias,
                selected,
                output,
                narrow(self.spec.tokens)?,
                narrow(self.spec.top_k)?,
                narrow(self.spec.hidden)?,
                narrow(self.spec.intermediate)?,
                self.spec.swiglu_limit,
            ),
        )?)
    }

    #[allow(clippy::too_many_arguments)]
    pub(crate) fn down_mlx(
        &self,
        stream: &Stream,
        input: &DeviceBuffer<bf16>,
        blocks: &DeviceBuffer<u32>,
        scales: &DeviceBuffer<u8>,
        bias: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        routing: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        Ok(self.down_mlx.launch(
            stream,
            self.down_launch()?,
            (
                input,
                blocks,
                scales,
                bias,
                selected,
                routing,
                output,
                narrow(self.spec.tokens)?,
                narrow(self.spec.top_k)?,
                narrow(self.spec.hidden)?,
                narrow(self.spec.intermediate)?,
            ),
        )?)
    }

    fn gate_launch(&self) -> Result<LaunchConfig> {
        Ok(LaunchConfig {
            grid: (narrow(self.spec.intermediate * self.spec.tokens * self.spec.top_k)?, 1, 1),
            block: (32, 1, 1),
            shared_memory_bytes: 0,
        })
    }

    fn down_launch(&self) -> Result<LaunchConfig> {
        Ok(LaunchConfig {
            grid: (narrow(self.spec.hidden * self.spec.tokens)?, 1, 1),
            block: (32, 1, 1),
            shared_memory_bytes: 0,
        })
    }
}