mistralrs-quant 0.8.1

Fast, flexible LLM inference.
Documentation 
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

macro_rules! dequant_kernel {
    ($wq:ty, $scalar:ty, $postfix:tt) => {
        paste! {
            pub(crate) fn [< dequantize_ $postfix >](
                wq_packed: *const $wq,
                scale: *const $scalar,
                zero: *const $scalar,
                out: *const $scalar,
                h: i32,
                w: i32
            );
        }
    };
}

pub mod eight_bit {
    use half::{bf16, f16};
    use paste::paste;

    #[allow(dead_code)]
    extern "C" {
        dequant_kernel!(u8, f32, 8bit_u8_kernel_f32);
        dequant_kernel!(u8, f16, 8bit_u8_kernel_f16);
        dequant_kernel!(u8, bf16, 8bit_u8_kernel_bf16);
    }
}

pub mod four_bit {
    use half::{bf16, f16};
    use paste::paste;

    #[allow(dead_code)]
    extern "C" {
        dequant_kernel!(u8, f32, 4bit_u8_kernel_f32);
        dequant_kernel!(u8, f16, 4bit_u8_kernel_f16);
        dequant_kernel!(u8, bf16, 4bit_u8_kernel_bf16);
    }
}

pub mod three_bit {
    use half::{bf16, f16};
    use paste::paste;

    #[allow(dead_code)]
    extern "C" {
        dequant_kernel!(i32, f32, 3bit_32_kernel_f32);
        dequant_kernel!(i32, f16, 3bit_32_kernel_f16);
        dequant_kernel!(i32, bf16, 3bit_32_kernel_bf16);
    }
}

pub mod two_bit {
    use half::{bf16, f16};
    use paste::paste;

    #[allow(dead_code)]
    extern "C" {
        dequant_kernel!(u8, f32, 2bit_u8_kernel_f32);
        dequant_kernel!(u8, f16, 2bit_u8_kernel_f16);
        dequant_kernel!(u8, bf16, 2bit_u8_kernel_bf16);
    }
}

pub mod one_bit {
    use half::{bf16, f16};
    use paste::paste;

    #[allow(dead_code)]
    extern "C" {
        dequant_kernel!(u8, f32, 1bit_u8_kernel_f32);
        dequant_kernel!(u8, f16, 1bit_u8_kernel_f16);
        dequant_kernel!(u8, bf16, 1bit_u8_kernel_bf16);
    }
}