From fbaefe0eaaccbb6ab9d733ad84580428f18db410 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Mon, 24 Aug 2026 10:46:57 +0200 Subject: [PATCH] revert load reordering and scale pre-loading --- .../vulkan-shaders/mul_mmq_cm1.comp | 58 +++++-------------- 1 file changed, 15 insertions(+), 43 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp index 7185a3a686..44fadf33de 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp @@ -350,61 +350,33 @@ void main() { USE_MAGIC_BIAS ? ACC_BIAS_BITS : 0); } - const uint K_SUB = BK / TK; - coopmat all_a[cms_per_row * K_SUB]; - coopmat all_b[cms_per_col * K_SUB]; + coopmat cache_a; + coopmat cache_b; - [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { - [[unroll]] for (uint h = 0; h < K_SUB; h++) { - coopMatLoad(all_a[cm_row * K_SUB + h], buf_a_qs, (warp_r * WM + cm_row * TM) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutRowMajor); - } - } - [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { - [[unroll]] for (uint h = 0; h < K_SUB; h++) { - coopMatLoad(all_b[cm_col * K_SUB + h], buf_b_qs, (warp_c * WN + cm_col * TN) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutColumnMajor); - } - } + [[unroll]] for (uint i = 0; i < BK; i += TK) { + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + coopMatLoad(cache_a, buf_a_qs, (warp_r * WM + cm_row * TM) * QPITCH + ks * (BK / 4) + i / 4, QPITCH, gl_CooperativeMatrixLayoutRowMajor); - [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { - [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { - [[unroll]] for (uint h = 0; h < K_SUB; h++) { - cm_result[cm_col * cms_per_row + cm_row] = coopMatMulAdd(all_a[cm_row * K_SUB + h], all_b[cm_col * K_SUB + h], cm_result[cm_col * cms_per_row + cm_row]); + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + coopMatLoad(cache_b, buf_b_qs, (warp_c * WN + cm_col * TN) * QPITCH + ks * (BK / 4) + i / 4, QPITCH, gl_CooperativeMatrixLayoutColumnMajor); + + cm_result[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, cm_result[cm_col * cms_per_row + cm_row]); } } } - // Pre-load scales into registers - float scale_a[cms_per_row * CM_ELEMS]; - float nbias_a[cms_per_row * CM_ELEMS]; - float scale_b[cms_per_col * CM_ELEMS]; - [[unroll]] for (uint r = 0; r < cms_per_row; r++) { - [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { - scale_a[r * CM_ELEMS + e] = buf_a_d[ks * BM + warp_r * WM + r * TM + elem_row[e]]; - if (USE_MAGIC_BIAS) { - nbias_a[r * CM_ELEMS + e] = -ACC_BIAS_F * scale_a[r * CM_ELEMS + e]; - } - } - } - [[unroll]] for (uint c = 0; c < cms_per_col; c++) { - [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { - scale_b[c * CM_ELEMS + e] = buf_b_d[ks * BN + warp_c * WN + c * TN + elem_col[e]]; - } - } - - // Apply scales from registers + // Apply scales directly from coopmat elements [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { const uint tile_idx = cm_col * cms_per_row + cm_row; [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { + const float da = buf_a_d[ks * BM + warp_r * WM + cm_row * TM + elem_row[e]]; + const float db = buf_b_d[ks * BN + warp_c * WN + cm_col * TN + elem_col[e]]; if (USE_MAGIC_BIAS) { - const float t = fma(intBitsToFloat(int(cm_result[tile_idx][e])), - scale_a[cm_row * CM_ELEMS + e], - nbias_a[cm_row * CM_ELEMS + e]); - sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, scale_b[cm_col * CM_ELEMS + e], - float(sums[tile_idx * CM_ELEMS + e]))); + const float t = fma(intBitsToFloat(int(cm_result[tile_idx][e])), da, -ACC_BIAS_F * da); + sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, db, float(sums[tile_idx * CM_ELEMS + e]))); } else { - sums[tile_idx * CM_ELEMS + e] += ACC_TYPE(float(cm_result[tile_idx][e]) - * scale_a[cm_row * CM_ELEMS + e] * scale_b[cm_col * CM_ELEMS + e]); + sums[tile_idx * CM_ELEMS + e] += ACC_TYPE(float(cm_result[tile_idx][e]) * da * db); } } }