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

use models::weights::{BlockQuantization, RoutedExpertBindings, TensorBinding, TensorStorage};

pub(super) fn routed_is_dense(bindings: RoutedExpertBindings<'_>) -> bool {
    all(bindings, |binding| matches!(binding.storage, TensorStorage::Dense { .. }))
}

pub(super) fn routed_is_mxfp4(bindings: RoutedExpertBindings<'_>) -> bool {
    let RoutedExpertBindings::SeparateGateUp { gate, up, down } = bindings else {
        return false;
    };
    [gate, up, down].iter().all(|binding| {
        matches!(
            binding.storage,
            TensorStorage::BlockQuantized { format, .. } if format.is_mxfp4()
        )
    })
}

pub(super) fn routed_is_mxfp8(bindings: RoutedExpertBindings<'_>) -> bool {
    all(bindings, |binding| {
        matches!(
            binding.storage,
            TensorStorage::BlockQuantized { format, .. } if format == BlockQuantization::MXFP8
        )
    })
}

fn all(bindings: RoutedExpertBindings<'_>, predicate: impl Fn(&TensorBinding) -> bool) -> bool {
    match bindings {
        RoutedExpertBindings::SeparateGateUp { gate, up, down } => {
            [gate, up, down].into_iter().all(predicate)
        },
        RoutedExpertBindings::InterleavedGateUp { gate_up, down } => {
            [gate_up, down].into_iter().all(predicate)
        },
        RoutedExpertBindings::Individual { .. } => false,
    }
}