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
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72

use mircuda::{DeviceBuffer, bf16};
use runtime::kv::{BlockTable, KvWritePlan};

use super::super::{DecodeDenseSwiGlu, DenseSwiGluWeights, GateUpBuffers};
use crate::{Result, backend::attention::graph::Nodes};

impl DecodeDenseSwiGlu {
    pub(super) fn execute_captured(
        &mut self,
        input: &DeviceBuffer<bf16>,
        weights: DenseSwiGluWeights<'_>,
        write_plan: &KvWritePlan,
        table: &BlockTable,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<Nodes> {
        let nodes = self.attention.execute_captured(
            input,
            weights.attention,
            write_plan,
            table,
            &mut self.scratch.attention,
        )?;
        self.hidden_ops.add(
            &self.stream,
            input,
            &self.scratch.attention,
            &mut self.scratch.residual,
        )?;
        let separate = self.gate_up.execute(
            &self.scratch.residual,
            &self.post_attention_norm,
            weights.post_attention_norm,
            weights.gate_up,
            &mut GateUpBuffers {
                normalized: &mut self.scratch.normalized,
                packed: &mut self.scratch.gate_up,
                separate: &mut self.scratch.gate_up_separate,
            },
        )?;
        let fused_down = separate
            && self.down.execute_gated(
                &self.scratch.gate_up_separate[0],
                &self.scratch.gate_up_separate[1],
                self.config.activation,
                weights.down,
                &mut self.scratch.mlp,
            )?;
        if separate && !fused_down {
            self.activation.execute_separate(
                &self.stream,
                &self.scratch.gate_up_separate[0],
                &self.scratch.gate_up_separate[1],
                &mut self.scratch.activated,
                self.config.activation.into(),
            )?;
        } else if !fused_down {
            self.activation.execute(
                &self.stream,
                &self.scratch.gate_up,
                &mut self.scratch.activated,
                self.config.activation.into(),
            )?;
        }
        if !fused_down {
            self.down
                .execute(&self.scratch.activated, weights.down, &mut self.scratch.mlp)?;
        }
        self.hidden_ops
            .add(&self.stream, &self.scratch.residual, &self.scratch.mlp, output)?;
        Ok(nodes)
    }
}