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 crate::{
    BucketedNvFp4MoeBf16, CudaBackend, DirectNvFp4MoeBf16, GatedActivation, GroupedNvFp4MoeBf16,
    HybridNvFp4MoeBf16, MoeExecution, NvFp4ExpertBank, Result, SelectedNvFp4MoeBf16,
};

#[derive(Debug)]
pub(super) enum Experts {
    Bucketed(Box<BucketedNvFp4MoeBf16>),
    Direct(Box<DirectNvFp4MoeBf16>),
    Grouped(Box<GroupedNvFp4MoeBf16>),
    Hybrid(Box<HybridNvFp4MoeBf16>),
    WeightOnly(Box<SelectedNvFp4MoeBf16>),
}

impl Experts {
    #[allow(clippy::too_many_arguments)]
    pub(super) fn new(
        backend: &CudaBackend,
        tokens: usize,
        selected: usize,
        activation: GatedActivation,
        execution: MoeExecution,
        gate: NvFp4ExpertBank,
        up: NvFp4ExpertBank,
        down: NvFp4ExpertBank,
    ) -> Result<Self> {
        match execution {
            MoeExecution::DirectW4A4 => backend
                .prepare_direct_nvfp4_moe_bf16(tokens, selected, activation, gate, up, down)
                .map(Box::new)
                .map(Self::Direct),
            MoeExecution::HybridW4A4 => backend
                .prepare_hybrid_gate_nvfp4_moe_bf16(tokens, selected, activation, gate, up, down)
                .map(Box::new)
                .map(Self::Hybrid),
            MoeExecution::IndexedGrouped | MoeExecution::FusedIndexedGrouped => backend
                .prepare_grouped_nvfp4_moe_bf16(
                    tokens, selected, activation, execution, gate, up, down,
                )
                .map(Box::new)
                .map(Self::Grouped),
            MoeExecution::SelectedWeightOnly => backend
                .prepare_batched_selected_nvfp4_moe_bf16(
                    tokens, selected, activation, gate, up, down,
                )
                .map(Box::new)
                .map(Self::WeightOnly),
            MoeExecution::Bucketed => backend
                .prepare_bucketed_nvfp4_moe_bf16(tokens, selected, activation, gate, up, down)
                .map(Box::new)
                .map(Self::Bucketed),
        }
    }

    pub(super) fn execute(
        &mut self,
        input: &DeviceBuffer<bf16>,
        selected: &DeviceBuffer<u32>,
        routing: &DeviceBuffer<bf16>,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        match self {
            Self::Bucketed(experts) => experts.execute(input, selected, routing, output),
            Self::Direct(experts) => experts.execute(input, selected, routing, output),
            Self::Grouped(experts) => experts.execute(input, selected, routing, output),
            Self::Hybrid(experts) => experts.execute(input, selected, routing, output),
            Self::WeightOnly(experts) => experts.execute(input, selected, routing, output),
        }
    }
}