sycl: wide stores for the q4_K/q5_K weight dequant (#29696)
This commit is contained in:
Titaniumtown authored and 
GitHub committed
1 parent
88aaf325a7
commit
13de11e78f
2 files changed
+97
-118
No files matched your search
@@ -212,20 +212,10 @@ template <typename dst_t>
|
||||
static void dequantize_row_q4_K_sycl(const void *vx, dst_t *y, const int64_t k,
|
||||
dpct::queue_ptr stream) {
|
||||
const int64_t nb = k / QK_K;
|
||||
{
|
||||
dpct::has_capability_or_fail(stream->get_device(),
|
||||
{sycl::aspect::fp16});
|
||||
|
||||
stream->submit([&](sycl::handler &cgh) {
|
||||
sycl::local_accessor<uint8_t, 1> scale_local_acc(sycl::range<1>(12), cgh);
|
||||
cgh.parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) *
|
||||
sycl::range<3>(1, 1, 32),
|
||||
sycl::range<3>(1, 1, 32)),
|
||||
[=](sycl::nd_item<3> item_ct1) {
|
||||
dequantize_block_q4_K(vx, y, get_pointer(scale_local_acc), item_ct1);
|
||||
});
|
||||
});
|
||||
}
|
||||
dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
|
||||
stream->parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t) nb * SYCL_DEQK_WI_PER_BLOCK),
|
||||
sycl::range<1>((size_t) SYCL_DEQK_WI_PER_BLOCK)),
|
||||
[=](sycl::nd_item<1> it) { dequantize_block_q4_K_wide(vx, y, nb, it); });
|
||||
}
|
||||
|
||||
template <typename dst_t>
|
||||
@@ -250,32 +240,10 @@ template <typename dst_t>
|
||||
static void dequantize_row_q5_K_sycl(const void *vx, dst_t *y, const int64_t k,
|
||||
dpct::queue_ptr stream) {
|
||||
const int64_t nb = k / QK_K;
|
||||
#if QK_K == 256
|
||||
{
|
||||
dpct::has_capability_or_fail(stream->get_device(),
|
||||
{sycl::aspect::fp16});
|
||||
|
||||
stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) *
|
||||
sycl::range<3>(1, 1, 64),
|
||||
sycl::range<3>(1, 1, 64)),
|
||||
[=](sycl::nd_item<3> item_ct1) {
|
||||
dequantize_block_q5_K(vx, y, item_ct1);
|
||||
});
|
||||
}
|
||||
#else
|
||||
{
|
||||
dpct::has_capability_or_fail(stream->get_device(),
|
||||
{sycl::aspect::fp16});
|
||||
|
||||
stream->parallel_for(sycl::nd_range<3>(sycl::range<3>(1, 1, nb) *
|
||||
sycl::range<3>(1, 1, 32),
|
||||
sycl::range<3>(1, 1, 32)),
|
||||
[=](sycl::nd_item<3> item_ct1) {
|
||||
dequantize_block_q5_K(vx, y, item_ct1);
|
||||
});
|
||||
}
|
||||
|
||||
#endif
|
||||
dpct::has_capability_or_fail(stream->get_device(), { sycl::aspect::fp16 });
|
||||
stream->parallel_for(sycl::nd_range<1>(sycl::range<1>((size_t) nb * SYCL_DEQK_WI_PER_BLOCK),
|
||||
sycl::range<1>((size_t) SYCL_DEQK_WI_PER_BLOCK)),
|
||||
[=](sycl::nd_item<1> it) { dequantize_block_q5_K_wide(vx, y, nb, it); });
|
||||
}
|
||||
|
||||
template <typename dst_t>
|
||||
|
||||
@@ -1108,6 +1108,53 @@ static inline void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8
|
||||
}
|
||||
#endif
|
||||
|
||||
// Work-items that cover one 256-weight K-quant block: four 64-element chunks, 8 work-items
|
||||
// per chunk, 4 qs bytes per work-item.
|
||||
static constexpr int SYCL_DEQK_WI_PER_BLOCK = 32;
|
||||
|
||||
template <typename dst_t>
|
||||
static inline void dequantize_q4_K_wide_one(dst_t * __restrict__ y, const uint8_t * __restrict__ qs,
|
||||
const uint8_t * __restrict__ scales, const float dall,
|
||||
const float dmin, const int tid) {
|
||||
const int il = tid >> 3; // 64-element chunk, 0..3
|
||||
const int ir = (tid & 7) * 4; // 4 qs bytes within the chunk
|
||||
|
||||
uint8_t sc, m;
|
||||
get_scale_min_k4(2 * il + 0, scales, sc, m);
|
||||
const float d1 = dall * sc;
|
||||
const float m1 = dmin * m;
|
||||
get_scale_min_k4(2 * il + 1, scales, sc, m);
|
||||
const float d2 = dall * sc;
|
||||
const float m2 = dmin * m;
|
||||
|
||||
const sycl::vec<uint8_t, 4> q = vec_aligned_load<uint8_t, 4>(qs + 32 * il + ir);
|
||||
|
||||
sycl::vec<dst_t, 4> lo, hi;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
lo[l] = d1 * (q[l] & 0xF) - m1;
|
||||
hi[l] = d2 * (q[l] >> 4) - m2;
|
||||
}
|
||||
dst_t * yb = y + 64 * il + ir;
|
||||
*reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 0) = lo;
|
||||
*reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 32) = hi;
|
||||
}
|
||||
|
||||
template <typename dst_t>
|
||||
static void dequantize_block_q4_K_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb,
|
||||
const sycl::nd_item<1> & it) {
|
||||
static_assert(QK_K == 256, "wide K-quant dequant assumes QK_K == 256");
|
||||
const int64_t g = it.get_global_id(0);
|
||||
const int64_t b = g / SYCL_DEQK_WI_PER_BLOCK;
|
||||
const int tid = (int) (g % SYCL_DEQK_WI_PER_BLOCK);
|
||||
if (b >= nb) {
|
||||
return;
|
||||
}
|
||||
const block_q4_K * x = (const block_q4_K *) vx + b;
|
||||
dequantize_q4_K_wide_one(yy + b * QK_K, x->qs, x->scales, (float) x->dm[0], (float) x->dm[1], tid);
|
||||
}
|
||||
|
||||
|
||||
template <typename dst_t>
|
||||
inline void dequantize_q4_K_common(dst_t * __restrict__ y, const uint8_t * __restrict__ qs_ptr, const float dall,
|
||||
const float dmin, uint8_t * __restrict__ scales_local, int il, int ir) {
|
||||
@@ -1130,41 +1177,6 @@ inline void dequantize_q4_K_common(dst_t * __restrict__ y, const uint8_t * __res
|
||||
}
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
static void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy,
|
||||
uint8_t* scales_local, const sycl::nd_item<3> &item_ct1) {
|
||||
const block_q4_K * x = (const block_q4_K *) vx;
|
||||
|
||||
const int64_t i = item_ct1.get_group(2);
|
||||
|
||||
#if QK_K == 256
|
||||
const int64_t tid = item_ct1.get_local_id(2);
|
||||
const int64_t il = tid / 8;
|
||||
const int64_t ir = tid % 8;
|
||||
|
||||
dst_t * y = yy + i * QK_K + 64 * il + 4 * ir;
|
||||
|
||||
const sycl::half2 dm = x[i].dm;
|
||||
const float dall = dm[0];
|
||||
const float dmin = dm[1];
|
||||
|
||||
if (tid < 12) {
|
||||
scales_local[tid] = x[i].scales[tid];
|
||||
}
|
||||
|
||||
item_ct1.barrier(sycl::access::fence_space::local_space);
|
||||
dequantize_q4_K_common(y, x[i].qs, dall, dmin, scales_local, il, ir);
|
||||
#else
|
||||
const int64_t tid = item_ct1.get_local_id(2);
|
||||
const uint8_t * q = x[i].qs;
|
||||
dst_t * y = yy + i*QK_K;
|
||||
const float d = (float)x[i].dm[0];
|
||||
const float m = (float)x[i].dm[1];
|
||||
y[tid+ 0] = d * (x[i].scales[0] & 0xF) * (q[tid] & 0xF) - m * (x[i].scales[0] >> 4);
|
||||
y[tid+32] = d * (x[i].scales[1] & 0xF) * (q[tid] >> 4) - m * (x[i].scales[1] >> 4);
|
||||
#endif
|
||||
}
|
||||
|
||||
template <typename dst_t>
|
||||
static void dequantize_block_q4_K_reorder(const void * __restrict__ vx, dst_t * __restrict__ yy, uint8_t * scales_local,
|
||||
const sycl::nd_item<1> & item_ct1, int64_t nb) {
|
||||
@@ -1195,52 +1207,51 @@ static void dequantize_block_q4_K_reorder(const void * __restrict__ vx, dst_t *
|
||||
dequantize_q4_K_common(y, qs_ptr, dall, dmin, scales_local, il, ir);
|
||||
}
|
||||
|
||||
template<typename dst_t>
|
||||
static void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy,
|
||||
const sycl::nd_item<3> &item_ct1) {
|
||||
const block_q5_K * x = (const block_q5_K *) vx;
|
||||
|
||||
const int64_t i = item_ct1.get_group(2);
|
||||
|
||||
#if QK_K == 256
|
||||
// assume 64 threads - this is very slightly better than the one below
|
||||
const int64_t tid = item_ct1.get_local_id(2);
|
||||
const int64_t il = tid/16; // il is in 0...3
|
||||
const int64_t ir = tid%16; // ir is in 0...15
|
||||
const int64_t is = 2*il; // is is in 0...6
|
||||
|
||||
dst_t * y = yy + i*QK_K + 64*il + 2*ir;
|
||||
|
||||
const float dall = x[i].dm[0];
|
||||
const float dmin = x[i].dm[1];
|
||||
|
||||
const uint8_t * ql = x[i].qs + 32*il + 2*ir;
|
||||
const uint8_t * qh = x[i].qh + 2*ir;
|
||||
template <typename dst_t>
|
||||
static inline void dequantize_q5_K_wide_one(dst_t * __restrict__ y, const uint8_t * __restrict__ qs,
|
||||
const uint8_t * __restrict__ qh_ptr,
|
||||
const uint8_t * __restrict__ scales, const float dall,
|
||||
const float dmin, const int tid) {
|
||||
const int il = tid >> 3; // 64-element chunk, 0..3
|
||||
const int ir = (tid & 7) * 4; // 4 qs bytes within the chunk; qh holds the 5th bit
|
||||
|
||||
uint8_t sc, m;
|
||||
get_scale_min_k4(is + 0, x[i].scales, sc, m);
|
||||
const float d1 = dall * sc; const float m1 = dmin * m;
|
||||
get_scale_min_k4(is + 1, x[i].scales, sc, m);
|
||||
const float d2 = dall * sc; const float m2 = dmin * m;
|
||||
get_scale_min_k4(2 * il + 0, scales, sc, m);
|
||||
const float d1 = dall * sc;
|
||||
const float m1 = dmin * m;
|
||||
get_scale_min_k4(2 * il + 1, scales, sc, m);
|
||||
const float d2 = dall * sc;
|
||||
const float m2 = dmin * m;
|
||||
|
||||
uint8_t hm = 1 << (2*il);
|
||||
y[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1;
|
||||
y[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1;
|
||||
hm <<= 1;
|
||||
y[32] = d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2;
|
||||
y[33] = d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2;
|
||||
#else
|
||||
const int64_t tid = item_ct1.get_local_id(2);
|
||||
const uint8_t q = x[i].qs[tid];
|
||||
const int64_t im = tid/8; // 0...3
|
||||
const int64_t in = tid%8; // 0...7
|
||||
const int64_t is = tid/16; // 0 or 1
|
||||
const uint8_t h = x[i].qh[in] >> im;
|
||||
const float d = x[i].d;
|
||||
dst_t * y = yy + i*QK_K + tid;
|
||||
y[ 0] = d * x[i].scales[is+0] * ((q & 0xF) - ((h >> 0) & 1 ? 0 : 16));
|
||||
y[32] = d * x[i].scales[is+2] * ((q >> 4) - ((h >> 4) & 1 ? 0 : 16));
|
||||
#endif
|
||||
const sycl::vec<uint8_t, 4> q = vec_aligned_load<uint8_t, 4>(qs + 32 * il + ir);
|
||||
const sycl::vec<uint8_t, 4> h = vec_aligned_load<uint8_t, 4>(qh_ptr + ir);
|
||||
|
||||
const uint8_t hm1 = 1 << (2 * il);
|
||||
const uint8_t hm2 = hm1 << 1;
|
||||
|
||||
sycl::vec<dst_t, 4> lo, hi;
|
||||
#pragma unroll
|
||||
for (int l = 0; l < 4; ++l) {
|
||||
lo[l] = d1 * ((q[l] & 0xF) + ((h[l] & hm1) ? 16 : 0)) - m1;
|
||||
hi[l] = d2 * ((q[l] >> 4) + ((h[l] & hm2) ? 16 : 0)) - m2;
|
||||
}
|
||||
dst_t * yb = y + 64 * il + ir;
|
||||
*reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 0) = lo;
|
||||
*reinterpret_cast<sycl::vec<dst_t, 4> *>(yb + 32) = hi;
|
||||
}
|
||||
|
||||
template <typename dst_t>
|
||||
static void dequantize_block_q5_K_wide(const void * __restrict__ vx, dst_t * __restrict__ yy, int64_t nb,
|
||||
const sycl::nd_item<1> & it) {
|
||||
static_assert(QK_K == 256, "wide K-quant dequant assumes QK_K == 256");
|
||||
const int64_t g = it.get_global_id(0);
|
||||
const int64_t b = g / SYCL_DEQK_WI_PER_BLOCK;
|
||||
const int tid = (int) (g % SYCL_DEQK_WI_PER_BLOCK);
|
||||
if (b >= nb) {
|
||||
return;
|
||||
}
|
||||
const block_q5_K * x = (const block_q5_K *) vx + b;
|
||||
dequantize_q5_K_wide_one(yy + b * QK_K, x->qs, x->qh, x->scales, (float) x->dm[0], (float) x->dm[1], tid);
|
||||
}
|
||||
|
||||
template <typename dst_t>
|
||||
|
||||
Reference in new issue
Block a user