libmir-cuda 0.2.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

use std::collections::HashSet;

use models::weights::{
    BlockFormat, HybridMoeExpertBindings, HybridMoeLayerBindings, TensorBinding, TensorCatalog,
    TensorInfo, TensorStorage,
};

use crate::{Error, NvFp4ExpertSource, Result};

pub(super) struct LayerSource<'a> {
    pub tensors: Vec<&'a TensorInfo>,
    pub names: Vec<String>,
    pub gate: Vec<NvFp4ExpertSource<'a>>,
    pub up: Vec<NvFp4ExpertSource<'a>>,
    pub down: Vec<NvFp4ExpertSource<'a>>,
}

impl<'a> LayerSource<'a> {
    pub fn discover(
        catalog: &'a TensorCatalog,
        bindings: &HybridMoeLayerBindings<'_>,
    ) -> Result<Self> {
        let names = common(bindings).into_iter().map(|binding| binding.source.clone()).collect();
        let mut seen = HashSet::new();
        let tensors = common(bindings)
            .into_iter()
            .flat_map(TensorBinding::physical_sources)
            .filter(|name| seen.insert((*name).to_owned()))
            .map(|name| required(catalog, name))
            .collect::<Result<_>>()?;
        let HybridMoeExpertBindings::Individual { gate, up, down } = &bindings.experts else {
            return Err(Error::UnsupportedDecoderLayer(
                "CUDA hybrid MoE requires individual NVFP4 expert bindings".into(),
            ));
        };
        Ok(Self {
            tensors,
            names,
            gate: expert_sources(catalog, gate)?,
            up: expert_sources(catalog, up)?,
            down: expert_sources(catalog, down)?,
        })
    }
}

fn common<'a>(bindings: &'a HybridMoeLayerBindings<'a>) -> Vec<&'a TensorBinding> {
    vec![
        bindings.input_norm,
        bindings.attention.query,
        bindings.attention.key,
        bindings.attention.value.unwrap_or(bindings.attention.key),
        bindings.attention.query_norm,
        bindings.attention.key_norm,
        bindings.attention.output,
        bindings.post_attention_norm,
        bindings.pre_dense_norm,
        bindings.dense.gate,
        bindings.dense.up,
        bindings.dense.down,
        bindings.post_dense_norm,
        bindings.router.projection,
        bindings.router.norm_scale,
        bindings.router.expert_scale,
        bindings.pre_expert_norm,
        bindings.post_expert_norm,
        bindings.post_feed_forward_norm,
        bindings.layer_scale,
    ]
}

fn expert_sources<'a>(
    catalog: &'a TensorCatalog,
    bindings: &[&TensorBinding],
) -> Result<Vec<NvFp4ExpertSource<'a>>> {
    bindings.iter().map(|binding| expert_source(catalog, binding)).collect()
}

fn expert_source<'a>(
    catalog: &'a TensorCatalog,
    binding: &TensorBinding,
) -> Result<NvFp4ExpertSource<'a>> {
    let TensorStorage::BlockQuantized {
        format: BlockFormat::NvFp4,
        scales,
        global_scale: Some(global_scale),
        input_scale: Some(input_scale),
        ..
    } = &binding.storage
    else {
        return Err(Error::UnsupportedDecoderLayer(format!(
            "CUDA expert {} requires a complete NVFP4 binding",
            binding.source
        )));
    };
    Ok(NvFp4ExpertSource {
        weight: required(catalog, &binding.source)?,
        weight_scale: required(catalog, scales)?,
        weight_scale_2: required(catalog, global_scale)?,
        input_scale: required(catalog, input_scale)?,
    })
}

fn required<'a>(catalog: &'a TensorCatalog, name: &str) -> Result<&'a TensorInfo> {
    catalog
        .tensors
        .iter()
        .find(|tensor| tensor.name == name)
        .ok_or_else(|| Error::MissingTensor(name.into()))
}