singe-kernel 0.1.0-alpha.4

Reusable CPU and GPU kernels.
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
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

// @generated by cargo xtask gen-cutile-dtype-kernels
// Do not edit generated entries by hand.

#[cutile::module]
mod kernels {
    use cutile::core::*;
    #[cutile::entry()]
    pub fn quantize_f32_u8<const B: i32>(
        out: &mut Tensor<u8, { [B] }>,
        input: &Tensor<f32, { [-1] }>,
        scale: f32,
        zero_point: f32,
    ) {
        let values: Tile<f32, { [B] }> = load_tile_like(input, out);
        let raw = values / broadcast_scalar(scale, out.shape())
            + broadcast_scalar(zero_point, out.shape());
        let low = constant(0.0f32, out.shape());
        let high = constant(255.0f32, out.shape());
        let clamped = min_tile(max_tile(raw, low), high);
        let floored = floor(clamped);
        let fraction = clamped - floored;
        let half = constant(0.5f32, out.shape());
        let one = constant(1.0f32, out.shape());
        let two = constant(2.0f32, out.shape());
        let floor_half = floor(floored / two);
        let floor_is_odd = cmpf(
            floor_half * two,
            floored,
            predicate::NotEqual,
            cmp_ordering::Ordered,
        );
        let above_half = cmpf(
            fraction,
            half,
            predicate::GreaterThan,
            cmp_ordering::Ordered,
        );
        let tie = cmpf(fraction, half, predicate::Equal, cmp_ordering::Ordered);
        let rounded_up = floored + one;
        let rounded = select(above_half | (tie & floor_is_odd), rounded_up, floored);
        let quantized: Tile<u8, { [B] }> = convert_tile(rounded);
        out.store(quantized);
    }
    #[cutile::entry()]
    pub fn quantize_f32_i8<const B: i32>(
        out: &mut Tensor<i8, { [B] }>,
        input: &Tensor<f32, { [-1] }>,
        scale: f32,
        zero_point: f32,
    ) {
        let values: Tile<f32, { [B] }> = load_tile_like(input, out);
        let raw = values / broadcast_scalar(scale, out.shape())
            + broadcast_scalar(zero_point, out.shape());
        let low = constant(-128.0f32, out.shape());
        let high = constant(127.0f32, out.shape());
        let clamped = min_tile(max_tile(raw, low), high);
        let floored = floor(clamped);
        let fraction = clamped - floored;
        let half = constant(0.5f32, out.shape());
        let one = constant(1.0f32, out.shape());
        let two = constant(2.0f32, out.shape());
        let floor_half = floor(floored / two);
        let floor_is_odd = cmpf(
            floor_half * two,
            floored,
            predicate::NotEqual,
            cmp_ordering::Ordered,
        );
        let above_half = cmpf(
            fraction,
            half,
            predicate::GreaterThan,
            cmp_ordering::Ordered,
        );
        let tie = cmpf(fraction, half, predicate::Equal, cmp_ordering::Ordered);
        let rounded_up = floored + one;
        let rounded = select(above_half | (tie & floor_is_odd), rounded_up, floored);
        let quantized: Tile<i8, { [B] }> = convert_tile(rounded);
        out.store(quantized);
    }
}
pub use crate::cuda::cutile::kernel::f32::utility::dequantize_f8e4m3_f32_block;
pub use crate::cuda::cutile::kernel::f32::utility::dequantize_i8_f32;
pub use crate::cuda::cutile::kernel::f32::utility::dequantize_i8_f32_grouped;
pub use crate::cuda::cutile::kernel::f32::utility::dequantize_u8_f32;
pub use crate::cuda::cutile::kernel::f32::utility::dequantize_u8_f32_grouped;
pub use crate::cuda::cutile::kernel::f32::utility::matmul_f8e4m3_f8e4m3_block_scaled_bf16;
pub use crate::cuda::cutile::kernel::f32::utility::matmul_f8e4m3_f8e4m3_block_scaled_f16;
pub use crate::cuda::cutile::kernel::f32::utility::matmul_f8e4m3_f8e4m3_block_scaled_f32;
pub use crate::cuda::cutile::kernel::f32::utility::matmul_f32_i8_dequantize_f32;
pub use crate::cuda::cutile::kernel::f32::utility::matmul_i8_i8_f32;
pub use crate::cuda::cutile::kernel::f32::utility::per_token_group_quant_f8e4m3_f32;
pub use crate::cuda::cutile::kernel::f32::utility::per_token_group_quant_f8e4m3_f32_column_major_scales;
pub use crate::cuda::cutile::kernel::f32::utility::per_token_group_quant_f8e4m3_f32_column_major_ue8m0_scales;
pub use crate::cuda::cutile::kernel::f32::utility::per_token_group_quant_i8_f32;
pub use crate::cuda::cutile::kernel::f32::utility::per_token_group_quant_i8_f32_column_major_scales;
pub use crate::cuda::cutile::kernel::f32::utility::per_token_group_quant_i8_f32_column_major_ue8m0_scales;
pub use crate::cuda::cutile::kernel::f32::utility::qk_rope_quantize_f8e4m3_f32_interleaved_rope_only;
pub use crate::cuda::cutile::kernel::f32::utility::qk_rope_quantize_f8e4m3_f32_interleaved_single_kernel;
pub use crate::cuda::cutile::kernel::f32::utility::quantize_f32_f8e4m3_block;
pub use crate::cuda::cutile::kernel::f32::utility::quantize_f32_f8e4m3_scaled;
pub use crate::cuda::cutile::kernel::f32::utility::rope_quantize_f8e4m3_f32_interleaved;
pub use kernels::*;