1#extension GL_EXT_control_flow_attributes : enable
 2
 3layout (push_constant) uniform parameter
 4{
 5    uint KX;
 6    uint KY;
 7    uint ne00;
 8    uint ne01;
 9    uint ne02;
10    uint ne12;
11    uint ne13;
12    uint nb11;
13    uint nb12;
14    uint nb13;
15    float scale;
16    float max_bias;
17    float m0;
18    float m1;
19    uint n_head_log2;
20    uint nrows_x;
21    uint has_sinks;
22} p;
23
24#include "types.glsl"
25
26layout(constant_id = 0) const uint BLOCK_SIZE = 128;
27layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
28layout(constant_id = 1) const uint num_iters = 4;
29
30layout (binding = 0) readonly buffer X {A_TYPE data_a[];};
31layout (binding = 1) readonly buffer Y {B_TYPE data_b[];};
32layout (binding = 2) readonly buffer Z {float data_c[];};
33layout (binding = 3) buffer D {D_TYPE data_d[];};
34layout (binding = 4) buffer M {float data_m[];};
35layout (binding = 5) buffer S {float data_s[];};
36
37shared FLOAT_TYPE vals[BLOCK_SIZE];
38
39float get_slope(uint rowx) {
40    float slope = 1.0f;
41
42    // ALiBi
43    if (p.max_bias > 0.0f) {
44        const uint h = (rowx / p.ne01) % p.ne02; // head index
45
46        const float base = h < p.n_head_log2 ? p.m0 : p.m1;
47        const uint   exp = h < p.n_head_log2 ? h + 1 : 2*(h - p.n_head_log2) + 1;
48
49        slope = pow(base, exp);
50    }
51
52    return slope;
53}