diff --git a/ggml/src/ggml-sycl/CMakeLists.txt b/ggml/src/ggml-sycl/CMakeLists.txt index a8d9c0d804bf..bd4dc57101af 100644 --- a/ggml/src/ggml-sycl/CMakeLists.txt +++ b/ggml/src/ggml-sycl/CMakeLists.txt @@ -215,4 +215,22 @@ if (GGML_SYCL_DEVICE_ARCH) "SHELL:-Xsycl-target-backend=spir64_gen \"-device ${GGML_SYCL_DEVICE_ARCH}\"" -fsycl-max-parallel-link-jobs=${GGML_SYCL_MAX_PARALLEL_LINK_JOBS} ) + + # The PQ2_0/PTQ1_0 XMX kernels (pq2_xmx.cpp) need 16-wide DPAS and 2D block loads, and AOT compiles them for + # every listed device, so they are only built when all of them are Xe-HPC, Xe2 or later. Otherwise those + # types keep the existing paths. + string(REPLACE "," ";" _ggml_sycl_aot_devices "${GGML_SYCL_DEVICE_ARCH}") + set(_ggml_sycl_pq2_xmx ON) + foreach (_ggml_sycl_dev IN LISTS _ggml_sycl_aot_devices) + string(STRIP "${_ggml_sycl_dev}" _ggml_sycl_dev) + string(TOLOWER "${_ggml_sycl_dev}" _ggml_sycl_dev) + # the -vg parts of Xe-HPC have no XMX + if (NOT _ggml_sycl_dev MATCHES "^(pvc|bmg|lnl|ptl|wcl|nvl|cri|xe-hpc|xe2|xe3)" OR _ggml_sycl_dev MATCHES "-vg") + set(_ggml_sycl_pq2_xmx OFF) + endif() + endforeach() + if (NOT _ggml_sycl_pq2_xmx) + message(STATUS "GGML_SYCL_DEVICE_ARCH includes a device without 16-wide DPAS, not building the PQ2_0/PTQ1_0 XMX path") + set_property(SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/pq2_xmx.cpp APPEND PROPERTY COMPILE_DEFINITIONS GGML_SYCL_NO_PQ2_XMX) + endif() endif() diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp index 34de284d83ac..8f2106d83088 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common.hpp @@ -224,6 +224,7 @@ inline dpct::err0 ggml_sycl_set_device(const int device) try { ////////////////////// struct optimize_feature { bool reorder=false; + bool xmx_pq2=false; // PQ2_0 rewritten into the XMX layout (pq2_xmx.hpp); only that path can read it }; struct sycl_device_info { @@ -243,6 +244,7 @@ struct sycl_device_info { sycl_hw_info hw_info; optimize_feature opt_feature; bool usm_system_support; // support for USM system allocations + int dpas_exec_size; // XMX DPAS width (8 or 16) for int8, 0 without XMX }; diff --git a/ggml/src/ggml-sycl/getrows.cpp b/ggml/src/ggml-sycl/getrows.cpp index 36f840e6f5d1..187640446b48 100644 --- a/ggml/src/ggml-sycl/getrows.cpp +++ b/ggml/src/ggml-sycl/getrows.cpp @@ -277,10 +277,16 @@ void ggml_sycl_op_get_rows(ggml_backend_sycl_context & ctx, ggml_tensor * dst) { src1_i32, (float *)dst->data, ctx.stream()); break; case GGML_TYPE_PTQ1_0: + GGML_ASSERT(!(dst->src[0]->extra && + ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.xmx_pq2) && + "PTQ1_0 in the XMX layout reached get_rows"); get_rows_sycl(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); break; case GGML_TYPE_PQ2_0: + GGML_ASSERT(!(dst->src[0]->extra && + ((ggml_tensor_extra_gpu *) dst->src[0]->extra)->optimized_feature.xmx_pq2) && + "PQ2_0 in the XMX layout reached get_rows"); get_rows_sycl(ctx, dst->src[0], dst->src[1], dst, (const float *)dst->src[0]->data, src1_i32, (float *)dst->data, ctx.stream()); break; diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index bb29a075d4f0..97b5f35659f7 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -56,6 +56,7 @@ #include "ggml-sycl/add-id.hpp" #include "ggml-sycl/backend.hpp" +#include "ggml-sycl/pq2_xmx.hpp" #include "ggml-sycl/common.hpp" #include "ggml-sycl/element_wise.hpp" #include "ggml-sycl/fwht.hpp" @@ -106,6 +107,23 @@ int g_ggml_sycl_dev2dev_memcpy = DEV2DEV_MEMCPY_SYCL; int g_ggml_sycl_usm_system = 0; int g_ggml_sycl_enable_host_pinned_mem = 1; +// int8 DPAS execution size of the device's XMX units (8 or 16), as the runtime reports it; 0 without XMX +static int ggml_sycl_dpas_exec_size(const sycl::device & device) { + namespace matrix = syclex::matrix; + if (!device.has(sycl::aspect::ext_intel_matrix)) { + return 0; + } + try { + for (const matrix::combination & c : device.get_info()) { + if (c.atype == matrix::matrix_type::sint8 && c.btype == matrix::matrix_type::sint8) { + return (int) c.nsize; + } + } + } catch (const sycl::exception &) { + } + return 0; +} + static ggml_sycl_device_info ggml_sycl_init() { ggml_sycl_device_info info = {}; @@ -170,6 +188,7 @@ static ggml_sycl_device_info ggml_sycl_init() { info.max_work_group_sizes[i] = prop.get_max_work_group_size(); info.devices[i].max_wg_per_cu = info.max_work_group_sizes[i] / prop.get_max_compute_units(); info.devices[i].hw_info = get_device_hw_info(&device); + info.devices[i].dpas_exec_size = ggml_sycl_dpas_exec_size(device); // Only check GPU devices; CPU devices use OpenCL and would otherwise // disable Level Zero for the GPUs on systems without ONEAPI_DEVICE_SELECTOR set. @@ -593,7 +612,9 @@ ggml_backend_sycl_buffer_init_tensor(ggml_backend_buffer_t buffer, case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K:{ + case GGML_TYPE_Q6_K: + case GGML_TYPE_PQ2_0: + case GGML_TYPE_PTQ1_0:{ ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{}; tensor->extra = extra; ctx->tensor_extras.push_back(extra); @@ -954,6 +975,26 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ GGML_UNUSED(buft); } +static bool ggml_sycl_device_has_dpas16(int device); + +// GGML_SYCL_DISABLE_XMX=1 keeps PQ2_0/PTQ1_0 off the XMX path (and its layout) +static bool ggml_sycl_xmx_disabled() { + static const bool disabled = ggml_sycl_get_env("GGML_SYCL_DISABLE_XMX", 0); + return disabled; +} + +// PTQ1_0 weights are expanded into the 34-byte PQ2_0 XMX blocks on first use (pq2_xmx.hpp), so on devices +// that run that path their allocation reserves room for the expanded form +static bool ggml_sycl_ptq1_xmx_expands(int device, const ggml_tensor * tensor) { + return tensor->type == GGML_TYPE_PTQ1_0 && g_ggml_sycl_enable_optimize && !ggml_sycl_xmx_disabled() && + tensor->ne[2] == 1 && tensor->ne[3] == 1 && ggml_sycl_pq2_xmx_supports_ne0(tensor->ne[0]) && + ggml_sycl_device_has_dpas16(device); +} + +static size_t ggml_sycl_ptq1_xmx_bytes(const ggml_tensor * tensor) { + return (size_t) (ggml_nelements(tensor) / QK_PTQ1_0) * sizeof(block_pq2_0); +} + static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { size_t size = ggml_nbytes(tensor); int64_t ne0 = tensor->ne[0]; @@ -964,9 +1005,12 @@ static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_t } } - return size; + const auto * buft_ctx = (const ggml_backend_sycl_buffer_type_context *) buft->context; + if (ggml_sycl_ptq1_xmx_expands(buft_ctx->device, tensor)) { + size = std::max(size, ggml_sycl_ptq1_xmx_bytes(tensor)); + } - GGML_UNUSED(buft); + return size; } static const ggml_backend_buffer_type_i ggml_backend_sycl_buffer_type_interface = { @@ -3764,6 +3808,13 @@ inline bool ggml_sycl_supports_mmq(enum ggml_type type) { return false; } +// The PQ2_0/PTQ1_0 XMX path feeds 2-bit weights to ESIMD DPAS at execution size 16 through 2D block loads, which +// every XMX device with 16-wide DPAS has (Xe-HPC, Xe2 and later). 8-wide XMX (Xe-HPG, Arrow Lake-H) keeps the +// existing paths. +static bool ggml_sycl_device_has_dpas16(int device) { + return ggml_sycl_info().devices[device].dpas_exec_size == 16; +} + inline bool ggml_sycl_supports_reorder_mul_mat_sycl(enum ggml_type type) { switch (type) { case GGML_TYPE_Q1_0: @@ -4529,6 +4580,50 @@ static bool can_use_mul_mat_vec_q(const ggml_tensor * src0, const ggml_tensor * src1->ne[1] <= MMVQ_MAX_BATCH_SIZE; } +// PQ2_0/PTQ1_0 weights on 16-wide DPAS devices are rewritten into the XMX layout on first use. From then on every +// mul_mat on them has to take that path, so the layout flag alone decides once it is set. +static bool ggml_sycl_pq2_xmx_use(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + const ggml_tensor * dst) { + if (src0->type != GGML_TYPE_PQ2_0 && src0->type != GGML_TYPE_PTQ1_0) { + return false; + } + // MUL_MAT_ID passes each expert as a 2D copy of the 3D weight that shares its extra, and a view shares + // its parent's data: rewriting either in place would corrupt the rest of the tensor + if (dst->op != GGML_OP_MUL_MAT || src0->view_src != nullptr) { + return false; + } + ggml_tensor_extra_gpu * extra = static_cast(src0->extra); + if (extra && extra->optimized_feature.xmx_pq2) { + return true; + } + if (!g_ggml_sycl_enable_optimize || ggml_sycl_xmx_disabled() || !ggml_sycl_device_has_dpas16(ctx.device)) { + return false; + } + // op offload refills COMPUTE buffers from host memory every time, so an in-place layout there goes stale + if (!extra || ggml_backend_buffer_is_sycl_split(src0->buffer) || + src0->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + return false; + } + if (src0->ne[2] != 1 || src0->ne[3] != 1 || !ggml_is_contiguous(src0) || + !ggml_sycl_pq2_xmx_supports_ne0(src0->ne[0]) || (uintptr_t) src0->data % 64 != 0) { + return false; + } + if (src1->type != GGML_TYPE_F32 || src1->nb[0] != sizeof(float) || dst->type != GGML_TYPE_F32 || + !ggml_is_contiguous(dst)) { + return false; + } + // PTQ1_0 expands to 34 bytes a block, which only fits where the buffer reserved room for it + if (src0->type == GGML_TYPE_PTQ1_0 && + ggml_backend_buft_get_alloc_size(src0->buffer->buft, src0) < ggml_sycl_ptq1_xmx_bytes(src0)) { + return false; + } + if (!ggml_sycl_pq2_xmx_reorder(const_cast(src0), ctx.stream())) { + return false; + } + extra->optimized_feature.xmx_pq2 = true; + return true; +} + static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { scope_op_debug_print scope_dbg_print(__func__, dst, /*num_src=*/2); @@ -4543,6 +4638,11 @@ static void ggml_sycl_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor return; } + if (ggml_sycl_pq2_xmx_use(ctx, src0, src1, dst)) { + ggml_sycl_pq2_xmx_mul_mat(ctx, src0, src1, dst); + return; + } + const bool split = ggml_backend_buffer_is_sycl_split(src0->buffer); int64_t min_compute_capability = INT_MAX; diff --git a/ggml/src/ggml-sycl/pq2_xmx.cpp b/ggml/src/ggml-sycl/pq2_xmx.cpp new file mode 100644 index 000000000000..20fbf60f4ae9 --- /dev/null +++ b/ggml/src/ggml-sycl/pq2_xmx.cpp @@ -0,0 +1,382 @@ +#include "pq2_xmx.hpp" +#include "dequantize.hpp" + +// GGML_SYCL_NO_PQ2_XMX: an AOT build for a device without 16-wide DPAS, which cannot compile these kernels +#if defined(__INTEL_LLVM_COMPILER) && !defined(GGML_SYCL_NO_PQ2_XMX) +#include + +namespace { + +namespace esimd = sycl::ext::intel::esimd; +namespace xmx = sycl::ext::intel::esimd::xmx; + +constexpr int PQ2_XMX_QS_BYTES = QK_PQ2_0 / 4; // 32 bytes of 2-bit codes per block +constexpr int PQ2_XMX_QS_DW = PQ2_XMX_QS_BYTES / 4; +constexpr int PQ2_XMX_WG = 16; // independent tiles per work-group when K is not split + +// thread count the K split aims for: the B50 runs 1024 hardware threads, and decode needs several in flight +// per EU to keep enough loads outstanding +constexpr int PQ2_XMX_TARGET_THREADS = 4096; +constexpr int PQ2_XMX_PREFILL_TARGET_THREADS = 512; + +static_assert(QK_PQ2_0 == 128 && sizeof(block_pq2_0) == 34, "PQ2_0 layout changed"); +static_assert(QK_PTQ1_0 == QK_PQ2_0, "PTQ1_0 expands block for block into PQ2_0 codes"); + +// PQ2_0 packs value+1 per 2-bit field, lowest first: as a little-endian dword that is already DPAS 2-bit packing. +// Subtracting 1 per field without borrow gives the s2 values DPAS multiplies. +template +ESIMD_INLINE esimd::simd pq2_codes_to_s2(esimd::simd x) { + constexpr uint32_t H = 0xAAAAAAAAu; + constexpr uint32_t L = 0x55555555u; + return ((x | H) - L) ^ (~x & H); +} + +// Each thread computes an (8*MR) token x (16*NR) row tile over its share of the K blocks. With S > 1 the S +// threads of a work-group split K for the same tile and reduce through SLM, so a mat-vec keeps enough threads +// streaming weights. One DPAS covers 8 tokens x 16 rows x 32 k; the four of a block accumulate in int32. +template +ESIMD_INLINE void pq2_xmx_thread(const uint32_t * wq, const uint16_t * wd, const uint32_t * a8, const float * as, + float * dst, int K, int nrows, int ncols, int nrows_dst, int n_tiles_n, + int tile, int ks) { + using namespace esimd; + constexpr int TM = 8 * MR; + constexpr int TN = 16 * NR; + + if constexpr (S > 1) { + slm_init(); + } + + const int m0 = (tile / n_tiles_n) * TM; + const int n0 = (tile % n_tiles_n) * TN; + const int nb = K / QK_PQ2_0; + const int b0 = (int) ((int64_t) ks * nb / S); + const int b1 = (int) ((int64_t) (ks + 1) * nb / S); + + // weights: nrows rows of nb*32 bytes; activations: ncols rows of K bytes. Rows past either read as zeros. + const uint32_t wsurf_w = (uint32_t) (nb * PQ2_XMX_QS_BYTES) - 1; + const uint32_t wsurf_h = (uint32_t) nrows - 1; + const uint32_t asurf_w = (uint32_t) K - 1; + const uint32_t asurf_h = (uint32_t) ncols - 1; + + // scale gathers clamp to a valid row/token; the clamped lanes are never stored + simd d_off[NR]; +#pragma unroll + for (int g = 0; g < NR; ++g) { + simd r(n0 + 16 * g, 1); + r.merge(simd(nrows - 1), r >= (uint32_t) nrows); + d_off[g] = r * (uint32_t) (nb * sizeof(uint16_t)); + } + simd s_off[MR]; +#pragma unroll + for (int s = 0; s < MR; ++s) { + simd m(m0 + 8 * s, 1); + m.merge(simd(ncols - 1), m >= (uint32_t) ncols); + s_off[s] = m * (uint32_t) (nb * sizeof(float)); + } + + simd acc[MR][NR]; +#pragma unroll + for (int s = 0; s < MR; ++s) { +#pragma unroll + for (int g = 0; g < NR; ++g) { + acc[s][g] = 0.0f; + } + } + + for (int b = b0; b < b1; ++b) { + // transposed load: w[g][j*16 + n] = dword j (k = 16j..16j+15) of row n, the DPAS B layout for 2 dwords per k32 + simd w[NR]; + simd dw[NR]; +#pragma unroll + for (int g = 0; g < NR; ++g) { + w[g] = pq2_codes_to_s2<128>(load_2d( + wq, wsurf_w, wsurf_h, wsurf_w, b * PQ2_XMX_QS_DW, n0 + 16 * g)); + simd dbits = gather(wd, d_off[g] + (uint32_t) (b * sizeof(uint16_t))); + simd dh = dbits.template bit_cast_view(); + dw[g] = convert(dh); + } + +#pragma unroll + for (int s = 0; s < MR; ++s) { + const int yrow = m0 + 8 * s; + + simd ci[NR]; +#pragma unroll + for (int g = 0; g < NR; ++g) { + ci[g] = 0; + } +#pragma unroll + for (int c = 0; c < QK_PQ2_0 / 32; ++c) { + // A operand: token t's 32 int8 values at dwords t*8 .. t*8+7 + simd ad = load_2d(a8, asurf_w, asurf_h, asurf_w, + b * 32 + 8 * c, yrow); + simd am = ad.template bit_cast_view(); +#pragma unroll + for (int g = 0; g < NR; ++g) { + simd bd = w[g].template select<32, 1>(32 * c); + simd bm = bd.template bit_cast_view(); + ci[g] = xmx::dpas<8, 8, int, int, signed char, signed char, xmx::dpas_argument_type::s2, + xmx::dpas_argument_type::s8>(ci[g], bm, am); + } + } + + const simd da = gather(as, s_off[s] + (uint32_t) (b * sizeof(float))); +#pragma unroll + for (int g = 0; g < NR; ++g) { +#pragma unroll + for (int t = 0; t < 8; ++t) { + const simd cit = ci[g].template select<16, 1>(16 * t); + acc[s][g].template select<16, 1>(16 * t) += convert(cit) * (dw[g] * da[t]); + } + } + } + } + + if constexpr (S > 1) { + // every thread parks its partial tile in SLM; thread 0 sums them and stores +#pragma unroll + for (int s = 0; s < MR; ++s) { +#pragma unroll + for (int g = 0; g < NR; ++g) { +#pragma unroll + for (int q = 0; q < 8; ++q) { + const uint32_t off = (uint32_t) ((((ks * MR + s) * NR + g) * 128 + 16 * q) * sizeof(float)); + slm_block_store(off, acc[s][g].template select<16, 1>(16 * q)); + } + } + } + barrier(); + if (ks != 0) { + return; + } + for (int o = 1; o < S; ++o) { +#pragma unroll + for (int s = 0; s < MR; ++s) { +#pragma unroll + for (int g = 0; g < NR; ++g) { +#pragma unroll + for (int q = 0; q < 8; ++q) { + const uint32_t off = (uint32_t) ((((o * MR + s) * NR + g) * 128 + 16 * q) * sizeof(float)); + acc[s][g].template select<16, 1>(16 * q) += slm_block_load(off); + } + } + } + } + } + + const simd lane(0, 1); +#pragma unroll + for (int s = 0; s < MR; ++s) { +#pragma unroll + for (int t = 0; t < 8; ++t) { + const int m = m0 + 8 * s + t; + if (m >= ncols) { + continue; + } + // row base in 64 bits: an output head at a large ubatch passes 4 GB + float * drow = dst + (size_t) m * nrows_dst; +#pragma unroll + for (int g = 0; g < NR; ++g) { + const simd n = lane + (uint32_t) (n0 + 16 * g); + const simd_mask<16> ok = n < (uint32_t) nrows; + scatter(drow, n * (uint32_t) sizeof(float), acc[s][g].template select<16, 1>(16 * t), ok); + } + } + } +} + +template +static void launch_pq2_xmx(const uint32_t * wq, const uint16_t * wd, const uint32_t * a8, const float * as, + float * dst, int K, int nrows, int ncols, int nrows_dst, dpct::queue_ptr stream) { + constexpr int TM = 8 * MR; + constexpr int TN = 16 * NR; + + const int n_tiles_m = (ncols + TM - 1) / TM; + const int n_tiles_n = (nrows + TN - 1) / TN; + const int n_tiles = n_tiles_m * n_tiles_n; + + stream->submit([&](sycl::handler & h) { + if constexpr (S > 1) { + const sycl::nd_range<1> nd{ sycl::range<1>((size_t) n_tiles * S), sycl::range<1>(S) }; + h.parallel_for(nd, [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] { + pq2_xmx_thread(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, n_tiles_n, + (int) it.get_group(0), (int) it.get_local_id(0)); + }); + } else { + const size_t global = (size_t) ((n_tiles + PQ2_XMX_WG - 1) / PQ2_XMX_WG) * PQ2_XMX_WG; + const sycl::nd_range<1> nd{ sycl::range<1>(global), sycl::range<1>(PQ2_XMX_WG) }; + h.parallel_for(nd, [=](sycl::nd_item<1> it) [[intel::sycl_explicit_simd]] { + const int tile = (int) it.get_global_id(0); + if (tile < n_tiles) { + pq2_xmx_thread(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, n_tiles_n, tile, 0); + } + }); + } + }); +} + +// K split for a tile count: the smallest power of two reaching the thread target, at least two blocks per thread +template +static void launch_pq2_xmx_split(const uint32_t * wq, const uint16_t * wd, const uint32_t * a8, const float * as, + float * dst, int K, int nrows, int ncols, int nrows_dst, dpct::queue_ptr stream) { + const int nb = K / QK_PQ2_0; + const int n_tiles = ((ncols + 8 * MR - 1) / (8 * MR)) * ((nrows + 16 * NR - 1) / (16 * NR)); + // 32 token tiles (prefill) carry a large SLM reduction, so they split at most in two and only while short + const int target = MR >= 4 ? PQ2_XMX_PREFILL_TARGET_THREADS : PQ2_XMX_TARGET_THREADS; + const int max_split = MR >= 4 ? 2 : 16; + int split = 1; + while (split < max_split && n_tiles * split < target && nb >= 4 * split) { + split *= 2; + } + switch (split) { + case 1: launch_pq2_xmx(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, stream); break; + case 2: launch_pq2_xmx(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, stream); break; + case 4: launch_pq2_xmx(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, stream); break; + case 8: launch_pq2_xmx(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, stream); break; + default: launch_pq2_xmx(wq, wd, a8, as, dst, K, nrows, ncols, nrows_dst, stream); break; + } +} + +} // namespace + +bool ggml_sycl_pq2_xmx_supports_ne0(int64_t ne0) { + return ne0 % QK_PQ2_0 == 0 && (ne0 / QK_PQ2_0) * PQ2_XMX_QS_BYTES >= 64; +} + +bool ggml_sycl_pq2_xmx_reorder(ggml_tensor * src0, dpct::queue_ptr stream) { + GGML_ASSERT((src0->type == GGML_TYPE_PQ2_0 || src0->type == GGML_TYPE_PTQ1_0) && ggml_is_contiguous(src0)); + + const size_t size = ggml_nbytes(src0); + const size_t nblk = (size_t) ggml_nelements(src0) / QK_PQ2_0; + uint8_t * data = (uint8_t *) src0->data; + + void * tmp = sycl::malloc_device(size, *stream); + if (!tmp) { + GGML_LOG_WARN("%s: failed to allocate %zu bytes for the PQ2_0 XMX reorder, skipping it\n", __func__, size); + return false; + } + stream->memcpy(tmp, data, size).wait(); + + uint8_t * qs = data; + sycl::half * d = (sycl::half *) (data + nblk * PQ2_XMX_QS_BYTES); + if (src0->type == GGML_TYPE_PQ2_0) { + stream->parallel_for(sycl::range<1>(nblk), [=](sycl::id<1> i) { + const block_pq2_0 * x = (const block_pq2_0 *) tmp + i; +#pragma unroll + for (int j = 0; j < PQ2_XMX_QS_BYTES; ++j) { + qs[i * PQ2_XMX_QS_BYTES + j] = x->qs[j]; + } + d[i] = x->d; + }).wait(); + } else { + // base-3 trits (value -1..1) become PQ2_0 codes (value + 1), four to a byte, lowest first; + // the caller made sure the buffer holds 34 bytes a block + stream->parallel_for(sycl::range<1>(nblk), [=](sycl::id<1> i) { + const block_ptq1_0 * x = (const block_ptq1_0 *) tmp + i; + for (int j = 0; j < PQ2_XMX_QS_BYTES; ++j) { + uint8_t byte = 0; +#pragma unroll + for (int k = 0; k < 4; ++k) { + byte |= (uint8_t) ((ptq1_0_trit(x, 4 * j + k) + 1) << (2 * k)); + } + qs[i * PQ2_XMX_QS_BYTES + j] = byte; + } + d[i] = x->d; + }).wait(); + } + + sycl::free(tmp, *stream); + return true; +} + +void ggml_sycl_pq2_xmx_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + ggml_tensor * dst) { + GGML_ASSERT((src0->type == GGML_TYPE_PQ2_0 || src0->type == GGML_TYPE_PTQ1_0) && src0->ne[2] == 1 && + src0->ne[3] == 1); + GGML_ASSERT(src1->type == GGML_TYPE_F32 && src1->nb[0] == sizeof(float)); + GGML_ASSERT(dst->type == GGML_TYPE_F32 && ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_sycl_pq2_xmx_supports_ne0(src0->ne[0])); + + const int K = (int) src0->ne[0]; + const int nrows = (int) src0->ne[1]; + const int nb = K / QK_PQ2_0; + const int ne11 = (int) src1->ne[1]; + const int ne12 = (int) src1->ne[2]; + const int ncols = (int) (src1->ne[1] * src1->ne[2] * src1->ne[3]); + + dpct::queue_ptr stream = ctx.stream(); + + // int8 activations (ncols rows of K bytes, 64-byte aligned for 2D loads) and one float scale per 128 values + ggml_sycl_pool_alloc a8_alloc(ctx.pool(), (size_t) ncols * K + 64); + ggml_sycl_pool_alloc as_alloc(ctx.pool(), (size_t) ncols * nb); + int8_t * a8 = (int8_t *) GGML_PAD((uintptr_t) a8_alloc.get(), 64); + float * as = as_alloc.get(); + + { + const char * src1_d = (const char *) src1->data; + const size_t nb11 = src1->nb[1], nb12 = src1->nb[2], nb13 = src1->nb[3]; + // one sub-group per (token, 128-block), 8 values per work-item + stream->parallel_for( + sycl::nd_range<1>(sycl::range<1>((size_t) ncols * nb * 16), sycl::range<1>(16)), + [=](sycl::nd_item<1> it) [[sycl::reqd_sub_group_size(16)]] { + const int grp = (int) it.get_group(0); + const int j = grp / nb; + const int b = grp % nb; + const int l = (int) it.get_local_id(0); + const int i1 = j % ne11; + const int i2 = (j / ne11) % ne12; + const int i3 = j / (ne11 * ne12); + + const float * x = (const float *) (src1_d + i1 * nb11 + i2 * nb12 + i3 * nb13) + b * QK_PQ2_0 + l * 8; + float v[8]; + float amax = 0.0f; +#pragma unroll + for (int i = 0; i < 8; ++i) { + v[i] = x[i]; + amax = sycl::fmax(amax, sycl::fabs(v[i])); + } + amax = sycl::reduce_over_group(it.get_sub_group(), amax, sycl::maximum()); + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + + sycl::vec q; +#pragma unroll + for (int i = 0; i < 8; ++i) { + q[i] = (int8_t) sycl::round(v[i] * id); + } + *(sycl::vec *) (a8 + (size_t) j * K + b * QK_PQ2_0 + l * 8) = q; + if (l == 0) { + as[(size_t) j * nb + b] = d; + } + }); + } + + const uint32_t * wq = (const uint32_t *) src0->data; + const uint16_t * wd = (const uint16_t *) ((const uint8_t *) src0->data + (size_t) nrows * nb * PQ2_XMX_QS_BYTES); + float * dd = (float *) dst->data; + const int nrows_dst = (int) dst->ne[0]; + + if (ncols <= 8) { + launch_pq2_xmx_split<1, 2>(wq, wd, (const uint32_t *) a8, as, dd, K, nrows, ncols, nrows_dst, stream); + } else if (ncols <= 16) { + launch_pq2_xmx_split<2, 2>(wq, wd, (const uint32_t *) a8, as, dd, K, nrows, ncols, nrows_dst, stream); + } else { + launch_pq2_xmx_split<4, 2>(wq, wd, (const uint32_t *) a8, as, dd, K, nrows, ncols, nrows_dst, stream); + } +} + +#else + +bool ggml_sycl_pq2_xmx_supports_ne0(int64_t) { + return false; +} + +bool ggml_sycl_pq2_xmx_reorder(ggml_tensor *, dpct::queue_ptr) { + return false; +} + +void ggml_sycl_pq2_xmx_mul_mat(ggml_backend_sycl_context &, const ggml_tensor *, const ggml_tensor *, ggml_tensor *) { + GGML_ABORT("PQ2_0 XMX path is not built in"); +} + +#endif // __INTEL_LLVM_COMPILER && !GGML_SYCL_NO_PQ2_XMX diff --git a/ggml/src/ggml-sycl/pq2_xmx.hpp b/ggml/src/ggml-sycl/pq2_xmx.hpp new file mode 100644 index 000000000000..4ad3c7493f0b --- /dev/null +++ b/ggml/src/ggml-sycl/pq2_xmx.hpp @@ -0,0 +1,20 @@ +#pragma once + +#include "common.hpp" + +// PQ2_0 in the XMX layout: the weight tensor is rewritten in place, once, into a plane of 32-byte qs blocks +// (row pitch nb * 32 bytes, so every row is 2D-block-load aligned) followed by a plane of fp16 block scales. +// PTQ1_0 weights take the same layout: their base-3 trits are expanded to PQ2_0 codes, which needs 34 bytes a +// block instead of 28, so the buffer type reserves that room for them on devices that use this path. +// Activations are quantized to int8 with one float scale per 128 values, so the four DPAS of a PQ2_0 block +// accumulate in integers before a single float rescale. + +// ne[0] of a PQ2_0 weight the XMX path accepts: a 2D surface needs a row of at least 64 bytes +bool ggml_sycl_pq2_xmx_supports_ne0(int64_t ne0); + +// rewrite src0 (PQ2_0 or PTQ1_0, AoS blocks) into the XMX layout in place +bool ggml_sycl_pq2_xmx_reorder(ggml_tensor * src0, dpct::queue_ptr stream); + +// dst = src0 * src1 for a src0 already in the XMX layout; src1 is f32 with contiguous rows, dst is contiguous +void ggml_sycl_pq2_xmx_mul_mat(ggml_backend_sycl_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, + ggml_tensor * dst);