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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

use mircuda::{DeviceBuffer, bf16};
use models::weights::{TensorBinding, TensorStorage};

use super::{
    CompressedInt8Bf16Linear, CompressedInt8Embedding, CompressedInt8Weight,
    awq::{AwqBf16Linear, AwqWeight},
    bitsandbytes::{BitsAndBytes4BitBf16Linear, BitsAndBytes4BitWeight},
    gptq::{GptqBf16Linear, GptqWeight},
};
use crate::{CudaBackend, CudaTensorSet, Error, Result};

#[derive(Clone, Debug)]
pub enum PackedIntegerWeight {
    Compressed(CompressedInt8Weight),
    Awq(AwqWeight),
    Gptq(GptqWeight),
    BitsAndBytes4Bit(BitsAndBytes4BitWeight),
}

impl PackedIntegerWeight {
    pub(crate) fn load_binding(
        tensors: &CudaTensorSet,
        binding: &TensorBinding,
        input: usize,
        output: usize,
    ) -> Result<Self> {
        match binding.storage {
            TensorStorage::PackedInt8 { .. } | TensorStorage::PackedInt4 { .. } => {
                CompressedInt8Weight::load_binding(tensors, binding, input, output)
                    .map(Self::Compressed)
            },
            TensorStorage::Awq { .. } => {
                AwqWeight::load_binding(tensors, binding, input, output).map(Self::Awq)
            },
            TensorStorage::Gptq { .. } => {
                GptqWeight::load_binding(tensors, binding, input, output).map(Self::Gptq)
            },
            TensorStorage::BitsAndBytes4Bit { .. } => {
                BitsAndBytes4BitWeight::load_binding(tensors, binding, input, output)
                    .map(Self::BitsAndBytes4Bit)
            },
            _ => Err(Error::InvalidQuantizedGemv("binding is not a packed integer weight")),
        }
    }

    pub(in crate::backend) fn validate(&self, input: usize, output: usize) -> Result<()> {
        match self {
            Self::Compressed(weight) => weight.validate(input, output),
            Self::Awq(weight) => weight.validate(input, output),
            Self::Gptq(weight) => weight.validate(input, output),
            Self::BitsAndBytes4Bit(weight) => weight.validate(),
        }
    }
}

#[derive(Clone, Debug)]
pub enum PackedIntegerBf16Linear {
    Compressed(CompressedInt8Bf16Linear),
    Awq(AwqBf16Linear),
    Gptq(GptqBf16Linear),
    BitsAndBytes4Bit(BitsAndBytes4BitBf16Linear),
}

impl PackedIntegerBf16Linear {
    pub(in crate::backend) fn new(
        backend: &CudaBackend,
        tokens: usize,
        input: usize,
        output: usize,
        weight: &PackedIntegerWeight,
    ) -> Result<Self> {
        match weight {
            PackedIntegerWeight::Compressed(weight) => {
                CompressedInt8Bf16Linear::new(backend, tokens, input, output, weight)
                    .map(Self::Compressed)
            },
            PackedIntegerWeight::Awq(weight) => {
                AwqBf16Linear::new(backend, tokens, input, output, weight).map(Self::Awq)
            },
            PackedIntegerWeight::Gptq(weight) => {
                GptqBf16Linear::new(backend, tokens, input, output, weight).map(Self::Gptq)
            },
            PackedIntegerWeight::BitsAndBytes4Bit(weight) => {
                BitsAndBytes4BitBf16Linear::new(backend, tokens, weight).map(Self::BitsAndBytes4Bit)
            },
        }
    }

    pub(in crate::backend) fn execute(
        &self,
        input: &DeviceBuffer<bf16>,
        weight: &PackedIntegerWeight,
        output: &mut DeviceBuffer<bf16>,
    ) -> Result<()> {
        match (self, weight) {
            (Self::Compressed(operation), PackedIntegerWeight::Compressed(weight)) => {
                operation.execute(input, weight, output)
            },
            (Self::Awq(operation), PackedIntegerWeight::Awq(weight)) => {
                operation.execute(input, weight, output)
            },
            (Self::Gptq(operation), PackedIntegerWeight::Gptq(weight)) => {
                operation.execute(input, weight, output)
            },
            (Self::BitsAndBytes4Bit(operation), PackedIntegerWeight::BitsAndBytes4Bit(weight)) => {
                operation.execute(input, weight, output)
            },
            _ => Err(Error::InvalidExecutionPlan("packed integer operation/weight mismatch")),
        }
    }
}

pub(in crate::backend) fn embedding(
    backend: &CudaBackend,
    vocab: usize,
    hidden: usize,
    output_scale: f32,
    weight: PackedIntegerWeight,
) -> Result<CompressedInt8Embedding> {
    match weight {
        PackedIntegerWeight::Compressed(weight) => {
            CompressedInt8Embedding::new(backend, vocab, hidden, output_scale, weight)
        },
        PackedIntegerWeight::Awq(_)
        | PackedIntegerWeight::Gptq(_)
        | PackedIntegerWeight::BitsAndBytes4Bit(_) => Err(Error::InvalidQuantizedGemv(
            "AWQ/GPTQ storage cannot be used as an embedding table",
        )),
    }
}