* Add mat-vec fast path of MUL_MAT_ID.
* Add shared accumulation vec logic and the other types supports.
* Add i-quant mat-mat for MUL_MAT_ID and fix some parts
* Remove n_experts from shader_lib_context.
}
const size_t base_q_bytes = (key.head_dim_qk + key.head_dim_v) * q_tile * GGML_WEBGPU_F16_SIZE_BYTES +
2 * q_tile * GGML_WEBGPU_F32_SIZE_BYTES;
- size_t bytes_per_kv = 0;
+ size_t bytes_per_kv = 0;
if (!key.kv_direct) {
bytes_per_kv += std::max(key.head_dim_qk, key.head_dim_v);
}
(v_offset_elems % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0u);
const bool kv_vec_type_supported =
K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q4_0 || K->type == GGML_TYPE_Q8_0;
- const bool use_vec = context.supports_subgroups && (context.src0->ne[1] < 20) && (context.src0->ne[0] % 32 == 0) &&
- (context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) &&
- kv_vec_type_supported && (K->type != GGML_TYPE_F16 || f16_vec4_aligned) &&
- (context.src2->type == K->type);
+ const bool use_vec = context.supports_subgroups && (context.src0->ne[1] < 20) && (context.src0->ne[0] % 32 == 0) &&
+ (context.src2->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) &&
+ kv_vec_type_supported && (K->type != GGML_TYPE_F16 || f16_vec4_aligned) &&
+ (context.src2->type == K->type);
const bool use_tile = context.supports_subgroups && !context.supports_subgroup_matrix && K->type == GGML_TYPE_F16 &&
V->type == GGML_TYPE_F16 && f16_vec4_aligned &&
(context.src0->ne[0] % GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH == 0) &&
struct ggml_webgpu_mul_mat_id_pipeline_key {
ggml_type src0_type;
ggml_type src1_type;
+ uint32_t n_experts;
+ int vectorized;
bool operator==(const ggml_webgpu_mul_mat_id_pipeline_key & other) const {
- return src0_type == other.src0_type && src1_type == other.src1_type;
+ return src0_type == other.src0_type && src1_type == other.src1_type && n_experts == other.n_experts &&
+ vectorized == other.vectorized;
}
};
size_t seed = 0;
ggml_webgpu_hash_combine(seed, key.src0_type);
ggml_webgpu_hash_combine(seed, key.src1_type);
+ ggml_webgpu_hash_combine(seed, key.n_experts);
+ ggml_webgpu_hash_combine(seed, key.vectorized);
return seed;
}
};
std::unordered_map<int, webgpu_pipeline> mul_mat_id_gather_pipelines; // key is fixed
std::unordered_map<ggml_webgpu_mul_mat_id_pipeline_key, webgpu_pipeline, ggml_webgpu_mul_mat_id_pipeline_key_hash>
mul_mat_id_pipelines; // src0_type/src1_type
+ std::unordered_map<ggml_webgpu_mul_mat_id_pipeline_key, webgpu_pipeline, ggml_webgpu_mul_mat_id_pipeline_key_hash>
+ mul_mat_id_vec_pipelines; // src0_type/src1_type
std::unordered_map<ggml_webgpu_set_rows_pipeline_key, webgpu_pipeline, ggml_webgpu_set_rows_pipeline_key_hash>
set_rows_pipelines;
key.type = context.dst->type;
key.d_state = (int) context.src0->ne[0];
key.xbc_overlap = ggml_webgpu_tensor_overlap(context.src1, context.src4) &&
- ggml_webgpu_tensor_overlap(context.src1, context.src5);
+ ggml_webgpu_tensor_overlap(context.src1, context.src5);
auto it = ssm_scan_pipelines.find(key);
if (it != ssm_scan_pipelines.end()) {
ggml_webgpu_mul_mat_vec_pipeline_key key = {};
key.src0_type = context.src0->type;
key.src1_type = context.src1->type;
- key.vectorized = (context.src0->ne[0] % 4 == 0 &&
+ key.vectorized = (context.src0->ne[0] % 4 == 0 &&
(context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ?
- 1 :
- 0;
+ 1 :
+ 0;
auto it = mul_mat_vec_pipelines.find(key);
if (it != mul_mat_vec_pipelines.end()) {
ggml_webgpu_mul_mat_id_pipeline_key key = {};
key.src0_type = context.src0->type;
key.src1_type = context.src1->type;
+ key.n_experts = context.src0->ne[2];
+ key.vectorized = (context.src0->ne[0] % 4 == 0 && context.src0->ne[1] % 4 == 0 &&
+ (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ?
+ 1 :
+ 0;
auto it = mul_mat_id_pipelines.find(key);
if (it != mul_mat_id_pipelines.end()) {
switch (context.src0->type) {
case GGML_TYPE_F32:
defines.push_back("SRC0_INNER_TYPE=f32");
- defines.push_back("FLOAT");
defines.push_back("INIT_SRC0_SHMEM_FLOAT");
defines.push_back("INIT_SRC1_SHMEM_FLOAT");
variant += "_f32";
break;
case GGML_TYPE_F16:
defines.push_back("SRC0_INNER_TYPE=f16");
- defines.push_back("FLOAT");
defines.push_back("INIT_SRC0_SHMEM_FLOAT");
defines.push_back("INIT_SRC1_SHMEM_FLOAT");
variant += "_f16";
defines.push_back("U32_DEQUANT_HELPERS");
defines.push_back("SRC0_INNER_TYPE=u32");
+ switch (context.src0->type) {
+ case GGML_TYPE_IQ1_S:
+ case GGML_TYPE_IQ1_M:
+ case GGML_TYPE_IQ4_NL:
+ case GGML_TYPE_IQ4_XS:
+ defines.push_back(type_upper + "_GRID");
+ break;
+ case GGML_TYPE_IQ2_XXS:
+ case GGML_TYPE_IQ2_XS:
+ case GGML_TYPE_IQ2_S:
+ case GGML_TYPE_IQ3_XXS:
+ case GGML_TYPE_IQ3_S:
+ defines.push_back(type_upper + "_GRID");
+ defines.push_back(type_upper + "_TABLES");
+ break;
+ default:
+ break;
+ }
+
variant += std::string("_") + src0_name;
break;
}
}
- defines.push_back("SCALAR");
+ // VEC/SCALAR controls
+ defines.push_back(key.vectorized ? "VEC" : "SCALAR");
// mul_mat_id is register-tile only.
const uint32_t tile_k =
return mul_mat_id_pipelines[key];
}
+ webgpu_pipeline get_mul_mat_id_vec_pipeline(const ggml_webgpu_shader_lib_context & context) {
+ ggml_webgpu_mul_mat_id_pipeline_key key = {};
+ key.src0_type = context.src0->type;
+ key.src1_type = context.src1->type;
+ key.n_experts = context.src0->ne[2];
+ key.vectorized = (context.src0->ne[0] % 4 == 0 &&
+ (context.src0->type == GGML_TYPE_F32 || context.src0->type == GGML_TYPE_F16)) ?
+ 1 :
+ 0;
+
+ auto it = mul_mat_id_vec_pipelines.find(key);
+ if (it != mul_mat_id_vec_pipelines.end()) {
+ return it->second;
+ }
+
+ std::vector<std::string> defines;
+ std::string variant = "mul_mat_id_vec";
+ const char * shader_src = wgsl_mul_mat_id_vec;
+
+ // src1 type
+ switch (context.src1->type) {
+ case GGML_TYPE_F32:
+ defines.push_back("SRC1_INNER_TYPE=f32");
+ break;
+ case GGML_TYPE_F16:
+ defines.push_back("SRC1_INNER_TYPE=f16");
+ break;
+ default:
+ GGML_ABORT("Unsupported src1 type for mul_mat fast shader");
+ }
+
+ // src0 type
+ switch (context.src0->type) {
+ case GGML_TYPE_F32:
+ defines.push_back("SRC0_INNER_TYPE=f32");
+ defines.push_back("MUL_ACC_FLOAT");
+ variant += "_f32";
+ break;
+ case GGML_TYPE_F16:
+ defines.push_back("SRC0_INNER_TYPE=f16");
+ defines.push_back("MUL_ACC_FLOAT");
+ variant += "_f16";
+ break;
+ default:
+ {
+ // Quantized types: use helpers but accumulate in f16
+ const struct ggml_type_traits * src0_traits = ggml_get_type_traits(context.src0->type);
+ std::string src0_name = src0_traits->type_name;
+ std::string type_upper = src0_name;
+ variant += "_" + src0_name;
+ std::transform(type_upper.begin(), type_upper.end(), type_upper.begin(), ::toupper);
+
+ defines.push_back("BYTE_HELPERS");
+ defines.push_back("MUL_ACC_" + type_upper);
+ defines.push_back("U32_DEQUANT_HELPERS");
+ defines.push_back("SRC0_INNER_TYPE=u32");
+ switch (context.src0->type) {
+ case GGML_TYPE_IQ1_S:
+ case GGML_TYPE_IQ1_M:
+ case GGML_TYPE_IQ2_S:
+ case GGML_TYPE_IQ3_S:
+ case GGML_TYPE_IQ4_NL:
+ case GGML_TYPE_IQ4_XS:
+ defines.push_back(type_upper + "_GRID");
+ break;
+ case GGML_TYPE_IQ2_XXS:
+ case GGML_TYPE_IQ2_XS:
+ case GGML_TYPE_IQ3_XXS:
+ defines.push_back(type_upper + "_GRID");
+ defines.push_back(type_upper + "_TABLES");
+ break;
+ default:
+ break;
+ }
+ break;
+ }
+ }
+
+ // VEC/SCALAR controls
+ defines.push_back(key.vectorized ? "VEC" : "SCALAR");
+
+ uint32_t wg_size = WEBGPU_MUL_MAT_VEC_WG_SIZE;
+ uint32_t outputs_per_wg = WEBGPU_MUL_MAT_VEC_FLOAT_OUTPUTS_PER_WG;
+
+ if (key.src0_type == GGML_TYPE_Q1_0) {
+ outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG;
+ } else if (key.src0_type >= GGML_TYPE_Q2_K) {
+ outputs_per_wg = WEBGPU_MUL_MAT_VEC_K_Q_OUTPUTS_PER_WG;
+ } else if (key.src0_type >= GGML_TYPE_Q4_0) {
+ outputs_per_wg = WEBGPU_MUL_MAT_VEC_LEGACY_Q_OUTPUTS_PER_WG;
+ }
+
+ // variant suffix for src1 type
+ variant += std::string("_") + (context.src1->type == GGML_TYPE_F32 ? "f32" : "f16");
+
+ defines.push_back(std::string("WG_SIZE=") + std::to_string(wg_size));
+ defines.push_back(std::string("OUTPUTS_PER_WG=") + std::to_string(outputs_per_wg));
+ defines.push_back(context.supports_subgroups ? "USE_SUBGROUP_REDUCTION" : "USE_WORKGROUP_REDUCTION");
+ variant += context.supports_subgroups ? "_sg_reduce" : "_wg_reduce";
+ if (key.vectorized) {
+ variant += "_vectorized";
+ }
+
+ defines.push_back(std::string("N_EXPERTS=") + std::to_string(key.n_experts));
+
+ auto processed = preprocessor.preprocess(shader_src, defines);
+
+ auto decisions = std::make_shared<ggml_webgpu_mul_mat_vec_shader_decisions>();
+ decisions->wg_size = wg_size;
+ decisions->outputs_per_wg = outputs_per_wg;
+
+ webgpu_pipeline pipeline = ggml_webgpu_create_pipeline(device, processed, variant);
+ pipeline.context = decisions;
+ mul_mat_id_vec_pipelines[key] = pipeline;
+ return mul_mat_id_vec_pipelines[key];
+ }
+
webgpu_pipeline get_unary_pipeline(const ggml_webgpu_shader_lib_context & context) {
const bool is_unary = context.dst->op == GGML_OP_UNARY;
const int op = is_unary ? (int) ggml_get_unary_op(context.dst) : context.dst->op;
case GGML_TYPE_Q5_0:
case GGML_TYPE_Q5_1:
case GGML_TYPE_Q8_0:
- case GGML_TYPE_Q8_1:
case GGML_TYPE_Q6_K:
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
}
+static webgpu_encoded_op ggml_webgpu_mul_mat_id_vec(webgpu_context & ctx,
+ ggml_tensor * src0,
+ ggml_tensor * src1,
+ ggml_tensor * src2,
+ ggml_tensor * dst) {
+ const uint32_t param_n_expert = (uint32_t) src0->ne[2];
+ const uint32_t param_n_expert_used = (uint32_t) dst->ne[1];
+
+ ggml_webgpu_shader_lib_context shader_lib_ctx = {};
+ shader_lib_ctx.src0 = src0;
+ shader_lib_ctx.src1 = src1;
+ shader_lib_ctx.src2 = src2;
+ shader_lib_ctx.dst = dst;
+ shader_lib_ctx.supports_subgroups = ctx->global_ctx->capabilities.supports_subgroups;
+ shader_lib_ctx.max_wg_size = ctx->global_ctx->capabilities.limits.maxComputeInvocationsPerWorkgroup;
+
+ webgpu_pipeline pipeline = ctx->shader_lib->get_mul_mat_id_vec_pipeline(shader_lib_ctx);
+
+ std::vector<uint32_t> params = {
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src0) / ggml_type_size(src0->type)),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src1) / ggml_type_size(src1->type)),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, src2) / ggml_type_size(src2->type)),
+ (uint32_t) (ggml_webgpu_tensor_misalignment(ctx, dst) / ggml_type_size(dst->type)),
+ (uint32_t) src0->ne[0],
+ (uint32_t) src0->ne[1],
+ param_n_expert,
+ param_n_expert_used,
+ (uint32_t) src1->ne[1],
+ (uint32_t) (src0->nb[1] / ggml_type_size(src0->type)),
+ (uint32_t) (src1->nb[1] / ggml_type_size(src1->type)),
+ (uint32_t) (src0->nb[2] / ggml_type_size(src0->type)),
+ (uint32_t) (src1->nb[2] / ggml_type_size(src1->type)),
+ };
+
+ std::vector<wgpu::BindGroupEntry> entries = {
+ ggml_webgpu_make_bind_group_entry(0, ggml_webgpu_tensor_buf(src0), ggml_webgpu_tensor_align_offset(ctx, src0),
+ ggml_webgpu_tensor_binding_size(ctx, src0)),
+ ggml_webgpu_make_bind_group_entry(1, ggml_webgpu_tensor_buf(src1), ggml_webgpu_tensor_align_offset(ctx, src1),
+ ggml_webgpu_tensor_binding_size(ctx, src1)),
+ ggml_webgpu_make_bind_group_entry(2, ggml_webgpu_tensor_buf(src2), ggml_webgpu_tensor_align_offset(ctx, src2),
+ ggml_webgpu_tensor_binding_size(ctx, src2)),
+ ggml_webgpu_make_bind_group_entry(3, ggml_webgpu_tensor_buf(dst), ggml_webgpu_tensor_align_offset(ctx, dst),
+ ggml_webgpu_tensor_binding_size(ctx, dst)),
+ };
+
+ uint32_t wg_x = 1;
+ uint32_t wg_y = 1;
+
+ auto * decisions = static_cast<ggml_webgpu_mul_mat_vec_shader_decisions *>(pipeline.context.get());
+
+ const uint32_t max_wg_per_dim = ctx->global_ctx->capabilities.limits.maxComputeWorkgroupsPerDimension;
+ uint32_t output_groups = CEIL_DIV(dst->ne[0], decisions->outputs_per_wg);
+ uint32_t total_wg = output_groups * param_n_expert_used;
+ compute_2d_workgroups(total_wg, max_wg_per_dim, wg_x, wg_y);
+
+ return ggml_backend_webgpu_build(ctx, pipeline, params, entries, wg_x, wg_y);
+}
+
static webgpu_encoded_op ggml_webgpu_mul_mat_id(webgpu_context & ctx,
ggml_tensor * src0,
ggml_tensor * src1,
ggml_tensor * src2,
ggml_tensor * dst) {
+ // we can use mat-vec fast path
+ if (dst->ne[2] == 1) {
+ return ggml_webgpu_mul_mat_id_vec(ctx, src0, src1, src2, dst);
+ }
+
ggml_webgpu_shader_lib_context shader_lib_ctx = {};
shader_lib_ctx.src0 = src0;
shader_lib_ctx.src1 = src1;
case GGML_TYPE_Q4_K:
case GGML_TYPE_Q5_K:
case GGML_TYPE_Q6_K:
+ case GGML_TYPE_IQ1_S:
+ case GGML_TYPE_IQ1_M:
+ case GGML_TYPE_IQ2_XXS:
+ case GGML_TYPE_IQ2_XS:
+ case GGML_TYPE_IQ2_S:
+ case GGML_TYPE_IQ3_XXS:
+ case GGML_TYPE_IQ3_S:
+ case GGML_TYPE_IQ4_NL:
+ case GGML_TYPE_IQ4_XS:
supports_op = true;
break;
default:
--- /dev/null
+#ifdef USE_SUBGROUP_REDUCTION
+enable subgroups;
+#endif
+enable f16;
+
+#define DECLARE_BYTE_LOADERS_SRC0
+#include "common_decls.tmpl"
+
+#include "mul_mat_vec_acc.tmpl"
+
+struct MulMatIdVecParams {
+ offset_src0: u32,
+ offset_src1: u32,
+ offset_ids: u32,
+ offset_dst: u32,
+
+ k: u32,
+ m: u32,
+ n_expert: u32,
+ n_expert_used: u32,
+ b_ne1: u32,
+
+ stride_01: u32,
+ stride_11: u32,
+ stride_02: u32,
+ stride_12: u32,
+};
+
+@group(0) @binding(0) var<storage, read_write> src0: array<SRC0_TYPE>; // [cols, rows, n_expert]
+@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>; // [cols, b_ne1, n_tokens(1)]
+@group(0) @binding(2) var<storage, read_write> ids: array<u32>; // [n_experd_used, n_tokens(1)]
+@group(0) @binding(3) var<storage, read_write> dst: array<f32>; // [rows, n_expert_used, n_tokens(1)]
+
+// "mul_mat_vec_acc.tmpl" requires params.k, params.m, params.stride_01
+@group(0) @binding(4) var<uniform> params: MulMatIdVecParams;
+
+// Flattened as [row][thread] to keep each row's reduction contiguous in memory.
+var<workgroup> partial_sums: array<f32, OUTPUTS_PER_WG * WG_SIZE>;
+
+fn partial_index(row: u32, thread: u32) -> u32 {
+ return row * WG_SIZE + thread;
+}
+
+var<workgroup> gathered_count_ids: array<u32, N_EXPERTS>;
+var<workgroup> gathered_expert_used: array<u32, N_EXPERTS>;
+
+@compute @workgroup_size(WG_SIZE)
+fn main(
+ @builtin(local_invocation_id) local_id: vec3<u32>,
+ @builtin(workgroup_id) wg_id: vec3<u32>,
+ @builtin(num_workgroups) num_wg: vec3<u32>
+#ifdef USE_SUBGROUP_REDUCTION
+ , @builtin(subgroup_id) subgroup_id: u32,
+ @builtin(subgroup_invocation_id) subgroup_invocation_id: u32,
+ @builtin(num_subgroups) num_subgroups: u32,
+ @builtin(subgroup_size) subgroup_size: u32
+#endif
+) {
+
+ let thread_id = local_id.x;
+
+ for (var i = thread_id;i < params.n_expert;i += WG_SIZE) {
+ gathered_count_ids[i] = 0;
+ }
+
+ workgroupBarrier();
+
+ // gather the selected experts for the target token.
+ for (var col = thread_id;col < params.n_expert_used;col += WG_SIZE) {
+ let expert = ids[params.offset_ids + col];
+ gathered_count_ids[expert] = 1;
+ gathered_expert_used[expert] = col;
+ }
+
+ workgroupBarrier();
+
+ let output_groups:u32 = (params.m + OUTPUTS_PER_WG - 1u) / OUTPUTS_PER_WG;
+ let wg_linear = wg_id.y * num_wg.x + wg_id.x;
+
+ var own_expert:u32 = 0;
+ var wg_in_batch:u32 = 0;
+ var wg_sum:u32 = 0;
+
+ for (var i = 0u;i < params.n_expert;i += 1) {
+ let wg_vec_count = gathered_count_ids[i]; // 1 or 0
+ let wg_per_matrix = output_groups * wg_vec_count;
+ if (wg_sum <= wg_linear && wg_linear < wg_sum + wg_per_matrix) {
+ own_expert = i;
+ wg_in_batch = wg_linear - wg_sum;
+ break;
+ }
+ wg_sum += wg_per_matrix;
+ }
+
+ let row_base = (wg_linear % output_groups) * OUTPUTS_PER_WG;
+ let dst1_stride = params.m;
+
+ let src0_batch_offset = params.offset_src0 + own_expert * params.stride_02;
+ let src1_idx_base = params.offset_src1 + (gathered_expert_used[own_expert] % params.b_ne1) * params.stride_11;
+ let dst_idx_base = params.offset_dst + gathered_expert_used[own_expert] * dst1_stride + row_base;
+
+ let acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);
+
+#ifdef USE_SUBGROUP_REDUCTION
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let subgroup_total = subgroupAdd(acc[row]);
+ if (subgroup_invocation_id == 0u) {
+ partial_sums[partial_index(row, subgroup_id)] = subgroup_total;
+ }
+ }
+
+ workgroupBarrier();
+
+ for (var row = subgroup_id; (row < OUTPUTS_PER_WG) && (row_base + row < params.m); row += num_subgroups) {
+ let output_row = row_base + row;
+ var row_acc = 0.0f;
+ for (var k = subgroup_invocation_id; k < num_subgroups; k += subgroup_size) {
+ row_acc += partial_sums[partial_index(row, k)];
+ }
+ let row_total = subgroupAdd(row_acc);
+ if (subgroup_invocation_id == 0) {
+ dst[dst_idx_base + row] = row_total;
+ }
+ }
+#endif
+
+#ifdef USE_WORKGROUP_REDUCTION
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ partial_sums[partial_index(row, thread_id)] = acc[row];
+ }
+
+ workgroupBarrier();
+
+ var stride:u32 = WG_SIZE / 2u;
+
+ while (stride > 0) {
+ if (thread_id < stride) {
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ partial_sums[partial_index(row, thread_id)] += partial_sums[partial_index(row, thread_id + stride)];
+ }
+ }
+
+ workgroupBarrier();
+ stride = stride / 2;
+ }
+
+ if (thread_id < OUTPUTS_PER_WG) {
+ let output_row = row_base + thread_id;
+ if (output_row < params.m) {
+ dst[dst_idx_base + thread_id] = partial_sums[partial_index(thread_id, 0)];
+ }
+ }
+#endif
+}
#define DECLARE_BYTE_LOADERS_SRC0
#include "common_decls.tmpl"
-#ifdef U32_DEQUANT_HELPERS
-#define SRC0_TYPE u32
-
-fn byte_of(v: u32, b: u32) -> u32 {
- return (v >> (b * 8u)) & 0xFFu;
-}
-
-fn sbyte_of(v: u32, b: u32) -> i32 {
- let raw = i32((v >> (b * 8u)) & 0xFFu);
- return select(raw, raw - 256, raw >= 128);
-}
-#endif
-
-#ifdef VEC
-#define VEC_SIZE 4u
-#define SRC0_TYPE vec4<SRC0_INNER_TYPE>
-#define SRC1_TYPE vec4<SRC1_INNER_TYPE>
-
-fn inner_dot(src0_val: SRC0_TYPE, src1_val: SRC1_TYPE) -> f32 {
- return f32(dot(SRC1_TYPE(src0_val), src1_val));
-}
-#endif
-
-#ifdef SCALAR
-#define VEC_SIZE 1u
-#define SRC0_TYPE SRC0_INNER_TYPE
-#define SRC1_TYPE SRC1_INNER_TYPE
-
-fn inner_dot(src0_val: SRC0_TYPE, src1_val: SRC1_TYPE) -> f32 {
- return f32(src0_val) * f32(src1_val);
-}
-#endif
+#include "mul_mat_vec_acc.tmpl"
struct MulMatParams {
offset_src0: u32,
@group(0) @binding(1) var<storage, read_write> src1: array<SRC1_TYPE>;
@group(0) @binding(2) var<storage, read_write> dst: array<f32>;
+// "mul_mat_vec_acc.tmpl" requires params.k, params.m, params.stride_01
@group(0) @binding(3) var<uniform> params: MulMatParams;
// Flattened as [row][thread] to keep each row's reduction contiguous in memory.
let src1_idx_base = params.offset_src1 + src13_idx * params.stride_13 + src12_idx * params.stride_12;
let dst_idx_base = params.offset_dst + dst3_idx * dst3_stride + dst2_idx * dst2_stride + row_base;
- var acc: array<f32, OUTPUTS_PER_WG>;
-
-#ifdef MUL_ACC_FLOAT
- let k_vec = params.k / VEC_SIZE;
- let src1_idx_base_vec = src1_idx_base / VEC_SIZE;
-
- // Each thread walks K, loads from the vector, and updates
- // a small block of output rows held in registers.
- for (var k = thread_id; k < k_vec; k += WG_SIZE) {
- let x = src1[src1_idx_base_vec + k];
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let src0_idx = (src0_batch_offset + output_row * params.stride_01) / VEC_SIZE + k;
- acc[row] += inner_dot(src0[src0_idx], x);
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q1_0
-#define BLOCK_SIZE 128
-#define BLOCK_SIZE_BYTES 18
-#define THREADS_PER_BLOCK 16
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let q_byte = load_u32_at_src0(block_byte_base + 2u + thread_within_block) & 0xFFu;
- var row_sum = 0.0;
- for (var bit = 0u; bit < 8u; bit++) {
- let w = select(-d, d, ((q_byte >> bit) & 1u) != 0u);
- row_sum += w * x_block[bit];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q4_0
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 18
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % 4;
- for (var block = thread_id/THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE/THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4] = f32(src1[x_base + i + 16]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- var row_sum = 0.0;
-
- let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * thread_within_block);
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_byte = get_byte(q_packed, byte_idx);
- let q_lo = (f32(q_byte & 0xFu) - 8.0) * d;
- let q_hi = (f32((q_byte >> 4u) & 0xFu) - 8.0) * d;
- row_sum += q_lo * x_block[byte_idx];
- row_sum += q_hi * x_block[byte_idx + 4u];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q4_1
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 20
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4] = f32(src1[x_base + i + 16]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let m = f32(load_f16_at_src0(block_byte_base + 2u));
- var row_sum = 0.0;
-
- let q_packed = load_u32_at_src0(block_byte_base + 4u + 4u * thread_within_block);
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_byte = get_byte(q_packed, byte_idx);
- let q_lo = f32(q_byte & 0xFu) * d + m;
- let q_hi = f32((q_byte >> 4u) & 0xFu) * d + m;
- row_sum += q_lo * x_block[byte_idx];
- row_sum += q_hi * x_block[byte_idx + 4u];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q5_0
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 22
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4] = f32(src1[x_base + i + 16]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let qh_packed = load_u32_at_src0(block_byte_base + 2u);
- let q_packed = load_u32_at_src0(block_byte_base + 6u + 4u * thread_within_block);
- let qh_shift = thread_within_block * 4u;
- var row_sum = 0.0;
-
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_byte = get_byte(q_packed, byte_idx);
- let qh_lo = ((qh_packed >> (qh_shift + byte_idx)) << 4u) & 0x10u;
- let qh_hi = (qh_packed >> (qh_shift + byte_idx + 12u)) & 0x10u;
- let q_lo = (f32((q_byte & 0xFu) | qh_lo) - 16.0) * d;
- let q_hi = (f32(((q_byte >> 4u) & 0xFu) | qh_hi) - 16.0) * d;
- row_sum += q_lo * x_block[byte_idx];
- row_sum += q_hi * x_block[byte_idx + 4u];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q5_1
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 24
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4] = f32(src1[x_base + i + 16]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let m = f32(load_f16_at_src0(block_byte_base + 2u));
- let qh_packed = load_u32_at_src0(block_byte_base + 4u);
- let q_packed = load_u32_at_src0(block_byte_base + 8u + 4u * thread_within_block);
- let qh_shift = thread_within_block * 4u;
- var row_sum = 0.0;
-
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_byte = get_byte(q_packed, byte_idx);
- let qh_lo = ((qh_packed >> (qh_shift + byte_idx)) << 4u) & 0x10u;
- let qh_hi = (qh_packed >> (qh_shift + byte_idx + 12u)) & 0x10u;
- let q_lo = f32((q_byte & 0xFu) | qh_lo) * d + m;
- let q_hi = f32(((q_byte >> 4u) & 0xFu) | qh_hi) * d + m;
- row_sum += q_lo * x_block[byte_idx];
- row_sum += q_hi * x_block[byte_idx + 4u];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q8_0
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 34
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- var row_sum = 0.0;
-
- for (var packed_idx = 0u; packed_idx < ELEMS_PER_THREAD / 4u; packed_idx++) {
- let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * (thread_within_block * 2u + packed_idx));
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_val = f32(get_byte_i32(q_packed, byte_idx)) * d;
- row_sum += q_val * x_block[packed_idx * 4u + byte_idx];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q8_1
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 36
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let m = f32(load_f16_at_src0(block_byte_base + 2u));
- var row_sum = 0.0;
-
- for (var packed_idx = 0u; packed_idx < ELEMS_PER_THREAD / 4u; packed_idx++) {
- let q_packed = load_u32_at_src0(block_byte_base + 4u + 4u * (thread_within_block * 2u + packed_idx));
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_val = f32(get_byte_i32(q_packed, byte_idx)) * d + m;
- row_sum += q_val * x_block[packed_idx * 4u + byte_idx];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q2_K
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 84
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let lane = tid / 2u;
- let phase = tid % 2u;
- let iq = lane / 4u;
- let ir = lane % 4u;
- let is = ir / 2u;
-
- let y_offset = 128u * iq + 8u * ir + 4u * phase;
- let sc0_byte = 8u * iq + is;
- let sc2_byte = 8u * iq + is + 2u;
- let sc4_byte = 8u * iq + is + 4u;
- let sc6_byte = 8u * iq + is + 6u;
- let qs_byte = 16u + (16u * iq + 4u * ir) * 2u + 4u * phase;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 4u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4u] = f32(src1[x_base + 32u + i]);
- x_block[i + 8u] = f32(src1[x_base + 64u + i]);
- x_block[i + 12u] = f32(src1[x_base + 96u + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let dall = f32(load_f16_at_src0(block_byte_base + 80u));
- let dmin = f32(load_f16_at_src0(block_byte_base + 82u)) * (1.0 / 16.0);
-
- let sc0 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc0_byte), sc0_byte & 3u);
- let sc2 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc2_byte), sc2_byte & 3u);
- let sc4 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc4_byte), sc4_byte & 3u);
- let sc6 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc6_byte), sc6_byte & 3u);
-
- let q_u32 = load_u32_at_src0_aligned(block_byte_base + qs_byte);
- let qs0 = q_u32 & 0xFFFFu;
- let qs1 = q_u32 >> 16u;
-
- var sumy = vec4<f32>(0.0, 0.0, 0.0, 0.0);
- var acc1 = vec4<f32>(0.0, 0.0, 0.0, 0.0);
- var acc2 = vec4<f32>(0.0, 0.0, 0.0, 0.0);
-
- sumy[0] = x_block[0] + x_block[1] + x_block[2] + x_block[3];
- sumy[1] = x_block[4] + x_block[5] + x_block[6] + x_block[7];
- sumy[2] = x_block[8] + x_block[9] + x_block[10] + x_block[11];
- sumy[3] = x_block[12] + x_block[13] + x_block[14] + x_block[15];
-
- acc1[0] = x_block[0] * f32(qs0 & 0x0003u) + x_block[2] * f32(qs1 & 0x0003u);
- acc2[0] = x_block[1] * f32(qs0 & 0x0300u) + x_block[3] * f32(qs1 & 0x0300u);
- acc1[1] = x_block[4] * f32(qs0 & 0x000Cu) + x_block[6] * f32(qs1 & 0x000Cu);
- acc2[1] = x_block[5] * f32(qs0 & 0x0C00u) + x_block[7] * f32(qs1 & 0x0C00u);
- acc1[2] = x_block[8] * f32(qs0 & 0x0030u) + x_block[10] * f32(qs1 & 0x0030u);
- acc2[2] = x_block[9] * f32(qs0 & 0x3000u) + x_block[11] * f32(qs1 & 0x3000u);
- acc1[3] = x_block[12] * f32(qs0 & 0x00C0u) + x_block[14] * f32(qs1 & 0x00C0u);
- acc2[3] = x_block[13] * f32(qs0 & 0xC000u) + x_block[15] * f32(qs1 & 0xC000u);
-
- acc[row] += dall * ((acc1[0] + (1.0/256.0) * acc2[0]) * f32(sc0 & 0xFu) +
- (acc1[1] + (1.0/256.0) * acc2[1]) * f32(sc2 & 0xFu) / 4.0 +
- (acc1[2] + (1.0/256.0) * acc2[2]) * f32(sc4 & 0xFu) / 16.0 +
- (acc1[3] + (1.0/256.0) * acc2[3]) * f32(sc6 & 0xFu) / 64.0)
- - dmin * (sumy[0] * f32(sc0 & 0xF0u) + sumy[1] * f32(sc2 & 0xF0u) +
- sumy[2] * f32(sc4 & 0xF0u) + sumy[3] * f32(sc6 & 0xF0u));
- }
- }
- }
-#endif
-
-
-#ifdef MUL_ACC_Q3_K
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 110
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let lane = tid / 2u;
- let phase = tid % 2u;
- let ip = lane / 4u;
- let il = 2u * ((lane % 4u) / 2u);
- let ir = lane % 2u;
- let l0 = 8u * ir;
-
- let q_byte = 32u + 32u * ip + l0 + 16u * phase;
- let h_byte = l0 + 16u * phase;
- let y_offset = 128u * ip + 32u * il + l0 + 16u * phase;
-
- let s_shift1 = 4u * ip;
- let s_shift2 = s_shift1 + il;
-
- let v1 = select(64.0, 4.0, il == 0u);
- let v2 = 4.0 * v1;
- let shift = 2u * il;
-
- var qm0: u32; var qm1: u32; var qm2: u32; var qm3: u32;
- if (il == 0u) {
- qm0 = 0x0003u; qm1 = 0x0300u; qm2 = 0x000Cu; qm3 = 0x0C00u;
- } else {
- qm0 = 0x0030u; qm1 = 0x3000u; qm2 = 0x00C0u; qm3 = 0xC000u;
- }
-
- let mm_idx = 2u * ip + il / 2u;
- var hm0: u32; var hm1: u32; var hm2: u32; var hm3: u32;
- switch (mm_idx) {
- case 0u: { hm0=0x0001u; hm1=0x0100u; hm2=0x0002u; hm3=0x0200u; }
- case 1u: { hm0=0x0004u; hm1=0x0400u; hm2=0x0008u; hm3=0x0800u; }
- case 2u: { hm0=0x0010u; hm1=0x1000u; hm2=0x0020u; hm3=0x2000u; }
- default: { hm0=0x0040u; hm1=0x4000u; hm2=0x0080u; hm3=0x8000u; }
- }
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 8u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 8u] = f32(src1[x_base + 32u + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let d = f32(load_f16_at_src0(block_byte_base + 108u));
- let a_base = 96u;
- let a_il0 = load_u16_at_src0(block_byte_base + a_base + il * 2u);
- let a_il1 = load_u16_at_src0(block_byte_base + a_base + (il + 1u) * 2u);
- let a_4 = load_u16_at_src0(block_byte_base + a_base + 8u);
- let a_5 = load_u16_at_src0(block_byte_base + a_base + 10u);
-
- var scales32 = a_4 | (a_5 << 16u);
- let aux32 = ((scales32 >> s_shift2) << 4u) & 0x30303030u;
- scales32 = a_il0 | (a_il1 << 16u);
- scales32 = ((scales32 >> s_shift1) & 0x0F0F0F0Fu) | aux32;
-
- let scale0 = f32(i32(byte_of(scales32, phase + 0u)) - 32);
- let scale1 = f32(i32(byte_of(scales32, phase + 2u)) - 32);
-
- let q_u32_0 = load_u32_at_src0(block_byte_base + q_byte + 0u);
- let q_u32_1 = load_u32_at_src0(block_byte_base + q_byte + 4u);
- let h_u32_0 = load_u32_at_src0(block_byte_base + h_byte + 0u);
- let h_u32_1 = load_u32_at_src0(block_byte_base + h_byte + 4u);
-
- var s1 = 0.0; var s2 = 0.0; var s3 = 0.0;
- var s4 = 0.0; var s5 = 0.0; var s6 = 0.0;
-
- for (var l = 0u; l < 8u; l += 2u) {
- let q_u32 = select(q_u32_0, q_u32_1, l >= 4u);
- let qs = select(q_u32 & 0xFFFFu, q_u32 >> 16u, (l & 2u) != 0u);
- let h_u32 = select(h_u32_0, h_u32_1, l >= 4u);
- let hv = select(h_u32 & 0xFFFFu, h_u32 >> 16u, (l & 2u) != 0u);
-
- s1 += x_block[l + 0u] * f32(qs & qm0);
- s2 += x_block[l + 1u] * f32(qs & qm1);
- s3 += select(0.0, x_block[l + 0u], (hv & hm0) == 0u) +
- select(0.0, x_block[l + 1u], (hv & hm1) == 0u);
- s4 += x_block[l + 8u] * f32(qs & qm2);
- s5 += x_block[l + 9u] * f32(qs & qm3);
- s6 += select(0.0, x_block[l + 8u], (hv & hm2) == 0u) +
- select(0.0, x_block[l + 9u], (hv & hm3) == 0u);
- }
-
- let d1 = d * (s1 + (1.0/256.0) * s2 - s3 * v1);
- let d2 = d * (s4 + (1.0/256.0) * s5 - s6 * v2);
- acc[row] += (d1 * scale0 + 0.25 * d2 * scale1) / f32(1u << shift);
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q4_K
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 144
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let il = tid / 4u;
- let ir = tid % 4u;
- let im = il / 2u;
- let in = il % 2u;
- let l0 = 4u * (2u * ir + in);
-
- let y_offset = 64u * im + l0;
- let q_offset = 32u * im + l0;
- let sc0_byte = 4u + im * 2u;
- let sc2_byte = 4u + (im + 2u) * 2u;
- let sc4_byte = 4u + (im + 4u) * 2u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 4u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4u] = f32(src1[x_base + 32u + i]);
- x_block[i + 8u] = f32(src1[x_base + 128u + i]);
- x_block[i + 12u] = f32(src1[x_base + 160u + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let d = f32(load_f16_at_src0(block_byte_base + 0u));
- let dmin = f32(load_f16_at_src0(block_byte_base + 2u));
-
- let sc0_u32 = load_u32_at_src0_aligned(block_byte_base + sc0_byte);
- let sc0 = select(sc0_u32 & 0xFFFFu, sc0_u32 >> 16u, (sc0_byte & 2u) != 0u);
- let sc2_u32 = load_u32_at_src0_aligned(block_byte_base + sc2_byte);
- let sc2 = select(sc2_u32 & 0xFFFFu, sc2_u32 >> 16u, (sc2_byte & 2u) != 0u);
- let sc4_u32 = load_u32_at_src0_aligned(block_byte_base + sc4_byte);
- let sc4 = select(sc4_u32 & 0xFFFFu, sc4_u32 >> 16u, (sc4_byte & 2u) != 0u);
-
- let sc16_0 = sc0 & 0x3F3Fu;
- let sc16_1 = sc2 & 0x3F3Fu;
- let sc16_2 = (sc4 & 0x0F0Fu) | ((sc0 & 0xC0C0u) >> 2u);
- let sc16_3 = ((sc4 >> 4u) & 0x0F0Fu) | ((sc2 & 0xC0C0u) >> 2u);
-
- let scale0 = f32(sc16_0 & 0xFFu);
- let scale1 = f32((sc16_0 >> 8u) & 0xFFu);
- let min0 = f32(sc16_1 & 0xFFu);
- let min1 = f32((sc16_1 >> 8u) & 0xFFu);
- let scale2 = f32(sc16_2 & 0xFFu);
- let scale3 = f32((sc16_2 >> 8u) & 0xFFu);
- let min2 = f32(sc16_3 & 0xFFu);
- let min3 = f32((sc16_3 >> 8u) & 0xFFu);
-
- let q1_u32 = load_u32_at_src0_aligned(block_byte_base + 16u + q_offset);
- let q2_u32 = load_u32_at_src0_aligned(block_byte_base + 80u + q_offset);
-
- var dot = vec4<f32>(0.0, 0.0, 0.0, 0.0);
- var sumx = vec4<f32>(0.0, 0.0, 0.0, 0.0);
- for (var i = 0u; i < 4u; i++) {
- let q1b = byte_of(q1_u32, i);
- let q2b = byte_of(q2_u32, i);
- dot[0] += x_block[i] * f32(q1b & 0x0Fu);
- dot[1] += x_block[i + 4u] * f32(q1b >> 4u);
- dot[2] += x_block[i + 8u] * f32(q2b & 0x0Fu);
- dot[3] += x_block[i + 12u] * f32(q2b >> 4u);
- sumx[0] += x_block[i];
- sumx[1] += x_block[i + 4u];
- sumx[2] += x_block[i + 8u];
- sumx[3] += x_block[i + 12u];
- }
-
- acc[row] += d * (dot[0] * scale0 + dot[1] * scale1 + dot[2] * scale2 + dot[3] * scale3)
- - dmin * (sumx[0] * min0 + sumx[1] * min1 + sumx[2] * min2 + sumx[3] * min3);
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q5_K
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 176
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let il = tid / 4u;
- let ir = tid % 4u;
- let im = il / 2u;
- let in = il % 2u;
- let l0 = 4u * (2u * ir + in);
-
- let y_offset = 64u * im + l0;
- let q_offset = 48u + 32u * im + l0;
- let qh_offset = 16u + 8u * ir + 4u * in;
- let sc0_byte = 4u + im * 2u;
- let sc2_byte = 4u + (im + 2u) * 2u;
- let sc4_byte = 4u + (im + 4u) * 2u;
-
- let hm1 = 1u << (2u * im);
- let hm2 = hm1 << 1u;
- let hm3 = hm1 << 4u;
- let hm4 = hm2 << 4u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 4u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4u] = f32(src1[x_base + 32u + i]);
- x_block[i + 8u] = f32(src1[x_base + 128u + i]);
- x_block[i + 12u] = f32(src1[x_base + 160u + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let d = f32(load_f16_at_src0(block_byte_base + 0u));
- let dmin = f32(load_f16_at_src0(block_byte_base + 2u));
-
- let sc0_u32 = load_u32_at_src0_aligned(block_byte_base + sc0_byte);
- let sc0 = select(sc0_u32 & 0xFFFFu, sc0_u32 >> 16u, (sc0_byte & 2u) != 0u);
- let sc2_u32 = load_u32_at_src0_aligned(block_byte_base + sc2_byte);
- let sc2 = select(sc2_u32 & 0xFFFFu, sc2_u32 >> 16u, (sc2_byte & 2u) != 0u);
- let sc4_u32 = load_u32_at_src0_aligned(block_byte_base + sc4_byte);
- let sc4 = select(sc4_u32 & 0xFFFFu, sc4_u32 >> 16u, (sc4_byte & 2u) != 0u);
-
- let sc16_0 = sc0 & 0x3F3Fu;
- let sc16_1 = sc2 & 0x3F3Fu;
- let sc16_2 = (sc4 & 0x0F0Fu) | ((sc0 & 0xC0C0u) >> 2u);
- let sc16_3 = ((sc4 >> 4u) & 0x0F0Fu) | ((sc2 & 0xC0C0u) >> 2u);
-
- let f0 = f32(sc16_0 & 0xFFu);
- let f1 = f32((sc16_0 >> 8u) & 0xFFu);
- let m0 = f32(sc16_1 & 0xFFu);
- let m1 = f32((sc16_1 >> 8u) & 0xFFu);
- let f4 = f32(sc16_2 & 0xFFu);
- let f5 = f32((sc16_2 >> 8u) & 0xFFu);
- let m4 = f32(sc16_3 & 0xFFu);
- let m5 = f32((sc16_3 >> 8u) & 0xFFu);
-
- let q1_u32 = load_u32_at_src0_aligned(block_byte_base + q_offset);
- let q2_u32 = load_u32_at_src0_aligned(block_byte_base + q_offset + 64u);
- let qh_u32 = load_u32_at_src0_aligned(block_byte_base + qh_offset);
-
- var vals = vec4<f32>(0.0, 0.0, 0.0, 0.0);
- var sumy = vec4<f32>(0.0, 0.0, 0.0, 0.0);
- for (var i = 0u; i < 4u; i++) {
- let q1b = byte_of(q1_u32, i);
- let q2b = byte_of(q2_u32, i);
- let qhb = byte_of(qh_u32, i);
-
- let yl0 = x_block[i];
- let yl8 = x_block[i + 4u];
- let yh0 = x_block[i + 8u];
- let yh8 = x_block[i + 12u];
-
- sumy[0] += yl0;
- sumy[1] += yl8;
- sumy[2] += yh0;
- sumy[3] += yh8;
-
- let q0 = f32((q1b & 0x0Fu) | select(0u, 0x10u, (qhb & hm1) != 0u));
- let q1 = f32((q1b >> 4u) | select(0u, 0x10u, (qhb & hm2) != 0u));
- let q2 = f32((q2b & 0x0Fu) | select(0u, 0x10u, (qhb & hm3) != 0u));
- let q3 = f32((q2b >> 4u) | select(0u, 0x10u, (qhb & hm4) != 0u));
-
- vals[0] += yl0 * q0;
- vals[1] += yl8 * q1;
- vals[2] += yh0 * q2;
- vals[3] += yh8 * q3;
- }
-
- acc[row] += d * (f0 * vals[0] + f1 * vals[1] + f4 * vals[2] + f5 * vals[3])
- - dmin * (sumy[0] * m0 + sumy[1] * m1 +
- sumy[2] * m4 + sumy[3] * m5);
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_Q6_K
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 210
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let ip = tid / 8u;
- let il = tid % 8u;
- let l0 = 4u * il;
- let is = 8u * ip + l0 / 16u;
-
- let y_offset = 128u * ip + l0;
- let q_offset_l = 64u * ip + l0;
- let q_offset_h = 32u * ip + l0;
-
- let num_blocks = params.k / BLOCK_SIZE;
- let sc_base_byte = 192u + (is & ~3u);
- let sc_byte_pos = is & 3u;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var l = 0u; l < 4u; l++) {
- x_block[l] = f32(src1[x_base + l]);
- x_block[l + 4u] = f32(src1[x_base + 32u + l]);
- x_block[l + 8u] = f32(src1[x_base + 64u + l]);
- x_block[l + 12u] = f32(src1[x_base + 96u + l]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let d = f32(load_f16_at_src0(block_byte_base + 208u));
- let ql1_u32 = load_u32_at_src0(block_byte_base + q_offset_l);
- let ql2_u32 = load_u32_at_src0(block_byte_base + q_offset_l + 32u);
- let qh_u32 = load_u32_at_src0(block_byte_base + 128u + q_offset_h);
- let sc_u32_0 = load_u32_at_src0(block_byte_base + sc_base_byte);
- let sc_u32_1 = load_u32_at_src0(block_byte_base + sc_base_byte + 4u);
-
- let sc0 = sbyte_of(sc_u32_0, sc_byte_pos);
- let sc2 = sbyte_of(sc_u32_0, sc_byte_pos + 2u);
- let sc4 = sbyte_of(sc_u32_1, sc_byte_pos);
- let sc6 = sbyte_of(sc_u32_1, sc_byte_pos + 2u);
-
- var sums = vec4<f32>(0.0, 0.0, 0.0, 0.0);
-
- for (var l = 0u; l < 4u; l++) {
- let q1b = byte_of(ql1_u32, l);
- let q2b = byte_of(ql2_u32, l);
- let qhb = byte_of(qh_u32, l);
-
- let dq0 = f32(i32((q1b & 0x0Fu) | ((qhb & 0x03u) << 4u)) - 32);
- let dq1 = f32(i32((q2b & 0x0Fu) | ((qhb & 0x0Cu) << 2u)) - 32);
- let dq2 = f32(i32((q1b >> 4u) | (qhb & 0x30u)) - 32);
- let dq3 = f32(i32((q2b >> 4u) | ((qhb & 0xC0u) >> 2u)) - 32);
-
- sums[0] += x_block[l] * dq0;
- sums[1] += x_block[l + 4u] * dq1;
- sums[2] += x_block[l + 8u] * dq2;
- sums[3] += x_block[l + 12u] * dq3;
- }
-
- acc[row] += d * (sums[0] * f32(sc0) + sums[1] * f32(sc2) +
- sums[2] * f32(sc4) + sums[3] * f32(sc6));
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ1_S
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 50
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let d = f32(load_f16_at_src0(block_byte_base));
- let qh = load_u32_at_src0(block_byte_base + 34u + sub_blk * 2u) & 0xFFFFu;
- let dl = d * f32(2u * ((qh >> 12u) & 7u) + 1u);
- let delta = select(IQ1_DELTA, -IQ1_DELTA, (qh & 0x8000u) != 0u);
- let qs_w = load_u32_at_src0(block_byte_base + 2u + sub_blk * 4u);
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let qs_byte = get_byte(qs_w, l);
- let ig = (qs_byte | (((qh >> (3u * l)) & 7u) << 8u)) * 8u;
- let gw = iq1_grid[ig / 16u];
- let bit_base = (ig % 16u) * 2u;
- for (var j = 0u; j < 8u; j++) {
- let g = (gw >> (bit_base + j * 2u)) & 3u;
- let gs = select(f32(g), f32(g) - 4.0, (g & 2u) != 0u);
- row_sum += dl * (gs + delta) * x_block[ll * 8u + j];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ1_M
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 56
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
-
- let sc_lo = load_u32_at_src0(block_byte_base + 48u);
- let sc_hi = load_u32_at_src0(block_byte_base + 52u);
- let sc0 = sc_lo & 0xFFFFu;
- let sc1 = (sc_lo >> 16u) & 0xFFFFu;
- let sc2 = sc_hi & 0xFFFFu;
- let sc3 = (sc_hi >> 16u) & 0xFFFFu;
- let d_bits = (sc0 >> 12u) | ((sc1 >> 8u) & 0xF0u) | ((sc2 >> 4u) & 0xF00u) | (sc3 & 0xF000u);
- let d = f32(bitcast<vec2<f16>>(d_bits)[0]);
-
- let sc_u16 = select(select(sc2, sc3, sub_blk >= 6u),
- select(sc0, sc1, sub_blk >= 2u),
- sub_blk < 4u);
-
- let qs_w = load_u32_at_src0(block_byte_base + sub_blk * 4u);
- let qh = load_u32_at_src0(block_byte_base + 32u + sub_blk * 2u) & 0xFFFFu;
- let qh_lo = qh & 0xFFu;
- let qh_hi = (qh >> 8u) & 0xFFu;
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let bit_off = 6u * (sub_blk % 2u) + 3u * (l / 2u);
- let sub_scale = (sc_u16 >> bit_off) & 0x7u;
- let dl = d * f32(2u * sub_scale + 1u);
- let qh_byte = select(qh_lo, qh_hi, l >= 2u);
- let ll2 = l % 2u;
- let grid_idx = get_byte(qs_w, l) | (((qh_byte >> (4u * ll2)) & 7u) << 8u);
- let delta = select(IQ1_DELTA, -IQ1_DELTA, ((qh_byte >> (3u + 4u * ll2)) & 1u) != 0u);
- let ig = grid_idx * 8u;
- let gw = iq1_grid[ig / 16u];
- let bit_base = (ig % 16u) * 2u;
- for (var j = 0u; j < 8u; j++) {
- let g = (gw >> (bit_base + j * 2u)) & 3u;
- let gs = select(f32(g), f32(g) - 4.0, (g & 2u) != 0u);
- row_sum += dl * (gs + delta) * x_block[ll * 8u + j];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ2_XXS
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 66
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let aux_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
- let aux_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
- let ls = aux_hi >> 28u;
- let db = d * (0.5 + f32(ls)) * 0.25;
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let grid_idx = (aux_lo >> (8u * l)) & 0xFFu;
- let signs_idx = (aux_hi >> (7u * l)) & 0x7Fu;
- let signs = (ksigns_iq2xs[signs_idx / 4u] >> ((signs_idx % 4u) * 8u)) & 0xFFu;
- let gw_lo = iq2xxs_grid[grid_idx * 2u];
- let gw_hi = iq2xxs_grid[grid_idx * 2u + 1u];
- for (var j = 0u; j < 8u; j++) {
- let gw = select(gw_hi, gw_lo, j < 4u);
- let b = f32((gw >> ((j & 3u) * 8u)) & 0xFFu);
- let s = select(1.0, -1.0, ((signs >> j) & 1u) != 0u);
- row_sum += db * b * s * x_block[ll * 8u + j];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ2_XS
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 74
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let qs_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
- let qs_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
- let scales_word = load_u32_at_src0(block_byte_base + 66u + (sub_blk / 4u) * 4u);
- let scales_byte = get_byte(scales_word, sub_blk % 4u);
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let qs_word = select(qs_hi, qs_lo, l < 2u);
- let half2 = (l % 2u) * 16u;
- let qs_val = (qs_word >> half2) & 0xFFFFu;
- let grid_idx = qs_val & 0x1FFu;
- let signs_idx = (qs_val >> 9u) & 0x7Fu;
- let sub_scale = (scales_byte >> (4u * (l / 2u))) & 0xFu;
- let db = d * (0.5 + f32(sub_scale)) * 0.25;
- let signs = (ksigns_iq2xs[signs_idx / 4u] >> ((signs_idx % 4u) * 8u)) & 0xFFu;
- let gw_lo = iq2xs_grid[grid_idx * 2u];
- let gw_hi = iq2xs_grid[grid_idx * 2u + 1u];
- for (var j = 0u; j < 8u; j++) {
- let gw = select(gw_hi, gw_lo, j < 4u);
- let b = f32((gw >> ((j & 3u) * 8u)) & 0xFFu);
- let s = select(1.0, -1.0, ((signs >> j) & 1u) != 0u);
- row_sum += db * b * s * x_block[ll * 8u + j];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ2_S
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 82
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let qs_w = load_u32_at_src0(block_byte_base + 2u + sub_blk * 4u);
- let sg_w = load_u32_at_src0(block_byte_base + 34u + sub_blk * 4u);
- let qh_word = load_u32_at_src0(block_byte_base + 66u + (sub_blk / 4u) * 4u);
- let qh_byte = get_byte(qh_word, sub_blk % 4u);
- let sc_word = load_u32_at_src0(block_byte_base + 74u + (sub_blk / 4u) * 4u);
- let scales_byte = get_byte(sc_word, sub_blk % 4u);
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let qs_byte = get_byte(qs_w, l);
- let sign_byte = get_byte(sg_w, l);
- let grid_idx = qs_byte | (((qh_byte >> (2u * l)) & 3u) << 8u);
- let sub_scale = (scales_byte >> (4u * (l / 2u))) & 0xFu;
- let db = d * (0.5 + f32(sub_scale)) * 0.25;
- let gw_lo = iq2s_grid[grid_idx * 2u];
- let gw_hi = iq2s_grid[grid_idx * 2u + 1u];
- for (var j = 0u; j < 8u; j++) {
- let gw = select(gw_hi, gw_lo, j < 4u);
- let b = f32((gw >> ((j & 3u) * 8u)) & 0xFFu);
- let s = select(1.0, -1.0, ((sign_byte >> j) & 1u) != 0u);
- row_sum += db * b * s * x_block[ll * 8u + j];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ3_XXS
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 98
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let qs_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
- let qs_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
- let aux = load_u32_at_src0(block_byte_base + 66u + sub_blk * 4u);
- let ls = aux >> 28u;
- let db = d * (0.5 + f32(ls)) * 0.5;
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let qs_word = select(qs_hi, qs_lo, l < 2u);
- let byte_pos = (l % 2u) * 2u;
- let grid_idx_0 = (qs_word >> (byte_pos * 8u)) & 0xFFu;
- let grid_idx_1 = (qs_word >> ((byte_pos + 1u) * 8u)) & 0xFFu;
- let signs_idx = (aux >> (7u * l)) & 0x7Fu;
- let signs = (ksigns_iq2xs[signs_idx / 4u] >> ((signs_idx % 4u) * 8u)) & 0xFFu;
- let grid1 = iq3xxs_grid[grid_idx_0];
- let grid2 = iq3xxs_grid[grid_idx_1];
- for (var j = 0u; j < 4u; j++) {
- let b1 = f32((grid1 >> (j * 8u)) & 0xFFu);
- let b2 = f32((grid2 >> (j * 8u)) & 0xFFu);
- let s1 = select(1.0, -1.0, ((signs >> j) & 1u) != 0u);
- let s2 = select(1.0, -1.0, ((signs >> (j + 4u)) & 1u) != 0u);
- row_sum += db * b1 * s1 * x_block[ll * 8u + j];
- row_sum += db * b2 * s2 * x_block[ll * 8u + j + 4u];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ3_S
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 110
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let slot0 = half * 2u;
- let y_offset = sub_blk * 32u + slot0 * 8u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let qs_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
- let qs_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
- let qh_word = load_u32_at_src0(block_byte_base + 66u + (sub_blk / 4u) * 4u);
- let qh_byte = get_byte(qh_word, sub_blk % 4u);
- let sg_w = load_u32_at_src0(block_byte_base + 74u + sub_blk * 4u);
- let sc_word = load_u32_at_src0(block_byte_base + 106u);
- let scales_byte = get_byte(sc_word, sub_blk / 2u);
- let sub_scale = (scales_byte >> (4u * (sub_blk % 2u))) & 0xFu;
- let db = d * (1.0 + 2.0 * f32(sub_scale));
-
- var row_sum = 0.0;
- for (var ll = 0u; ll < 2u; ll++) {
- let l = slot0 + ll;
- let qs_word = select(qs_hi, qs_lo, l < 2u);
- let byte_pos = (l % 2u) * 2u;
- let qs0 = (qs_word >> (byte_pos * 8u)) & 0xFFu;
- let qs1 = (qs_word >> ((byte_pos + 1u) * 8u)) & 0xFFu;
- let grid_idx_1 = qs0 | (((qh_byte >> (2u * l)) & 1u) << 8u);
- let grid_idx_2 = qs1 | (((qh_byte >> (2u * l + 1u)) & 1u) << 8u);
- let sign_byte = get_byte(sg_w, l);
- let grid1 = iq3s_grid[grid_idx_1];
- let grid2 = iq3s_grid[grid_idx_2];
- for (var j = 0u; j < 4u; j++) {
- let b1 = f32((grid1 >> (j * 8u)) & 0xFFu);
- let b2 = f32((grid2 >> (j * 8u)) & 0xFFu);
- let s1 = select(1.0, -1.0, ((sign_byte >> j) & 1u) != 0u);
- let s2 = select(1.0, -1.0, ((sign_byte >> (j + 4u)) & 1u) != 0u);
- row_sum += db * b1 * s1 * x_block[ll * 8u + j];
- row_sum += db * b2 * s2 * x_block[ll * 8u + j + 4u];
- }
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ4_NL
-#define BLOCK_SIZE 32
-#define BLOCK_SIZE_BYTES 18
-#define THREADS_PER_BLOCK 4
-#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
-
- let num_blocks = params.k / BLOCK_SIZE;
- let thread_within_block = thread_id % THREADS_PER_BLOCK;
- for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4u;
- var x_block: array<f32, ELEMS_PER_THREAD>;
- for (var i = 0u; i < ELEMS_PER_THREAD / 2u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- x_block[i + 4u] = f32(src1[x_base + i + 16u]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- var row_sum = 0.0;
-
- let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * thread_within_block);
- for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
- let q_byte = get_byte(q_packed, byte_idx);
- let q_lo = f32(kvalues_iq4nl[q_byte & 0xFu]) * d;
- let q_hi = f32(kvalues_iq4nl[(q_byte >> 4u) & 0xFu]) * d;
- row_sum += q_lo * x_block[byte_idx];
- row_sum += q_hi * x_block[byte_idx + 4u];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
-
-#ifdef MUL_ACC_IQ4_XS
-#define BLOCK_SIZE 256
-#define BLOCK_SIZE_BYTES 136
-#define THREADS_PER_BLOCK 16
-
- let tid = thread_id % THREADS_PER_BLOCK;
- let block_group = thread_id / THREADS_PER_BLOCK;
- let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
-
- let sub_blk = tid / 2u;
- let half = tid % 2u;
- let y_offset = sub_blk * 32u + half * 16u;
-
- let num_blocks = params.k / BLOCK_SIZE;
-
- for (var block = block_group; block < num_blocks; block += num_block_groups) {
- let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
- var x_block: array<f32, 16>;
- for (var i = 0u; i < 16u; i++) {
- x_block[i] = f32(src1[x_base + i]);
- }
-
- for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
- let output_row = row_base + row;
- if (output_row < params.m) {
- let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
- let d = f32(load_f16_at_src0(block_byte_base));
- let scales_h = load_u16_at_src0(block_byte_base + 2u);
- let scales_l_word = load_u32_at_src0(block_byte_base + 4u);
- let sl_byte = get_byte(scales_l_word, sub_blk / 2u);
- let sl = (sl_byte >> (4u * (sub_blk % 2u))) & 0xFu;
- let sh_bits = (scales_h >> (2u * sub_blk)) & 3u;
- let ls = i32(sl | (sh_bits << 4u));
- let dl = d * f32(ls - 32);
-
- let qs_byte_off = 8u + sub_blk * 16u;
- let q_w0 = load_u32_at_src0(block_byte_base + qs_byte_off);
- let q_w1 = load_u32_at_src0(block_byte_base + qs_byte_off + 4u);
- let q_w2 = load_u32_at_src0(block_byte_base + qs_byte_off + 8u);
- let q_w3 = load_u32_at_src0(block_byte_base + qs_byte_off + 12u);
-
- var row_sum = 0.0;
- for (var i = 0u; i < 16u; i++) {
- let q_word = select(
- select(q_w0, q_w1, i >= 4u),
- select(q_w2, q_w3, i >= 12u),
- i >= 8u);
- let q_byte = get_byte(q_word, i % 4u);
- let nib = select(q_byte & 0xFu, (q_byte >> 4u) & 0xFu, half == 1u);
- row_sum += f32(kvalues_iq4nl[nib]) * dl * x_block[i];
- }
- acc[row] += row_sum;
- }
- }
- }
-#endif
+ let acc = accumulate_vec_dot(thread_id, row_base, src0_batch_offset, src1_idx_base);
#ifdef USE_SUBGROUP_REDUCTION
for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
--- /dev/null
+#ifdef U32_DEQUANT_HELPERS
+#define SRC0_TYPE u32
+
+fn byte_of(v: u32, b: u32) -> u32 {
+ return (v >> (b * 8u)) & 0xFFu;
+}
+
+fn sbyte_of(v: u32, b: u32) -> i32 {
+ let raw = i32((v >> (b * 8u)) & 0xFFu);
+ return select(raw, raw - 256, raw >= 128);
+}
+#endif
+
+#ifdef VEC
+#define VEC_SIZE 4u
+#define SRC0_TYPE vec4<SRC0_INNER_TYPE>
+#define SRC1_TYPE vec4<SRC1_INNER_TYPE>
+
+fn inner_dot(src0_val: SRC0_TYPE, src1_val: SRC1_TYPE) -> f32 {
+ return f32(dot(SRC1_TYPE(src0_val), src1_val));
+}
+#endif
+
+#ifdef SCALAR
+#define VEC_SIZE 1u
+#define SRC0_TYPE SRC0_INNER_TYPE
+#define SRC1_TYPE SRC1_INNER_TYPE
+
+fn inner_dot(src0_val: SRC0_TYPE, src1_val: SRC1_TYPE) -> f32 {
+ return f32(src0_val) * f32(src1_val);
+}
+#endif
+
+#ifdef MUL_ACC_FLOAT
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let k_vec = params.k / VEC_SIZE;
+ let src1_idx_base_vec = src1_idx_base / VEC_SIZE;
+
+ // Each thread walks K, loads from the vector, and updates
+ // a small block of output rows held in registers.
+ for (var k = thread_id; k < k_vec; k += WG_SIZE) {
+ let x = src1[src1_idx_base_vec + k];
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let src0_idx = (src0_batch_offset + output_row * params.stride_01) / VEC_SIZE + k;
+ acc[row] += inner_dot(src0[src0_idx], x);
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q1_0
+#define BLOCK_SIZE 128
+#define BLOCK_SIZE_BYTES 18
+#define THREADS_PER_BLOCK 16
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let q_byte = load_u32_at_src0(block_byte_base + 2u + thread_within_block) & 0xFFu;
+ var row_sum = 0.0;
+ for (var bit = 0u; bit < 8u; bit++) {
+ let w = select(-d, d, ((q_byte >> bit) & 1u) != 0u);
+ row_sum += w * x_block[bit];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q4_0
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 18
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % 4;
+ for (var block = thread_id/THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE/THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4] = f32(src1[x_base + i + 16]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ var row_sum = 0.0;
+
+ let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * thread_within_block);
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_byte = get_byte(q_packed, byte_idx);
+ let q_lo = (f32(q_byte & 0xFu) - 8.0) * d;
+ let q_hi = (f32((q_byte >> 4u) & 0xFu) - 8.0) * d;
+ row_sum += q_lo * x_block[byte_idx];
+ row_sum += q_hi * x_block[byte_idx + 4u];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q4_1
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 20
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4] = f32(src1[x_base + i + 16]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let m = f32(load_f16_at_src0(block_byte_base + 2u));
+ var row_sum = 0.0;
+
+ let q_packed = load_u32_at_src0(block_byte_base + 4u + 4u * thread_within_block);
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_byte = get_byte(q_packed, byte_idx);
+ let q_lo = f32(q_byte & 0xFu) * d + m;
+ let q_hi = f32((q_byte >> 4u) & 0xFu) * d + m;
+ row_sum += q_lo * x_block[byte_idx];
+ row_sum += q_hi * x_block[byte_idx + 4u];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q5_0
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 22
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4] = f32(src1[x_base + i + 16]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let qh_packed = load_u32_at_src0(block_byte_base + 2u);
+ let q_packed = load_u32_at_src0(block_byte_base + 6u + 4u * thread_within_block);
+ let qh_shift = thread_within_block * 4u;
+ var row_sum = 0.0;
+
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_byte = get_byte(q_packed, byte_idx);
+ let qh_lo = ((qh_packed >> (qh_shift + byte_idx)) << 4u) & 0x10u;
+ let qh_hi = (qh_packed >> (qh_shift + byte_idx + 12u)) & 0x10u;
+ let q_lo = (f32((q_byte & 0xFu) | qh_lo) - 16.0) * d;
+ let q_hi = (f32(((q_byte >> 4u) & 0xFu) | qh_hi) - 16.0) * d;
+ row_sum += q_lo * x_block[byte_idx];
+ row_sum += q_hi * x_block[byte_idx + 4u];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q5_1
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 24
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD / 2; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4] = f32(src1[x_base + i + 16]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let m = f32(load_f16_at_src0(block_byte_base + 2u));
+ let qh_packed = load_u32_at_src0(block_byte_base + 4u);
+ let q_packed = load_u32_at_src0(block_byte_base + 8u + 4u * thread_within_block);
+ let qh_shift = thread_within_block * 4u;
+ var row_sum = 0.0;
+
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_byte = get_byte(q_packed, byte_idx);
+ let qh_lo = ((qh_packed >> (qh_shift + byte_idx)) << 4u) & 0x10u;
+ let qh_hi = (qh_packed >> (qh_shift + byte_idx + 12u)) & 0x10u;
+ let q_lo = f32((q_byte & 0xFu) | qh_lo) * d + m;
+ let q_hi = f32(((q_byte >> 4u) & 0xFu) | qh_hi) * d + m;
+ row_sum += q_lo * x_block[byte_idx];
+ row_sum += q_hi * x_block[byte_idx + 4u];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q8_0
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 34
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ var row_sum = 0.0;
+
+ for (var packed_idx = 0u; packed_idx < ELEMS_PER_THREAD / 4u; packed_idx++) {
+ let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * (thread_within_block * 2u + packed_idx));
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_val = f32(get_byte_i32(q_packed, byte_idx)) * d;
+ row_sum += q_val * x_block[packed_idx * 4u + byte_idx];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q8_1
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 36
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * ELEMS_PER_THREAD;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let m = f32(load_f16_at_src0(block_byte_base + 2u));
+ var row_sum = 0.0;
+
+ for (var packed_idx = 0u; packed_idx < ELEMS_PER_THREAD / 4u; packed_idx++) {
+ let q_packed = load_u32_at_src0(block_byte_base + 4u + 4u * (thread_within_block * 2u + packed_idx));
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_val = f32(get_byte_i32(q_packed, byte_idx)) * d + m;
+ row_sum += q_val * x_block[packed_idx * 4u + byte_idx];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q2_K
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 84
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let lane = tid / 2u;
+ let phase = tid % 2u;
+ let iq = lane / 4u;
+ let ir = lane % 4u;
+ let is = ir / 2u;
+
+ let y_offset = 128u * iq + 8u * ir + 4u * phase;
+ let sc0_byte = 8u * iq + is;
+ let sc2_byte = 8u * iq + is + 2u;
+ let sc4_byte = 8u * iq + is + 4u;
+ let sc6_byte = 8u * iq + is + 6u;
+ let qs_byte = 16u + (16u * iq + 4u * ir) * 2u + 4u * phase;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 4u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4u] = f32(src1[x_base + 32u + i]);
+ x_block[i + 8u] = f32(src1[x_base + 64u + i]);
+ x_block[i + 12u] = f32(src1[x_base + 96u + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let dall = f32(load_f16_at_src0(block_byte_base + 80u));
+ let dmin = f32(load_f16_at_src0(block_byte_base + 82u)) * (1.0 / 16.0);
+
+ let sc0 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc0_byte), sc0_byte & 3u);
+ let sc2 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc2_byte), sc2_byte & 3u);
+ let sc4 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc4_byte), sc4_byte & 3u);
+ let sc6 = byte_of(load_u32_at_src0_aligned(block_byte_base + sc6_byte), sc6_byte & 3u);
+
+ let q_u32 = load_u32_at_src0_aligned(block_byte_base + qs_byte);
+ let qs0 = q_u32 & 0xFFFFu;
+ let qs1 = q_u32 >> 16u;
+
+ var sumy = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+ var acc1 = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+ var acc2 = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+
+ sumy[0] = x_block[0] + x_block[1] + x_block[2] + x_block[3];
+ sumy[1] = x_block[4] + x_block[5] + x_block[6] + x_block[7];
+ sumy[2] = x_block[8] + x_block[9] + x_block[10] + x_block[11];
+ sumy[3] = x_block[12] + x_block[13] + x_block[14] + x_block[15];
+
+ acc1[0] = x_block[0] * f32(qs0 & 0x0003u) + x_block[2] * f32(qs1 & 0x0003u);
+ acc2[0] = x_block[1] * f32(qs0 & 0x0300u) + x_block[3] * f32(qs1 & 0x0300u);
+ acc1[1] = x_block[4] * f32(qs0 & 0x000Cu) + x_block[6] * f32(qs1 & 0x000Cu);
+ acc2[1] = x_block[5] * f32(qs0 & 0x0C00u) + x_block[7] * f32(qs1 & 0x0C00u);
+ acc1[2] = x_block[8] * f32(qs0 & 0x0030u) + x_block[10] * f32(qs1 & 0x0030u);
+ acc2[2] = x_block[9] * f32(qs0 & 0x3000u) + x_block[11] * f32(qs1 & 0x3000u);
+ acc1[3] = x_block[12] * f32(qs0 & 0x00C0u) + x_block[14] * f32(qs1 & 0x00C0u);
+ acc2[3] = x_block[13] * f32(qs0 & 0xC000u) + x_block[15] * f32(qs1 & 0xC000u);
+
+ acc[row] += dall * ((acc1[0] + (1.0/256.0) * acc2[0]) * f32(sc0 & 0xFu) +
+ (acc1[1] + (1.0/256.0) * acc2[1]) * f32(sc2 & 0xFu) / 4.0 +
+ (acc1[2] + (1.0/256.0) * acc2[2]) * f32(sc4 & 0xFu) / 16.0 +
+ (acc1[3] + (1.0/256.0) * acc2[3]) * f32(sc6 & 0xFu) / 64.0)
+ - dmin * (sumy[0] * f32(sc0 & 0xF0u) + sumy[1] * f32(sc2 & 0xF0u) +
+ sumy[2] * f32(sc4 & 0xF0u) + sumy[3] * f32(sc6 & 0xF0u));
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+
+#ifdef MUL_ACC_Q3_K
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 110
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let lane = tid / 2u;
+ let phase = tid % 2u;
+ let ip = lane / 4u;
+ let il = 2u * ((lane % 4u) / 2u);
+ let ir = lane % 2u;
+ let l0 = 8u * ir;
+
+ let q_byte = 32u + 32u * ip + l0 + 16u * phase;
+ let h_byte = l0 + 16u * phase;
+ let y_offset = 128u * ip + 32u * il + l0 + 16u * phase;
+
+ let s_shift1 = 4u * ip;
+ let s_shift2 = s_shift1 + il;
+
+ let v1 = select(64.0, 4.0, il == 0u);
+ let v2 = 4.0 * v1;
+ let shift = 2u * il;
+
+ var qm0: u32; var qm1: u32; var qm2: u32; var qm3: u32;
+ if (il == 0u) {
+ qm0 = 0x0003u; qm1 = 0x0300u; qm2 = 0x000Cu; qm3 = 0x0C00u;
+ } else {
+ qm0 = 0x0030u; qm1 = 0x3000u; qm2 = 0x00C0u; qm3 = 0xC000u;
+ }
+
+ let mm_idx = 2u * ip + il / 2u;
+ var hm0: u32; var hm1: u32; var hm2: u32; var hm3: u32;
+ switch (mm_idx) {
+ case 0u: { hm0=0x0001u; hm1=0x0100u; hm2=0x0002u; hm3=0x0200u; }
+ case 1u: { hm0=0x0004u; hm1=0x0400u; hm2=0x0008u; hm3=0x0800u; }
+ case 2u: { hm0=0x0010u; hm1=0x1000u; hm2=0x0020u; hm3=0x2000u; }
+ default: { hm0=0x0040u; hm1=0x4000u; hm2=0x0080u; hm3=0x8000u; }
+ }
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 8u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 8u] = f32(src1[x_base + 32u + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let d = f32(load_f16_at_src0(block_byte_base + 108u));
+ let a_base = 96u;
+ let a_il0 = load_u16_at_src0(block_byte_base + a_base + il * 2u);
+ let a_il1 = load_u16_at_src0(block_byte_base + a_base + (il + 1u) * 2u);
+ let a_4 = load_u16_at_src0(block_byte_base + a_base + 8u);
+ let a_5 = load_u16_at_src0(block_byte_base + a_base + 10u);
+
+ var scales32 = a_4 | (a_5 << 16u);
+ let aux32 = ((scales32 >> s_shift2) << 4u) & 0x30303030u;
+ scales32 = a_il0 | (a_il1 << 16u);
+ scales32 = ((scales32 >> s_shift1) & 0x0F0F0F0Fu) | aux32;
+
+ let scale0 = f32(i32(byte_of(scales32, phase + 0u)) - 32);
+ let scale1 = f32(i32(byte_of(scales32, phase + 2u)) - 32);
+
+ let q_u32_0 = load_u32_at_src0(block_byte_base + q_byte + 0u);
+ let q_u32_1 = load_u32_at_src0(block_byte_base + q_byte + 4u);
+ let h_u32_0 = load_u32_at_src0(block_byte_base + h_byte + 0u);
+ let h_u32_1 = load_u32_at_src0(block_byte_base + h_byte + 4u);
+
+ var s1 = 0.0; var s2 = 0.0; var s3 = 0.0;
+ var s4 = 0.0; var s5 = 0.0; var s6 = 0.0;
+
+ for (var l = 0u; l < 8u; l += 2u) {
+ let q_u32 = select(q_u32_0, q_u32_1, l >= 4u);
+ let qs = select(q_u32 & 0xFFFFu, q_u32 >> 16u, (l & 2u) != 0u);
+ let h_u32 = select(h_u32_0, h_u32_1, l >= 4u);
+ let hv = select(h_u32 & 0xFFFFu, h_u32 >> 16u, (l & 2u) != 0u);
+
+ s1 += x_block[l + 0u] * f32(qs & qm0);
+ s2 += x_block[l + 1u] * f32(qs & qm1);
+ s3 += select(0.0, x_block[l + 0u], (hv & hm0) == 0u) +
+ select(0.0, x_block[l + 1u], (hv & hm1) == 0u);
+ s4 += x_block[l + 8u] * f32(qs & qm2);
+ s5 += x_block[l + 9u] * f32(qs & qm3);
+ s6 += select(0.0, x_block[l + 8u], (hv & hm2) == 0u) +
+ select(0.0, x_block[l + 9u], (hv & hm3) == 0u);
+ }
+
+ let d1 = d * (s1 + (1.0/256.0) * s2 - s3 * v1);
+ let d2 = d * (s4 + (1.0/256.0) * s5 - s6 * v2);
+ acc[row] += (d1 * scale0 + 0.25 * d2 * scale1) / f32(1u << shift);
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q4_K
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 144
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let il = tid / 4u;
+ let ir = tid % 4u;
+ let im = il / 2u;
+ let in = il % 2u;
+ let l0 = 4u * (2u * ir + in);
+
+ let y_offset = 64u * im + l0;
+ let q_offset = 32u * im + l0;
+ let sc0_byte = 4u + im * 2u;
+ let sc2_byte = 4u + (im + 2u) * 2u;
+ let sc4_byte = 4u + (im + 4u) * 2u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 4u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4u] = f32(src1[x_base + 32u + i]);
+ x_block[i + 8u] = f32(src1[x_base + 128u + i]);
+ x_block[i + 12u] = f32(src1[x_base + 160u + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let d = f32(load_f16_at_src0(block_byte_base + 0u));
+ let dmin = f32(load_f16_at_src0(block_byte_base + 2u));
+
+ let sc0_u32 = load_u32_at_src0_aligned(block_byte_base + sc0_byte);
+ let sc0 = select(sc0_u32 & 0xFFFFu, sc0_u32 >> 16u, (sc0_byte & 2u) != 0u);
+ let sc2_u32 = load_u32_at_src0_aligned(block_byte_base + sc2_byte);
+ let sc2 = select(sc2_u32 & 0xFFFFu, sc2_u32 >> 16u, (sc2_byte & 2u) != 0u);
+ let sc4_u32 = load_u32_at_src0_aligned(block_byte_base + sc4_byte);
+ let sc4 = select(sc4_u32 & 0xFFFFu, sc4_u32 >> 16u, (sc4_byte & 2u) != 0u);
+
+ let sc16_0 = sc0 & 0x3F3Fu;
+ let sc16_1 = sc2 & 0x3F3Fu;
+ let sc16_2 = (sc4 & 0x0F0Fu) | ((sc0 & 0xC0C0u) >> 2u);
+ let sc16_3 = ((sc4 >> 4u) & 0x0F0Fu) | ((sc2 & 0xC0C0u) >> 2u);
+
+ let scale0 = f32(sc16_0 & 0xFFu);
+ let scale1 = f32((sc16_0 >> 8u) & 0xFFu);
+ let min0 = f32(sc16_1 & 0xFFu);
+ let min1 = f32((sc16_1 >> 8u) & 0xFFu);
+ let scale2 = f32(sc16_2 & 0xFFu);
+ let scale3 = f32((sc16_2 >> 8u) & 0xFFu);
+ let min2 = f32(sc16_3 & 0xFFu);
+ let min3 = f32((sc16_3 >> 8u) & 0xFFu);
+
+ let q1_u32 = load_u32_at_src0_aligned(block_byte_base + 16u + q_offset);
+ let q2_u32 = load_u32_at_src0_aligned(block_byte_base + 80u + q_offset);
+
+ var dot = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+ var sumx = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+ for (var i = 0u; i < 4u; i++) {
+ let q1b = byte_of(q1_u32, i);
+ let q2b = byte_of(q2_u32, i);
+ dot[0] += x_block[i] * f32(q1b & 0x0Fu);
+ dot[1] += x_block[i + 4u] * f32(q1b >> 4u);
+ dot[2] += x_block[i + 8u] * f32(q2b & 0x0Fu);
+ dot[3] += x_block[i + 12u] * f32(q2b >> 4u);
+ sumx[0] += x_block[i];
+ sumx[1] += x_block[i + 4u];
+ sumx[2] += x_block[i + 8u];
+ sumx[3] += x_block[i + 12u];
+ }
+
+ acc[row] += d * (dot[0] * scale0 + dot[1] * scale1 + dot[2] * scale2 + dot[3] * scale3)
+ - dmin * (sumx[0] * min0 + sumx[1] * min1 + sumx[2] * min2 + sumx[3] * min3);
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q5_K
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 176
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let il = tid / 4u;
+ let ir = tid % 4u;
+ let im = il / 2u;
+ let in = il % 2u;
+ let l0 = 4u * (2u * ir + in);
+
+ let y_offset = 64u * im + l0;
+ let q_offset = 48u + 32u * im + l0;
+ let qh_offset = 16u + 8u * ir + 4u * in;
+ let sc0_byte = 4u + im * 2u;
+ let sc2_byte = 4u + (im + 2u) * 2u;
+ let sc4_byte = 4u + (im + 4u) * 2u;
+
+ let hm1 = 1u << (2u * im);
+ let hm2 = hm1 << 1u;
+ let hm3 = hm1 << 4u;
+ let hm4 = hm2 << 4u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 4u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4u] = f32(src1[x_base + 32u + i]);
+ x_block[i + 8u] = f32(src1[x_base + 128u + i]);
+ x_block[i + 12u] = f32(src1[x_base + 160u + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let d = f32(load_f16_at_src0(block_byte_base + 0u));
+ let dmin = f32(load_f16_at_src0(block_byte_base + 2u));
+
+ let sc0_u32 = load_u32_at_src0_aligned(block_byte_base + sc0_byte);
+ let sc0 = select(sc0_u32 & 0xFFFFu, sc0_u32 >> 16u, (sc0_byte & 2u) != 0u);
+ let sc2_u32 = load_u32_at_src0_aligned(block_byte_base + sc2_byte);
+ let sc2 = select(sc2_u32 & 0xFFFFu, sc2_u32 >> 16u, (sc2_byte & 2u) != 0u);
+ let sc4_u32 = load_u32_at_src0_aligned(block_byte_base + sc4_byte);
+ let sc4 = select(sc4_u32 & 0xFFFFu, sc4_u32 >> 16u, (sc4_byte & 2u) != 0u);
+
+ let sc16_0 = sc0 & 0x3F3Fu;
+ let sc16_1 = sc2 & 0x3F3Fu;
+ let sc16_2 = (sc4 & 0x0F0Fu) | ((sc0 & 0xC0C0u) >> 2u);
+ let sc16_3 = ((sc4 >> 4u) & 0x0F0Fu) | ((sc2 & 0xC0C0u) >> 2u);
+
+ let f0 = f32(sc16_0 & 0xFFu);
+ let f1 = f32((sc16_0 >> 8u) & 0xFFu);
+ let m0 = f32(sc16_1 & 0xFFu);
+ let m1 = f32((sc16_1 >> 8u) & 0xFFu);
+ let f4 = f32(sc16_2 & 0xFFu);
+ let f5 = f32((sc16_2 >> 8u) & 0xFFu);
+ let m4 = f32(sc16_3 & 0xFFu);
+ let m5 = f32((sc16_3 >> 8u) & 0xFFu);
+
+ let q1_u32 = load_u32_at_src0_aligned(block_byte_base + q_offset);
+ let q2_u32 = load_u32_at_src0_aligned(block_byte_base + q_offset + 64u);
+ let qh_u32 = load_u32_at_src0_aligned(block_byte_base + qh_offset);
+
+ var vals = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+ var sumy = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+ for (var i = 0u; i < 4u; i++) {
+ let q1b = byte_of(q1_u32, i);
+ let q2b = byte_of(q2_u32, i);
+ let qhb = byte_of(qh_u32, i);
+
+ let yl0 = x_block[i];
+ let yl8 = x_block[i + 4u];
+ let yh0 = x_block[i + 8u];
+ let yh8 = x_block[i + 12u];
+
+ sumy[0] += yl0;
+ sumy[1] += yl8;
+ sumy[2] += yh0;
+ sumy[3] += yh8;
+
+ let q0 = f32((q1b & 0x0Fu) | select(0u, 0x10u, (qhb & hm1) != 0u));
+ let q1 = f32((q1b >> 4u) | select(0u, 0x10u, (qhb & hm2) != 0u));
+ let q2 = f32((q2b & 0x0Fu) | select(0u, 0x10u, (qhb & hm3) != 0u));
+ let q3 = f32((q2b >> 4u) | select(0u, 0x10u, (qhb & hm4) != 0u));
+
+ vals[0] += yl0 * q0;
+ vals[1] += yl8 * q1;
+ vals[2] += yh0 * q2;
+ vals[3] += yh8 * q3;
+ }
+
+ acc[row] += d * (f0 * vals[0] + f1 * vals[1] + f4 * vals[2] + f5 * vals[3])
+ - dmin * (sumy[0] * m0 + sumy[1] * m1 +
+ sumy[2] * m4 + sumy[3] * m5);
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_Q6_K
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 210
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let ip = tid / 8u;
+ let il = tid % 8u;
+ let l0 = 4u * il;
+ let is = 8u * ip + l0 / 16u;
+
+ let y_offset = 128u * ip + l0;
+ let q_offset_l = 64u * ip + l0;
+ let q_offset_h = 32u * ip + l0;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let sc_base_byte = 192u + (is & ~3u);
+ let sc_byte_pos = is & 3u;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var l = 0u; l < 4u; l++) {
+ x_block[l] = f32(src1[x_base + l]);
+ x_block[l + 4u] = f32(src1[x_base + 32u + l]);
+ x_block[l + 8u] = f32(src1[x_base + 64u + l]);
+ x_block[l + 12u] = f32(src1[x_base + 96u + l]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let d = f32(load_f16_at_src0(block_byte_base + 208u));
+ let ql1_u32 = load_u32_at_src0(block_byte_base + q_offset_l);
+ let ql2_u32 = load_u32_at_src0(block_byte_base + q_offset_l + 32u);
+ let qh_u32 = load_u32_at_src0(block_byte_base + 128u + q_offset_h);
+ let sc_u32_0 = load_u32_at_src0(block_byte_base + sc_base_byte);
+ let sc_u32_1 = load_u32_at_src0(block_byte_base + sc_base_byte + 4u);
+
+ let sc0 = sbyte_of(sc_u32_0, sc_byte_pos);
+ let sc2 = sbyte_of(sc_u32_0, sc_byte_pos + 2u);
+ let sc4 = sbyte_of(sc_u32_1, sc_byte_pos);
+ let sc6 = sbyte_of(sc_u32_1, sc_byte_pos + 2u);
+
+ var sums = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+
+ for (var l = 0u; l < 4u; l++) {
+ let q1b = byte_of(ql1_u32, l);
+ let q2b = byte_of(ql2_u32, l);
+ let qhb = byte_of(qh_u32, l);
+
+ let dq0 = f32(i32((q1b & 0x0Fu) | ((qhb & 0x03u) << 4u)) - 32);
+ let dq1 = f32(i32((q2b & 0x0Fu) | ((qhb & 0x0Cu) << 2u)) - 32);
+ let dq2 = f32(i32((q1b >> 4u) | (qhb & 0x30u)) - 32);
+ let dq3 = f32(i32((q2b >> 4u) | ((qhb & 0xC0u) >> 2u)) - 32);
+
+ sums[0] += x_block[l] * dq0;
+ sums[1] += x_block[l + 4u] * dq1;
+ sums[2] += x_block[l + 8u] * dq2;
+ sums[3] += x_block[l + 12u] * dq3;
+ }
+
+ acc[row] += d * (sums[0] * f32(sc0) + sums[1] * f32(sc2) +
+ sums[2] * f32(sc4) + sums[3] * f32(sc6));
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ1_S
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 50
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let qh = load_u32_at_src0(block_byte_base + 34u + sub_blk * 2u) & 0xFFFFu;
+ let dl = d * f32(2u * ((qh >> 12u) & 7u) + 1u);
+ let delta = select(IQ1_DELTA, -IQ1_DELTA, (qh & 0x8000u) != 0u);
+ let qs_w = load_u32_at_src0(block_byte_base + 2u + sub_blk * 4u);
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let qs_byte = get_byte(qs_w, l);
+ let ig = (qs_byte | (((qh >> (3u * l)) & 7u) << 8u)) * 8u;
+ let gw = iq1_grid[ig / 16u];
+ let bit_base = (ig % 16u) * 2u;
+ for (var j = 0u; j < 8u; j++) {
+ let g = (gw >> (bit_base + j * 2u)) & 3u;
+ let gs = select(f32(g), f32(g) - 4.0, (g & 2u) != 0u);
+ row_sum += dl * (gs + delta) * x_block[ll * 8u + j];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ1_M
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 56
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+
+ let sc_lo = load_u32_at_src0(block_byte_base + 48u);
+ let sc_hi = load_u32_at_src0(block_byte_base + 52u);
+ let sc0 = sc_lo & 0xFFFFu;
+ let sc1 = (sc_lo >> 16u) & 0xFFFFu;
+ let sc2 = sc_hi & 0xFFFFu;
+ let sc3 = (sc_hi >> 16u) & 0xFFFFu;
+ let d_bits = (sc0 >> 12u) | ((sc1 >> 8u) & 0xF0u) | ((sc2 >> 4u) & 0xF00u) | (sc3 & 0xF000u);
+ let d = f32(bitcast<vec2<f16>>(d_bits)[0]);
+
+ let sc_u16 = select(select(sc2, sc3, sub_blk >= 6u),
+ select(sc0, sc1, sub_blk >= 2u),
+ sub_blk < 4u);
+
+ let qs_w = load_u32_at_src0(block_byte_base + sub_blk * 4u);
+ let qh = load_u32_at_src0(block_byte_base + 32u + sub_blk * 2u) & 0xFFFFu;
+ let qh_lo = qh & 0xFFu;
+ let qh_hi = (qh >> 8u) & 0xFFu;
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let bit_off = 6u * (sub_blk % 2u) + 3u * (l / 2u);
+ let sub_scale = (sc_u16 >> bit_off) & 0x7u;
+ let dl = d * f32(2u * sub_scale + 1u);
+ let qh_byte = select(qh_lo, qh_hi, l >= 2u);
+ let ll2 = l % 2u;
+ let grid_idx = get_byte(qs_w, l) | (((qh_byte >> (4u * ll2)) & 7u) << 8u);
+ let delta = select(IQ1_DELTA, -IQ1_DELTA, ((qh_byte >> (3u + 4u * ll2)) & 1u) != 0u);
+ let ig = grid_idx * 8u;
+ let gw = iq1_grid[ig / 16u];
+ let bit_base = (ig % 16u) * 2u;
+ for (var j = 0u; j < 8u; j++) {
+ let g = (gw >> (bit_base + j * 2u)) & 3u;
+ let gs = select(f32(g), f32(g) - 4.0, (g & 2u) != 0u);
+ row_sum += dl * (gs + delta) * x_block[ll * 8u + j];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ2_XXS
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 66
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let aux_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
+ let aux_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
+ let ls = aux_hi >> 28u;
+ let db = d * (0.5 + f32(ls)) * 0.25;
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let grid_idx = (aux_lo >> (8u * l)) & 0xFFu;
+ let signs_idx = (aux_hi >> (7u * l)) & 0x7Fu;
+ let signs = (ksigns_iq2xs[signs_idx / 4u] >> ((signs_idx % 4u) * 8u)) & 0xFFu;
+ let gw_lo = iq2xxs_grid[grid_idx * 2u];
+ let gw_hi = iq2xxs_grid[grid_idx * 2u + 1u];
+ for (var j = 0u; j < 8u; j++) {
+ let gw = select(gw_hi, gw_lo, j < 4u);
+ let b = f32((gw >> ((j & 3u) * 8u)) & 0xFFu);
+ let s = select(1.0, -1.0, ((signs >> j) & 1u) != 0u);
+ row_sum += db * b * s * x_block[ll * 8u + j];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ2_XS
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 74
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let qs_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
+ let qs_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
+ let scales_word = load_u32_at_src0(block_byte_base + 66u + (sub_blk / 4u) * 4u);
+ let scales_byte = get_byte(scales_word, sub_blk % 4u);
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let qs_word = select(qs_hi, qs_lo, l < 2u);
+ let half2 = (l % 2u) * 16u;
+ let qs_val = (qs_word >> half2) & 0xFFFFu;
+ let grid_idx = qs_val & 0x1FFu;
+ let signs_idx = (qs_val >> 9u) & 0x7Fu;
+ let sub_scale = (scales_byte >> (4u * (l / 2u))) & 0xFu;
+ let db = d * (0.5 + f32(sub_scale)) * 0.25;
+ let signs = (ksigns_iq2xs[signs_idx / 4u] >> ((signs_idx % 4u) * 8u)) & 0xFFu;
+ let gw_lo = iq2xs_grid[grid_idx * 2u];
+ let gw_hi = iq2xs_grid[grid_idx * 2u + 1u];
+ for (var j = 0u; j < 8u; j++) {
+ let gw = select(gw_hi, gw_lo, j < 4u);
+ let b = f32((gw >> ((j & 3u) * 8u)) & 0xFFu);
+ let s = select(1.0, -1.0, ((signs >> j) & 1u) != 0u);
+ row_sum += db * b * s * x_block[ll * 8u + j];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ2_S
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 82
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let qs_w = load_u32_at_src0(block_byte_base + 2u + sub_blk * 4u);
+ let sg_w = load_u32_at_src0(block_byte_base + 34u + sub_blk * 4u);
+ let qh_word = load_u32_at_src0(block_byte_base + 66u + (sub_blk / 4u) * 4u);
+ let qh_byte = get_byte(qh_word, sub_blk % 4u);
+ let sc_word = load_u32_at_src0(block_byte_base + 74u + (sub_blk / 4u) * 4u);
+ let scales_byte = get_byte(sc_word, sub_blk % 4u);
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let qs_byte = get_byte(qs_w, l);
+ let sign_byte = get_byte(sg_w, l);
+ let grid_idx = qs_byte | (((qh_byte >> (2u * l)) & 3u) << 8u);
+ let sub_scale = (scales_byte >> (4u * (l / 2u))) & 0xFu;
+ let db = d * (0.5 + f32(sub_scale)) * 0.25;
+ let gw_lo = iq2s_grid[grid_idx * 2u];
+ let gw_hi = iq2s_grid[grid_idx * 2u + 1u];
+ for (var j = 0u; j < 8u; j++) {
+ let gw = select(gw_hi, gw_lo, j < 4u);
+ let b = f32((gw >> ((j & 3u) * 8u)) & 0xFFu);
+ let s = select(1.0, -1.0, ((sign_byte >> j) & 1u) != 0u);
+ row_sum += db * b * s * x_block[ll * 8u + j];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ3_XXS
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 98
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let qs_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
+ let qs_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
+ let aux = load_u32_at_src0(block_byte_base + 66u + sub_blk * 4u);
+ let ls = aux >> 28u;
+ let db = d * (0.5 + f32(ls)) * 0.5;
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let qs_word = select(qs_hi, qs_lo, l < 2u);
+ let byte_pos = (l % 2u) * 2u;
+ let grid_idx_0 = (qs_word >> (byte_pos * 8u)) & 0xFFu;
+ let grid_idx_1 = (qs_word >> ((byte_pos + 1u) * 8u)) & 0xFFu;
+ let signs_idx = (aux >> (7u * l)) & 0x7Fu;
+ let signs = (ksigns_iq2xs[signs_idx / 4u] >> ((signs_idx % 4u) * 8u)) & 0xFFu;
+ let grid1 = iq3xxs_grid[grid_idx_0];
+ let grid2 = iq3xxs_grid[grid_idx_1];
+ for (var j = 0u; j < 4u; j++) {
+ let b1 = f32((grid1 >> (j * 8u)) & 0xFFu);
+ let b2 = f32((grid2 >> (j * 8u)) & 0xFFu);
+ let s1 = select(1.0, -1.0, ((signs >> j) & 1u) != 0u);
+ let s2 = select(1.0, -1.0, ((signs >> (j + 4u)) & 1u) != 0u);
+ row_sum += db * b1 * s1 * x_block[ll * 8u + j];
+ row_sum += db * b2 * s2 * x_block[ll * 8u + j + 4u];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ3_S
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 110
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let slot0 = half * 2u;
+ let y_offset = sub_blk * 32u + slot0 * 8u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let qs_lo = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u);
+ let qs_hi = load_u32_at_src0(block_byte_base + 2u + sub_blk * 8u + 4u);
+ let qh_word = load_u32_at_src0(block_byte_base + 66u + (sub_blk / 4u) * 4u);
+ let qh_byte = get_byte(qh_word, sub_blk % 4u);
+ let sg_w = load_u32_at_src0(block_byte_base + 74u + sub_blk * 4u);
+ let sc_word = load_u32_at_src0(block_byte_base + 106u);
+ let scales_byte = get_byte(sc_word, sub_blk / 2u);
+ let sub_scale = (scales_byte >> (4u * (sub_blk % 2u))) & 0xFu;
+ let db = d * (1.0 + 2.0 * f32(sub_scale));
+
+ var row_sum = 0.0;
+ for (var ll = 0u; ll < 2u; ll++) {
+ let l = slot0 + ll;
+ let qs_word = select(qs_hi, qs_lo, l < 2u);
+ let byte_pos = (l % 2u) * 2u;
+ let qs0 = (qs_word >> (byte_pos * 8u)) & 0xFFu;
+ let qs1 = (qs_word >> ((byte_pos + 1u) * 8u)) & 0xFFu;
+ let grid_idx_1 = qs0 | (((qh_byte >> (2u * l)) & 1u) << 8u);
+ let grid_idx_2 = qs1 | (((qh_byte >> (2u * l + 1u)) & 1u) << 8u);
+ let sign_byte = get_byte(sg_w, l);
+ let grid1 = iq3s_grid[grid_idx_1];
+ let grid2 = iq3s_grid[grid_idx_2];
+ for (var j = 0u; j < 4u; j++) {
+ let b1 = f32((grid1 >> (j * 8u)) & 0xFFu);
+ let b2 = f32((grid2 >> (j * 8u)) & 0xFFu);
+ let s1 = select(1.0, -1.0, ((sign_byte >> j) & 1u) != 0u);
+ let s2 = select(1.0, -1.0, ((sign_byte >> (j + 4u)) & 1u) != 0u);
+ row_sum += db * b1 * s1 * x_block[ll * 8u + j];
+ row_sum += db * b2 * s2 * x_block[ll * 8u + j + 4u];
+ }
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ4_NL
+#define BLOCK_SIZE 32
+#define BLOCK_SIZE_BYTES 18
+#define THREADS_PER_BLOCK 4
+#define ELEMS_PER_THREAD (BLOCK_SIZE/THREADS_PER_BLOCK)
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+ let thread_within_block = thread_id % THREADS_PER_BLOCK;
+ for (var block = thread_id / THREADS_PER_BLOCK; block < num_blocks; block += WG_SIZE / THREADS_PER_BLOCK) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + thread_within_block * 4u;
+ var x_block: array<f32, ELEMS_PER_THREAD>;
+ for (var i = 0u; i < ELEMS_PER_THREAD / 2u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ x_block[i + 4u] = f32(src1[x_base + i + 16u]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ var row_sum = 0.0;
+
+ let q_packed = load_u32_at_src0(block_byte_base + 2u + 4u * thread_within_block);
+ for (var byte_idx = 0u; byte_idx < 4u; byte_idx++) {
+ let q_byte = get_byte(q_packed, byte_idx);
+ let q_lo = f32(kvalues_iq4nl[q_byte & 0xFu]) * d;
+ let q_hi = f32(kvalues_iq4nl[(q_byte >> 4u) & 0xFu]) * d;
+ row_sum += q_lo * x_block[byte_idx];
+ row_sum += q_hi * x_block[byte_idx + 4u];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
+
+#ifdef MUL_ACC_IQ4_XS
+#define BLOCK_SIZE 256
+#define BLOCK_SIZE_BYTES 136
+#define THREADS_PER_BLOCK 16
+fn accumulate_vec_dot(thread_id: u32, row_base: u32, src0_batch_offset: u32, src1_idx_base: u32) -> array<f32, OUTPUTS_PER_WG> {
+ var acc: array<f32, OUTPUTS_PER_WG>;
+
+ let tid = thread_id % THREADS_PER_BLOCK;
+ let block_group = thread_id / THREADS_PER_BLOCK;
+ let num_block_groups: u32 = WG_SIZE / THREADS_PER_BLOCK;
+
+ let sub_blk = tid / 2u;
+ let half = tid % 2u;
+ let y_offset = sub_blk * 32u + half * 16u;
+
+ let num_blocks = params.k / BLOCK_SIZE;
+
+ for (var block = block_group; block < num_blocks; block += num_block_groups) {
+ let x_base = src1_idx_base + block * BLOCK_SIZE + y_offset;
+ var x_block: array<f32, 16>;
+ for (var i = 0u; i < 16u; i++) {
+ x_block[i] = f32(src1[x_base + i]);
+ }
+
+ for (var row = 0u; row < OUTPUTS_PER_WG; row++) {
+ let output_row = row_base + row;
+ if (output_row < params.m) {
+ let block_byte_base = (src0_batch_offset + output_row * params.stride_01 + block) * BLOCK_SIZE_BYTES;
+ let d = f32(load_f16_at_src0(block_byte_base));
+ let scales_h = load_u16_at_src0(block_byte_base + 2u);
+ let scales_l_word = load_u32_at_src0(block_byte_base + 4u);
+ let sl_byte = get_byte(scales_l_word, sub_blk / 2u);
+ let sl = (sl_byte >> (4u * (sub_blk % 2u))) & 0xFu;
+ let sh_bits = (scales_h >> (2u * sub_blk)) & 3u;
+ let ls = i32(sl | (sh_bits << 4u));
+ let dl = d * f32(ls - 32);
+
+ let qs_byte_off = 8u + sub_blk * 16u;
+ let q_w0 = load_u32_at_src0(block_byte_base + qs_byte_off);
+ let q_w1 = load_u32_at_src0(block_byte_base + qs_byte_off + 4u);
+ let q_w2 = load_u32_at_src0(block_byte_base + qs_byte_off + 8u);
+ let q_w3 = load_u32_at_src0(block_byte_base + qs_byte_off + 12u);
+
+ var row_sum = 0.0;
+ for (var i = 0u; i < 16u; i++) {
+ let q_word = select(
+ select(q_w0, q_w1, i >= 4u),
+ select(q_w2, q_w3, i >= 12u),
+ i >= 8u);
+ let q_byte = get_byte(q_word, i % 4u);
+ let nib = select(q_byte & 0xFu, (q_byte >> 4u) & 0xFu, half == 1u);
+ row_sum += f32(kvalues_iq4nl[nib]) * dl * x_block[i];
+ }
+ acc[row] += row_sum;
+ }
+ }
+ }
+
+ return acc;
+}
+#endif
overview / pkg / ggml / sources / llama.cpp / commitdiff