char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Ruben Ortlamvulkan: use spec constant for matrix matrix multiplication A-type (#25773)91f6a6cf3

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
10.3 KiB146 linesraw
1// Asymmetric K/V flash attention: aliased SSBO views of bindings 1 (K) and 2 (V)
2// covering every supported FA element type, plus an uber dequantize4() that
3// switches on FaTypeK / FaTypeV. After spec-constant specialization the driver
4// folds away every path except the one matching the K/V type for this pipeline.
5//
6// Included by flash_attn.comp and flash_attn_cm1.comp. Not included by
7// flash_attn_cm2.comp, which has its own buffer_reference-based decode path.
8//
9// We use macros (rather than per-quant decode functions taking a struct) on
10// purpose: the FA shaders don't enable GL_EXT_shader_explicit_arithmetic_types_float16
11// when FLOAT16 isn't defined, which makes float16-containing struct values
12// illegal to return from / pass to functions. Macros expand inline where the
13// float16 stays in storage and is converted to FLOAT_TYPE at use.
14
15// F32 is fed as a vec4 "block" (4 floats), matching what dequant_funcs_cm2.glsl
16// does for F32 in the cm2 shader. FaBlockBytesK/V == 16 for F32.
17layout (binding = 1) readonly buffer K_PACKED_F32  { vec4 data[]; }                k_packed_f32;
18layout (binding = 2) readonly buffer V_PACKED_F32  { vec4 data[]; }                v_packed_f32;
19
20layout (binding = 1) readonly buffer K_PACKED_Q4_0 { block_q4_0_packed16 data[]; } k_packed_q4_0;
21layout (binding = 2) readonly buffer V_PACKED_Q4_0 { block_q4_0_packed16 data[]; } v_packed_q4_0;
22layout (binding = 1) readonly buffer K_PACKED_Q4_1 { block_q4_1_packed16 data[]; } k_packed_q4_1;
23layout (binding = 2) readonly buffer V_PACKED_Q4_1 { block_q4_1_packed16 data[]; } v_packed_q4_1;
24layout (binding = 1) readonly buffer K_PACKED_Q5_0 { block_q5_0_packed16 data[]; } k_packed_q5_0;
25layout (binding = 2) readonly buffer V_PACKED_Q5_0 { block_q5_0_packed16 data[]; } v_packed_q5_0;
26layout (binding = 1) readonly buffer K_PACKED_Q5_1 { block_q5_1_packed16 data[]; } k_packed_q5_1;
27layout (binding = 2) readonly buffer V_PACKED_Q5_1 { block_q5_1_packed16 data[]; } v_packed_q5_1;
28layout (binding = 1) readonly buffer K_PACKED_Q8_0 { block_q8_0_packed16 data[]; } k_packed_q8_0;
29layout (binding = 2) readonly buffer V_PACKED_Q8_0 { block_q8_0_packed16 data[]; } v_packed_q8_0;
30layout (binding = 1) readonly buffer K_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } k_packed_iq4_nl;
31layout (binding = 2) readonly buffer V_PACKED_IQ4_NL { block_iq4_nl_packed16 data[]; } v_packed_iq4_nl;
32
33layout (binding = 1) readonly buffer K_PACKED_BF16 { u16vec4 data[]; } k_packed_bf16;
34layout (binding = 2) readonly buffer V_PACKED_BF16 { u16vec4 data[]; } v_packed_bf16;
35
36// Q4_1 and Q5_1 packed32 views: aliased to the same memory as the packed16
37// views, used by the MMQ K-side hot path for fast 4-uint loads.
38layout (binding = 1) readonly buffer K_PACKED_Q4_1_P32 { block_q4_1_packed32 data[]; } k_packed_q4_1_p32;
39layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 data[]; } k_packed_q5_1_p32;
40
41// Per-quant decode bodies are expanded once for the K view set and once for
42// the V view set. The macros take the buffer name as a parameter.
43#define FA_DEQUANT4_F32(BUF) \
44    return FLOAT_TYPEV4(BUF.data[a_offset + ib]);
45
46#define FA_DEQUANT4_Q4_0(BUF) {                                                                   \
47    uint vui_lo = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0]);                          \
48    uint vui_hi = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1]);                          \
49    uint shift = (iqs & 0x10) >> 2;                                                               \
50    vui_lo >>= shift;                                                                             \
51    vui_hi >>= shift;                                                                             \
52    FLOAT_TYPEV4 nibbles = FLOAT_TYPEV4(vui_lo & 0xF, (vui_lo >> 8) & 0xF,                        \
53                                        vui_hi & 0xF, (vui_hi >> 8) & 0xF);                       \
54    return FLOAT_TYPE(BUF.data[a_offset + ib].d) * (nibbles - FLOAT_TYPE(8.0f));                  \
55}
56
57#define FA_DEQUANT4_Q4_1(BUF) {                                                                   \
58    uint vui_lo = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0]);                          \
59    uint vui_hi = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1]);                          \
60    uint shift = (iqs & 0x10) >> 2;                                                               \
61    vui_lo >>= shift;                                                                             \
62    vui_hi >>= shift;                                                                             \
63    FLOAT_TYPEV4 nibbles = FLOAT_TYPEV4(vui_lo & 0xF, (vui_lo >> 8) & 0xF,                        \
64                                        vui_hi & 0xF, (vui_hi >> 8) & 0xF);                       \
65    return FLOAT_TYPE(BUF.data[a_offset + ib].d) * nibbles                                        \
66         + FLOAT_TYPE(BUF.data[a_offset + ib].m);                                                 \
67}
68
69#define FA_DEQUANT4_Q5_0(BUF) {                                                                   \
70    uint vui_lo = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0]);                          \
71    uint vui_hi = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1]);                          \
72    uint shift = (iqs & 0x10) >> 2;                                                               \
73    vui_lo >>= shift;                                                                             \
74    vui_hi >>= shift;                                                                             \
75    uint qh = uint(BUF.data[a_offset + ib].qh[0])                                                 \
76            | (uint(BUF.data[a_offset + ib].qh[1]) << 16);                                        \
77    FLOAT_TYPEV4 hb = FLOAT_TYPEV4((qh >> iqs)       & 1, (qh >> (iqs + 1)) & 1,                  \
78                                   (qh >> (iqs + 2)) & 1, (qh >> (iqs + 3)) & 1)                  \
79                      * FLOAT_TYPE(16.0f);                                                        \
80    FLOAT_TYPEV4 nibbles = FLOAT_TYPEV4(vui_lo & 0xF, (vui_lo >> 8) & 0xF,                        \
81                                        vui_hi & 0xF, (vui_hi >> 8) & 0xF);                       \
82    return FLOAT_TYPE(BUF.data[a_offset + ib].d) * (nibbles + hb - FLOAT_TYPE(16.0f));            \
83}
84
85#define FA_DEQUANT4_Q5_1(BUF) {                                                                   \
86    uint vui_lo = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0]);                          \
87    uint vui_hi = uint(BUF.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1]);                          \
88    uint shift = (iqs & 0x10) >> 2;                                                               \
89    vui_lo >>= shift;                                                                             \
90    vui_hi >>= shift;                                                                             \
91    uint qh = BUF.data[a_offset + ib].qh;                                                         \
92    FLOAT_TYPEV4 hb = FLOAT_TYPEV4((qh >> iqs)       & 1, (qh >> (iqs + 1)) & 1,                  \
93                                   (qh >> (iqs + 2)) & 1, (qh >> (iqs + 3)) & 1)                  \
94                      * FLOAT_TYPE(16.0f);                                                        \
95    FLOAT_TYPEV4 nibbles = FLOAT_TYPEV4(vui_lo & 0xF, (vui_lo >> 8) & 0xF,                        \
96                                        vui_hi & 0xF, (vui_hi >> 8) & 0xF);                       \
97    return FLOAT_TYPE(BUF.data[a_offset + ib].d) * (nibbles + hb)                                 \
98         + FLOAT_TYPE(BUF.data[a_offset + ib].m);                                                 \
99}
100
101#define FA_DEQUANT4_Q8_0(BUF) {                                                                   \
102    const i8vec2 v0 = unpack8(int32_t(BUF.data[a_offset + ib].qs[iqs / 2    ])).xy;               \
103    const i8vec2 v1 = unpack8(int32_t(BUF.data[a_offset + ib].qs[iqs / 2 + 1])).xy;               \
104    return FLOAT_TYPE(BUF.data[a_offset + ib].d) * FLOAT_TYPEV4(v0.x, v0.y, v1.x, v1.y);          \
105}
106
107#define FA_DEQUANT4_IQ4_NL(BUF) {                                                                 \
108    const uint shift = (iqs & 0x10) >> 2;                                                         \
109    const uint qs_i = (iqs & 0xC) >> 1;                                                           \
110    const uint qsw = uint(BUF.data[a_offset + ib].qs[qs_i])                                       \
111                   | (uint(BUF.data[a_offset + ib].qs[qs_i + 1u]) << 16);                         \
112    const FLOAT_TYPE d = FLOAT_TYPE(BUF.data[a_offset + ib].d);                                   \
113    const u8vec4 q = unpack8((qsw >> shift) & 0x0F0F0F0Fu);                                       \
114    return d * FLOAT_TYPEV4(kvalues_iq4nl[q.x], kvalues_iq4nl[q.y],                               \
115                            kvalues_iq4nl[q.z], kvalues_iq4nl[q.w]);                              \
116}
117
118#define FA_DEQUANT4_BF16(BUF) \
119    return FLOAT_TYPEV4(bf16_to_fp32(uvec4(BUF.data[(a_offset + ib) / 4])));
120
121FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) {
122    if (binding_idx == BINDING_IDX_K) {
123        switch (FaTypeK) {
124            case GGML_TYPE_F32:  FA_DEQUANT4_F32 (k_packed_f32)
125            case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0)
126            case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1)
127            case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0)
128            case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1)
129            case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0)
130            case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl)
131            case GGML_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16)
132        }
133    } else {
134        switch (FaTypeV) {
135            case GGML_TYPE_F32:  FA_DEQUANT4_F32 (v_packed_f32)
136            case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0)
137            case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1)
138            case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0)
139            case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1)
140            case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0)
141            case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl)
142            case GGML_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16)
143        }
144    }
145    return FLOAT_TYPEV4(0);
146}