From 6421fbe532a927e2ad080e1860d578ae043fc3af Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Fri, 28 Aug 2026 16:02:01 +0200 Subject: [PATCH] clean up --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1 + .../vulkan-shaders/mul_mmq_cm1.comp | 87 ++++++++----------- .../vulkan-shaders/mul_mmq_cm1_funcs.glsl | 54 +++++------- 3 files changed, 60 insertions(+), 82 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index cf76db9a09..8534e9e232 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4888,6 +4888,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } else #endif { + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4], matmul_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4], matmul_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); } 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 bd41b1c186..8c195e0f12 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp @@ -10,7 +10,6 @@ #extension GL_KHR_memory_scope_semantics : enable #if defined(MUL_MAT_ID_USE_SUBGROUPS) -#extension GL_KHR_shader_subgroup_basic : enable #extension GL_KHR_shader_subgroup_ballot : enable #endif @@ -91,10 +90,9 @@ layout (constant_id = 11) const uint CM_ARCH = 0; // vk_device_architecture (gg const uint QPITCH = BK_STEP * (BK / 4) + 4; -// Shared memory cache shared uint32_t buf_a_qs[BM * QPITCH]; #if defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1) || defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K) -shared vec2 buf_a_dm[BM * BK_STEP]; // .x = d (== old buf_a_d), .y = m (== old buf_a_m) +shared vec2 buf_a_dm[BM * BK_STEP]; // .x = d, .y = m #else shared float buf_a_d[BM * BK_STEP * KSCALES]; #endif @@ -122,13 +120,21 @@ const uint CM_ELEMS = (TM * TN) / WARP; #define ACC_BIAS_F 12582912.0f const bool USE_MAGIC_BIAS = WARP != 32; -// Coopmat accumulator layout (RDNA): col constant per lane, 8 distinct rows. -// Row order differs by arch: RDNA4 blocked, RDNA3/3.5 interleaved. +// Accumulator row for element e: RDNA4 blocked, RDNA3/3.5 interleaved. uint cm_elem_row(uint e) { const uint row_half = gl_SubgroupInvocationID / TN; return (CM_ARCH == CM_ARCH_AMD_RDNA4) ? (e + row_half * CM_ELEMS) : (row_half + 2u * e); } +// min_term = asymmetric-quant min*b_sum correction (0 for symmetric types). +ACC_TYPE cm1_accumulate(ACC_TYPE prev, int acc_e, float scale_a, float nbias_a, float scale_b, float min_term) { + if (USE_MAGIC_BIAS) { + const float t = fma(intBitsToFloat(acc_e), scale_a, nbias_a); + return ACC_TYPE(fma(t, scale_b, float(prev) + min_term)); + } + return prev + ACC_TYPE(fma(float(acc_e) * scale_a, scale_b, min_term)); +} + #ifdef MUL_MAT_ID #define NUM_WARPS (BLOCK_SIZE / WARP) #include "mul_mm_id_funcs.glsl" @@ -237,7 +243,6 @@ void main() { barrier(); #endif - // Workgroup has no work if (ic * BN >= _ne1) return; #endif @@ -274,7 +279,6 @@ void main() { block_a_prefetch pre_a[A_LOADS * BK_STEP]; block_b_prefetch pre_b[B_LOADS * BK_STEP]; - // Prefetch first block if (start_k < end_k) { PREFETCH_BLOCK(start_k) } @@ -290,7 +294,6 @@ void main() { barrier(); for (uint block = start_k; block < end_k; block += BK * BK_STEP) { - // Store prefetched data to shmem STORE_BLOCK_TO_LDS(block) barrier(); @@ -298,14 +301,12 @@ void main() { pos_a_ib += BK_STEP; pos_b_ib += BK_STEP; - // Prefetch next block (overlaps with compute) const uint next_block = block + BK * BK_STEP; if (next_block < end_k) { PREFETCH_BLOCK(next_block) } if (active_col_tile) { - // Compute from shmem [[unroll]] for (uint ks = 0; ks < BK_STEP; ks++) { const uint K_SUB = BK / TK; @@ -337,35 +338,37 @@ void main() { const uint tile_idx = r * cms_per_col + c; [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { - if (USE_MAGIC_BIAS) { - const float t = fma(intBitsToFloat(int(acc[e])), - scale_a[e], - nbias_a[e]); - sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, scale_b_v, - float(sums[tile_idx * CM_ELEMS + e]))); - } else { - sums[tile_idx * CM_ELEMS + e] += ACC_TYPE(float(acc[e]) - * scale_a[e] * scale_b_v); - } + sums[tile_idx * CM_ELEMS + e] = cm1_accumulate( + sums[tile_idx * CM_ELEMS + e], int(acc[e]), + scale_a[e], nbias_a[e], scale_b_v, 0.0); } } } } #elif defined(DATA_A_Q4_1) || defined(DATA_A_Q5_1) || defined(DATA_A_Q4_K) || defined(DATA_A_Q5_K) - // Preload all B fragments for this ks to avoid re-loading them per r (halves B LDS reads). + // Preload all A/B fragments up front (ILP). + coopmat cache_a[cms_per_row * K_SUB]; coopmat cache_b[cms_per_col * K_SUB]; + + [[unroll]] for (uint r = 0; r < cms_per_row; r++) { + [[unroll]] for (uint h = 0; h < K_SUB; h++) { + coopMatLoad(cache_a[r * K_SUB + h], buf_a_qs, (a_row0 + r * TM) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutRowMajor); + } + } [[unroll]] for (uint c = 0; c < cms_per_col; c++) { [[unroll]] for (uint h = 0; h < K_SUB; h++) { coopMatLoad(cache_b[c * K_SUB + h], buf_b_qs, (b_col0 + c * TN) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutColumnMajor); } } - [[unroll]] for (uint r = 0; r < cms_per_row; r++) { - coopmat cache_a[K_SUB]; - [[unroll]] for (uint h = 0; h < K_SUB; h++) { - coopMatLoad(cache_a[h], buf_a_qs, (a_row0 + r * TM) * QPITCH + ks * (BK / 4) + h * (TK / 4), QPITCH, gl_CooperativeMatrixLayoutRowMajor); - } + float scale_b[cms_per_col]; + float bs[cms_per_col]; + [[unroll]] for (uint c = 0; c < cms_per_col; c++) { + scale_b[c] = buf_b_d[ks * BN + b_col0 + c * TN + elem_col0]; + bs[c] = float(buf_b_s[ks * BN + b_col0 + c * TN + elem_col0]); + } + [[unroll]] for (uint r = 0; r < cms_per_row; r++) { float scale_a[CM_ELEMS]; float nbias_a[CM_ELEMS]; float ma[CM_ELEMS]; @@ -384,28 +387,19 @@ void main() { USE_MAGIC_BIAS ? ACC_BIAS_BITS : 0); [[unroll]] for (uint h = 0; h < K_SUB; h++) { - acc = coopMatMulAdd(cache_a[h], cache_b[c * K_SUB + h], acc); + acc = coopMatMulAdd(cache_a[r * K_SUB + h], cache_b[c * K_SUB + h], acc); } const uint tile_idx = r * cms_per_col + c; - const float sb = buf_b_d[ks * BN + b_col0 + c * TN + elem_col0]; - const float bs = float(buf_b_s[ks * BN + b_col0 + c * TN + elem_col0]); [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { - if (USE_MAGIC_BIAS) { - const float t = fma(intBitsToFloat(int(acc[e])), - scale_a[e], nbias_a[e]); - sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, sb, - fma(ma[e], bs, - float(sums[tile_idx * CM_ELEMS + e])))); - } else { - sums[tile_idx * CM_ELEMS + e] += ACC_TYPE( - fma(float(acc[e]) * scale_a[e], sb, ma[e] * bs)); - } + sums[tile_idx * CM_ELEMS + e] = cm1_accumulate( + sums[tile_idx * CM_ELEMS + e], int(acc[e]), + scale_a[e], nbias_a[e], scale_b[c], ma[e] * bs[c]); } } } #else - // Simple types: pre-load all caches for better data reuse + // Preload all A/B fragments up front (ILP). coopmat cache_a[cms_per_row * K_SUB]; coopmat cache_b[cms_per_col * K_SUB]; @@ -446,16 +440,9 @@ void main() { const uint tile_idx = r * cms_per_col + c; [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { - if (USE_MAGIC_BIAS) { - const float t = fma(intBitsToFloat(int(acc[e])), - scale_a[e], - nbias_a[e]); - sums[tile_idx * CM_ELEMS + e] = ACC_TYPE(fma(t, scale_b[c], - float(sums[tile_idx * CM_ELEMS + e]))); - } else { - sums[tile_idx * CM_ELEMS + e] += ACC_TYPE(float(acc[e]) - * scale_a[e] * scale_b[c]); - } + sums[tile_idx * CM_ELEMS + e] = cm1_accumulate( + sums[tile_idx * CM_ELEMS + e], int(acc[e]), + scale_a[e], nbias_a[e], scale_b[c], 0.0); } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1_funcs.glsl index c9d2da0fef..3b07934f98 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1_funcs.glsl @@ -47,16 +47,14 @@ block_a_prefetch block_a_load(uint ib, uint loadr) { } void block_a_to_shmem(block_a_prefetch blk, uint buf_ib, uint ks, uint loadr) { + // Store raw unsigned nibbles; the -8 offset is absorbed by the min term. uint32_t lo4 = blk.qs & 0x0F0F0F0F; uint32_t hi4 = (blk.qs >> 4) & 0x0F0F0F0F; - lo4 = ((lo4 | 0x80808080) - 0x08080808) ^ 0x80808080; - hi4 = ((hi4 | 0x80808080) - 0x08080808) ^ 0x80808080; buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr ] = lo4; buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr + 4] = hi4; if (loadr == 0) { - const float d = float(blk.dm.x); - buf_a_dm[ks * BM + buf_ib] = vec2(d, 8.0 * d + float(blk.dm.y)); + buf_a_dm[ks * BM + buf_ib] = vec2(float(blk.dm.x), float(blk.dm.y)); } } @@ -109,18 +107,16 @@ block_a_prefetch block_a_load(uint ib, uint loadr) { } void block_a_to_shmem(block_a_prefetch blk, uint buf_ib, uint ks, uint loadr) { + // Store raw unsigned 5-bit values; the -16 offset is absorbed by the min term. uint32_t lo4 = blk.qs & 0x0F0F0F0F; uint32_t hi4 = (blk.qs >> 4) & 0x0F0F0F0F; lo4 |= ((blk.qh >> (4u * loadr )) & 0xFu) * 0x02040810u & 0x10101010u; hi4 |= ((blk.qh >> (4u * loadr + 16u )) & 0xFu) * 0x02040810u & 0x10101010u; - lo4 = ((lo4 | 0x80808080) - 0x10101010) ^ 0x80808080; - hi4 = ((hi4 | 0x80808080) - 0x10101010) ^ 0x80808080; buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr ] = lo4; buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr + 4] = hi4; if (loadr == 0) { - const float d = float(blk.dm.x); - buf_a_dm[ks * BM + buf_ib] = vec2(d, 16.0 * d + float(blk.dm.y)); + buf_a_dm[ks * BM + buf_ib] = vec2(float(blk.dm.x), float(blk.dm.y)); } } @@ -240,25 +236,22 @@ block_a_prefetch block_a_load(uint ib, uint loadr) { } void block_a_to_shmem(block_a_prefetch blk, uint buf_ib, uint ks, uint loadr) { - uint32_t v0 = ((blk.qs0 | 0x80808080) - 0x08080808) ^ 0x80808080; - uint32_t v1 = ((blk.qs1 | 0x80808080) - 0x08080808) ^ 0x80808080; - buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr * 2 ] = v0; - buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr * 2 + 1] = v1; + // Store raw unsigned nibbles (blk.qs already masked); no -8 recentering needed. + buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr * 2 ] = blk.qs0; + buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr * 2 + 1] = blk.qs1; if (loadr == 0) { const uint ib_k = blk.ib / 8; const uint sub = blk.ib % 8; - uint sc_val, mn_val; - if (sub < 4) { - sc_val = uint(data_a[ib_k].scales[sub]) & 0x3Fu; - mn_val = uint(data_a[ib_k].scales[sub + 4]) & 0x3Fu; - } else { - sc_val = (uint(data_a[ib_k].scales[sub + 4]) & 0xFu) | ((uint(data_a[ib_k].scales[sub - 4]) & 0xC0u) >> 2); - mn_val = (uint(data_a[ib_k].scales[sub + 4]) >> 4) | ((uint(data_a[ib_k].scales[sub]) & 0xC0u) >> 2); - } + const uint j = sub & 3u; + const uint s_j = uint(data_a[ib_k].scales[j]); + const uint s_j4 = uint(data_a[ib_k].scales[j + 4]); + const uint s_j8 = uint(data_a[ib_k].scales[j + 8]); + const uint sc_val = (sub < 4) ? (s_j & 0x3Fu) : ((s_j8 & 0x0Fu) | ((s_j >> 6) << 4)); + const uint mn_val = (sub < 4) ? (s_j4 & 0x3Fu) : ((s_j8 >> 4) | ((s_j4 >> 6) << 4)); vec2 dm = vec2(data_a_packed32[ib_k].dm); float d_scaled = dm.x * float(sc_val); - buf_a_dm[ks * BM + buf_ib] = vec2(d_scaled, 8.0 * d_scaled - (dm.y * float(mn_val))); + buf_a_dm[ks * BM + buf_ib] = vec2(d_scaled, -(dm.y * float(mn_val))); } } @@ -295,27 +288,24 @@ block_a_prefetch block_a_load(uint ib, uint loadr) { } void block_a_to_shmem(block_a_prefetch blk, uint buf_ib, uint ks, uint loadr) { + // Store raw unsigned 5-bit values (qs nibble | qh bit); no -16 recentering needed. uint32_t v0 = blk.qs0 | blk.qh0; uint32_t v1 = blk.qs1 | blk.qh1; - v0 = ((v0 | 0x80808080) - 0x10101010) ^ 0x80808080; - v1 = ((v1 | 0x80808080) - 0x10101010) ^ 0x80808080; buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr * 2 ] = v0; buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr * 2 + 1] = v1; if (loadr == 0) { const uint ib_k = blk.ib / 8; const uint sub = blk.ib % 8; - uint sc_val, mn_val; - if (sub < 4) { - sc_val = uint(data_a[ib_k].scales[sub]) & 0x3Fu; - mn_val = uint(data_a[ib_k].scales[sub + 4]) & 0x3Fu; - } else { - sc_val = (uint(data_a[ib_k].scales[sub + 4]) & 0xFu) | ((uint(data_a[ib_k].scales[sub - 4]) & 0xC0u) >> 2); - mn_val = (uint(data_a[ib_k].scales[sub + 4]) >> 4) | ((uint(data_a[ib_k].scales[sub]) & 0xC0u) >> 2); - } + const uint j = sub & 3u; + const uint s_j = uint(data_a[ib_k].scales[j]); + const uint s_j4 = uint(data_a[ib_k].scales[j + 4]); + const uint s_j8 = uint(data_a[ib_k].scales[j + 8]); + const uint sc_val = (sub < 4) ? (s_j & 0x3Fu) : ((s_j8 & 0x0Fu) | ((s_j >> 6) << 4)); + const uint mn_val = (sub < 4) ? (s_j4 & 0x3Fu) : ((s_j8 >> 4) | ((s_j4 >> 6) << 4)); vec2 dm = vec2(data_a_packed32[ib_k].dm); float d_scaled = dm.x * float(sc_val); - buf_a_dm[ks * BM + buf_ib] = vec2(d_scaled, 16.0 * d_scaled - (dm.y * float(mn_val))); + buf_a_dm[ks * BM + buf_ib] = vec2(d_scaled, -(dm.y * float(mn_val))); } }