Commit 411d54851 for llama.cpp
commit 411d548515017ca75c2a12c39659b6e7e1551dd7
Author: losewayy <1134285612@qq.com>
Date: Tue Oct 13 19:03:47 2026 +0800
vulkan : split matmul dispatch when it would exceed maxComputeWorkGroupCount (#29533)
* vulkan : split matmul dispatch when it would exceed maxComputeWorkGroupCount
* vulkan : fix narrowing in chunked matmul dispatch
* vulkan : zero fill the A side of mul_mmq for k steps past the end
The A side skipped the shared memory write for the k steps past end_k, while the B side was already told the bound and zero filled. The A registers are read for all BK_STEP steps, so the stale scale went into the dot product and produced NaN at k = 32.
Pass the bound to block_a_to_shmem and write zeros when it is false, for every DATA_A type, mirroring the B side.
Also assert that the B and D subbuffer sizes are not VK_WHOLE_SIZE before they are adjusted per chunk, and add two batched cases of the chunked shape.
Signed-off-by: losewayy <104015127+losewayy@users.noreply.github.com>
* vulkan : remove redundant matmul tile adjustment
---------
Signed-off-by: losewayy <104015127+losewayy@users.noreply.github.com>
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-common.h b/ggml/src/ggml-vulkan/ggml-vulkan-common.h
index 134078bd2..8922025ed 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-common.h
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-common.h
@@ -94,7 +94,7 @@ void deferred_memset(void * dst, uint32_t val, size_t size, std::vector<vk_stagi
// matmul
vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type type);
-void ggml_vk_matmul(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3, uint32_t padded_n);
+void ggml_vk_matmul(ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, vk_subbuffer&& a, vk_subbuffer&& b, vk_subbuffer&& d, vk_subbuffer&& split_k_buffer, uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3, uint32_t padded_n, uint64_t b_col_bytes, uint64_t d_col_bytes);
bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor);
vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src, const ggml_tensor * dst, ggml_type to);
vk_pipeline ggml_vk_get_quantize_pipeline(ggml_backend_vk_context * ctx, ggml_type type);
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-debug.cpp b/ggml/src/ggml-vulkan/ggml-vulkan-debug.cpp
index 8d8fe6592..97747e866 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan-debug.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan-debug.cpp
@@ -173,7 +173,8 @@ void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t n, size
ctx, subctx, p, ggml_vk_subbuffer(ctx, d_X), ggml_vk_subbuffer(ctx, d_Y), ggml_vk_subbuffer(ctx, d_D), ggml_vk_subbuffer(ctx, ctx->prealloc_split_k),
m, n, k,
k, k, m, k*m, k*n, m*n,
- split_k, batch, batch, batch, 1, 1, n
+ split_k, batch, batch, batch, 1, 1, n,
+ k * sizeof(Y_TYPE), m * sizeof(float)
);
}
ggml_vk_ctx_end(subctx);
@@ -640,7 +641,8 @@ void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t
ctx, subctx, p, { qx_buf, 0, qx_sz }, { qy_buf, 0, qy_sz }, { d_buf, 0, d_sz }, { ctx->prealloc_split_k, 0, ctx->prealloc_size_split_k },
m, n, k,
k, k, m, k*m, k*n, m*n,
- split_k, batch, batch, batch, 1, 1, n
+ split_k, batch, batch, batch, 1, 1, n,
+ (uint64_t)(k / ggml_blck_size(GGML_TYPE_Q8_1)) * ggml_type_size(GGML_TYPE_Q8_1), m * sizeof(float)
);
}
} else {
@@ -649,7 +651,8 @@ void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, size_t
ctx, subctx, p, { qx_buf, 0, qx_sz }, { y_buf, 0, y_sz }, { d_buf, 0, d_sz }, { ctx->prealloc_split_k, 0, ctx->prealloc_size_split_k },
m, n, k,
k, k, m, k*m, k*n, m*n,
- split_k, batch, batch, batch, 1, 1, n
+ split_k, batch, batch, batch, 1, 1, n,
+ k * sizeof(float), m * sizeof(float)
);
}
}
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 05c63b64e..5435171f1 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -5907,18 +5907,34 @@ void ggml_vk_matmul(
uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d,
uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d,
uint32_t split_k, uint32_t batch, uint32_t ne02, uint32_t ne12, uint32_t broadcast2, uint32_t broadcast3,
- uint32_t padded_n) {
+ uint32_t padded_n, uint64_t b_col_bytes, uint64_t d_col_bytes) {
VK_LOG_DEBUG("ggml_vk_matmul(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), split_k: (" << (split_k_buffer.buffer != nullptr ? split_k_buffer.buffer->buffer : VK_NULL_HANDLE) << ", " << split_k_buffer.offset << ", " << split_k_buffer.size << "), m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", split_k: " << split_k << ", batch: " << batch << ", ne02: " << ne02 << ", ne12: " << ne12 << ", broadcast2: " << broadcast2 << ", broadcast3: " << broadcast3 << ", padded_n: " << padded_n << ")");
if (split_k == 1) {
- ggml_pipeline_request_descriptor_sets(ctx, pipeline, CEIL_DIV(batch, ctx->device->properties.limits.maxComputeWorkGroupCount[2]));
-
- uint32_t base_work_group_z = 0;
- while (base_work_group_z < batch) {
- uint32_t groups_z = std::min(batch - base_work_group_z, ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
-
- const vk_mat_mat_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, base_work_group_z, batch, k, ne02, ne12, broadcast2, broadcast3, padded_n };
- ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d }, pc, { m, n, groups_z });
- base_work_group_z += groups_z;
+ // a small tile can still produce more workgroups than the device allows on the n axis:
+ // split the dispatch into n-chunks, shifting the B and D column windows
+ const uint64_t max_n = (uint64_t)ctx->device->properties.limits.maxComputeWorkGroupCount[1] * pipeline->wg_denoms[1];
+ const uint64_t n_chunks = CEIL_DIV(n, max_n);
+ ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_chunks * CEIL_DIV(batch, ctx->device->properties.limits.maxComputeWorkGroupCount[2]));
+
+ for (uint64_t n_base = 0; n_base < n; n_base += max_n) {
+ const uint32_t n_chunk = (uint32_t) std::min(n - n_base, max_n);
+ GGML_ASSERT(n_base * b_col_bytes < b.size && n_base * d_col_bytes < d.size);
+ GGML_ASSERT(b.size != VK_WHOLE_SIZE && d.size != VK_WHOLE_SIZE);
+ vk_subbuffer b_chunk = b;
+ vk_subbuffer d_chunk = d;
+ b_chunk.offset += n_base * b_col_bytes;
+ b_chunk.size -= n_base * b_col_bytes;
+ d_chunk.offset += n_base * d_col_bytes;
+ d_chunk.size -= n_base * d_col_bytes;
+
+ uint32_t base_work_group_z = 0;
+ while (base_work_group_z < batch) {
+ uint32_t groups_z = std::min(batch - base_work_group_z, ctx->device->properties.limits.maxComputeWorkGroupCount[2]);
+
+ const vk_mat_mat_push_constants pc = { m, n_chunk, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, base_work_group_z, batch, k, ne02, ne12, broadcast2, broadcast3, uint32_t(padded_n - n_base) };
+ ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b_chunk, d_chunk }, pc, { m, n_chunk, groups_z });
+ base_work_group_z += groups_z;
+ }
}
return;
}
@@ -6539,6 +6555,11 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
stride_batch_y = src1->nb[0] / ggml_type_size(src1->type);
}
+ // column stride in bytes of the B buffer: raw src1 keeps its own stride, staging buffers are tightly packed
+ const uint64_t b_col_bytes = (y_non_contig || qy_needs_dequant || quantize_y)
+ ? (quantize_y ? (uint64_t)(ne10 / ggml_blck_size(GGML_TYPE_Q8_1)) * ggml_type_size(GGML_TYPE_Q8_1) : (uint64_t)ne10 * sizeof(ggml_fp16_t))
+ : (uint64_t)src1->nb[1];
+
// compute
ggml_vk_matmul(
ctx, subctx, pipeline,
@@ -6546,7 +6567,8 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub
ggml_vk_subbuffer(ctx, d_D, d_buf_offset), { ctx->prealloc_split_k, 0, d_sz * split_k },
ne01, ne11, ne10,
ne10, ne10, stride_d, stride_batch_x, stride_batch_y, stride_batch_d,
- split_k, ne12*ne13, ne02, ne12, r2, r3, padded_n
+ split_k, ne12*ne13, ne02, ne12, r2, r3, padded_n,
+ b_col_bytes, (uint64_t)dst->nb[1]
); // NOLINT
if (x_non_contig || qx_needs_dequant) {
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp
index 2819eb46e..1866386a7 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp
@@ -201,9 +201,7 @@ void main() {
const uint iqs = loadr_a;
[[unroll]] for (uint k_step = 0; k_step < BK_STEP; k_step++) {
- if (block + k_step * BK < end_k) {
- block_a_to_shmem(k_step * BM + buf_ib, ib + k_step, iqs);
- }
+ block_a_to_shmem(k_step * BM + buf_ib, ib + k_step, iqs, block + k_step * BK < end_k);
}
}
[[unroll]] for (uint l = 0; loadc_b + l < BN; l += loadstride_b) {
diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl
index 136d5b745..d9f2f5748 100644
--- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl
+++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl
@@ -7,18 +7,26 @@
// Each iqs value maps to a 32-bit integer
#if defined(DATA_A_Q2_0)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint block_idx = ib / 2;
- const uint byte_idx = (ib & 1u) * 8u + iqs;
- const uint bits = uint(data_a[block_idx].qs[byte_idx]);
- buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
- int8_t(bits & 3u),
- int8_t((bits >> 2u) & 3u),
- int8_t((bits >> 4u) & 3u),
- int8_t(bits >> 6u)));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint block_idx = ib / 2;
+ const uint byte_idx = (ib & 1u) * 8u + iqs;
+ const uint bits = uint(data_a[block_idx].qs[byte_idx]);
+ buf_a[buf_ib].qs[iqs] = pack32(i8vec4(
+ int8_t(bits & 3u),
+ int8_t((bits >> 2u) & 3u),
+ int8_t((bits >> 4u) & 3u),
+ int8_t(bits >> 6u)));
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPE(data_a[block_idx].d);
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+ }
}
}
@@ -43,21 +51,37 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_Q4_0) || defined(DATA_A_Q4_1)
// 2-byte loads for Q4_0 blocks (18 bytes)
// 4-byte loads for Q4_1 blocks (20 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ #ifdef DATA_A_Q4_0
+ buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
+ data_a_packed16[ib].qs[iqs * 2 + 1]));
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+ }
+ #else // DATA_A_Q4_1
+ buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
+ }
+ #endif
+ } else {
#ifdef DATA_A_Q4_0
- buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
- data_a_packed16[ib].qs[iqs * 2 + 1]));
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
- }
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+ }
#else // DATA_A_Q4_1
- buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
- }
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+ }
#endif
+ }
}
void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
@@ -93,23 +117,41 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_Q5_0) || defined(DATA_A_Q5_1)
// 2-byte loads for Q5_0 blocks (22 bytes)
// 4-byte loads for Q5_1 blocks (24 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ #ifdef DATA_A_Q5_0
+ buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
+ data_a_packed16[ib].qs[iqs * 2 + 1]));
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+ buf_a[buf_ib].qh = pack32(u16vec2(data_a_packed16[ib].qh[0], data_a_packed16[ib].qh[1]));
+ }
+ #else // DATA_A_Q5_1
+ buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
+ buf_a[buf_ib].qh = data_a_packed32[ib].qh;
+ }
+ #endif
+ } else {
#ifdef DATA_A_Q5_0
- buf_a[buf_ib].qs[iqs] = pack32(u16vec2(data_a_packed16[ib].qs[iqs * 2],
- data_a_packed16[ib].qs[iqs * 2 + 1]));
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
- buf_a[buf_ib].qh = pack32(u16vec2(data_a_packed16[ib].qh[0], data_a_packed16[ib].qh[1]));
- }
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+ buf_a[buf_ib].qh = 0;
+ }
#else // DATA_A_Q5_1
- buf_a[buf_ib].qs[iqs] = data_a_packed32[ib].qs[iqs];
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib].dm);
- buf_a[buf_ib].qh = data_a_packed32[ib].qh;
- }
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+ buf_a[buf_ib].qh = 0;
+ }
#endif
+ }
}
void block_a_to_registers(const uint reg_ib, const uint buf_ib) {
@@ -148,12 +190,20 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_Q8_0)
// 2-byte loads for Q8_0 blocks (34 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- buf_a[buf_ib].qs[iqs] = pack32(i16vec2(data_a_packed16[ib].qs[iqs * 2],
- data_a_packed16[ib].qs[iqs * 2 + 1]));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ buf_a[buf_ib].qs[iqs] = pack32(i16vec2(data_a_packed16[ib].qs[iqs * 2],
+ data_a_packed16[ib].qs[iqs * 2 + 1]));
+
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPE(data_a_packed16[ib].d);
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPE(0.0f);
+ }
}
}
@@ -180,20 +230,29 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_MXFP4)
// 1-byte loads for mxfp4 blocks (17 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4 ],
- data_a[ib].qs[iqs * 4 + 1],
- data_a[ib].qs[iqs * 4 + 2],
- data_a[ib].qs[iqs * 4 + 3]));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4 ],
+ data_a[ib].qs[iqs * 4 + 1],
+ data_a[ib].qs[iqs * 4 + 2],
+ data_a[ib].qs[iqs * 4 + 3]));
- const u8vec4 i_a0 = unpack8( qs & 0x0F0F0F0F);
- const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F);
+ const u8vec4 i_a0 = unpack8( qs & 0x0F0F0F0F);
+ const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F);
- buf_a[buf_ib].qs[iqs ] = pack32(i8vec4(kvalues_mxfp4[i_a0.x], kvalues_mxfp4[i_a0.y], kvalues_mxfp4[i_a0.z], kvalues_mxfp4[i_a0.w]));
- buf_a[buf_ib].qs[iqs + 4] = pack32(i8vec4(kvalues_mxfp4[i_a1.x], kvalues_mxfp4[i_a1.y], kvalues_mxfp4[i_a1.z], kvalues_mxfp4[i_a1.w]));
+ buf_a[buf_ib].qs[iqs ] = pack32(i8vec4(kvalues_mxfp4[i_a0.x], kvalues_mxfp4[i_a0.y], kvalues_mxfp4[i_a0.z], kvalues_mxfp4[i_a0.w]));
+ buf_a[buf_ib].qs[iqs + 4] = pack32(i8vec4(kvalues_mxfp4[i_a1.x], kvalues_mxfp4[i_a1.y], kvalues_mxfp4[i_a1.z], kvalues_mxfp4[i_a1.w]));
+
+ if (iqs == 0) {
+ buf_a[buf_ib].d = FLOAT_TYPE(e8m0_to_fp32(data_a[ib].e) * 0.5);
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
+ buf_a[buf_ib].qs[iqs + 4] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].d = FLOAT_TYPE(e8m0_to_fp32(data_a[ib].e) * 0.5);
+ if (iqs == 0) {
+ buf_a[buf_ib].d = FLOAT_TYPE(0.0f);
+ }
}
}
@@ -218,19 +277,28 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#endif
#if defined(DATA_A_IQ4_XS)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint ib_k = ib / 8;
- const uint ib32 = ib % 8;
- const uint32_t vui = data_a_packed32[ib_k].qs[4 * ib32 + iqs];
- const i32vec2 qs = iq4nl_to_i8x8(vui);
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint ib_k = ib / 8;
+ const uint ib32 = ib % 8;
+ const uint32_t vui = data_a_packed32[ib_k].qs[4 * ib32 + iqs];
+ const i32vec2 qs = iq4nl_to_i8x8(vui);
+
+ buf_a[buf_ib].qs[iqs ] = qs.x;
+ buf_a[buf_ib].qs[iqs + 4] = qs.y;
- buf_a[buf_ib].qs[iqs ] = qs.x;
- buf_a[buf_ib].qs[iqs + 4] = qs.y;
+ if (iqs == 0) {
+ const uint sl = (data_a_packed32[ib_k].scales_l >> (4 * ib32)) & 0xF;
+ const uint sh = (data_a_packed32[ib_k].scales_h >> (2 * ib32)) & 3;
+ buf_a[buf_ib].d = FLOAT_TYPE(float(data_a[ib_k].d) * float(int(sl | (sh << 4)) - 32));
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
+ buf_a[buf_ib].qs[iqs + 4] = 0;
- if (iqs == 0) {
- const uint sl = (data_a_packed32[ib_k].scales_l >> (4 * ib32)) & 0xF;
- const uint sh = (data_a_packed32[ib_k].scales_h >> (2 * ib32)) & 3;
- buf_a[buf_ib].d = FLOAT_TYPE(float(data_a[ib_k].d) * float(int(sl | (sh << 4)) - 32));
+ if (iqs == 0) {
+ buf_a[buf_ib].d = FLOAT_TYPE(0.0f);
+ }
}
}
@@ -256,24 +324,33 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
// iqs still refers to a 32-bit integer, meaning 0..7 for 32-wide quants
#if defined(DATA_A_Q2_K)
// 4-byte loads for Q2_K blocks (84 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint ib_k = ib / 8;
- const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint ib_k = ib / 8;
+ const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
+
+ const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
+ const uint qs_shift = ((iqs_k % 32) / 8) * 2;
- const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
- const uint qs_shift = ((iqs_k % 32) / 8) * 2;
+ // Repack 4x4 quants into one int
+ const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx ] >> qs_shift) & 0x03030303;
+ const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x03030303;
+ const uint32_t vals2 = (data_a_packed32[ib_k].qs[qs_idx + 2] >> qs_shift) & 0x03030303;
+ const uint32_t vals3 = (data_a_packed32[ib_k].qs[qs_idx + 3] >> qs_shift) & 0x03030303;
- // Repack 4x4 quants into one int
- const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx ] >> qs_shift) & 0x03030303;
- const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x03030303;
- const uint32_t vals2 = (data_a_packed32[ib_k].qs[qs_idx + 2] >> qs_shift) & 0x03030303;
- const uint32_t vals3 = (data_a_packed32[ib_k].qs[qs_idx + 3] >> qs_shift) & 0x03030303;
+ buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 2) | (vals2 << 4) | (vals3 << 6);
- buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 2) | (vals2 << 4) | (vals3 << 6);
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib_k].dm);
+ buf_a[buf_ib].scales = unpack8(uint32_t(data_a_packed16[ib_k].scales[iqs_k / 8])).xy; // vec4 used due to #12147
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
- if (iqs == 0) {
- buf_a[buf_ib].dm = FLOAT_TYPEV2(data_a_packed32[ib_k].dm);
- buf_a[buf_ib].scales = unpack8(uint32_t(data_a_packed16[ib_k].scales[iqs_k / 8])).xy; // vec4 used due to #12147
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+ buf_a[buf_ib].scales = u8vec2(0);
+ }
}
}
@@ -305,35 +382,43 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_Q3_K)
// 2-byte loads for Q3_K blocks (110 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint ib_k = ib / 8;
- const uint hm_idx = iqs * QUANT_R_MMQ;
- const uint iqs_k = (ib % 8) * 8 + hm_idx;
-
- const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
- const uint qs_shift = ((iqs_k % 32) / 8) * 2;
- const uint hm_shift = iqs_k / 8;
-
- // Repack 2x4 quants into one int
- // Add the 3rd bit instead of subtracting it to allow packing the quants
- // vec4 for unpack8 used due to #12147
- const i8vec2 vals00 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 ] >> qs_shift) & uint16_t(0x0303)))).xy |
- unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 ] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
- const i8vec2 vals01 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 1 ] >> qs_shift) & uint16_t(0x0303)))).xy |
- unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 1] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
- const i8vec2 vals10 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 2 ] >> qs_shift) & uint16_t(0x0303)))).xy |
- unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 2] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
- const i8vec2 vals11 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 3 ] >> qs_shift) & uint16_t(0x0303)))).xy |
- unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 3] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
- buf_a[buf_ib].qs[iqs] = pack32(u8vec4(vals00.x, vals00.y, vals01.x, vals01.y)) |
- (pack32(u8vec4(vals10.x, vals10.y, vals11.x, vals11.y)) << 4);
-
- if (iqs == 0) {
- const uint is = iqs_k / 4;
- const i8vec2 scales = i8vec2(unpack8(uint32_t(((data_a_packed16[ib_k].scales[(is % 8 ) / 2] >> (4 * (is / 8))) & 0x0F0F) |
- (((data_a_packed16[ib_k].scales[(8 + (is % 4)) / 2] >> (2 * (is / 4))) & 0x0303) << 4))).xy); // vec4 used due to #12147
-
- buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales - 32));
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint ib_k = ib / 8;
+ const uint hm_idx = iqs * QUANT_R_MMQ;
+ const uint iqs_k = (ib % 8) * 8 + hm_idx;
+
+ const uint qs_idx = (iqs_k / 32) * 8 + (iqs_k % 8);
+ const uint qs_shift = ((iqs_k % 32) / 8) * 2;
+ const uint hm_shift = iqs_k / 8;
+
+ // Repack 2x4 quants into one int
+ // Add the 3rd bit instead of subtracting it to allow packing the quants
+ // vec4 for unpack8 used due to #12147
+ const i8vec2 vals00 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 ] >> qs_shift) & uint16_t(0x0303)))).xy |
+ unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 ] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+ const i8vec2 vals01 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 1 ] >> qs_shift) & uint16_t(0x0303)))).xy |
+ unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 1] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+ const i8vec2 vals10 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 2 ] >> qs_shift) & uint16_t(0x0303)))).xy |
+ unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 2] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+ const i8vec2 vals11 = unpack8(int32_t(int16_t((data_a_packed16[ib_k].qs[qs_idx * 2 + 3 ] >> qs_shift) & uint16_t(0x0303)))).xy |
+ unpack8(int32_t(int16_t(((data_a_packed16[ib_k].hmask[hm_idx * 2 + 3] >> hm_shift) & uint16_t(0x0101))) << 2)).xy;
+ buf_a[buf_ib].qs[iqs] = pack32(u8vec4(vals00.x, vals00.y, vals01.x, vals01.y)) |
+ (pack32(u8vec4(vals10.x, vals10.y, vals11.x, vals11.y)) << 4);
+
+ if (iqs == 0) {
+ const uint is = iqs_k / 4;
+ const i8vec2 scales = i8vec2(unpack8(uint32_t(((data_a_packed16[ib_k].scales[(is % 8 ) / 2] >> (4 * (is / 8))) & 0x0F0F) |
+ (((data_a_packed16[ib_k].scales[(8 + (is % 4)) / 2] >> (2 * (is / 4))) & 0x0303) << 4))).xy); // vec4 used due to #12147
+
+ buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales - 32));
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
+
+ if (iqs == 0) {
+ buf_a[buf_ib].d_scales = FLOAT_TYPEV2(0.0f);
+ }
}
}
@@ -371,39 +456,47 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K)
// 4-byte loads for Q4_K blocks (144 bytes) and Q5_K blocks (176 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint ib_k = ib / 8;
- const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint ib_k = ib / 8;
+ const uint iqs_k = (ib % 8) * 8 + iqs * QUANT_R_MMQ;
- const uint qs_idx = (iqs_k / 16) * 8 + (iqs_k % 8);
- const uint qs_shift = ((iqs_k % 16) / 8) * 4;
+ const uint qs_idx = (iqs_k / 16) * 8 + (iqs_k % 8);
+ const uint qs_shift = ((iqs_k % 16) / 8) * 4;
- // Repack 2x4 quants into one int
-#if defined(DATA_A_Q4_K)
- const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx ] >> qs_shift) & 0x0F0F0F0F;
- const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x0F0F0F0F;
+ // Repack 2x4 quants into one int
+ #if defined(DATA_A_Q4_K)
+ const uint32_t vals0 = (data_a_packed32[ib_k].qs[qs_idx ] >> qs_shift) & 0x0F0F0F0F;
+ const uint32_t vals1 = (data_a_packed32[ib_k].qs[qs_idx + 1] >> qs_shift) & 0x0F0F0F0F;
- buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 4);
-#else // defined(DATA_A_Q5_K)
- const uint qh_idx = iqs * QUANT_R_MMQ;
- const uint qh_shift = iqs_k / 8;
+ buf_a[buf_ib].qs[iqs] = vals0 | (vals1 << 4);
+ #else // defined(DATA_A_Q5_K)
+ const uint qh_idx = iqs * QUANT_R_MMQ;
+ const uint qh_shift = iqs_k / 8;
- buf_a[buf_ib].qs[iqs] = int32_t(((data_a_packed32[ib_k].qs[qs_idx] >> qs_shift) & 0x0F0F0F0F) |
- (((data_a_packed32[ib_k].qh[qh_idx] >> qh_shift) & 0x01010101) << 4));
-#endif
+ buf_a[buf_ib].qs[iqs] = int32_t(((data_a_packed32[ib_k].qs[qs_idx] >> qs_shift) & 0x0F0F0F0F) |
+ (((data_a_packed32[ib_k].qh[qh_idx] >> qh_shift) & 0x01010101) << 4));
+ #endif
- if (iqs == 0) {
- // Scale index
- const uint is = iqs_k / 8;
- u8vec2 scale_dm;
- if (is < 4) {
- scale_dm = u8vec2(data_a[ib_k].scales[is] & 0x3F, data_a[ib_k].scales[is + 4] & 0x3F);
- } else {
- scale_dm = u8vec2((data_a[ib_k].scales[is+4] & 0xF) | ((data_a[ib_k].scales[is-4] & 0xC0) >> 2),
- (data_a[ib_k].scales[is+4] >> 4) | ((data_a[ib_k].scales[is ] & 0xC0) >> 2));
+ if (iqs == 0) {
+ // Scale index
+ const uint is = iqs_k / 8;
+ u8vec2 scale_dm;
+ if (is < 4) {
+ scale_dm = u8vec2(data_a[ib_k].scales[is] & 0x3F, data_a[ib_k].scales[is + 4] & 0x3F);
+ } else {
+ scale_dm = u8vec2((data_a[ib_k].scales[is+4] & 0xF) | ((data_a[ib_k].scales[is-4] & 0xC0) >> 2),
+ (data_a[ib_k].scales[is+4] >> 4) | ((data_a[ib_k].scales[is ] & 0xC0) >> 2));
+ }
+
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(vec2(data_a_packed32[ib_k].dm) * vec2(scale_dm));
}
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
- buf_a[buf_ib].dm = FLOAT_TYPEV2(vec2(data_a_packed32[ib_k].dm) * vec2(scale_dm));
+ if (iqs == 0) {
+ buf_a[buf_ib].dm = FLOAT_TYPEV2(0.0f);
+ }
}
}
@@ -434,27 +527,35 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_Q6_K)
// 2-byte loads for Q6_K blocks (210 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint ib_k = ib / 8;
- const uint iqs_k = (ib % 8) * 8 + iqs;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint ib_k = ib / 8;
+ const uint iqs_k = (ib % 8) * 8 + iqs;
+
+ const uint ql_idx = (iqs_k / 32) * 16 + iqs_k % 16;
+ const uint ql_shift = ((iqs_k % 32) / 16) * 4;
- const uint ql_idx = (iqs_k / 32) * 16 + iqs_k % 16;
- const uint ql_shift = ((iqs_k % 32) / 16) * 4;
+ const uint qh_idx = (iqs_k / 32) * 8 + iqs;
+ const uint qh_shift = ((iqs_k % 32) / 8) * 2;
- const uint qh_idx = (iqs_k / 32) * 8 + iqs;
- const uint qh_shift = ((iqs_k % 32) / 8) * 2;
+ const i8vec2 vals00 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 ] >> ql_shift) & uint16_t(0x0F0F))).xy |
+ unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 ] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
+ const i8vec2 vals01 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 + 1] >> ql_shift) & uint16_t(0x0F0F))).xy |
+ unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 + 1] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
+ buf_a[buf_ib].qs[iqs] = pack32(i8vec4(vals00.x, vals00.y, vals01.x, vals01.y));
- const i8vec2 vals00 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 ] >> ql_shift) & uint16_t(0x0F0F))).xy |
- unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 ] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
- const i8vec2 vals01 = (unpack8(int32_t((data_a_packed16[ib_k].ql[ql_idx * 2 + 1] >> ql_shift) & uint16_t(0x0F0F))).xy |
- unpack8(int32_t(((data_a_packed16[ib_k].qh[qh_idx * 2 + 1] >> qh_shift) & uint16_t(0x0303)) << 4)).xy) - int8_t(32);
- buf_a[buf_ib].qs[iqs] = pack32(i8vec4(vals00.x, vals00.y, vals01.x, vals01.y));
+ if (iqs == 0) {
+ const uint is = iqs_k / 4;
+ const i8vec2 scales = unpack8(int32_t(data_a_packed16[ib_k].scales[is / 2])).xy;
- if (iqs == 0) {
- const uint is = iqs_k / 4;
- const i8vec2 scales = unpack8(int32_t(data_a_packed16[ib_k].scales[is / 2])).xy;
+ buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales));
+ }
+ } else {
+ buf_a[buf_ib].qs[iqs] = 0;
- buf_a[buf_ib].d_scales = FLOAT_TYPEV2(float(data_a_packed16[ib_k].d) * vec2(scales));
+ if (iqs == 0) {
+ buf_a[buf_ib].d_scales = FLOAT_TYPEV2(0.0f);
+ }
}
}
@@ -491,32 +592,41 @@ ACC_TYPE mmq_dot_product(const uint ib_a) {
#if defined(DATA_A_IQ3_S)
// 2-byte loads for IQ3_S blocks (110 bytes)
-void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) {
- const uint ib_k = ib / 8;
- const uint ib32 = ib % 8;
+void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs, const bool is_in_bounds) {
+ if (is_in_bounds) {
+ const uint ib_k = ib / 8;
+ const uint ib32 = ib % 8;
- // grid indices for qs[2 * iqs] and qs[2 * iqs + 1]
- const uint qs = uint(data_a_packed16[ib_k].qs[ib32 * 4 + iqs]);
- // their two high index bits
- const uint qh = uint(data_a_packed16[ib_k].qh[ib32 / 2]) >> ((ib32 & 1) * 8 + 2 * iqs);
- // one sign bit per value, 8 values
- const uint signs = uint(data_a_packed16[ib_k].signs[ib32 * 2 + iqs / 2]) >> ((iqs & 1) * 8);
+ // grid indices for qs[2 * iqs] and qs[2 * iqs + 1]
+ const uint qs = uint(data_a_packed16[ib_k].qs[ib32 * 4 + iqs]);
+ // their two high index bits
+ const uint qh = uint(data_a_packed16[ib_k].qh[ib32 / 2]) >> ((ib32 & 1) * 8 + 2 * iqs);
+ // one sign bit per value, 8 values
+ const uint signs = uint(data_a_packed16[ib_k].signs[ib32 * 2 + iqs / 2]) >> ((iqs & 1) * 8);
- // grid holds 4 values of 1..15, one per byte
- const ivec4 vals0 = ivec4(unpack8(iq3s_grid[( qs & 0xFF) | ((qh & 1) << 8)]));
- const ivec4 vals1 = ivec4(unpack8(iq3s_grid[((qs >> 8) & 0xFF) | ((qh & 2) << 7)]));
+ // grid holds 4 values of 1..15, one per byte
+ const ivec4 vals0 = ivec4(unpack8(iq3s_grid[( qs & 0xFF) | ((qh & 1) << 8)]));
+ const ivec4 vals1 = ivec4(unpack8(iq3s_grid[((qs >> 8) & 0xFF) | ((qh & 2) << 7)]));
- // negate with (v ^ -s) - -s to avoid branches
- const ivec4 m0 = -(ivec4(signs, signs >> 1, signs >> 2, signs >> 3) & 1);
- const ivec4 m1 = -(ivec4(signs >> 4, signs >> 5, signs >> 6, signs >> 7) & 1);
+ // negate with (v ^ -s) - -s to avoid branches
+ const ivec4 m0 = -(ivec4(signs, signs >> 1, signs >> 2, signs >> 3) & 1);
+ const ivec4 m1 = -(ivec4(signs >> 4, signs >> 5, signs >> 6, signs >> 7) & 1);
- buf_a[buf_ib].qs[2 * iqs ] = pack32(i8vec4((vals0 ^ m0) - m0));
- buf_a[buf_ib].qs[2 * iqs + 1] = pack32(i8vec4((vals1 ^ m1) - m1));
+ buf_a[buf_ib].qs[2 * iqs ] = pack32(i8vec4((vals0 ^ m0) - m0));
+ buf_a[buf_ib].qs[2 * iqs + 1] = pack32(i8vec4((vals1 ^ m1) - m1));
- if (iqs == 0) {
- const uint scale = (uint(data_a_packed16[ib_k].scales[ib32 / 4]) >> ((ib32 & 3) * 4)) & 0xF;
+ if (iqs == 0) {
+ const uint scale = (uint(data_a_packed16[ib_k].scales[ib32 / 4]) >> ((ib32 & 3) * 4)) & 0xF;
+
+ buf_a[buf_ib].d = FLOAT_TYPE(float(data_a_packed16[ib_k].d) * float(1 + 2 * scale));
+ }
+ } else {
+ buf_a[buf_ib].qs[2 * iqs] = 0;
+ buf_a[buf_ib].qs[2 * iqs + 1] = 0;
- buf_a[buf_ib].d = FLOAT_TYPE(float(data_a_packed16[ib_k].d) * float(1 + 2 * scale));
+ if (iqs == 0) {
+ buf_a[buf_ib].d = FLOAT_TYPE(0.0f);
+ }
}
}
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index b54abbea8..0080c7e16 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -10723,6 +10723,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 1, 509, 2051, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 1, 509, 2051, {1, 1}, {1, 1}));
+ // testing very large N
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F32, 16, 2105344, 32, {1, 1}, {1, 1}));
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 16, 8400000, 32, {1, 1}, {1, 1}));
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 16, 8400000, 32, {1, 1}, {1, 1}));
+ // same shapes with a repeat in dim 3
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F16, GGML_TYPE_F16, 16, 8400000, 32, {1, 1}, {2, 1}));
+ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 16, 8400000, 32, {1, 1}, {2, 1}));
+
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 31, 509, 2051, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_F32, GGML_TYPE_F32, 32, 509, 2112, {1, 1}, {1, 1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 32, 509, 2112, {1, 1}, {1, 1}));