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 super::{ClampedRoutedConfig, weights::ClampedRoutedExpertWeights};
use crate::{CudaTensor, CudaTensorDType, Error, Result};

pub(super) fn validate_common(
    config: ClampedRoutedConfig,
    tensors: [&CudaTensor; 8],
) -> Result<()> {
    let q = config.query_heads * config.head_dim;
    let kv = config.kv_heads * config.head_dim;
    let shapes: [&[usize]; 8] = [
        &[config.hidden],
        &[q],
        &[kv],
        &[kv],
        &[config.hidden],
        &[config.query_heads],
        &[config.hidden],
        &[config.experts],
    ];
    for (tensor, shape) in tensors.into_iter().zip(shapes) {
        validate(tensor, shape, CudaTensorDType::Bf16)?;
    }
    Ok(())
}

pub(super) fn validate_native_experts(
    config: ClampedRoutedConfig,
    weights: &ClampedRoutedExpertWeights,
) -> Result<()> {
    let ClampedRoutedExpertWeights::Native(weights) = weights else {
        return Err(Error::InvalidExecutionPlan("expected native clamped-routed experts"));
    };
    let gate_rows = 2 * config.intermediate;
    validate(
        &weights.gate_up_blocks,
        &[config.experts, gate_rows, config.hidden / 32, 16],
        CudaTensorDType::U8,
    )?;
    validate(
        &weights.gate_up_scales,
        &[config.experts, gate_rows, config.hidden / 32],
        CudaTensorDType::U8,
    )?;
    validate(&weights.gate_up_bias, &[config.experts, gate_rows], CudaTensorDType::Bf16)?;
    validate(
        &weights.down_blocks,
        &[config.experts, config.hidden, config.intermediate / 32, 16],
        CudaTensorDType::U8,
    )?;
    validate(
        &weights.down_scales,
        &[config.experts, config.hidden, config.intermediate / 32],
        CudaTensorDType::U8,
    )?;
    validate(&weights.down_bias, &[config.experts, config.hidden], CudaTensorDType::Bf16)
}

pub(super) fn validate_mlx_experts(
    config: ClampedRoutedConfig,
    weights: &ClampedRoutedExpertWeights,
) -> Result<()> {
    let ClampedRoutedExpertWeights::Mlx(weights) = weights else {
        return Err(Error::InvalidExecutionPlan("expected MLX clamped-routed experts"));
    };
    let gate_shape = [config.experts, config.intermediate, config.hidden / 8];
    let gate_scale_shape = [config.experts, config.intermediate, config.hidden / 32];
    let gate_bias_shape = [config.experts, config.intermediate];
    for tensor in [&weights.gate_blocks, &weights.up_blocks] {
        validate(tensor, &gate_shape, CudaTensorDType::U32)?;
    }
    for tensor in [&weights.gate_scales, &weights.up_scales] {
        validate(tensor, &gate_scale_shape, CudaTensorDType::U8)?;
    }
    for tensor in [&weights.gate_bias, &weights.up_bias] {
        validate(tensor, &gate_bias_shape, CudaTensorDType::Bf16)?;
    }
    validate(
        &weights.down_blocks,
        &[config.experts, config.hidden, config.intermediate / 8],
        CudaTensorDType::U32,
    )?;
    validate(
        &weights.down_scales,
        &[config.experts, config.hidden, config.intermediate / 32],
        CudaTensorDType::U8,
    )?;
    validate(&weights.down_bias, &[config.experts, config.hidden], CudaTensorDType::Bf16)
}

fn validate(tensor: &CudaTensor, shape: &[usize], dtype: CudaTensorDType) -> Result<()> {
    if tensor.shape() != shape {
        return Err(Error::InvalidQuantizedTensor {
            name: tensor.name().into(),
            expected: shape.to_vec(),
            actual: tensor.shape().to_vec(),
        });
    }
    if tensor.dtype() != dtype {
        return Err(Error::DTypeMismatch {
            name: tensor.name().into(),
            expected: match dtype {
                CudaTensorDType::Bf16 => "BF16",
                CudaTensorDType::U32 => "U32",
                _ => "U8",
            },
        });
    }
    Ok(())
}