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
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

use mircuda::{DeviceBuffer, bf16};

use super::{
    AffineSharedExpertMoeConfig, AffineSharedExpertMoeWeights, routed::SharedRoutedExecution,
    scratch::AffineSharedMoeScratch,
};
use crate::{
    CudaBackend, DenseRole, Error, Result,
    backend::linear::CheckpointProjection,
    kernels::{ElementwiseBf16, SigmoidMultiplyBf16},
};

#[derive(Debug)]
pub struct CudaAffineSharedExpertMoeExecution {
    backend: CudaBackend,
    config: AffineSharedExpertMoeConfig,
    tokens: usize,
    routed: SharedRoutedExecution,
    shared_gate: CheckpointProjection,
    shared_up: CheckpointProjection,
    shared_down: CheckpointProjection,
    shared_output_gate: CheckpointProjection,
    shared_activation: ElementwiseBf16,
    output_add: ElementwiseBf16,
    sigmoid_multiply: SigmoidMultiplyBf16,
    weights: AffineSharedExpertMoeWeights,
    scratch: AffineSharedMoeScratch,
}

impl CudaAffineSharedExpertMoeExecution {
    pub(super) fn new(
        backend: &CudaBackend,
        config: AffineSharedExpertMoeConfig,
        weights: &AffineSharedExpertMoeWeights,
        tokens: usize,
    ) -> Result<Self> {
        let shared = |input, output, role, weight| {
            CheckpointProjection::new(backend, tokens, input, output, role, weight)
        };
        Ok(Self {
            backend: backend.clone(),
            config,
            tokens,
            routed: SharedRoutedExecution::new(backend, config, weights, tokens)?,
            shared_gate: shared(
                config.hidden_size,
                config.shared_intermediate_size,
                DenseRole::DenseGateUp,
                &weights.shared_gate,
            )?,
            shared_up: shared(
                config.hidden_size,
                config.shared_intermediate_size,
                DenseRole::DenseGateUp,
                &weights.shared_up,
            )?,
            shared_down: shared(
                config.shared_intermediate_size,
                config.hidden_size,
                DenseRole::DenseDown,
                &weights.shared_down,
            )?,
            shared_output_gate: CheckpointProjection::new(
                backend,
                tokens,
                config.hidden_size,
                1,
                DenseRole::Router,
                &weights.shared_output_gate,
            )?,
            shared_activation: ElementwiseBf16::compile(
                &backend.inner.compiler,
                tokens * config.shared_intermediate_size,
            )?,
            output_add: ElementwiseBf16::compile(
                &backend.inner.compiler,
                tokens * config.hidden_size,
            )?,
            sigmoid_multiply: SigmoidMultiplyBf16::compile(
                &backend.inner.compiler,
                tokens,
                config.hidden_size,
            )?,
            weights: weights.clone(),
            scratch: AffineSharedMoeScratch::new(backend, config, tokens)?,
        })
    }

    pub fn execute(
        &mut self,
        input: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        self.validate(input, output)?;
        self.routed.execute(&self.backend, &self.weights, input, &mut self.scratch)?;
        self.shared(input)?;
        self.output_add.add(
            &self.backend.inner.stream,
            &self.scratch.routed_output,
            &self.scratch.gated_shared_output,
            output,
        )
    }

    fn shared(&mut self, input: &DeviceBuffer<bf16>) -> Result<()> {
        let stream = &self.backend.inner.stream;
        self.shared_gate.execute(input, &mut self.scratch.shared_gate)?;
        self.shared_up.execute(input, &mut self.scratch.shared_up)?;
        self.shared_activation.gated(
            stream,
            &self.scratch.shared_gate,
            &self.scratch.shared_up,
            &mut self.scratch.shared_intermediate,
            self.config.activation.into(),
        )?;
        self.shared_down
            .execute(&self.scratch.shared_intermediate, &mut self.scratch.shared_output)?;
        self.shared_output_gate.execute(input, &mut self.scratch.shared_output_gate)?;
        self.sigmoid_multiply.execute(
            stream,
            &self.scratch.shared_output,
            &self.scratch.shared_output_gate,
            &mut self.scratch.gated_shared_output,
        )
    }

    fn validate(&self, input: &DeviceBuffer<bf16>, output: &DeviceBuffer<bf16>) -> Result<()> {
        let expected = self.tokens * self.config.hidden_size;
        if input.len() != expected || output.len() != expected {
            return Err(Error::InvalidDecoderKernel("affine shared-expert MoE buffer mismatch"));
        }
        Ok(())
    }
}