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

use mircuda::Stream;

use super::{SelectedDenseMoe, SelectedDenseReduceLaunch, launch};
use crate::{Result, kernels::geometry::narrow};

impl SelectedDenseMoe {
    pub fn reduce(
        &self,
        stream: &Stream,
        launch: &mut SelectedDenseReduceLaunch<'_>,
    ) -> Result<()> {
        self.validate_reduce(launch)?;
        if self.spec.down_transposed && self.spec.tokens == 1 {
            return self.reduce_split(stream, launch);
        }
        let spec = self.spec;
        Ok(self.reduce.launch(
            stream,
            launch::reduce(spec)?,
            (
                launch.input,
                launch.selected,
                launch.routing,
                launch.weight,
                launch.bias,
                &mut *launch.output,
                narrow(spec.output_features)?,
                narrow(spec.input_features)?,
                narrow(spec.expert_count)?,
                narrow(spec.selected_count)?,
                u32::from(spec.down_transposed),
                u32::from(spec.down_bias),
            ),
        )?)
    }

    fn reduce_split(
        &self,
        stream: &Stream,
        launch: &mut SelectedDenseReduceLaunch<'_>,
    ) -> Result<()> {
        let spec = self.spec;
        self.project.launch(
            stream,
            launch::project(spec)?,
            (
                launch.input,
                launch.selected,
                launch.routing,
                launch.weight,
                launch.bias,
                &mut *launch.partial,
                narrow(spec.output_features)?,
                narrow(spec.input_features)?,
                narrow(spec.expert_count)?,
                narrow(spec.selected_count)?,
                u32::from(spec.down_bias),
                u32::from(spec.down_transposed),
            ),
        )?;
        self.finalize(stream, launch)
    }

    pub(super) fn finalize(
        &self,
        stream: &Stream,
        launch: &mut SelectedDenseReduceLaunch<'_>,
    ) -> Result<()> {
        let spec = self.spec;
        Ok(self.finalize.launch(
            stream,
            launch::finalize(spec)?,
            (
                &*launch.partial,
                &mut *launch.output,
                narrow(spec.input_features)?,
                narrow(spec.selected_count)?,
            ),
        )?)
    }
}