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

use super::{DecodeMoeBlockBf16, DecodeMoeBlockWeights};
use crate::Result;

impl DecodeMoeBlockBf16 {
    pub(super) fn execute_dense(&mut self, weights: DecodeMoeBlockWeights<'_>) -> Result<()> {
        self.pre_dense_norm.execute(
            &self.scratch.hidden,
            weights.pre_dense_norm,
            &mut self.scratch.normalized,
        )?;
        self.dense_gate_up.execute(
            &self.scratch.normalized,
            weights.dense_gate_up,
            &mut self.scratch.dense_gate_up,
        )?;
        self.dense_activation.execute(
            &self.stream,
            &self.scratch.dense_gate_up,
            &mut self.scratch.dense_activated,
            self.config.activation.into(),
        )?;
        self.dense_down.execute(
            &self.scratch.dense_activated,
            weights.dense_down,
            &mut self.scratch.normalized,
        )?;
        self.post_dense_norm.execute(
            &self.scratch.normalized,
            weights.post_dense_norm,
            &mut self.scratch.dense,
        )
    }

    pub(super) fn execute_experts(&mut self, weights: DecodeMoeBlockWeights<'_>) -> Result<()> {
        let selection = self.router.execute(&self.scratch.hidden, weights.router)?;
        self.pre_expert_norm.execute(
            &self.scratch.hidden,
            weights.pre_expert_norm,
            &mut self.scratch.normalized,
        )?;
        self.experts.execute(
            &self.scratch.normalized,
            selection.indices,
            selection.weights,
            &mut self.scratch.expert,
        )?;
        self.post_expert_norm.execute(
            &self.scratch.expert,
            weights.post_expert_norm,
            &mut self.scratch.expert_norm,
        )
    }
}