rublas 0.1.78

Linear algebra and matrix multiplication for Ruda.
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
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
131
132
133
134
135
136

use ruda_kernel::dsl as kernel_dsl;
use ruda_kernel::dsl::prelude::*;

#[ruda(launch)]
pub(super) fn awq_gemm<F: Float>(
    input: &Array<F>,
    qweight: &Array<i32>,
    qzeros: &Array<i32>,
    scales: &Array<F>,
    bias: &Array<F>,
    output: &mut Array<F>,
    input_features: u32,
    output_features: u32,
    group_size: u32,
    has_bias: u32,
    #[define(F)] _dtype: StorageType,
) {
    let position = ABSOLUTE_POS;
    if position >= output.len() {
        terminate!();
    }
    let k = input_features as usize;
    let n = output_features as usize;
    let row = position / n;
    let column = position % n;
    let packed_columns = n / 8;
    let lane = column % 8;
    let shift = ((lane % 2) * 4 + lane / 2) as u32 * 4;
    let mut sum = 0.0f32;
    let mut inner = 0usize;
    while inner < k {
        let group = inner / group_size as usize;
        let word = u32::cast_from(qweight[inner * packed_columns + column / 8]);
        let zero_word = u32::cast_from(qzeros[group * packed_columns + column / 8]);
        let quant = (word >> shift) & 15u32;
        let zero = (zero_word >> shift) & 15u32;
        let weight = F::cast_from(
            (f32::cast_from(quant) - f32::cast_from(zero))
                * f32::cast_from(scales[group * n + column]),
        );
        sum += f32::cast_from(input[row * k + inner]) * f32::cast_from(weight);
        inner += 1;
    }
    let mut value = F::cast_from(sum);
    if has_bias != 0 {
        value += bias[column];
    }
    output[position] = value;
}

/// Mixed activation/weight storage with the same rounded AWQ coefficients.
#[ruda(launch)]
pub(super) fn awq_gemm_mixed<G: Float, W: Float>(
    input: &Array<G>,
    qweight: &Array<i32>,
    qzeros: &Array<i32>,
    scales: &Array<W>,
    bias: &Array<W>,
    output: &mut Array<G>,
    input_features: u32,
    output_features: u32,
    group_size: u32,
    has_bias: u32,
    #[define(G)] _input_dtype: StorageType,
    #[define(W)] _weight_dtype: StorageType,
) {
    let position = ABSOLUTE_POS;
    if position >= output.len() { terminate!(); }
    let k = input_features as usize;
    let n = output_features as usize;
    let row = position / n;
    let column = position % n;
    let packed_columns = n / 8;
    let lane = column % 8;
    let shift = ((lane % 2) * 4 + lane / 2) as u32 * 4;
    let mut sum = 0.0f32;
    let mut inner = 0usize;
    while inner < k {
        let group = inner / group_size as usize;
        let word = u32::cast_from(qweight[inner * packed_columns + column / 8]);
        let zero_word = u32::cast_from(qzeros[group * packed_columns + column / 8]);
        let quant = (word >> shift) & 15u32;
        let zero = (zero_word >> shift) & 15u32;
        let weight = W::cast_from(
            (f32::cast_from(quant) - f32::cast_from(zero))
                * f32::cast_from(scales[group * n + column]),
        );
        sum += f32::cast_from(input[row * k + inner]) * f32::cast_from(weight);
        inner += 1;
    }
    let mut value = G::cast_from(sum);
    if has_bias != 0 { value += G::cast_from(bias[column]); }
    output[position] = value;
}

/// Input VJP reads packed AWQ words directly, without materializing a dense weight.
#[ruda(launch)]
pub(super) fn awq_input_backward<G: Float, W: Float>(
    gradient: &Array<G>,
    qweight: &Array<i32>,
    qzeros: &Array<i32>,
    scales: &Array<W>,
    output: &mut Array<G>,
    input_features: u32,
    output_features: u32,
    group_size: u32,
    #[define(G)] _gradient_dtype: StorageType,
    #[define(W)] _weight_dtype: StorageType,
) {
    let position = ABSOLUTE_POS;
    if position >= output.len() { terminate!(); }
    let k = input_features as usize;
    let n = output_features as usize;
    let row = position / k;
    let inner = position % k;
    let group = inner / group_size as usize;
    let packed_columns = n / 8;
    let mut sum = 0.0f32;
    let mut column = 0usize;
    while column < n {
        let lane = column % 8;
        let shift = ((lane % 2) * 4 + lane / 2) as u32 * 4;
        let word = u32::cast_from(qweight[inner * packed_columns + column / 8]);
        let zero_word = u32::cast_from(qzeros[group * packed_columns + column / 8]);
        let quant = (word >> shift) & 15u32;
        let zero = (zero_word >> shift) & 15u32;
        // W rounding must match forward even when G is FP32 and W is FP16/BF16.
        let weight = W::cast_from(
            (f32::cast_from(quant) - f32::cast_from(zero))
                * f32::cast_from(scales[group * n + column]),
        );
        sum += f32::cast_from(gradient[row * n + column]) * f32::cast_from(weight);
        column += 1;
    }
    output[position] = G::cast_from(sum);
}