sycl: support reordered Q4_K/Q5_K/Q6_K MoE MUL_MAT_ID (#24452)

* sycl: support reordered Q4_K and Q5_K MoE MUL_MAT_ID

Extend reordered-weight handling to fused MoE MUL_MAT_ID for Q4_K and Q5_K expert tensors and add Q5_K reordered DMMV coverage. Unsupported 3D reorder cases now fall back instead of aborting.

* sycl: extend MoE reorder to Q6_K mul_mat_id
This commit is contained in:
Frosty40
2026-06-16 08:35:00 +03:00
committed by GitHub
parent fdd109883d
commit ac79caa7ce
4 changed files with 465 additions and 10 deletions
+200 -9
View File
@@ -3685,6 +3685,149 @@ static bool reorder_qw_q4_k(uint8_t * data_device, size_t size, size_t offset, d
return true;
}
// Reorder each expert slice into a self-contained SoA layout.
static bool reorder_qw_q4_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
GGML_ASSERT(expert_bytes % sizeof(block_q4_K) == 0);
const int blocks_per_expert = (int) (expert_bytes / sizeof(block_q4_K));
const size_t total_bytes = expert_bytes * (size_t) n_expert;
sycl_reorder_temp_buffer tmp(stream, total_bytes);
if (!tmp) {
GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
return false;
}
uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
sycl::event copy_event;
SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
if (!g_ggml_sycl_use_async_mem_op) {
copy_event.wait();
}
const int total_blocks = blocks_per_expert * (int) n_expert;
auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
const int gb = gb_;
const int e = gb / blocks_per_expert;
const int ib = gb % blocks_per_expert;
const block_q4_K * x = (const block_q4_K *) (tmp_buf + (size_t) e * expert_bytes);
uint8_t * base = data_device + (size_t) e * expert_bytes;
auto * qs_ptr = base;
auto * scales_ptr = qs_ptr + QK_K / 2 * blocks_per_expert;
auto * dm_ptr = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * blocks_per_expert);
for (int j = 0; j < QK_K / 2; ++j) {
qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
}
for (int j = 0; j < K_SCALE_SIZE; ++j) {
scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
}
dm_ptr[ib] = x[ib].dm;
});
if (!g_ggml_sycl_use_async_mem_op) {
reorder_event.wait_and_throw();
}
return true;
}
// Reorder each Q5_K expert slice into [qs][qh][scales][dm].
static bool reorder_qw_q5_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
GGML_ASSERT(expert_bytes % sizeof(block_q5_K) == 0);
const int blocks_per_expert = (int) (expert_bytes / sizeof(block_q5_K));
const size_t total_bytes = expert_bytes * (size_t) n_expert;
sycl_reorder_temp_buffer tmp(stream, total_bytes);
if (!tmp) {
GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
return false;
}
uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
sycl::event copy_event;
SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
if (!g_ggml_sycl_use_async_mem_op) {
copy_event.wait();
}
const int total_blocks = blocks_per_expert * (int) n_expert;
auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
const int gb = gb_;
const int e = gb / blocks_per_expert;
const int ib = gb % blocks_per_expert;
const block_q5_K * x = (const block_q5_K *) (tmp_buf + (size_t) e * expert_bytes);
uint8_t * base = data_device + (size_t) e * expert_bytes;
auto * qs_ptr = base;
auto * qh_ptr = qs_ptr + (QK_K / 2) * blocks_per_expert;
auto * scales_ptr = qh_ptr + (QK_K / 8) * blocks_per_expert;
auto * dm_ptr = (sycl::half2 *) (scales_ptr + K_SCALE_SIZE * blocks_per_expert);
for (int j = 0; j < QK_K / 2; ++j) {
qs_ptr[ib * (QK_K / 2) + j] = x[ib].qs[j];
}
for (int j = 0; j < QK_K / 8; ++j) {
qh_ptr[ib * (QK_K / 8) + j] = x[ib].qh[j];
}
for (int j = 0; j < K_SCALE_SIZE; ++j) {
scales_ptr[ib * K_SCALE_SIZE + j] = x[ib].scales[j];
}
dm_ptr[ib] = x[ib].dm;
});
if (!g_ggml_sycl_use_async_mem_op) {
reorder_event.wait_and_throw();
}
return true;
}
// Reorder each Q6_K expert slice into [ql][qh][scales][d].
static bool reorder_qw_q6_k_moe(uint8_t * data_device, size_t expert_bytes, int64_t n_expert, dpct::queue_ptr stream) {
GGML_ASSERT(expert_bytes % sizeof(block_q6_K) == 0);
const int blocks_per_expert = (int) (expert_bytes / sizeof(block_q6_K));
const size_t total_bytes = expert_bytes * (size_t) n_expert;
sycl_reorder_temp_buffer tmp(stream, total_bytes);
if (!tmp) {
GGML_LOG_WARN("%s: failed to allocate %zu bytes for reorder temp buffer, skipping reorder\n", __func__, total_bytes);
return false;
}
uint8_t * tmp_buf = static_cast<uint8_t *>(tmp.ptr);
sycl::event copy_event;
SYCL_CHECK(CHECK_TRY_ERROR(copy_event = stream->memcpy(tmp_buf, data_device, total_bytes)));
if (!g_ggml_sycl_use_async_mem_op) {
copy_event.wait();
}
const int total_blocks = blocks_per_expert * (int) n_expert;
auto reorder_event = stream->parallel_for(total_blocks, [=](auto gb_) {
const int gb = gb_;
const int e = gb / blocks_per_expert;
const int ib = gb % blocks_per_expert;
const block_q6_K * x = (const block_q6_K *) (tmp_buf + (size_t) e * expert_bytes);
uint8_t * base = data_device + (size_t) e * expert_bytes;
auto * ql_ptr = base;
auto * qh_ptr = ql_ptr + (QK_K / 2) * blocks_per_expert;
auto * scales_ptr = qh_ptr + (QK_K / 4) * blocks_per_expert;
auto * d_ptr = (sycl::half *) (scales_ptr + (QK_K / 16) * blocks_per_expert);
for (int j = 0; j < QK_K / 2; ++j) {
ql_ptr[ib * (QK_K / 2) + j] = x[ib].ql[j];
}
for (int j = 0; j < QK_K / 4; ++j) {
qh_ptr[ib * (QK_K / 4) + j] = x[ib].qh[j];
}
for (int j = 0; j < QK_K / 16; ++j) {
scales_ptr[ib * (QK_K / 16) + j] = x[ib].scales[j];
}
d_ptr[ib] = x[ib].d;
});
if (!g_ggml_sycl_use_async_mem_op) {
reorder_event.wait_and_throw();
}
return true;
}
static bool reorder_qw_q3_k(uint8_t * data_device, size_t size, size_t offset, dpct::queue_ptr stream) {
GGML_ASSERT(size % sizeof(block_q3_K) == 0);
GGML_ASSERT(offset % sizeof(block_q3_K) == 0);
@@ -3840,6 +3983,22 @@ static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) {
size_t nrows = src0->ne[1];
size_t size = ggml_nbytes(src0);
// MoE expert weights are addressed per expert via nb[2], so each slice must
// remain self-contained after reorder.
if (src0->ne[2] > 1) {
GGML_ASSERT((size_t) size == (size_t) src0->ne[2] * src0->nb[2]);
switch (src0->type) {
case GGML_TYPE_Q4_K:
return reorder_qw_q4_k_moe(data_device, src0->nb[2], src0->ne[2], stream);
case GGML_TYPE_Q5_K:
return reorder_qw_q5_k_moe(data_device, src0->nb[2], src0->ne[2], stream);
case GGML_TYPE_Q6_K:
return reorder_qw_q6_k_moe(data_device, src0->nb[2], src0->ne[2], stream);
default:
return false;
}
}
switch (src0->type) {
case GGML_TYPE_Q4_0:
return reorder_qw_q4_0(data_device, ncols, nrows, size, 0, stream);
@@ -3854,7 +4013,6 @@ static bool reorder_qw(const ggml_tensor * src0, dpct::queue_ptr stream) {
case GGML_TYPE_Q6_K:
return reorder_qw_q6_k(data_device, size, 0, stream);
default:
GGML_ABORT("reorder_qw() called with unsupported type");
return false;
}
}
@@ -3902,6 +4060,23 @@ static void opt_for_reorder(ggml_backend_sycl_context * ctx, const ggml_tensor *
}
}
// Lazily reorder supported MoE expert weights once their fused path is used.
static void opt_for_reorder_id(ggml_backend_sycl_context * ctx, const ggml_tensor * src0) {
if (g_ggml_sycl_disable_optimize || !ctx->opt_feature.reorder) {
return;
}
if (src0->type != GGML_TYPE_Q4_K && src0->type != GGML_TYPE_Q5_K && src0->type != GGML_TYPE_Q6_K) {
return;
}
ggml_tensor_extra_gpu * extra = static_cast<ggml_tensor_extra_gpu *>(src0->extra);
if (!extra || extra->optimized_feature.reorder) {
return;
}
if (reorder_qw(src0, ctx->stream())) {
extra->optimized_feature.reorder = true;
}
}
static bool can_use_dequantize_mul_mat_vec(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
// The F16/BF16 qk=1 kernel iterates with stride 2*DMMV_X, requiring ne[0] to be
@@ -4067,11 +4242,6 @@ static bool ggml_sycl_mul_mat_id_mmvq_fused(
if (ne10 != src0->ne[0] || ne10 % QK8_1 != 0) return false;
if (!ggml_is_contiguous(src1)) return false;
// Reorder layout not supported; fall back.
const ggml_tensor_extra_gpu * src0_extra =
static_cast<const ggml_tensor_extra_gpu *>(src0->extra);
if (src0_extra && src0_extra->optimized_feature.reorder) return false;
const int64_t n_ids_per_group = ids->ne[0];
if (ids->ne[1] != 1) return false;
if (ne11 != 1 && ne11 != n_ids_per_group) return false;
@@ -4081,16 +4251,37 @@ static bool ggml_sycl_mul_mat_id_mmvq_fused(
const int n_experts_used = (int) n_ids_per_group;
const int nrows = (int) src0->ne[1];
// Lazily reorder the (Q4_K) expert weights into a per-expert SoA layout, then run the reorder
// GEMV. Placed after the bail checks so a non-dispatchable op does not pay the reorder cost.
opt_for_reorder_id(&ctx, src0);
const ggml_tensor_extra_gpu * src0_extra =
static_cast<const ggml_tensor_extra_gpu *>(src0->extra);
const bool use_reorder = src0_extra && src0_extra->optimized_feature.reorder;
ggml_sycl_pool_alloc<char> src1_q8_alloc(ctx.pool(),
(size_t) ne11 * src1_padded_cols * sizeof(block_q8_1) / QK8_1);
char * src1_ddq = src1_q8_alloc.get();
quantize_row_q8_1_sycl<quantize_q8_1>(
(const float *) src1->data, src1_ddq, (int) ne10, (int) ne11,
src1_padded_cols, stream);
if (use_reorder) {
quantize_row_q8_1_sycl<quantize_and_reorder_q8_1_soa>(
(const float *) src1->data, src1_ddq, (int) ne10, (int) ne11,
src1_padded_cols, stream);
} else {
quantize_row_q8_1_sycl<quantize_q8_1>(
(const float *) src1->data, src1_ddq, (int) ne10, (int) ne11,
src1_padded_cols, stream);
}
const size_t bytes_per_qrow = (size_t) src1_padded_cols * sizeof(block_q8_1) / QK8_1;
const size_t src1_row_stride = (ne11 == 1) ? 0 : bytes_per_qrow;
if (use_reorder) {
return ggml_sycl_mul_mat_vec_q_id_reorder(
src0->type, src0->data, src1_ddq, (const int32_t *) ids->data,
(float *) dst->data, (int) ne10, nrows, n_experts_used,
/*expert_weight_stride=*/ src0->nb[2],
/*dst_row_stride=*/ dst->nb[1],
src1_row_stride, stream);
}
return ggml_sycl_mul_mat_vec_q_id(
src0->type, src0->data, src1_ddq, (const int32_t *) ids->data,
(float *) dst->data, (int) ne10, nrows, n_experts_used,