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

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

use crate::{
    CudaBackend, CudaTensorSet, Error, NvFp4Config, NvFp4LinearWeight, NvFp4Tensors,
    NvFp4WeightOnlyWeight, Result,
};

pub(super) fn load_native(
    backend: &CudaBackend,
    tensors: &CudaTensorSet,
    binding: &TensorBinding,
) -> Result<NvFp4LinearWeight> {
    let TensorStorage::BlockQuantized {
        format: BlockQuantization::NVFP4,
        scales,
        global_scale: Some(global_scale),
        input_scale: Some(input_scale),
        ..
    } = &binding.storage
    else {
        return Err(Error::InvalidNvFp4("projection has no complete NVFP4 binding"));
    };
    let [output, input] = binding.logical_shape.as_deref().ok_or_else(|| {
        Error::UnsupportedDecoderLayer(format!(
            "NVFP4 projection has no logical matrix shape: {}",
            binding.source
        ))
    })?
    else {
        return Err(Error::UnsupportedDecoderLayer(format!(
            "NVFP4 projection is not a matrix: {}",
            binding.source
        )));
    };
    let get = |name: &str| tensors.get(name).ok_or_else(|| Error::MissingTensor(name.into()));
    backend.prepare_nvfp4_linear_weight(
        NvFp4Config::new(*input, *output),
        NvFp4Tensors {
            weight: get(&binding.source)?,
            weight_scale: get(scales)?,
            weight_scale_2: get(global_scale)?,
            input_scale: get(input_scale)?,
        },
    )
}

pub(super) fn load_weight_only(
    tensor_set: &CudaTensorSet,
    binding: &TensorBinding,
) -> Result<NvFp4WeightOnlyWeight> {
    let (config, tensors) = tensors(binding, tensor_set)?;
    NvFp4WeightOnlyWeight::load(config, tensors)
}

fn tensors<'a>(
    binding: &TensorBinding,
    tensors: &'a CudaTensorSet,
) -> Result<(NvFp4Config, NvFp4Tensors<'a>)> {
    let TensorStorage::BlockQuantized {
        format:
            BlockQuantization {
                format: models::weights::BlockFormat::NvFp4,
                ..
            },
        scales,
        global_scale: Some(global_scale),
        input_scale: Some(input_scale),
        ..
    } = &binding.storage
    else {
        return Err(Error::InvalidNvFp4("projection has no complete NVFP4 binding"));
    };
    let [output, input] = binding.logical_shape.as_deref().ok_or_else(|| {
        Error::UnsupportedDecoderLayer(format!(
            "NVFP4 projection has no logical matrix shape: {}",
            binding.source
        ))
    })?
    else {
        return Err(Error::UnsupportedDecoderLayer(format!(
            "NVFP4 projection is not a matrix: {}",
            binding.source
        )));
    };
    let get = |name: &str| tensors.get(name).ok_or_else(|| Error::MissingTensor(name.into()));
    Ok((
        NvFp4Config::new(*input, *output),
        NvFp4Tensors {
            weight: get(&binding.source)?,
            weight_scale: get(scales)?,
            weight_scale_2: get(global_scale)?,
            input_scale: get(input_scale)?,
        },
    ))
}