char-fork/llama.cpp

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

Sigbjørn Skjæretsycl : add Q2_K reordered MMVQ and ESIMD kernels (again) (#27490)867229003

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
27.3 KiB589 linesraw
1#ifndef GGML_SYCL_ESIMD_HPP
2#define GGML_SYCL_ESIMD_HPP
3
4#include <sycl/ext/intel/esimd.hpp>
5
6#include "common.hpp"
7
8namespace ggml_sycl_esimd {
9
10constexpr int GGML_SYCL_DMMV_ESIMD_WG_SIZE = 4;
11
12//
13// Shared ESIMD building blocks for the reordered K-quant dequantize-matvec
14// kernels.
15//
16// The reordered K-quant ESIMD matvec kernels share one skeleton: per super-block,
17// load a 256-float activation slice, load one weight block, dequantize it into 8
18// chunks of 32 and MAC each chunk against the matching activation slice, then
19// reduce and run a lane-0 epilogue.
20//
21// Each K-quant kernel emits exactly 8 chunks of 32 mapping to activation slices
22// 0..7, so the per-block work is captured by esimd_reorder_q_traits<T>::mac_pair,
23// which dequantizes two weight blocks and MACs both against a shared activation
24// vector with the two FMA chains interleaved (co-scheduled to hide FMA latency).
25// The "pair" is the (row0,row1) row pair owned by one work-group, so the
26// layout+dequant is written once per quant type here.
27//
28
29template <ggml_type T> struct esimd_reorder_q_traits;
30
31// build a 32-lane vector whose low 16 lanes are `lo` and high 16 are `hi`
32// (a super-chunk splits into two 16-wide halves with distinct scale/min codes).
33static ESIMD_INLINE sycl::ext::intel::esimd::simd<float, 32> splat_lo_hi(float lo, float hi) {
34    using namespace sycl::ext::intel::esimd;
35    simd<float, 32> v;
36    v.select<16, 1>(0)  = lo;
37    v.select<16, 1>(16) = hi;
38    return v;
39}
40
41// unpack one block of Q4_K/Q5_K scale/min codes (get_scale_min_k4 layout) into 8
42// float scales (dall * sc) and 8 float mins (-dmin * m); the min carries the
43// negation so the dequant epilogue adds.
44static ESIMD_INLINE void unpack_scale_min_k4(
45        sycl::ext::intel::esimd::simd<uint8_t, 12> scales, float dall, float dmin,
46        sycl::ext::intel::esimd::simd<float, 8> & scale_f,
47        sycl::ext::intel::esimd::simd<float, 8> & min_f) {
48    using namespace sycl::ext::intel::esimd;
49    simd<uint8_t, 8> sc = 0;
50    simd<uint8_t, 8> m  = 0;
51    simd<uint8_t, 4> scale_lo = scales.select<4, 1>(0);
52    simd<uint8_t, 4> min_lo   = scales.select<4, 1>(4);
53    simd<uint8_t, 4> hi_bits  = scales.select<4, 1>(8);
54    sc.select<4, 1>(0) = scale_lo & simd<uint8_t, 4>(0x3F);
55    sc.select<4, 1>(4) = (hi_bits & simd<uint8_t, 4>(0x0F)) |
56                         ((scale_lo >> simd<uint8_t, 4>(6)) << simd<uint8_t, 4>(4));
57    m.select<4, 1>(0)  = min_lo & simd<uint8_t, 4>(0x3F);
58    m.select<4, 1>(4)  = (hi_bits >> simd<uint8_t, 4>(4)) |
59                         ((min_lo >> simd<uint8_t, 4>(6)) << simd<uint8_t, 4>(4));
60    scale_f = convert<float>(sc) * dall;
61    min_f   = convert<float>(m) * (-dmin);
62}
63
64// ---------------------------------------------------------------------------
65// Q2_K, SOA reorder layout produced by reorder_qw_q2_k:
66//   [qs: nb*(QK_K/4)] [scales: nb*(QK_K/16)] [dm: nb*sizeof(half2)]
67// with nb = nrows*num_blocks_per_row.
68//
69// 2 bits per weight. The 8 output chunks of 32 (matching dequantize_row_q2_K)
70// map to super-chunk s (0..7): byte base 32*(s/4) into the 64-byte qs array,
71// bit shift 2*(s%4); the low 16 lanes use scales[2s], the high 16 use
72// scales[2s+1], with dl = d*(sc & 0xF), ml = dmin*(sc >> 4), deq = dl*q - ml.
73// ---------------------------------------------------------------------------
74template <> struct esimd_reorder_q_traits<GGML_TYPE_Q2_K> {
75    struct ptrs {
76        const uint8_t *    qs;
77        const uint8_t *    scales;
78        const sycl::half * dm;
79    };
80
81    static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) {
82        const uint8_t * qs     = (const uint8_t *) vx;
83        const uint8_t * scales = qs + nb * (QK_K / 4);
84        const sycl::half * dm  = (const sycl::half *) (scales + nb * (QK_K / 16));
85        return { qs, scales, dm };
86    }
87
88    static ESIMD_INLINE void mac_pair(
89            const ptrs & pa, size_t bia,
90            const ptrs & pb, size_t bib, bool has_b,
91            sycl::ext::intel::esimd::simd<float, 256> & y_vec,
92            sycl::ext::intel::esimd::simd<float, 32> & acc_a,
93            sycl::ext::intel::esimd::simd<float, 32> & acc_b) {
94        using namespace sycl::ext::intel::esimd;
95
96        simd<uint8_t, 64> qs_a     = block_load<uint8_t, 64>(pa.qs + bia * (QK_K / 4));
97        simd<uint8_t, 64> qs_b     = 0;
98        simd<uint8_t, 16> scales_a = block_load<uint8_t, 16>(pa.scales + bia * (QK_K / 16));
99        simd<uint8_t, 16> scales_b = 0;
100
101        const float dall_a = (float) pa.dm[bia * 2 + 0];
102        const float dmin_a = (float) pa.dm[bia * 2 + 1];
103        float dall_b = 0.0f;
104        float dmin_b = 0.0f;
105        if (has_b) {
106            qs_b     = block_load<uint8_t, 64>(pb.qs + bib * (QK_K / 4));
107            scales_b = block_load<uint8_t, 16>(pb.scales + bib * (QK_K / 16));
108            dall_b = (float) pb.dm[bib * 2 + 0];
109            dmin_b = (float) pb.dm[bib * 2 + 1];
110        }
111
112        // per-chunk scale (d * (sc & 0xF)) and min (-dmin * (sc >> 4)), all 16 codes;
113        // min carries the negation so the dequant epilogue adds (matches Q4_K/Q5_K)
114        simd<float, 16> scale_f_a = convert<float>(scales_a & simd<uint8_t, 16>(0x0F)) * dall_a;
115        simd<float, 16> min_f_a   = convert<float>(scales_a >> simd<uint8_t, 16>(4))  * (-dmin_a);
116        simd<float, 16> scale_f_b = convert<float>(scales_b & simd<uint8_t, 16>(0x0F)) * dall_b;
117        simd<float, 16> min_f_b   = convert<float>(scales_b >> simd<uint8_t, 16>(4))  * (-dmin_b);
118
119#pragma unroll
120        for (int s = 0; s < 8; ++s) {
121            const int     byte_base = 32 * (s / 4);
122            const uint8_t shift     = (uint8_t) (2 * (s % 4));
123            simd<float, 32> y_s = y_vec.select<32, 1>(s * 32);
124
125            simd<uint8_t, 32> qa = (qs_a.select<32, 1>(byte_base) >> shift) & simd<uint8_t, 32>(3);
126            simd<uint8_t, 32> qb = (qs_b.select<32, 1>(byte_base) >> shift) & simd<uint8_t, 32>(3);
127
128            const float scale_a_lo = scale_f_a[2 * s + 0];
129            const float scale_a_hi = scale_f_a[2 * s + 1];
130            const float min_a_lo   = min_f_a[2 * s + 0];
131            const float min_a_hi   = min_f_a[2 * s + 1];
132            const float scale_b_lo = scale_f_b[2 * s + 0];
133            const float scale_b_hi = scale_f_b[2 * s + 1];
134            const float min_b_lo   = min_f_b[2 * s + 0];
135            const float min_b_hi   = min_f_b[2 * s + 1];
136
137            simd<float, 32> scale_vec_a = splat_lo_hi(scale_a_lo, scale_a_hi);
138            simd<float, 32> min_vec_a   = splat_lo_hi(min_a_lo, min_a_hi);
139            simd<float, 32> scale_vec_b = splat_lo_hi(scale_b_lo, scale_b_hi);
140            simd<float, 32> min_vec_b   = splat_lo_hi(min_b_lo, min_b_hi);
141
142            simd<float, 32> deq_a = convert<float>(qa) * scale_vec_a + min_vec_a;
143            simd<float, 32> deq_b = convert<float>(qb) * scale_vec_b + min_vec_b;
144
145            acc_a += y_s * deq_a;
146            acc_b += y_s * deq_b;
147        }
148    }
149};
150
151// ---------------------------------------------------------------------------
152// Q3_K, SOA reorder layout produced by reorder_qw_q3_k:
153//   [qs: nb*(QK_K/4)] [hmask: nb*(QK_K/8)] [scales: nb*12] [d: nb*sizeof(half)]
154// with nb = nrows*num_blocks_per_row. Single super-block scale d, no dmin.
155//
156// 3 bits per weight: 2 low bits in qs, 1 high bit in hmask. The 8 output chunks
157// of 32 (matching dequantize_row_q3_K) map to super-chunk s (0..7): byte base
158// 32*(s/4) into the 64-byte qs array, bit shift 2*(s%4); the low 16 lanes use
159// scale code 2s, the high 16 use 2s+1. hmask is a 32-byte array (like Q5_K's
160// qh) where chunk s uses bit s of the same 32 bytes, but INVERTED: the value is
161// (q & 3) - (hmask_bit_set ? 0 : 4), i.e. (q & 3) + 4*bit - 4.
162//
163// The 16 6-bit scale codes are packed into 12 bytes (get_scale_min layout for
164// Q3_K): low nibbles from bytes 0..7, high 2 bits from bytes 8..11 shifted by
165// 0/2/4/6; the dequant scale is d * (code - 32).
166// ---------------------------------------------------------------------------
167template <> struct esimd_reorder_q_traits<GGML_TYPE_Q3_K> {
168    struct ptrs {
169        const uint8_t *    qs;
170        const uint8_t *    hmask;
171        const uint8_t *    scales;
172        const sycl::half * d;
173    };
174
175    static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) {
176        const uint8_t * qs     = (const uint8_t *) vx;
177        const uint8_t * hmask  = qs + nb * (QK_K / 4);
178        const uint8_t * scales = hmask + nb * (QK_K / 8);
179        const sycl::half * d   = (const sycl::half *) (scales + nb * 12);
180        return { qs, hmask, scales, d };
181    }
182
183    // unpack the 12 packed bytes into 16 6-bit scale codes (dequantize_row_q3_K
184    // aux layout), returned as float scale = d * (code - 32).
185    // done with wide (8/16-lane) ops rather than four 4-lane groups.
186    static ESIMD_INLINE sycl::ext::intel::esimd::simd<float, 16> unpack_scales(
187            sycl::ext::intel::esimd::simd<uint8_t, 12> in, float d) {
188        using namespace sycl::ext::intel::esimd;
189
190        // low 6-bit part: codes 0..7 = low nibble of bytes 0..7,
191        //                 codes 8..15 = high nibble of bytes 0..7
192        simd<uint8_t, 8>  lo8 = in.select<8, 1>(0);
193        simd<uint8_t, 16> code;
194        code.select<8, 1>(0) = lo8 & simd<uint8_t, 8>(0x0F);
195        code.select<8, 1>(8) = lo8 >> simd<uint8_t, 8>(4);
196
197        // high 2-bit part: bytes 8..11 replicated 4x, group g (0..3) shifted 2*g
198        simd<uint8_t, 16> hib;
199        hib.select<4, 1>(0)  = in.select<4, 1>(8);
200        hib.select<4, 1>(4)  = in.select<4, 1>(8);
201        hib.select<4, 1>(8)  = in.select<4, 1>(8);
202        hib.select<4, 1>(12) = in.select<4, 1>(8);
203        simd<uint8_t, 16> hshift;
204        hshift.select<4, 1>(0)  = 0;
205        hshift.select<4, 1>(4)  = 2;
206        hshift.select<4, 1>(8)  = 4;
207        hshift.select<4, 1>(12) = 6;
208        hib = (hib >> hshift) & simd<uint8_t, 16>(0x03);
209
210        code = code | (hib << simd<uint8_t, 16>(4));
211        return (convert<float>(code) - 32.0f) * d;
212    }
213
214    static ESIMD_INLINE void mac_pair(
215            const ptrs & pa, size_t bia,
216            const ptrs & pb, size_t bib, bool has_b,
217            sycl::ext::intel::esimd::simd<float, 256> & y_vec,
218            sycl::ext::intel::esimd::simd<float, 32> & acc_a,
219            sycl::ext::intel::esimd::simd<float, 32> & acc_b) {
220        using namespace sycl::ext::intel::esimd;
221
222        simd<uint8_t, 64> qs_a     = block_load<uint8_t, 64>(pa.qs + bia * (QK_K / 4));
223        simd<uint8_t, 64> qs_b     = 0;
224        simd<uint8_t, 32> hmask_a  = block_load<uint8_t, 32>(pa.hmask + bia * (QK_K / 8));
225        simd<uint8_t, 32> hmask_b  = 0;
226        simd<uint8_t, 12> scales_a = block_load<uint8_t, 12>(pa.scales + bia * 12);
227        simd<uint8_t, 12> scales_b = 0;
228
229        const float d_a = (float) pa.d[bia];
230        float d_b = 0.0f;
231        if (has_b) {
232            qs_b     = block_load<uint8_t, 64>(pb.qs + bib * (QK_K / 4));
233            hmask_b  = block_load<uint8_t, 32>(pb.hmask + bib * (QK_K / 8));
234            scales_b = block_load<uint8_t, 12>(pb.scales + bib * 12);
235            d_b = (float) pb.d[bib];
236        }
237
238        simd<float, 16> scale_f_a = unpack_scales(scales_a, d_a);
239        simd<float, 16> scale_f_b = unpack_scales(scales_b, d_b);
240
241#pragma unroll
242        for (int s = 0; s < 8; ++s) {
243            const int     byte_base = 32 * (s / 4);
244            const uint8_t shift     = (uint8_t) (2 * (s % 4));
245            simd<float, 32> y_s = y_vec.select<32, 1>(s * 32);
246
247            // 2 low bits from qs, high bit from hmask (bit s of the same 32 bytes);
248            // value = (q & 3) + 4*bit - 4  (inverted hmask: subtract 4 when bit clear).
249            // merge in the integer domain: q3 = (q & 3) | (bit << 2) in {0..7},
250            // then a single convert + subtract yields q3 - 4 (one convert, not two)
251            simd<uint16_t, 32> q3_a = convert<uint16_t>(
252                    (qs_a.select<32, 1>(byte_base) >> shift) & simd<uint8_t, 32>(3));
253            q3_a |= convert<uint16_t>(
254                    ((hmask_a >> simd<uint8_t, 32>((uint8_t) s)) & simd<uint8_t, 32>(1)) << simd<uint8_t, 32>(2));
255            simd<uint16_t, 32> q3_b = convert<uint16_t>(
256                    (qs_b.select<32, 1>(byte_base) >> shift) & simd<uint8_t, 32>(3));
257            q3_b |= convert<uint16_t>(
258                    ((hmask_b >> simd<uint8_t, 32>((uint8_t) s)) & simd<uint8_t, 32>(1)) << simd<uint8_t, 32>(2));
259
260            simd<float, 32> qf_a = convert<float>(q3_a) - 4.0f;
261            simd<float, 32> qf_b = convert<float>(q3_b) - 4.0f;
262
263            const float scale_a_lo = scale_f_a[2 * s + 0];
264            const float scale_a_hi = scale_f_a[2 * s + 1];
265            const float scale_b_lo = scale_f_b[2 * s + 0];
266            const float scale_b_hi = scale_f_b[2 * s + 1];
267
268            simd<float, 32> scale_vec_a = splat_lo_hi(scale_a_lo, scale_a_hi);
269            simd<float, 32> scale_vec_b = splat_lo_hi(scale_b_lo, scale_b_hi);
270
271            simd<float, 32> deq_a = qf_a * scale_vec_a;
272            simd<float, 32> deq_b = qf_b * scale_vec_b;
273
274            acc_a += y_s * deq_a;
275            acc_b += y_s * deq_b;
276        }
277    }
278};
279
280// ---------------------------------------------------------------------------
281// Q4_K, SOA reorder layout produced by reorder_qw_q4_k:
282//   [qs: nb*(QK_K/2)] [scales: nb*K_SCALE_SIZE] [dm: nb*sizeof(half2)]
283// with nb = nrows*num_blocks_per_row.
284// ---------------------------------------------------------------------------
285template <> struct esimd_reorder_q_traits<GGML_TYPE_Q4_K> {
286    struct ptrs {
287        const uint8_t *    qs;
288        const uint8_t *    scales;
289        const sycl::half * dm;
290    };
291
292    static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) {
293        const uint8_t * qs     = (const uint8_t *) vx;
294        const uint8_t * scales = qs + nb * (QK_K / 2);
295        const sycl::half * dm  = (const sycl::half *) (scales + nb * K_SCALE_SIZE);
296        return { qs, scales, dm };
297    }
298
299    static ESIMD_INLINE void mac_pair(
300            const ptrs & pa, size_t bia,
301            const ptrs & pb, size_t bib, bool has_b,
302            sycl::ext::intel::esimd::simd<float, 256> & y_vec,
303            sycl::ext::intel::esimd::simd<float, 32> & acc_a,
304            sycl::ext::intel::esimd::simd<float, 32> & acc_b) {
305        using namespace sycl::ext::intel::esimd;
306
307        simd<uint8_t, 128> qs_a     = block_load<uint8_t, 128>(pa.qs + bia * (QK_K / 2));
308        simd<uint8_t, 128> qs_b     = 0;
309        simd<uint8_t, 12>  scales_a = block_load<uint8_t, 12>(pa.scales + bia * K_SCALE_SIZE);
310        simd<uint8_t, 12>  scales_b = 0;
311
312        const float dall_a = (float) pa.dm[bia * 2 + 0];
313        const float dmin_a = (float) pa.dm[bia * 2 + 1];
314        float dall_b = 0.0f;
315        float dmin_b = 0.0f;
316        if (has_b) {
317            qs_b     = block_load<uint8_t, 128>(pb.qs + bib * (QK_K / 2));
318            scales_b = block_load<uint8_t, 12>(pb.scales + bib * K_SCALE_SIZE);
319            dall_b = (float) pb.dm[bib * 2 + 0];
320            dmin_b = (float) pb.dm[bib * 2 + 1];
321        }
322
323        simd<float, 8> scale_f_a, min_f_a, scale_f_b, min_f_b;
324        unpack_scale_min_k4(scales_a, dall_a, dmin_a, scale_f_a, min_f_a);
325        unpack_scale_min_k4(scales_b, dall_b, dmin_b, scale_f_b, min_f_b);
326
327        simd<uint8_t, 128> qs_lo_a = qs_a & simd<uint8_t, 128>(0x0F);
328        simd<uint8_t, 128> qs_hi_a = qs_a >> simd<uint8_t, 128>(4);
329        simd<uint8_t, 128> qs_lo_b = qs_b & simd<uint8_t, 128>(0x0F);
330        simd<uint8_t, 128> qs_hi_b = qs_b >> simd<uint8_t, 128>(4);
331
332#pragma unroll
333        for (int sb = 0; sb < 8; sb += 2) {
334            const int q_offset = sb * 16;
335            simd<float, 32> y_lo = y_vec.select<32, 1>(sb * 32);
336            simd<float, 32> y_hi = y_vec.select<32, 1>((sb + 1) * 32);
337
338            const float scale_a_lo = scale_f_a[sb];
339            const float scale_a_hi = scale_f_a[sb + 1];
340            const float min_a_lo   = min_f_a[sb];
341            const float min_a_hi   = min_f_a[sb + 1];
342            const float scale_b_lo = scale_f_b[sb];
343            const float scale_b_hi = scale_f_b[sb + 1];
344            const float min_b_lo   = min_f_b[sb];
345            const float min_b_hi   = min_f_b[sb + 1];
346
347            simd<uint8_t, 32> qa_lo = qs_lo_a.select<32, 1>(q_offset);
348            simd<uint8_t, 32> qa_hi = qs_hi_a.select<32, 1>(q_offset);
349            simd<uint8_t, 32> qb_lo = qs_lo_b.select<32, 1>(q_offset);
350            simd<uint8_t, 32> qb_hi = qs_hi_b.select<32, 1>(q_offset);
351
352            simd<float, 32> deq_a_lo = convert<float>(qa_lo) * scale_a_lo + min_a_lo;
353            simd<float, 32> deq_a_hi = convert<float>(qa_hi) * scale_a_hi + min_a_hi;
354            simd<float, 32> deq_b_lo = convert<float>(qb_lo) * scale_b_lo + min_b_lo;
355            simd<float, 32> deq_b_hi = convert<float>(qb_hi) * scale_b_hi + min_b_hi;
356
357            acc_a += y_lo * deq_a_lo;
358            acc_b += y_lo * deq_b_lo;
359            acc_a += y_hi * deq_a_hi;
360            acc_b += y_hi * deq_b_hi;
361        }
362    }
363};
364
365// ---------------------------------------------------------------------------
366// Q5_K, SOA reorder layout produced by reorder_qw_q5_k:
367//   [qs: nb*(QK_K/2)] [qh: nb*(QK_K/8)] [scales: nb*K_SCALE_SIZE] [dm: nb*sizeof(half2)]
368// with nb = nrows*num_blocks_per_row.
369//
370// Identical to Q4_K except each 4-bit quant gains a 5th (high) bit from qh:
371// output chunk c (0..7) adds 16 when bit c of qh[l] is set, where qh[l] indexes
372// the same 32 bytes for every chunk (matches dequantize_row_q5_K).
373// ---------------------------------------------------------------------------
374template <> struct esimd_reorder_q_traits<GGML_TYPE_Q5_K> {
375    struct ptrs {
376        const uint8_t *    qs;
377        const uint8_t *    qh;
378        const uint8_t *    scales;
379        const sycl::half * dm;
380    };
381
382    static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) {
383        const uint8_t * qs     = (const uint8_t *) vx;
384        const uint8_t * qh     = qs + nb * (QK_K / 2);
385        const uint8_t * scales = qh + nb * (QK_K / 8);
386        const sycl::half * dm  = (const sycl::half *) (scales + nb * K_SCALE_SIZE);
387        return { qs, qh, scales, dm };
388    }
389
390    // extract bit `bit` (0..7) of each lane and move it to bit position 4,
391    // e.g. for the 4-bit base quant's 5th (high) bit. `bit` is always a
392    // compile-time-known unrolled loop constant at call sites, so this folds
393    // to a single mask (bit==4), mask+left-shift (bit<4), or mask+right-shift
394    // (bit>4) instead of the shift+mask+shift a naive `(qh>>bit & 1) << 4` emits.
395    static ESIMD_INLINE sycl::ext::intel::esimd::simd<uint16_t, 32> extract_bit_to_pos4(
396            sycl::ext::intel::esimd::simd<uint8_t, 32> qh, int bit) {
397        using namespace sycl::ext::intel::esimd;
398        simd<uint16_t, 32> masked = convert<uint16_t>(qh & simd<uint8_t, 32>((uint8_t) (1u << bit)));
399        if (bit < 4) {
400            return masked << simd<uint16_t, 32>((uint16_t) (4 - bit));
401        } else if (bit > 4) {
402            return masked >> simd<uint16_t, 32>((uint16_t) (bit - 4));
403        }
404        return masked;
405    }
406
407    static ESIMD_INLINE void mac_pair(
408            const ptrs & pa, size_t bia,
409            const ptrs & pb, size_t bib, bool has_b,
410            sycl::ext::intel::esimd::simd<float, 256> & y_vec,
411            sycl::ext::intel::esimd::simd<float, 32> & acc_a,
412            sycl::ext::intel::esimd::simd<float, 32> & acc_b) {
413        using namespace sycl::ext::intel::esimd;
414
415        simd<uint8_t, 128> qs_a     = block_load<uint8_t, 128>(pa.qs + bia * (QK_K / 2));
416        simd<uint8_t, 128> qs_b     = 0;
417        simd<uint8_t, 32>  qh_a     = block_load<uint8_t, 32>(pa.qh + bia * (QK_K / 8));
418        simd<uint8_t, 32>  qh_b     = 0;
419        simd<uint8_t, 12>  scales_a = block_load<uint8_t, 12>(pa.scales + bia * K_SCALE_SIZE);
420        simd<uint8_t, 12>  scales_b = 0;
421
422        const float dall_a = (float) pa.dm[bia * 2 + 0];
423        const float dmin_a = (float) pa.dm[bia * 2 + 1];
424        float dall_b = 0.0f;
425        float dmin_b = 0.0f;
426        if (has_b) {
427            qs_b     = block_load<uint8_t, 128>(pb.qs + bib * (QK_K / 2));
428            qh_b     = block_load<uint8_t, 32>(pb.qh + bib * (QK_K / 8));
429            scales_b = block_load<uint8_t, 12>(pb.scales + bib * K_SCALE_SIZE);
430            dall_b = (float) pb.dm[bib * 2 + 0];
431            dmin_b = (float) pb.dm[bib * 2 + 1];
432        }
433
434        simd<float, 8> scale_f_a, min_f_a, scale_f_b, min_f_b;
435        unpack_scale_min_k4(scales_a, dall_a, dmin_a, scale_f_a, min_f_a);
436        unpack_scale_min_k4(scales_b, dall_b, dmin_b, scale_f_b, min_f_b);
437
438        simd<uint8_t, 128> qs_lo_a = qs_a & simd<uint8_t, 128>(0x0F);
439        simd<uint8_t, 128> qs_hi_a = qs_a >> simd<uint8_t, 128>(4);
440        simd<uint8_t, 128> qs_lo_b = qs_b & simd<uint8_t, 128>(0x0F);
441        simd<uint8_t, 128> qs_hi_b = qs_b >> simd<uint8_t, 128>(4);
442
443#pragma unroll
444        for (int sb = 0; sb < 8; sb += 2) {
445            const int q_offset = sb * 16;
446            simd<float, 32> y_lo = y_vec.select<32, 1>(sb * 32);
447            simd<float, 32> y_hi = y_vec.select<32, 1>((sb + 1) * 32);
448
449            const float scale_a_lo = scale_f_a[sb];
450            const float scale_a_hi = scale_f_a[sb + 1];
451            const float min_a_lo   = min_f_a[sb];
452            const float min_a_hi   = min_f_a[sb + 1];
453            const float scale_b_lo = scale_f_b[sb];
454            const float scale_b_hi = scale_f_b[sb + 1];
455            const float min_b_lo   = min_f_b[sb];
456            const float min_b_hi   = min_f_b[sb + 1];
457
458            simd<uint8_t, 32> qa_lo_u8 = qs_lo_a.select<32, 1>(q_offset);
459            simd<uint8_t, 32> qa_hi_u8 = qs_hi_a.select<32, 1>(q_offset);
460            simd<uint8_t, 32> qb_lo_u8 = qs_lo_b.select<32, 1>(q_offset);
461            simd<uint8_t, 32> qb_hi_u8 = qs_hi_b.select<32, 1>(q_offset);
462            simd<uint16_t, 32> qa_lo = convert<uint16_t>(qa_lo_u8);
463            simd<uint16_t, 32> qa_hi = convert<uint16_t>(qa_hi_u8);
464            simd<uint16_t, 32> qb_lo = convert<uint16_t>(qb_lo_u8);
465            simd<uint16_t, 32> qb_hi = convert<uint16_t>(qb_hi_u8);
466
467            // add the 5th bit: chunk sb uses qh bit sb, chunk sb+1 uses qh bit sb+1;
468            // qh always indexes the same 32 bytes regardless of chunk
469            qa_lo += extract_bit_to_pos4(qh_a, sb);
470            qa_hi += extract_bit_to_pos4(qh_a, sb + 1);
471            qb_lo += extract_bit_to_pos4(qh_b, sb);
472            qb_hi += extract_bit_to_pos4(qh_b, sb + 1);
473
474            simd<float, 32> deq_a_lo = convert<float>(qa_lo) * scale_a_lo + min_a_lo;
475            simd<float, 32> deq_a_hi = convert<float>(qa_hi) * scale_a_hi + min_a_hi;
476            simd<float, 32> deq_b_lo = convert<float>(qb_lo) * scale_b_lo + min_b_lo;
477            simd<float, 32> deq_b_hi = convert<float>(qb_hi) * scale_b_hi + min_b_hi;
478
479            acc_a += y_lo * deq_a_lo;
480            acc_b += y_lo * deq_b_lo;
481            acc_a += y_hi * deq_a_hi;
482            acc_b += y_hi * deq_b_hi;
483        }
484    }
485};
486
487// ---------------------------------------------------------------------------
488// Q6_K, SOA reorder layout:
489//   [ql: nb*(QK_K/2)] [qh: nb*(QK_K/4)] [scales(int8): nb*(QK_K/16)] [d: nb*half]
490// ---------------------------------------------------------------------------
491template <> struct esimd_reorder_q_traits<GGML_TYPE_Q6_K> {
492    struct ptrs {
493        const uint8_t *    ql;
494        const uint8_t *    qh;
495        const int8_t *     scales;
496        const sycl::half * d;
497    };
498
499    static ESIMD_INLINE ptrs make_ptrs(const void * vx, size_t nb) {
500        const uint8_t *    ql     = (const uint8_t *) vx;
501        const uint8_t *    qh     = ql + nb * (QK_K / 2);
502        const int8_t *     scales = (const int8_t *) (qh + nb * (QK_K / 4));
503        const sycl::half * d      = (const sycl::half *) (scales + nb * (QK_K / 16));
504        return { ql, qh, scales, d };
505    }
506
507    static ESIMD_INLINE void mac_pair(
508            const ptrs & pa, size_t bia,
509            const ptrs & pb, size_t bib, bool has_b,
510            sycl::ext::intel::esimd::simd<float, 256> & y_vec,
511            sycl::ext::intel::esimd::simd<float, 32> & acc_a,
512            sycl::ext::intel::esimd::simd<float, 32> & acc_b) {
513        using namespace sycl::ext::intel::esimd;
514
515        simd<uint8_t, 128> ql_a     = block_load<uint8_t, 128>(pa.ql + bia * (QK_K / 2));
516        simd<uint8_t, 128> ql_b     = 0;
517        simd<uint8_t, 64>  qh_a     = block_load<uint8_t, 64>(pa.qh + bia * (QK_K / 4));
518        simd<uint8_t, 64>  qh_b     = 0;
519        simd<int8_t, 16>   scales_a = block_load<int8_t, 16>(pa.scales + bia * (QK_K / 16));
520        simd<int8_t, 16>   scales_b = 0;
521
522        const float d_a = (float) pa.d[bia];
523        float d_b = 0.0f;
524        if (has_b) {
525            ql_b     = block_load<uint8_t, 128>(pb.ql + bib * (QK_K / 2));
526            qh_b     = block_load<uint8_t, 64>(pb.qh + bib * (QK_K / 4));
527            scales_b = block_load<int8_t, 16>(pb.scales + bib * (QK_K / 16));
528            d_b = (float) pb.d[bib];
529        }
530
531        simd<float, 16> sc_a = convert<float>(scales_a);
532        simd<float, 16> sc_b = convert<float>(scales_b);
533
534#pragma unroll
535        for (int im = 0; im < 2; ++im) {
536            simd<uint8_t, 32> ql_lo_a   = ql_a.select<32, 1>(64 * im);
537            simd<uint8_t, 32> ql_hi_a   = ql_a.select<32, 1>(64 * im + 32);
538            simd<uint8_t, 32> qh_bits_a = qh_a.select<32, 1>(32 * im);
539            simd<uint8_t, 32> ql_lo_b   = ql_b.select<32, 1>(64 * im);
540            simd<uint8_t, 32> ql_hi_b   = ql_b.select<32, 1>(64 * im + 32);
541            simd<uint8_t, 32> qh_bits_b = qh_b.select<32, 1>(32 * im);
542
543            // reconstruct each 32-wide 6-bit group (matches dequantize_row_q6_K)
544#pragma unroll
545            for (int g = 0; g < 4; ++g) {
546                simd<float, 32> y_g = y_vec.select<32, 1>(32 * (4 * im + g));
547
548                const float scale_a_lo = sc_a[8 * im + 2 * g + 0] * d_a;
549                const float scale_a_hi = sc_a[8 * im + 2 * g + 1] * d_a;
550                const float scale_b_lo = sc_b[8 * im + 2 * g + 0] * d_b;
551                const float scale_b_hi = sc_b[8 * im + 2 * g + 1] * d_b;
552
553                simd<float, 32> scale_vec_a = splat_lo_hi(scale_a_lo, scale_a_hi);
554                simd<float, 32> scale_vec_b = splat_lo_hi(scale_b_lo, scale_b_hi);
555
556                simd<uint8_t, 32> qa;
557                simd<uint8_t, 32> qb;
558                switch (g) {
559                    case 0:
560                        qa = (ql_lo_a & simd<uint8_t, 32>(0x0F)) | ((qh_bits_a & simd<uint8_t, 32>(0x03)) << simd<uint8_t, 32>(4));
561                        qb = (ql_lo_b & simd<uint8_t, 32>(0x0F)) | ((qh_bits_b & simd<uint8_t, 32>(0x03)) << simd<uint8_t, 32>(4));
562                        break;
563                    case 1:
564                        qa = (ql_hi_a & simd<uint8_t, 32>(0x0F)) | ((qh_bits_a & simd<uint8_t, 32>(0x0C)) << simd<uint8_t, 32>(2));
565                        qb = (ql_hi_b & simd<uint8_t, 32>(0x0F)) | ((qh_bits_b & simd<uint8_t, 32>(0x0C)) << simd<uint8_t, 32>(2));
566                        break;
567                    case 2:
568                        qa = (ql_lo_a >> simd<uint8_t, 32>(4)) | (qh_bits_a & simd<uint8_t, 32>(0x30));
569                        qb = (ql_lo_b >> simd<uint8_t, 32>(4)) | (qh_bits_b & simd<uint8_t, 32>(0x30));
570                        break;
571                    default:
572                        qa = (ql_hi_a >> simd<uint8_t, 32>(4)) | ((qh_bits_a & simd<uint8_t, 32>(0xC0)) >> simd<uint8_t, 32>(2));
573                        qb = (ql_hi_b >> simd<uint8_t, 32>(4)) | ((qh_bits_b & simd<uint8_t, 32>(0xC0)) >> simd<uint8_t, 32>(2));
574                        break;
575                }
576
577                simd<float, 32> deq_a = (convert<float>(qa) - 32.0f) * scale_vec_a;
578                simd<float, 32> deq_b = (convert<float>(qb) - 32.0f) * scale_vec_b;
579
580                acc_a += y_g * deq_a;
581                acc_b += y_g * deq_b;
582            }
583        }
584    }
585};
586
587} // namespace ggml_sycl_esimd
588
589#endif // GGML_SYCL_ESIMD_HPP