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

use mircuda::{DeviceBuffer, bf16};
use models::weights::BlockActivationMode;

use self::nvfp4::AutoNvFp4Experts;
use crate::{
    CudaBackend, Error, ExecutionPhase, GatedActivation, NvFp4ExpertBank, Result,
    backend::linear::{DenseExpertWeights, SelectedDenseMoeBf16},
};

mod nvfp4;

#[derive(Clone, Debug)]
pub(in crate::backend) enum ExpertWeights {
    NvFp4 {
        gate: NvFp4ExpertBank,
        up: NvFp4ExpertBank,
        down: NvFp4ExpertBank,
        activation_mode: BlockActivationMode,
    },
    Dense(DenseExpertWeights),
}

#[derive(Debug)]
pub(in crate::backend) enum Experts {
    NvFp4(Box<AutoNvFp4Experts>),
    Dense {
        operation: Box<SelectedDenseMoeBf16>,
        intermediate: DeviceBuffer<bf16>,
    },
}

impl Experts {
    #[allow(clippy::too_many_arguments)]
    pub(in crate::backend) fn new(
        backend: &CudaBackend,
        phase: ExecutionPhase,
        tokens: usize,
        selected: usize,
        activation: GatedActivation,
        weights: &ExpertWeights,
    ) -> Result<Self> {
        match weights {
            ExpertWeights::Dense(weights) => Ok(Self::Dense {
                operation: Box::new(SelectedDenseMoeBf16::new(
                    backend,
                    tokens,
                    selected,
                    weights,
                    activation.into(),
                )?),
                intermediate: backend
                    .pool()
                    .allocate(backend.stream(), weights.intermediate_elements(tokens, selected)?)?,
            }),
            ExpertWeights::NvFp4 { gate, up, down, activation_mode } => AutoNvFp4Experts::new(
                backend,
                phase,
                tokens,
                selected,
                activation,
                gate.clone(),
                up.clone(),
                down.clone(),
                *activation_mode,
            )
            .map(Box::new)
            .map(Self::NvFp4),
        }
    }

    pub(in crate::backend) fn execute(
        &mut self,
        input: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        routing: &DeviceBuffer<bf16>,
        weights: &ExpertWeights,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        match (self, weights) {
            (Self::NvFp4(experts), ExpertWeights::NvFp4 { .. }) => {
                experts.execute(input, selected, routing, output)
            },
            (Self::Dense { operation, intermediate }, ExpertWeights::Dense(weights)) => {
                operation.execute(input, selected, routing, weights, intermediate, output)
            },
            _ => Err(Error::InvalidExecutionPlan(
                "routed-expert plan differs from layer weight storage",
            )),
        }
    }
}