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

use models::weights::{
    BlockActivationMode, BlockFormat, ExpertProjectionRole, RoutedExpertBindings, TensorBinding,
    TensorCatalog, TensorStorage,
};

use crate::{
    CudaBackend, Error, NvFp4ExpertBankConfig, NvFp4ExpertSource, Result,
    backend::block::experts::ExpertWeights,
};

pub(super) fn load(
    backend: &CudaBackend,
    catalog: &TensorCatalog,
    bindings: RoutedExpertBindings<'_>,
    experts: usize,
    hidden: usize,
    intermediate: usize,
) -> Result<ExpertWeights> {
    let activation_mode = mode(bindings)?;
    let gate = sources(catalog, bindings.individual(ExpertProjectionRole::Gate))?;
    let up = sources(catalog, bindings.individual(ExpertProjectionRole::Up))?;
    let down = sources(catalog, bindings.individual(ExpertProjectionRole::Down))?;
    let bank = |input, output, sources: &[NvFp4ExpertSource<'_>]| {
        backend.prepare_nvfp4_expert_bank(
            NvFp4ExpertBankConfig {
                experts,
                input_features: input,
                output_features: output,
            },
            sources,
        )
    };
    Ok(ExpertWeights::NvFp4 {
        gate: bank(hidden, intermediate, &gate)?,
        up: bank(hidden, intermediate, &up)?,
        down: bank(intermediate, hidden, &down)?,
        activation_mode,
    })
}

fn sources<'a>(
    catalog: &'a TensorCatalog,
    bindings: Vec<&TensorBinding>,
) -> Result<Vec<NvFp4ExpertSource<'a>>> {
    bindings.into_iter().map(|binding| source(catalog, binding)).collect()
}

fn source<'a>(
    catalog: &'a TensorCatalog,
    binding: &TensorBinding,
) -> Result<NvFp4ExpertSource<'a>> {
    let TensorStorage::BlockQuantized {
        format,
        scales,
        global_scale: Some(global_scale),
        input_scale: Some(input_scale),
        ..
    } = &binding.storage
    else {
        return Err(Error::UnsupportedDecoderLayer(format!(
            "CUDA individual expert is not complete NVFP4: {}",
            binding.source
        )));
    };
    if format.format != BlockFormat::NvFp4 {
        return Err(Error::UnsupportedDecoderLayer(format!(
            "CUDA individual expert is not NVFP4: {}",
            binding.source
        )));
    }
    let get = |name: &str| catalog.get(name).ok_or_else(|| Error::MissingTensor(name.into()));
    Ok(NvFp4ExpertSource {
        weight: get(&binding.source)?,
        weight_scale: get(scales)?,
        weight_scale_2: get(global_scale)?,
        input_scale: get(input_scale)?,
    })
}

fn mode(bindings: RoutedExpertBindings<'_>) -> Result<BlockActivationMode> {
    let mut mode = None;
    let bindings =
        [ExpertProjectionRole::Gate, ExpertProjectionRole::Up, ExpertProjectionRole::Down]
            .into_iter()
            .flat_map(|projection| bindings.individual(projection));
    for binding in bindings {
        let TensorStorage::BlockQuantized { format, .. } = binding.storage else {
            return Err(Error::UnsupportedDecoderLayer(
                "CUDA individual experts mix non-block storage".into(),
            ));
        };
        if format.format != BlockFormat::NvFp4
            || mode
                .replace(format.activation_mode)
                .is_some_and(|value| value != format.activation_mode)
        {
            return Err(Error::UnsupportedDecoderLayer(
                "CUDA individual experts mix NVFP4 activation contracts".into(),
            ));
        }
    }
    mode.ok_or_else(|| Error::UnsupportedDecoderLayer("CUDA has no individual experts".into()))
}