From 1a63d545a381f7573f8b2f501e83a6c394c77e1e Mon Sep 17 00:00:00 2001 From: "Jiang, Fish" Date: Wed, 10 Jun 2026 15:44:37 +0800 Subject: [PATCH 1/5] vulkan: GEMM/Group GEMM optimizations for Intel Xe (3/3, Xe-LPG Plus/Xe2/Xe3) Co-authored-by: Xia, Jie Co-authored-by: Liu, Russell --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 288 +++++++++++++++--- .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 40 ++- .../vulkan-shaders/mul_mm_funcs.glsl | 215 ++++++++++--- .../vulkan-shaders/vulkan-shaders-gen.cpp | 93 ++++++ 4 files changed, 540 insertions(+), 96 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c0ab9f1c658..a3c166680eb 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -198,6 +198,8 @@ static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); struct vk_matmul_pipeline_struct { vk_pipeline l, m, s; vk_pipeline a_l, a_m, a_s; + // Selected at dispatch by ggml_vk_guess_matmul_pipeline under some condition. + vk_pipeline l_alt, a_l_alt; // Returns true when all unaligned pipelines are null. // We only check for unaligned variants since one of the unaligned pipelines must exist // while aligned pipelines are optional @@ -784,6 +786,7 @@ struct vk_device_struct { vk_matmul_pipeline2 pipeline_matmul_id_f16_f32; vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_COUNT]; // f16 B-type variant (coopmat1 only) vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_COUNT]; vk_pipeline pipeline_matmul_split_k_reduce; @@ -793,6 +796,8 @@ struct vk_device_struct { vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT]; + // Q4_K f32_f32 alt with NUM_ROWS=4 register-tile on M; selected when ne00<=1024. + vk_pipeline pipeline_dequant_mul_mat_vec_q4_k_4rows_f32_f32[DMMV_WG_SIZE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_q8_1_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_id_q8_1_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT]; @@ -3950,6 +3955,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { m_warptile_mmqid_int_k = { 128, 64, 64, 32, mul_mat_subgroup_size_16, 32, 1, 2, 2, 1, mul_mat_subgroup_size_16 }; s_warptile_mmqid_int_k = { mul_mat_subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 1, 2, 1, 1, mul_mat_subgroup_size_16 }; + l_mmq_wg_denoms = l_wg_denoms = { 128, 128, 1 }; + m_mmq_wg_denoms = m_wg_denoms = { 64, 64, 1 }; + s_mmq_wg_denoms = s_wg_denoms = { 32, 32, 1 }; + l_align = 128; + m_align = 64; + s_align = 32; + // chip specific tuning if ((device->architecture == AMD_GCN) && (device->driver_id != vk::DriverId::eAmdProprietary)) { m_warptile_mmq = m_warptile_mmq_int = { 256, 64, 64, 32, 16, 16, 2, 2, 2, 1, 16 }; @@ -3960,18 +3972,19 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; l_warptile_mmq_int_k = { 256, 128, 128, 32, subgroup_size_16, 64, 1, 4, 2, 1, subgroup_size_16 }; } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { - // Xe2/Xe3 with coopmat enabled - warptile performance tuning + // Xe1/Xe2/Xe3 with coopmat enabled - warptile performance tuning l_warptile = { 512, 128, 128, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - l_warptile_mmq = { 512, 128, 128, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; + if (device->architecture == INTEL_PRE_XE2) { + l_warptile_mmq = { 512, 256, 128, 32, 32, 32, 2, 8, 8, 16, 16 }; + l_mmq_wg_denoms = { 256, 128, 1 }; + l_align = 32; //set as BK + } else { + l_warptile_mmq = { 512, 128, 256, 32, 32, 32, 2, 8, 16, 16, 16 }; + l_mmq_wg_denoms = { 128, 256, 1 }; + l_align = 32; //set as BK + } } - l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; - m_mmq_wg_denoms = m_wg_denoms = { 64, 64, 1 }; - s_mmq_wg_denoms = s_wg_denoms = { 32, 32, 1 }; - l_align = 128; - m_align = 64; - s_align = 32; - for (uint32_t i = 0; i < GGML_TYPE_COUNT; ++i) { ggml_type t = (ggml_type)i; // Disable medium and large matrix multiplication if not enough shared memory is available @@ -4349,7 +4362,33 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->coopmat_support) { + // set_warp_tile: save current tile state and apply override. + // restore_warp_tile: revert to the previously saved state. + // The two lambdas share saved state; callers must not nest two active overrides simultaneously. + std::vector saved_wt; + std::array saved_wgd{}; + uint32_t saved_al = 0; + + auto set_warp_tile = [&](std::vector wt, + std::array wgd, + uint32_t al) { + saved_wt = l_warptile_mmq; + saved_wgd = l_mmq_wg_denoms; + saved_al = l_align; + l_warptile_mmq = std::move(wt); + l_mmq_wg_denoms = wgd; + l_align = al; + }; + + auto restore_warp_tile = [&]() { + l_warptile_mmq = std::move(saved_wt); + l_mmq_wg_denoms = saved_wgd; + l_align = saved_al; + }; + // Create 6 variants, {s,m,l}x{unaligned,aligned} + // required_subgroup_size is derived from the WARP element (last) of the warptile, + // so it is always consistent with the tile and no separate tracking variable is needed. #define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ if (device->mul_mat ## ID ## _l[TYPE]) \ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true); \ @@ -4373,6 +4412,22 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(TYPE, PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ } \ + // build the alt l/a_l pipelines into the ->l_alt / ->a_l_alt slots. + // Same shader SPV as ->l/->a_l, but compiled with a different warp tile via the saved l_warptile_mmq. +#define CREATE_MM_ALT(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l_alt, #NAMELC #F16ACC "_l_alt", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, true, (l_ ## WARPTILE).back()); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l_alt, #NAMELC #F16ACC "_aligned_l_alt", NAMELC ## _aligned ## F16ACC ## _cm1_len, NAMELC ## _aligned ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, true, (l_ ## WARPTILE).back()); + +#define CREATE_MM2_ALT(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->coopmat_acc_f16_support) { \ + CREATE_MM_ALT(TYPE, PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + if (device->coopmat_acc_f32_support) { \ + CREATE_MM_ALT(TYPE, PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); @@ -4433,6 +4488,81 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4].f32acc, matmul_nvfp4_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); } + // f16 B-type dense GEMM pipelines for coopmat1 (used when y_non_contig auto-converts f32→f16) + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q1_0], matmul_q1_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + // Xe3 LPG: q8_0 uses BM=BN=128 BK=64 tile; restore brings back the 128x256 xe2 base used by K-quants/IQ*. + if (device->architecture == INTEL_XE2 && device->coopmat_support && device->shader_core_count == 12) set_warp_tile({256, 128, 128, 64, 32, 32, 2, 8, 16, 16, 16}, {128, 128, 1}, 64); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + if (device->architecture == INTEL_XE2 && device->coopmat_support && device->shader_core_count == 12) restore_warp_tile(); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + // Xe1: Q4_K/Q6_K use sgs=32 with a 128x128 tile; IQ* reverts to the xe1 standard tile. + if (device->architecture == INTEL_PRE_XE2) set_warp_tile({512, 128, 128, 32, 32, 32, 2, 8, 8, 16, 32}, {128, 128, 1}, l_align); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K], matmul_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + if (device->architecture == INTEL_PRE_XE2) restore_warp_tile(); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_S], matmul_iq1_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_M], matmul_iq1_m_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_S], matmul_iq2_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_S], matmul_iq3_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_NVFP4], matmul_nvfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // build the ALT l/a_l pipelines into ->l_alt / ->a_l_alt for f16-B-type + // dense quant matmul (pipeline_dequant_mul_mat_mat_f16[*]). The dispatch in + // ggml_vk_guess_matmul_pipeline picks these when n<256 && m<4096. + // Uses the BM=128/WM=16 warp tile. + if (device->vendor_id == VK_VENDOR_ID_INTEL) { + // Intel XE2 coopmat shape: TM=8, TN=16, TK=16 (tm_m/tn_m/tk_m are out of scope here) + set_warp_tile({ 512, 128, 128, 32, 16, 32, 2, 8, 16, 16, 16 }, { 128, 128, 1 }, 32); + + CREATE_MM2_ALT(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K], matmul_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_S], matmul_iq1_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_M], matmul_iq1_m_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_S], matmul_iq2_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_S], matmul_iq3_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2_ALT(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + restore_warp_tile(); + } + + // Intel matmul_id warptile tuning + if (device->vendor_id == VK_VENDOR_ID_INTEL) { + if (device->architecture == INTEL_PRE_XE2) { + l_warptile_mmq = { 512, 128, 128, 32, 32, 32, 2, 8, 8, 16, 32 }; + l_mmq_wg_denoms = { 128, 128, 1 }; + } + else if (device->architecture == INTEL_XE2) { + l_warptile_mmq = { 512, 128, 128, 32, 32, 32, 2, 8, 16, 16, 32 }; + l_mmq_wg_denoms = { 128, 128, 1 }; + } + l_align = 32; //set as BK + } + GGML_ASSERT(device->subgroup_ballot); CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_subgroup_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); @@ -4464,10 +4594,16 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); - CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); #undef CREATE_MM2 #undef CREATE_MM +#undef CREATE_MM2_ALT +#undef CREATE_MM_ALT } else #endif // defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->fp16) { @@ -4869,6 +5005,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32", arr_dmmv_q2_k_f32_f32_len[reduc16], arr_dmmv_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32", arr_dmmv_q3_k_f32_f32_len[reduc16], arr_dmmv_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32", arr_dmmv_q4_k_f32_f32_len[reduc16], arr_dmmv_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + // Q4_K f32_f32 alt: NUM_ROWS=4 register-tile on M. Selected at dispatch when k<=1024. + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q4_k_4rows_f32_f32[w][i], "mul_mat_vec_q4_k_4rows_f32_f32", arr_dmmv_q4_k_f32_f32_len[reduc16], arr_dmmv_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {4, 1, 1}, {wg_size_subgroup16, 4, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32", arr_dmmv_q5_k_f32_f32_len[reduc16], arr_dmmv_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f32_f32", arr_dmmv_q6_k_f32_f32_len[reduc16], arr_dmmv_q6_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ1_S][i], "mul_mat_vec_iq1_s_f32_f32", arr_dmmv_iq1_s_f32_f32_len[reduc16], arr_dmmv_iq1_s_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -7190,7 +7328,7 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte return pipelines; } - if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2) { + if (src1_type != GGML_TYPE_F32 && src1_type != GGML_TYPE_F16 && !ctx->device->coopmat2) { return nullptr; } @@ -7227,6 +7365,9 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_conte return prec == GGML_PREC_DEFAULT ? ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f32acc; } if (ctx->device->coopmat_support) { + if (src1_type == GGML_TYPE_F16) { + return (ctx->device->fp16 && ctx->device->coopmat_acc_f16_support && prec == GGML_PREC_DEFAULT) ? ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f32acc; + } return (ctx->device->fp16 && ctx->device->coopmat_acc_f16_support && prec == GGML_PREC_DEFAULT) ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; } return (ctx->device->fp16 && prec == GGML_PREC_DEFAULT) ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; @@ -7313,6 +7454,14 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * return ctx->device->pipeline_dequant_mul_mat_vec_q8_1_f32[dmmv_wg][a_type][num_cols-1]; } + // Q4_K f32_f32 with small K wins from NUM_ROWS=4 register-tiling on M (Qwen3-0.6B shapes). + if (a_type == GGML_TYPE_Q4_K && b_type == GGML_TYPE_F32 && k <= 1024) { + vk_pipeline alt = ctx->device->pipeline_dequant_mul_mat_vec_q4_k_4rows_f32_f32[dmmv_wg][num_cols-1]; + if (alt) { + return alt; + } + } + return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32[dmmv_wg][a_type][num_cols-1] : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32[dmmv_wg][a_type][num_cols-1]; } @@ -7351,19 +7500,16 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co return pipelines; } - GGML_ASSERT(src1_type == GGML_TYPE_F32 || (ctx->device->coopmat2 && src1_type == GGML_TYPE_F16)); + GGML_ASSERT(src1_type == GGML_TYPE_F32 || src1_type == GGML_TYPE_F16); switch (src0_type) { case GGML_TYPE_Q1_0: case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: - 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: @@ -7374,7 +7520,27 @@ static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_co case GGML_TYPE_IQ3_S: case GGML_TYPE_IQ4_XS: case GGML_TYPE_IQ4_NL: + if (src1_type == GGML_TYPE_F16 && !ctx->device->coopmat2) { + return nullptr; // no f16 B pipeline for these types on coopmat1 + } + break; + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q5_1: case GGML_TYPE_MXFP4: + if (src1_type == GGML_TYPE_F16 && !ctx->device->coopmat2) { + // Use the dedicated f16 B pipeline + vk_matmul_pipeline2& mmp_f16b = ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0_type]; + bool prefer_fp16acc = ctx->device->fp16; + bool support_fp16acc = !mmp_f16b.f16acc->is_empty(); + bool support_fp32acc = !mmp_f16b.f32acc->is_empty(); + if (support_fp16acc && (prefer_fp16acc || !support_fp32acc)) { + return mmp_f16b.f16acc; + } else if (support_fp32acc) { + return mmp_f16b.f32acc; + } + return nullptr; + } case GGML_TYPE_NVFP4: break; default: @@ -8254,6 +8420,13 @@ static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, if ((mm_m && (m <= 64 || n <= 64)) || !mm_l) { return aligned ? mmp->a_m : mmp->m; } + + if (ctx->device->vendor_id == VK_VENDOR_ID_INTEL && + n < 256 && m < 4096 && + (aligned ? mmp->a_l_alt : mmp->l_alt) != nullptr) { + return aligned ? mmp->a_l_alt : mmp->l_alt; + } + return aligned ? mmp->a_l : mmp->l; } @@ -8645,6 +8818,10 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub const bool x_non_contig = (ctx->device->coopmat2 && src0->type == GGML_TYPE_F32) || !ggml_vk_dim01_contiguous(src0); const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || + // Intel coopmat1: force f32->f16 conversion so the f16-B-type pipeline is used. + (ctx->device->coopmat_support && !ctx->device->coopmat2 && + ctx->device->vendor_id == VK_VENDOR_ID_INTEL && + ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32) || (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); @@ -8921,15 +9098,21 @@ static bool ggml_vk_should_use_mmvq(const vk_device& device, uint32_t m, uint32_ } } - if (device->driver_id == vk::DriverId::eIntelProprietaryWindows) { - // Intel Windows proprietary driver MMVQ performance for !Q2/Q3/Q6 is worse than fp16, - // see https://github.com/ggml-org/llama.cpp/issues/17628 and - // https://github.com/ggml-org/llama.cpp/pull/23056 + if (k < 2048) { return false; } - if (k < 2048) { - return false; + if (device->driver_id == vk::DriverId::eIntelProprietaryWindows) { + // Intel Windows proprietary driver tuning + switch (src0_type) { + case GGML_TYPE_MXFP4: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q8_0: // MMVQ tg-slower than float path on Xe3 (n>1 pp returns true earlier) + return false; + default: + return true; + } } switch (src0_type) { @@ -8982,6 +9165,11 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& GGML_ASSERT(ne11 == 1 || ne12 * ne13 == 1); bool batch_n = ne11 > 1; + // For large batch_n (N > max_cols), dispatch one workgroup per column with NUM_COLS=1 + // instead of processing all columns in a single workgroup + const bool large_batch_n = batch_n && ne11 > mul_mat_vec_max_cols; + const uint32_t effective_num_cols = large_batch_n ? 1 : static_cast(ne11); + const bool x_non_contig = !ggml_vk_dim01_contiguous(src0); const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); @@ -9000,12 +9188,12 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& } // Check for mmq first - vk_pipeline dmmv = quantize_y ? ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, GGML_TYPE_Q8_1, ne11, ne20, ne00) : nullptr; + vk_pipeline dmmv = quantize_y ? ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, GGML_TYPE_Q8_1, effective_num_cols, ne20, ne00) : nullptr; vk_pipeline to_q8_1 = nullptr; if (dmmv == nullptr) { // Fall back to f16 dequant mul mat - dmmv = ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, src1->type, ne11, ne20, ne00); + dmmv = ggml_vk_get_dequantize_mul_mat_vec(ctx, src0->type, src1->type, effective_num_cols, ne20, ne00); quantize_y = false; } @@ -9158,12 +9346,15 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& fusion_flags |= MAT_VEC_FUSION_FLAGS_BIAS1; } - ggml_pipeline_request_descriptor_sets(ctx, dmmv, CEIL_DIV(ne12 * ne13, ctx->device->properties.limits.maxComputeWorkGroupCount[1])); + // For large_batch_n, dispatch ne11 workgroups in Y (one per column) instead of ne12*ne13 + const uint32_t total_y = large_batch_n ? static_cast(ne11) : static_cast(ne12 * ne13); + + ggml_pipeline_request_descriptor_sets(ctx, dmmv, CEIL_DIV(total_y, ctx->device->properties.limits.maxComputeWorkGroupCount[1])); uint32_t base_work_group_y = 0; - while (base_work_group_y < ne12 * ne13) { + while (base_work_group_y < total_y) { - uint32_t groups_y = std::min((uint32_t)(ne12 * ne13) - base_work_group_y, ctx->device->properties.limits.maxComputeWorkGroupCount[1]); + uint32_t groups_y = std::min(total_y - base_work_group_y, ctx->device->properties.limits.maxComputeWorkGroupCount[1]); const vk_mat_vec_push_constants pc = { (uint32_t)ne00, (uint32_t)ne10, (uint32_t)ne10, (uint32_t)ne01, stride_batch_x, stride_batch_y, stride_batch_d, @@ -9493,8 +9684,10 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c src1->ne[2] <= ctx->device->properties.limits.maxComputeWorkGroupCount[2]) { ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, cgraph, node_idx); // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) - // when ne12 and ne13 are one. - } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && + // when ne12 and ne13 are one. For small M (<=4) with large N, dispatch one workgroup per column. + } else if ((dst->ne[1] == 1 || + (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1) || + (dst->ne[0] <= 4 && src1->ne[2] * src1->ne[3] == 1)) && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || ggml_is_quantized(src0->type))) { ggml_vk_mul_mat_vec_q_f16(ctx, subctx, cgraph, node_idx); } else { @@ -9582,6 +9775,10 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& #endif const bool y_non_contig = y_decode_vector_staging || (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || + // Intel coopmat1: force f32->f16 conversion so the f16-B-type pipeline is used. + (ctx->device->coopmat_support && !ctx->device->coopmat2 && + ctx->device->vendor_id == VK_VENDOR_ID_INTEL && + (src0->type == GGML_TYPE_MXFP4 || src0->type == GGML_TYPE_Q4_K || src0->type == GGML_TYPE_Q5_K || src0->type == GGML_TYPE_Q5_1) && src1->type == GGML_TYPE_F32) || (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); @@ -10097,7 +10294,8 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx ggml_tensor * src1 = dst->src[1]; ggml_tensor * src2 = dst->src[2]; VK_LOG_DEBUG("ggml_vk_mul_mat_id(" << src0 << ", " << src1 << ", " << src2 << ", " << dst << ")"); - if (ggml_vk_use_mul_mat_vec_id(cgraph, node_idx)) { + // vec path only supports f32/Q8_1 B; fall through to matrix path for f16 B + if (ggml_vk_use_mul_mat_vec_id(cgraph, node_idx) && src1->type != GGML_TYPE_F16) { ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, cgraph, node_idx); } else { ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst); @@ -10315,7 +10513,6 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type); vk_pipeline pipeline = nullptr; - { std::lock_guard guard(ctx->device->compile_mutex); auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16; @@ -10474,10 +10671,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx pc, { dispatch_x, workgroups_y, workgroups_z }); ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, split_k, (sinks != nullptr) }; + const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, static_cast(ne1), static_cast(ne2), static_cast(ne3), split_k, (sinks != nullptr) }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, {split_k_buf, sinks_buf, dst_buf}, - pc2, { (uint32_t)ne1, HSV, (uint32_t)(ne2 * ne3) }); + pc2, { static_cast(ne1), HSV, static_cast(ne2 * ne3) }); ctx->prealloc_split_k_need_sync = true; } else { if (gqa_ratio > 1) { @@ -16482,13 +16679,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg bool need_disable = false; - // topk_moe often overwrites the source, but for a given row all the src values are - // loaded before anything is stored. If there's only one row, this is safe, so treat - // this as a special case. - bool is_topk_moe_single_row = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT && - ggml_nrows(cgraph->nodes[i]->src[0]) == 1; + // Note: topk_moe handles src/dst overlap internally for all batch sizes, + // so skip the overlap check entirely for topk_moe fusions. + bool is_topk_moe = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT; - if (!is_topk_moe_single_row) { + if (!is_topk_moe) { for (int j = 0; j < 2; ++j) { ggml_tensor *dst = output_nodes[j]; if (!dst) { @@ -16997,6 +17192,12 @@ void ggml_backend_vk_get_device_description(int device, char * description, size ggml_vk_get_device_description(dev_idx, description, description_size); } +static bool ggml_backend_vk_is_intel_xe2(int device) { + GGML_ASSERT(device < (int) vk_instance.device_indices.size()); + int dev_idx = vk_instance.device_indices[device]; + return ggml_vk_get_device(dev_idx)->architecture == INTEL_XE2; +} + void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total) { GGML_ASSERT(device < (int) vk_instance.device_indices.size()); GGML_ASSERT(device < (int) vk_instance.device_supports_membudget.size()); @@ -17917,11 +18118,20 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, return devices[device]; } +static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { + if (std::strcmp(name, "ggml_backend_vk_is_intel_xe2") == 0) { + return (void *) ggml_backend_vk_is_intel_xe2; + } + return nullptr; + + GGML_UNUSED(reg); +} + static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = { /* .get_name = */ ggml_backend_vk_reg_get_name, /* .get_device_count = */ ggml_backend_vk_reg_get_device_count, /* .get_device = */ ggml_backend_vk_reg_get_device, - /* .get_proc_address = */ NULL, + /* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address, }; ggml_backend_reg_t ggml_backend_vk_reg() { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 57c0410e455..a163061c8d8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -118,7 +118,9 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit #define BK_STEP 2 #endif -#ifdef COOPMAT +#if defined(COOPMAT) && defined(LOAD_A_OPT) +#define SHMEM_STRIDE (BK / 2) +#elif defined(COOPMAT) #define SHMEM_STRIDE (BK / 2 + 4) #else #define SHMEM_STRIDE (BK / 2 + 1) @@ -235,6 +237,10 @@ void main() { // Workgroup has no work if (ic * BN >= _ne1) return; + + uint required_work_items = (_ne1 - ic * BN) * BK / LOAD_VEC_B; + uint required_warp_c = (_ne1 - ic * BN + WN - 1) / WN ; + #endif #ifdef MUL_MAT_ID @@ -259,8 +265,8 @@ void main() { #endif #ifdef COOPMAT - coopmat cache_a; - coopmat cache_b; + coopmat cache_a[cms_per_row]; + coopmat cache_b[cms_per_col]; coopmat sums[cms_per_row * cms_per_col]; [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { @@ -285,6 +291,9 @@ void main() { [[unroll]] for (uint l = 0; l < BM; l += loadstride_a) { load_a_to_shmem(pos_a, loadr_a, loadc_a + l, ir * BM + loadc_a + l, block, end_k); } + #ifdef MUL_MAT_ID + if (gl_LocalInvocationID.x < required_work_items) { + #endif [[unroll]] for (uint l = 0; l < BN; l += loadstride_b) { #if !defined(MUL_MAT_ID) load_b_to_shmem(pos_b, loadr_b, loadc_b + l, ic * BN + loadc_b + l, block, end_k); @@ -292,6 +301,9 @@ void main() { load_b_to_shmem(pos_b, loadr_b, loadc_b + l, ic, _ne1, block, end_k); #endif } + #ifdef MUL_MAT_ID + } + #endif barrier(); @@ -299,18 +311,32 @@ void main() { pos_b += BK / LOAD_VEC_B_EFF; #ifdef COOPMAT +#ifdef MUL_MAT_ID + if (warp_c < required_warp_c) { +#endif [[unroll]] for (uint i = 0; i < BK; i += TK) { [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { // Load from shared into cache - coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); +#ifndef LOAD_A_OPT + coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); +#else + coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * TK / 2 + i * BM / 2, TK / 2, gl_CooperativeMatrixLayoutRowMajor); +#endif + } - [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { - coopMatLoad(cache_b, buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + coopMatLoad(cache_b[cm_col], buf_b, (warp_c * WN + cm_col * TN) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + } - sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a, cache_b, sums[cm_col * cms_per_row + cm_row]); + [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { + [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { + sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a[cm_row], cache_b[cm_col], sums[cm_col * cms_per_row + cm_row]); } } } +#ifdef MUL_MAT_ID + } +#endif #else [[unroll]] for (uint i = 0; i < BK / BK_STEP; i++) { // Load from shared into cache diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index 56a8a0f187f..0ecd1d5c7eb 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -11,6 +11,19 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 3] = aa[1].zw; return; } + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; +#else + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; +#endif + FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); + buf_a[buf_idx ] = aa[0].xy; + buf_a[buf_idx + 1] = aa[0].zw; + buf_a[buf_idx + 2] = aa[1].xy; + buf_a[buf_idx + 3] = aa[1].zw; #elif LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -20,9 +33,26 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = aa.zw; return; } + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; +#else + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; #endif + FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); + buf_a[buf_idx ] = aa.xy; + buf_a[buf_idx + 1] = aa.zw; +#else // LOAD_VEC_BATCH_A == 2 const uint idx = pos_a + col * p.stride_a + row * 2; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row; +#else + const uint buf_idx_outer = (row) / (TK / 2); + const uint buf_idx_inner = (row) % (TK / 2); + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner; +#endif if (idx_m < p.M && block + row * 2 + 1 < end_k) { buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], data_a_scalar[idx + 1]); @@ -86,7 +116,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw); #elif defined(DATA_A_Q5_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; +#ifndef LOAD_A_OPT + const uint buf_idx = col * SHMEM_STRIDE + row; +#else + const uint buf_idx_outer = row / (TK); + const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; + const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; +#endif const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -98,12 +134,24 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint vui = uint(data_a_packed16[ib].qs[iqs]); const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d; - +#ifndef LOAD_A_OPT buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz); buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw); +#else + buf_a[buf_idx_0] = FLOAT_TYPEV2(v.xz); + buf_a[buf_idx_1] = FLOAT_TYPEV2(v.yw); +#endif #elif defined(DATA_A_Q5_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; +#else + // Follow Q5_0 LOAD_A_OPT pattern, adapted for LOAD_VEC_A=8 (2 VEC2 per half) + const uint eff_row = row * LOAD_VEC_A / 4; // 2 VEC2 positions per load + const uint buf_idx_outer = eff_row / TK; + const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + eff_row; + const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; +#endif const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -119,13 +167,26 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y; const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y; +#ifndef LOAD_A_OPT buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xz); buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz); buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw); buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw); +#else + buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xz); + buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v1.xz); + buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v0.yw); + buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.yw); +#endif #elif defined(DATA_A_Q8_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; +#else + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; +#endif const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -196,43 +257,70 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin dl * (qs.y - hm.y)); #elif defined(DATA_A_Q4_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; +#else + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; +#endif - const uint ib = idx / 64; // 4 values per idx - const uint iqs = (idx % 64) * 2; // 0,2,4..126 - - const uint n = iqs / 32; // 0,1,2,3 - const uint b = (iqs % 32) / 16; // 0,1 - const uint is = 2 * n + b; // 0..7 - const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 - - const vec2 loadd = vec2(data_a[ib].dm); - - const uvec3 scales = uvec3(data_a_packed32[ib].scales[0], - data_a_packed32[ib].scales[1], - data_a_packed32[ib].scales[2]); - const uint scalesoffs = (is & 3) * 8; - - const uint scidx0 = (is < 4) ? 0 : 2; - const uint scidxshift0 = scalesoffs; - const uint scidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - const uint mbidx0 = (is < 4) ? 1 : 2; - const uint mbidxshift0 = (is < 4) ? scalesoffs : scalesoffs + 4; - const uint mbidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - - const uint8_t sc = uint8_t(((scales[scidx0] >> scidxshift0) & 0xF) | ((scales[0] >> scidxshift1) & 0x30)); - const uint8_t mbyte = uint8_t(((scales[mbidx0] >> mbidxshift0) & 0xF) | ((scales[1] >> mbidxshift1) & 0x30)); + const uint ib = idx / 64; // 4 values per idx + const uint iqs = (idx % 64) * 2; // 0,2,4..126 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2u * n + b; // 0..7 + const uint j = is & 3u; // low 2 bits: 0..3 + const int jsh = int(j * 8u); // byte bit-offset: 0,8,16,24 + const uint qsi = n * 32u + (iqs % 16u) * 2u; // 0,2,4..126 + + const vec2 loadd = vec2(data_a_packed32[ib].dm); + + // Q4_K scales: 12 bytes packed as 3 uint32 words. + // Decode 6-bit sc and mbyte via bitfieldExtract — avoids 4 scattered byte reads and + // 10 runtime ternaries. The condition (is < 4u) is warp-uniform (stride_a/LOAD_VEC_A + // is a multiple of 64), so this branch compiles to a uniform predicate with no divergence. + // sw0 = scales[0..3], sw1 = scales[4..7], sw2 = scales[8..11] + // is<4: sc = sw0[j] bits[0:5]; mbyte = sw1[j] bits[0:5] + // is≥4: sc = sw2[j] bits[0:3] | sw0[j] bits[6:7]<<4 + // mbyte = sw2[j] bits[4:7] | sw1[j] bits[6:7]<<4 + const uint sw0 = data_a_packed32[ib].scales[0]; + const uint sw1 = data_a_packed32[ib].scales[1]; + const uint sw2 = data_a_packed32[ib].scales[2]; + + uint sc, mbyte; + if (is < 4u) { + sc = bitfieldExtract(sw0, jsh, 6); + mbyte = bitfieldExtract(sw1, jsh, 6); + } else { + sc = bitfieldExtract(sw2, jsh, 4) | (bitfieldExtract(sw0, jsh + 6, 2) << 4u); + mbyte = bitfieldExtract(sw2, jsh + 4, 4) | (bitfieldExtract(sw1, jsh + 6, 2) << 4u); + } - const float d = loadd.x * sc; - const float m = -loadd.y * mbyte; + const float d = loadd.x * float(sc); + const float m = -loadd.y * float(mbyte); - const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F)); + // Nibble decode: 4× bitfieldExtract avoids byte-register scatter/gather chain + // (shr + and 0x0F0F0F0F + 4×mov:b + 4×mov:ub→uw + 4×itof → 4×bfe + 4×itof) + const uint qs_word = data_a_packed32[ib].qs[qsi / 4]; + const int base = int(b * 4u); // 0 or 4, per-lane + const vec4 q = vec4(float(bitfieldExtract(qs_word, base, 4)), + float(bitfieldExtract(qs_word, base + 8, 4)), + float(bitfieldExtract(qs_word, base + 16, 4)), + float(bitfieldExtract(qs_word, base + 24, 4))); buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)); buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); #elif defined(DATA_A_Q5_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; +#else + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; +#endif const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -243,25 +331,26 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 const uint qhi = (iqs % 16) * 2; // 0,2,4..30 - const vec2 loadd = vec2(data_a[ib].dm); + // OPT: bitfieldExtract scale decode (same layout as Q4_K, avoids 10 ternaries + 4 byte reads) + const vec2 loadd = vec2(data_a_packed32[ib].dm); - const uvec3 scales = uvec3(data_a_packed32[ib].scales[0], - data_a_packed32[ib].scales[1], - data_a_packed32[ib].scales[2]); - const uint scalesoffs = (is & 3) * 8; + const uint j = is & 3u; + const int jsh = int(j * 8u); + const uint sw0 = data_a_packed32[ib].scales[0]; + const uint sw1 = data_a_packed32[ib].scales[1]; + const uint sw2 = data_a_packed32[ib].scales[2]; - const uint scidx0 = (is < 4) ? 0 : 2; - const uint scidxshift0 = scalesoffs; - const uint scidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - const uint mbidx0 = (is < 4) ? 1 : 2; - const uint mbidxshift0 = (is < 4) ? scalesoffs : scalesoffs + 4; - const uint mbidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - - const uint8_t sc = uint8_t(((scales[scidx0] >> scidxshift0) & 0xF) | ((scales[0] >> scidxshift1) & 0x30)); - const uint8_t mbyte = uint8_t(((scales[mbidx0] >> mbidxshift0) & 0xF) | ((scales[1] >> mbidxshift1) & 0x30)); + uint sc, mbyte; + if (is < 4u) { + sc = bitfieldExtract(sw0, jsh, 6); + mbyte = bitfieldExtract(sw1, jsh, 6); + } else { + sc = bitfieldExtract(sw2, jsh, 4) | (bitfieldExtract(sw0, jsh + 6, 2) << 4u); + mbyte = bitfieldExtract(sw2, jsh + 4, 4) | (bitfieldExtract(sw1, jsh + 6, 2) << 4u); + } - const float d = loadd.x * sc; - const float m = -loadd.y * mbyte; + const float d = loadd.x * float(sc); + const float m = -loadd.y * float(mbyte); const uint qs = (data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F; const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4; @@ -271,7 +360,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); #elif defined(DATA_A_Q6_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; +#else + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; +#endif const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -497,7 +592,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin kvalues_iq4nl[vui >> 12]); #elif defined(DATA_A_MXFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; +#ifndef LOAD_A_OPT const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; +#else + const uint buf_idx_outer = row / (TK); + const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; + const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; +#endif const uint ib = idx / 8; const uint iqs = (idx & 0x07) * 2; @@ -505,11 +606,25 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float d = e8m0_to_fp32(data_a[ib].e) * 0.5; const uint vui = uint(data_a[ib].qs[iqs]); const uint vui2 = uint(data_a[ib].qs[iqs+1]); - - buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d); + // --------------------------------------------------------------------------- + // MXFP4 magnitude table: nibble bits[2:0] → {0,1,2,3,4,6,8,12} packed as + // 4-bit fields in a uint32 (nibble 0 in bits [3:0], nibble 7 in bits [31:28]). + // Decode: mag = bitfieldExtract(MXFP4_LUT, int((nibble & 7u) << 2), 4) + // val = float(mag) * ((nibble >= 8u) ? -d : d) + // --------------------------------------------------------------------------- + const uint MXFP4_LUT = 0xC8643210u; +#define MXFP4_VAL(nibble, scale) (float(bitfieldExtract(MXFP4_LUT, int(((nibble) & 7u) << 2), 4)) * (((nibble) >= 8u) ? -(scale) : (scale))) +#ifndef LOAD_A_OPT + buf_a[buf_idx ] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), + MXFP4_VAL(vui2 & 0xF, d)); + buf_a[buf_idx + 8] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), + MXFP4_VAL(vui2 >> 4, d)); +#else + buf_a[buf_idx_0] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), + MXFP4_VAL(vui2 & 0xF, d)); + buf_a[buf_idx_1] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), + MXFP4_VAL(vui2 >> 4, d)); +#endif #elif defined(DATA_A_NVFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; // lo and hi nibbles are 8 elements apart, which doesn't quite line up with diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 1925582ffed..03a797dc155 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -74,6 +74,92 @@ const std::vector type_names = { "bf16", }; +const std::vector load_a_opt_kernels = { + "matmul_id_subgroup_mxfp4_f32_cm1", + "matmul_f16_cm1", + "matmul_f16_f32_cm1", + "matmul_f32_f32_cm1", + "matmul_q4_k_f32_cm1", + "matmul_q6_k_f32_cm1", + "matmul_q5_0_f32_cm1", + "matmul_q8_0_f32_cm1", + "matmul_id_subgroup_mxfp4_f32_aligned_cm1", + "matmul_f16_aligned_cm1", + "matmul_f16_f32_aligned_cm1", + "matmul_f32_f32_aligned_cm1", + "matmul_q4_k_f32_aligned_cm1", + "matmul_q6_k_f32_aligned_cm1", + "matmul_q5_0_f32_aligned_cm1", + "matmul_q8_0_f32_aligned_cm1", + // f16acc variants (Step 6.2) + "matmul_q4_k_f32_f16acc_cm1", + "matmul_q6_k_f32_f16acc_cm1", + "matmul_q5_0_f32_f16acc_cm1", + "matmul_q8_0_f32_f16acc_cm1", + "matmul_id_subgroup_mxfp4_f32_f16acc_cm1", + "matmul_q4_k_f32_aligned_f16acc_cm1", + "matmul_q6_k_f32_aligned_f16acc_cm1", + "matmul_q5_0_f32_aligned_f16acc_cm1", + "matmul_q8_0_f32_aligned_f16acc_cm1", + "matmul_id_subgroup_mxfp4_f32_aligned_f16acc_cm1", + // f16 B-type MXFP4 variants (Step 6.2/6.3) + "matmul_id_subgroup_mxfp4_f16_cm1", + "matmul_id_subgroup_mxfp4_f16_aligned_cm1", + "matmul_id_subgroup_mxfp4_f16_f16acc_cm1", + "matmul_id_subgroup_mxfp4_f16_aligned_f16acc_cm1", + // Q4_K / Q5_K MUL_MAT_ID subgroup variants + "matmul_id_subgroup_q4_k_f32_cm1", + "matmul_id_subgroup_q4_k_f32_aligned_cm1", + "matmul_id_subgroup_q4_k_f32_f16acc_cm1", + "matmul_id_subgroup_q4_k_f32_aligned_f16acc_cm1", + "matmul_id_subgroup_q4_k_f16_cm1", + "matmul_id_subgroup_q4_k_f16_aligned_cm1", + "matmul_id_subgroup_q4_k_f16_f16acc_cm1", + "matmul_id_subgroup_q4_k_f16_aligned_f16acc_cm1", + "matmul_id_subgroup_q5_k_f32_cm1", + "matmul_id_subgroup_q5_k_f32_aligned_cm1", + "matmul_id_subgroup_q5_k_f32_f16acc_cm1", + "matmul_id_subgroup_q5_k_f32_aligned_f16acc_cm1", + "matmul_id_subgroup_q5_k_f16_cm1", + "matmul_id_subgroup_q5_k_f16_aligned_cm1", + "matmul_id_subgroup_q5_k_f16_f16acc_cm1", + "matmul_id_subgroup_q5_k_f16_aligned_f16acc_cm1", + // Q5_1 dense GEMM and MUL_MAT_ID variants + "matmul_q5_1_f32_cm1", + "matmul_q5_1_f32_aligned_cm1", + "matmul_q5_1_f32_f16acc_cm1", + "matmul_q5_1_f32_aligned_f16acc_cm1", + "matmul_q5_1_f16_cm1", + "matmul_q5_1_f16_aligned_cm1", + "matmul_q5_1_f16_f16acc_cm1", + "matmul_q5_1_f16_aligned_f16acc_cm1", + "matmul_id_subgroup_q5_1_f32_cm1", + "matmul_id_subgroup_q5_1_f32_aligned_cm1", + "matmul_id_subgroup_q5_1_f32_f16acc_cm1", + "matmul_id_subgroup_q5_1_f32_aligned_f16acc_cm1", + "matmul_id_subgroup_q5_1_f16_cm1", + "matmul_id_subgroup_q5_1_f16_aligned_cm1", + "matmul_id_subgroup_q5_1_f16_f16acc_cm1", + "matmul_id_subgroup_q5_1_f16_aligned_f16acc_cm1", + // f16 B-type dense GEMM variants (Step 6.3) + "matmul_q5_0_f16_cm1", + "matmul_q5_0_f16_aligned_cm1", + "matmul_q5_0_f16_f16acc_cm1", + "matmul_q5_0_f16_aligned_f16acc_cm1", + "matmul_q8_0_f16_cm1", + "matmul_q8_0_f16_aligned_cm1", + "matmul_q8_0_f16_f16acc_cm1", + "matmul_q8_0_f16_aligned_f16acc_cm1", + "matmul_q4_k_f16_cm1", + "matmul_q4_k_f16_aligned_cm1", + "matmul_q4_k_f16_f16acc_cm1", + "matmul_q4_k_f16_aligned_f16acc_cm1", + "matmul_q6_k_f16_cm1", + "matmul_q6_k_f16_aligned_cm1", + "matmul_q6_k_f16_f16acc_cm1", + "matmul_q6_k_f16_aligned_f16acc_cm1", +}; + enum MatMulIdType { NONE, DEFAULT, @@ -110,6 +196,13 @@ int execute_command(std::vector& command, std::string& stdout_str, cmd += part + " "; } + for (int32_t ii = 0; ii < load_a_opt_kernels.size(); ii++) { + if (cmd.find(load_a_opt_kernels.at(ii)) != std::string::npos) { + cmd.append(" -DLOAD_A_OPT"); + std::cout << cmd << std::endl; + } + } + if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { throw std::runtime_error("Failed to create process"); } From 4d54c3aa38bccb86dc08d96a1f18d4311e9710f6 Mon Sep 17 00:00:00 2001 From: "Jiang, Fish" Date: Wed, 17 Jun 2026 10:53:42 +0800 Subject: [PATCH 2/5] Add a new spec constant for SHMEM_STRIDE_PAD --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 3 +++ ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp | 9 ++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index a3c166680eb..9a535c6a8a1 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -4119,6 +4119,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { claimed_task.parameter_count = parameter_count; claimed_task.wg_denoms = wg_denoms; claimed_task.specialization_constants = specialization_constants; + if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { + claimed_task.specialization_constants.push_back(0u); + } claimed_task.disable_robustness = disable_robustness; claimed_task.require_full_subgroups = require_full_subgroups; claimed_task.required_subgroup_size = required_subgroup_size; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index a163061c8d8..950207f7a5a 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -118,13 +118,12 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit #define BK_STEP 2 #endif -#if defined(COOPMAT) && defined(LOAD_A_OPT) -#define SHMEM_STRIDE (BK / 2) -#elif defined(COOPMAT) -#define SHMEM_STRIDE (BK / 2 + 4) +#ifdef COOPMAT +layout(constant_id = 11) const uint SHMEM_STRIDE_PAD = 4; #else -#define SHMEM_STRIDE (BK / 2 + 1) +const uint SHMEM_STRIDE_PAD = 1; #endif +#define SHMEM_STRIDE (BK / 2 + SHMEM_STRIDE_PAD) shared FLOAT_TYPEV2 buf_a[BM * SHMEM_STRIDE]; shared FLOAT_TYPEV2 buf_b[BN * SHMEM_STRIDE]; From cb11307ddc04bf8e054f23b2d9d365338dcab3f3 Mon Sep 17 00:00:00 2001 From: "Jiang, Fish" Date: Wed, 17 Jun 2026 20:23:49 +0800 Subject: [PATCH 3/5] Move coopmat cache_a/cache_b declarations inside loop body --- ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 950207f7a5a..2201f2983ce 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -264,8 +264,6 @@ void main() { #endif #ifdef COOPMAT - coopmat cache_a[cms_per_row]; - coopmat cache_b[cms_per_col]; coopmat sums[cms_per_row * cms_per_col]; [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { @@ -313,6 +311,8 @@ void main() { #ifdef MUL_MAT_ID if (warp_c < required_warp_c) { #endif + coopmat cache_a[cms_per_row]; + coopmat cache_b[cms_per_col]; [[unroll]] for (uint i = 0; i < BK; i += TK) { [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { // Load from shared into cache From 4c62a3ae67f233cc1f9d0b39a1799cb4eb67cc02 Mon Sep 17 00:00:00 2001 From: "Jiang, Fish" Date: Tue, 23 Jun 2026 12:24:34 +0800 Subject: [PATCH 4/5] Rename LOAD_A_OPT to APPLY_SLM_A_RESHAPE, and add a new spec constant --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 73 ++- .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 16 +- .../vulkan-shaders/mul_mm_funcs.glsl | 421 +++++++++++------- .../vulkan-shaders/vulkan-shaders-gen.cpp | 93 ---- 4 files changed, 315 insertions(+), 288 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 9a535c6a8a1..0df95908be3 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -3846,6 +3846,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { (device->subgroup_size_control && device->subgroup_max_size >= 16); // mulmat + // Warptile layout (indices match mul_mm.comp constantIDs): + // [0..9] : BLOCK_SIZE, BM, BN, BK, WM, WN, WMITER, TM, TN, TK + // [10] : WARP / required_subgroup_size (read via WARPTILE_SUBGROUP_SIZE_IDX) + // [11] : ALIGNED (appended by ggml_vk_mul_mm_spec) + // [12,13] : SHMEM_STRIDE_PAD, APPLY_SLM_A_RESHAPE (Intel coopmat only, appended by ggml_vk_mul_mm_spec) + static constexpr size_t WARPTILE_SUBGROUP_SIZE_IDX = 10; std::vector l_warptile, m_warptile, s_warptile, l_warptile_id, m_warptile_id, s_warptile_id, l_warptile_mmq, m_warptile_mmq, s_warptile_mmq, @@ -3974,7 +3980,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { // Xe1/Xe2/Xe3 with coopmat enabled - warptile performance tuning l_warptile = { 512, 128, 128, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; - if (device->architecture == INTEL_PRE_XE2) { + if (device->architecture == INTEL_XE1) { l_warptile_mmq = { 512, 256, 128, 32, 32, 32, 2, 8, 8, 16, 16 }; l_mmq_wg_denoms = { 256, 128, 1 }; l_align = 32; //set as BK @@ -4119,9 +4125,6 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { claimed_task.parameter_count = parameter_count; claimed_task.wg_denoms = wg_denoms; claimed_task.specialization_constants = specialization_constants; - if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { - claimed_task.specialization_constants.push_back(0u); - } claimed_task.disable_robustness = disable_robustness; claimed_task.require_full_subgroups = require_full_subgroups; claimed_task.required_subgroup_size = required_subgroup_size; @@ -4265,8 +4268,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } #endif - auto const &ggml_vk_mul_mm_spec = [](std::vector spec, bool aligned) { - spec.push_back(aligned ? 1u : 0u); + auto const &ggml_vk_mul_mm_spec = [&device](std::vector spec, bool aligned) { + spec.push_back(aligned ? 1u : 0u); // constantID=11: ALIGNED + if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { + spec.push_back(0u); // constantID=12: SHMEM_STRIDE_PAD = 0 + spec.push_back(1u); // constantID=13: APPLY_SLM_A_RESHAPE = true + } return spec; }; @@ -4394,17 +4401,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // so it is always consistent with the tile and no separate tracking variable is needed. #define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ if (device->mul_mat ## ID ## _l[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true, (l_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ if (device->mul_mat ## ID ## _m[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true, (m_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ if (device->mul_mat ## ID ## _s[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true, (s_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ if (device->mul_mat ## ID ## _l[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true, (l_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ if (device->mul_mat ## ID ## _m[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true, (m_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ if (device->mul_mat ## ID ## _s[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true, (s_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ // Create 2 variants, {f16,f32} accumulator #define CREATE_MM2(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ @@ -4419,9 +4426,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // Same shader SPV as ->l/->a_l, but compiled with a different warp tile via the saved l_warptile_mmq. #define CREATE_MM_ALT(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ if (device->mul_mat ## ID ## _l[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l_alt, #NAMELC #F16ACC "_l_alt", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, true, (l_ ## WARPTILE).back()); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l_alt, #NAMELC #F16ACC "_l_alt", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true, (l_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); \ if (device->mul_mat ## ID ## _l[TYPE]) \ - ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l_alt, #NAMELC #F16ACC "_aligned_l_alt", NAMELC ## _aligned ## F16ACC ## _cm1_len, NAMELC ## _aligned ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, true, (l_ ## WARPTILE).back()); + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l_alt, #NAMELC #F16ACC "_aligned_l_alt", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true, (l_ ## WARPTILE)[WARPTILE_SUBGROUP_SIZE_IDX]); #define CREATE_MM2_ALT(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ if (device->coopmat_acc_f16_support) { \ @@ -4505,10 +4512,10 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); // Xe1: Q4_K/Q6_K use sgs=32 with a 128x128 tile; IQ* reverts to the xe1 standard tile. - if (device->architecture == INTEL_PRE_XE2) set_warp_tile({512, 128, 128, 32, 32, 32, 2, 8, 8, 16, 32}, {128, 128, 1}, l_align); + if (device->architecture == INTEL_XE1) set_warp_tile({512, 128, 128, 32, 32, 32, 2, 8, 8, 16, 32}, {128, 128, 1}, l_align); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K], matmul_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); - if (device->architecture == INTEL_PRE_XE2) restore_warp_tile(); + if (device->architecture == INTEL_XE1) restore_warp_tile(); CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_S], matmul_iq1_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_M], matmul_iq1_m_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); @@ -4555,7 +4562,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // Intel matmul_id warptile tuning if (device->vendor_id == VK_VENDOR_ID_INTEL) { - if (device->architecture == INTEL_PRE_XE2) { + if (device->architecture == INTEL_XE1) { l_warptile_mmq = { 512, 128, 128, 32, 32, 32, 2, 8, 8, 16, 32 }; l_mmq_wg_denoms = { 128, 128, 1 }; } @@ -10516,6 +10523,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type); vk_pipeline pipeline = nullptr; + { std::lock_guard guard(ctx->device->compile_mutex); auto &pipelines = ctx->device->pipeline_flash_attn_f32_f16; @@ -10674,10 +10682,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx pc, { dispatch_x, workgroups_y, workgroups_z }); ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, static_cast(ne1), static_cast(ne2), static_cast(ne3), split_k, (sinks != nullptr) }; + const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, split_k, (sinks != nullptr) }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, {split_k_buf, sinks_buf, dst_buf}, - pc2, { static_cast(ne1), HSV, static_cast(ne2 * ne3) }); + pc2, { (uint32_t)ne1, HSV, (uint32_t)(ne2 * ne3) }); ctx->prealloc_split_k_need_sync = true; } else { if (gqa_ratio > 1) { @@ -16682,11 +16690,13 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg bool need_disable = false; - // Note: topk_moe handles src/dst overlap internally for all batch sizes, - // so skip the overlap check entirely for topk_moe fusions. - bool is_topk_moe = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT; + // topk_moe often overwrites the source, but for a given row all the src values are + // loaded before anything is stored. If there's only one row, this is safe, so treat + // this as a special case. + bool is_topk_moe_single_row = ctx->fused_topk_moe_mode != TOPK_MOE_COUNT && + ggml_nrows(cgraph->nodes[i]->src[0]) == 1; - if (!is_topk_moe) { + if (!is_topk_moe_single_row) { for (int j = 0; j < 2; ++j) { ggml_tensor *dst = output_nodes[j]; if (!dst) { @@ -17195,12 +17205,6 @@ void ggml_backend_vk_get_device_description(int device, char * description, size ggml_vk_get_device_description(dev_idx, description, description_size); } -static bool ggml_backend_vk_is_intel_xe2(int device) { - GGML_ASSERT(device < (int) vk_instance.device_indices.size()); - int dev_idx = vk_instance.device_indices[device]; - return ggml_vk_get_device(dev_idx)->architecture == INTEL_XE2; -} - void ggml_backend_vk_get_device_memory(int device, size_t * free, size_t * total) { GGML_ASSERT(device < (int) vk_instance.device_indices.size()); GGML_ASSERT(device < (int) vk_instance.device_supports_membudget.size()); @@ -18121,20 +18125,11 @@ static ggml_backend_dev_t ggml_backend_vk_reg_get_device(ggml_backend_reg_t reg, return devices[device]; } -static void * ggml_backend_vk_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { - if (std::strcmp(name, "ggml_backend_vk_is_intel_xe2") == 0) { - return (void *) ggml_backend_vk_is_intel_xe2; - } - return nullptr; - - GGML_UNUSED(reg); -} - static const struct ggml_backend_reg_i ggml_backend_vk_reg_i = { /* .get_name = */ ggml_backend_vk_reg_get_name, /* .get_device_count = */ ggml_backend_vk_reg_get_device_count, /* .get_device = */ ggml_backend_vk_reg_get_device, - /* .get_proc_address = */ ggml_backend_vk_reg_get_proc_address, + /* .get_proc_address = */ NULL, }; ggml_backend_reg_t ggml_backend_vk_reg() { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 2201f2983ce..bfb7b5f42cc 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -119,9 +119,11 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit #endif #ifdef COOPMAT -layout(constant_id = 11) const uint SHMEM_STRIDE_PAD = 4; +layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4; +layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false; #else const uint SHMEM_STRIDE_PAD = 1; +const bool APPLY_SLM_A_RESHAPE = false; #endif #define SHMEM_STRIDE (BK / 2 + SHMEM_STRIDE_PAD) @@ -237,7 +239,7 @@ void main() { // Workgroup has no work if (ic * BN >= _ne1) return; - uint required_work_items = (_ne1 - ic * BN) * BK / LOAD_VEC_B; + uint required_work_items = (_ne1 - ic * BN) * BK / LOAD_VEC_B_EFF / LOAD_VEC_BATCH_B; uint required_warp_c = (_ne1 - ic * BN + WN - 1) / WN ; #endif @@ -316,11 +318,11 @@ void main() { [[unroll]] for (uint i = 0; i < BK; i += TK) { [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { // Load from shared into cache -#ifndef LOAD_A_OPT - coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); -#else - coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * TK / 2 + i * BM / 2, TK / 2, gl_CooperativeMatrixLayoutRowMajor); -#endif + if (!APPLY_SLM_A_RESHAPE) { + coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } else { + coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * TK / 2 + i * BM / 2, TK / 2, gl_CooperativeMatrixLayoutRowMajor); + } } [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index 0ecd1d5c7eb..ce3767f8991 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -3,7 +3,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin #if LOAD_VEC_A == 8 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; + } FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); buf_a[buf_idx ] = aa[0].xy; buf_a[buf_idx + 1] = aa[0].zw; @@ -11,48 +18,32 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 3] = aa[1].zw; return; } - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; -#else - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; -#endif - FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); - buf_a[buf_idx ] = aa[0].xy; - buf_a[buf_idx + 1] = aa[0].zw; - buf_a[buf_idx + 2] = aa[1].xy; - buf_a[buf_idx + 3] = aa[1].zw; #elif LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; + } FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); buf_a[buf_idx ] = aa.xy; buf_a[buf_idx + 1] = aa.zw; return; } - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; -#else - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; #endif - FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); - buf_a[buf_idx ] = aa.xy; - buf_a[buf_idx + 1] = aa.zw; -#else // LOAD_VEC_BATCH_A == 2 const uint idx = pos_a + col * p.stride_a + row * 2; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row; -#else - const uint buf_idx_outer = (row) / (TK / 2); - const uint buf_idx_inner = (row) % (TK / 2); - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner; -#endif + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row; + } else { + const uint buf_idx_outer = row / (TK / 2); + const uint buf_idx_inner = row % (TK / 2); + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner; + } if (idx_m < p.M && block + row * 2 + 1 < end_k) { buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], data_a_scalar[idx + 1]); @@ -65,7 +56,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin #if LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; + } FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx])); buf_a[buf_idx ] = aa.xy; buf_a[buf_idx + 1] = aa.zw; @@ -73,7 +71,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin } #endif const uint idx = pos_a + col * p.stride_a + row * 2; - const uint buf_idx = col * SHMEM_STRIDE + row; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row; + } else { + const uint outer = row / (TK / 2); + const uint inner = row % (TK / 2); + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner; + } if (idx_m < p.M && block + row * 2 + 1 < end_k) { buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), TO_FLOAT_TYPE(data_a_scalar[idx + 1])); @@ -84,7 +89,6 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin } #elif defined(DATA_A_Q4_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -94,13 +98,24 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d; const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.zw); + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; + buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); + buf_a[buf_idx + 1] = FLOAT_TYPEV2(v0.zw); + buf_a[buf_idx + 8] = FLOAT_TYPEV2(v1.xy); + buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.zw); + } else { + const uint eff_row = row * LOAD_VEC_A / 4; + const uint outer = eff_row / TK; + const uint buf_idx_0 = outer * BM * TK + col * TK / 2 + eff_row; + const uint buf_idx_1 = outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; + buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xy); + buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v0.zw); + buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v1.xy); + buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.zw); + } #elif defined(DATA_A_Q4_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -110,19 +125,24 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y; const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx + 1 ] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx + 8 ] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw); + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; + buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); + buf_a[buf_idx + 1 ] = FLOAT_TYPEV2(v0.zw); + buf_a[buf_idx + 8 ] = FLOAT_TYPEV2(v1.xy); + buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw); + } else { + const uint eff_row = row * LOAD_VEC_A / 4; + const uint outer = eff_row / TK; + const uint buf_idx_0 = outer * BM * TK + col * TK / 2 + eff_row; + const uint buf_idx_1 = outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; + buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xy); + buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v0.zw); + buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v1.xy); + buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.zw); + } #elif defined(DATA_A_Q5_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row; -#else - const uint buf_idx_outer = row / (TK); - const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; - const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; -#endif const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -134,24 +154,19 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint vui = uint(data_a_packed16[ib].qs[iqs]); const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d; -#ifndef LOAD_A_OPT - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw); -#else - buf_a[buf_idx_0] = FLOAT_TYPEV2(v.xz); - buf_a[buf_idx_1] = FLOAT_TYPEV2(v.yw); -#endif + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + row; + buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz); + buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw); + } else { + const uint buf_idx_outer = row / (TK); + const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; + const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; + buf_a[buf_idx_0] = FLOAT_TYPEV2(v.xz); + buf_a[buf_idx_1] = FLOAT_TYPEV2(v.yw); + } #elif defined(DATA_A_Q5_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; -#else - // Follow Q5_0 LOAD_A_OPT pattern, adapted for LOAD_VEC_A=8 (2 VEC2 per half) - const uint eff_row = row * LOAD_VEC_A / 4; // 2 VEC2 positions per load - const uint buf_idx_outer = eff_row / TK; - const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + eff_row; - const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; -#endif const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -167,26 +182,32 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y; const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y; -#ifndef LOAD_A_OPT - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xz); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw); - buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw); -#else - buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xz); - buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v1.xz); - buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v0.yw); - buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.yw); -#endif + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; + buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xz); + buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz); + buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw); + buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw); + } else { + const uint eff_row = row * LOAD_VEC_A / 4; + const uint buf_idx_outer = eff_row / TK; + const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + eff_row; + const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; + buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xz); + buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v1.xz); + buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v0.yw); + buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.yw); + } #elif defined(DATA_A_Q8_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; -#else - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; -#endif + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -200,7 +221,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); #elif defined(DATA_A_Q1_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 16; const uint iqs = idx & 0xfu; @@ -214,7 +242,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d); #elif defined(DATA_A_Q2_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -233,7 +268,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); #elif defined(DATA_A_Q3_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -257,13 +299,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin dl * (qs.y - hm.y)); #elif defined(DATA_A_Q4_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; -#else - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; -#endif + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; + } const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -314,13 +357,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); #elif defined(DATA_A_Q5_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; -#else - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; -#endif + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; + } const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -360,13 +404,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); #elif defined(DATA_A_Q6_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; -#else - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - const uint buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; -#endif + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; + const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; + buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; + } const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -388,7 +433,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx] = FLOAT_TYPEV2(q.x, q.y); #elif defined(DATA_A_IQ1_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -407,7 +459,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin } #elif defined(DATA_A_IQ1_M) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; @@ -429,7 +488,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin } #elif defined(DATA_A_IQ2_XXS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -460,7 +526,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin (sign & 128) != 0 ? -grid1.w : grid1.w); #elif defined(DATA_A_IQ2_XS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -486,7 +559,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin (sign & 128) != 0 ? -grid1.w : grid1.w); #elif defined(DATA_A_IQ2_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; // 0..31 @@ -514,7 +594,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin (sign & 128) != 0 ? -grid1.w : grid1.w); #elif defined(DATA_A_IQ3_XXS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -538,7 +625,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin (sign & 8) != 0 ? -v.w : v.w); #elif defined(DATA_A_IQ3_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -560,7 +654,14 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin (sign & 8) != 0 ? -v.w : v.w); #elif defined(DATA_A_IQ4_XS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + uint buf_idx; + if (!APPLY_SLM_A_RESHAPE) { + buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; + } else { + const uint outer = (row * LOAD_VEC_A) / TK; + const uint inner = (row * LOAD_VEC_A) % TK; + buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; + } const uint ib = idx / 64; // 4 values per idx const uint ib32 = (idx % 64) / 8; // 0..7 @@ -578,7 +679,6 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); #elif defined(DATA_A_IQ4_NL) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -586,19 +686,24 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d); const uint vui = uint(data_a_packed16[ib].qs[iqs]); - buf_a[buf_idx ] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], - kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]); - buf_a[buf_idx + 8] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], - kvalues_iq4nl[vui >> 12]); + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; + buf_a[buf_idx ] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], + kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]); + buf_a[buf_idx + 8] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], + kvalues_iq4nl[vui >> 12]); + } else { + const uint eff_row = row * LOAD_VEC_A / 4; + const uint outer = eff_row / TK; + const uint buf_idx_0 = outer * BM * TK + col * TK / 2 + eff_row; + const uint buf_idx_1 = outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; + buf_a[buf_idx_0] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], + kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]); + buf_a[buf_idx_1] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], + kvalues_iq4nl[vui >> 12]); + } #elif defined(DATA_A_MXFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; -#ifndef LOAD_A_OPT - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; -#else - const uint buf_idx_outer = row / (TK); - const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; - const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; -#endif const uint ib = idx / 8; const uint iqs = (idx & 0x07) * 2; @@ -614,23 +719,23 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin // --------------------------------------------------------------------------- const uint MXFP4_LUT = 0xC8643210u; #define MXFP4_VAL(nibble, scale) (float(bitfieldExtract(MXFP4_LUT, int(((nibble) & 7u) << 2), 4)) * (((nibble) >= 8u) ? -(scale) : (scale))) -#ifndef LOAD_A_OPT - buf_a[buf_idx ] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), - MXFP4_VAL(vui2 & 0xF, d)); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), - MXFP4_VAL(vui2 >> 4, d)); -#else - buf_a[buf_idx_0] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), - MXFP4_VAL(vui2 & 0xF, d)); - buf_a[buf_idx_1] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), - MXFP4_VAL(vui2 >> 4, d)); -#endif + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; + buf_a[buf_idx ] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), + MXFP4_VAL(vui2 & 0xF, d)); + buf_a[buf_idx + 8] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), + MXFP4_VAL(vui2 >> 4, d)); + } else { + const uint buf_idx_outer = row / (TK); + const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; + const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; + buf_a[buf_idx_0] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), + MXFP4_VAL(vui2 & 0xF, d)); + buf_a[buf_idx_1] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), + MXFP4_VAL(vui2 >> 4, d)); + } #elif defined(DATA_A_NVFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - // lo and hi nibbles are 8 elements apart, which doesn't quite line up with - // how the thread mapping and buf_idx calculation works for other types. - const uint buf_idx = col * SHMEM_STRIDE + (row & 3) + (row & ~3) * 2; - const uint ib = idx / 16u; const uint sub = (idx & 0xC) >> 2; const uint iqs = (idx & 0xF) * 2; @@ -638,10 +743,28 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint vui = uint(data_a[ib].qs[iqs]); const uint vui2 = uint(data_a[ib].qs[iqs+1]); - buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d); - buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d); + // lo and hi nibbles are 8 elements apart, which doesn't quite line up with + // how the thread mapping and buf_idx calculation works for other types. + const uint eff_row = (row & 3) + (row & ~3) * 2; + if (!APPLY_SLM_A_RESHAPE) { + const uint buf_idx = col * SHMEM_STRIDE + eff_row; + buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d); + buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d); + } else { + // Reshape [BM, BK/2] -> [BK/TK, BM, TK/2] (FLOAT_TYPEV2). Because the + // hi pair is only 8 K-elements (= 4 FLOAT_TYPEV2) away, it stays inside + // the same outer K-tile as the lo pair (TK = 16 > 8) — no need for the + // +BM*TK/2 split MXFP4 uses. + const uint buf_idx_outer = eff_row / (TK / 2); + const uint buf_idx_0 = buf_idx_outer * (BM * TK / 2) + col * (TK / 2) + (eff_row % (TK / 2)); + const uint buf_idx_1 = buf_idx_0 + 4; + buf_a[buf_idx_0] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d); + buf_a[buf_idx_1] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d); + } #endif } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 03a797dc155..1925582ffed 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -74,92 +74,6 @@ const std::vector type_names = { "bf16", }; -const std::vector load_a_opt_kernels = { - "matmul_id_subgroup_mxfp4_f32_cm1", - "matmul_f16_cm1", - "matmul_f16_f32_cm1", - "matmul_f32_f32_cm1", - "matmul_q4_k_f32_cm1", - "matmul_q6_k_f32_cm1", - "matmul_q5_0_f32_cm1", - "matmul_q8_0_f32_cm1", - "matmul_id_subgroup_mxfp4_f32_aligned_cm1", - "matmul_f16_aligned_cm1", - "matmul_f16_f32_aligned_cm1", - "matmul_f32_f32_aligned_cm1", - "matmul_q4_k_f32_aligned_cm1", - "matmul_q6_k_f32_aligned_cm1", - "matmul_q5_0_f32_aligned_cm1", - "matmul_q8_0_f32_aligned_cm1", - // f16acc variants (Step 6.2) - "matmul_q4_k_f32_f16acc_cm1", - "matmul_q6_k_f32_f16acc_cm1", - "matmul_q5_0_f32_f16acc_cm1", - "matmul_q8_0_f32_f16acc_cm1", - "matmul_id_subgroup_mxfp4_f32_f16acc_cm1", - "matmul_q4_k_f32_aligned_f16acc_cm1", - "matmul_q6_k_f32_aligned_f16acc_cm1", - "matmul_q5_0_f32_aligned_f16acc_cm1", - "matmul_q8_0_f32_aligned_f16acc_cm1", - "matmul_id_subgroup_mxfp4_f32_aligned_f16acc_cm1", - // f16 B-type MXFP4 variants (Step 6.2/6.3) - "matmul_id_subgroup_mxfp4_f16_cm1", - "matmul_id_subgroup_mxfp4_f16_aligned_cm1", - "matmul_id_subgroup_mxfp4_f16_f16acc_cm1", - "matmul_id_subgroup_mxfp4_f16_aligned_f16acc_cm1", - // Q4_K / Q5_K MUL_MAT_ID subgroup variants - "matmul_id_subgroup_q4_k_f32_cm1", - "matmul_id_subgroup_q4_k_f32_aligned_cm1", - "matmul_id_subgroup_q4_k_f32_f16acc_cm1", - "matmul_id_subgroup_q4_k_f32_aligned_f16acc_cm1", - "matmul_id_subgroup_q4_k_f16_cm1", - "matmul_id_subgroup_q4_k_f16_aligned_cm1", - "matmul_id_subgroup_q4_k_f16_f16acc_cm1", - "matmul_id_subgroup_q4_k_f16_aligned_f16acc_cm1", - "matmul_id_subgroup_q5_k_f32_cm1", - "matmul_id_subgroup_q5_k_f32_aligned_cm1", - "matmul_id_subgroup_q5_k_f32_f16acc_cm1", - "matmul_id_subgroup_q5_k_f32_aligned_f16acc_cm1", - "matmul_id_subgroup_q5_k_f16_cm1", - "matmul_id_subgroup_q5_k_f16_aligned_cm1", - "matmul_id_subgroup_q5_k_f16_f16acc_cm1", - "matmul_id_subgroup_q5_k_f16_aligned_f16acc_cm1", - // Q5_1 dense GEMM and MUL_MAT_ID variants - "matmul_q5_1_f32_cm1", - "matmul_q5_1_f32_aligned_cm1", - "matmul_q5_1_f32_f16acc_cm1", - "matmul_q5_1_f32_aligned_f16acc_cm1", - "matmul_q5_1_f16_cm1", - "matmul_q5_1_f16_aligned_cm1", - "matmul_q5_1_f16_f16acc_cm1", - "matmul_q5_1_f16_aligned_f16acc_cm1", - "matmul_id_subgroup_q5_1_f32_cm1", - "matmul_id_subgroup_q5_1_f32_aligned_cm1", - "matmul_id_subgroup_q5_1_f32_f16acc_cm1", - "matmul_id_subgroup_q5_1_f32_aligned_f16acc_cm1", - "matmul_id_subgroup_q5_1_f16_cm1", - "matmul_id_subgroup_q5_1_f16_aligned_cm1", - "matmul_id_subgroup_q5_1_f16_f16acc_cm1", - "matmul_id_subgroup_q5_1_f16_aligned_f16acc_cm1", - // f16 B-type dense GEMM variants (Step 6.3) - "matmul_q5_0_f16_cm1", - "matmul_q5_0_f16_aligned_cm1", - "matmul_q5_0_f16_f16acc_cm1", - "matmul_q5_0_f16_aligned_f16acc_cm1", - "matmul_q8_0_f16_cm1", - "matmul_q8_0_f16_aligned_cm1", - "matmul_q8_0_f16_f16acc_cm1", - "matmul_q8_0_f16_aligned_f16acc_cm1", - "matmul_q4_k_f16_cm1", - "matmul_q4_k_f16_aligned_cm1", - "matmul_q4_k_f16_f16acc_cm1", - "matmul_q4_k_f16_aligned_f16acc_cm1", - "matmul_q6_k_f16_cm1", - "matmul_q6_k_f16_aligned_cm1", - "matmul_q6_k_f16_f16acc_cm1", - "matmul_q6_k_f16_aligned_f16acc_cm1", -}; - enum MatMulIdType { NONE, DEFAULT, @@ -196,13 +110,6 @@ int execute_command(std::vector& command, std::string& stdout_str, cmd += part + " "; } - for (int32_t ii = 0; ii < load_a_opt_kernels.size(); ii++) { - if (cmd.find(load_a_opt_kernels.at(ii)) != std::string::npos) { - cmd.append(" -DLOAD_A_OPT"); - std::cout << cmd << std::endl; - } - } - if (!CreateProcessA(NULL, cmd.data(), NULL, NULL, TRUE, 0, NULL, NULL, &si, &pi)) { throw std::runtime_error("Failed to create process"); } From 4cdf2db0e2717bfcf1f682316627c15cd35921bf Mon Sep 17 00:00:00 2001 From: "Jiang, Fish" Date: Fri, 3 Jul 2026 12:55:56 +0800 Subject: [PATCH 5/5] vulkan: refactor SLM A addressing via a_shmem_index/store_a helpers; move coopmat cache declarations inside inner loop --- .../ggml-vulkan/vulkan-shaders/mul_mm.comp | 10 +- .../vulkan-shaders/mul_mm_funcs.glsl | 472 +++++------------- 2 files changed, 137 insertions(+), 345 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index bfb7b5f42cc..bcfb10af0b7 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -313,16 +313,12 @@ void main() { #ifdef MUL_MAT_ID if (warp_c < required_warp_c) { #endif - coopmat cache_a[cms_per_row]; - coopmat cache_b[cms_per_col]; [[unroll]] for (uint i = 0; i < BK; i += TK) { + coopmat cache_a[cms_per_row]; + coopmat cache_b[cms_per_col]; [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) { // Load from shared into cache - if (!APPLY_SLM_A_RESHAPE) { - coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, gl_CooperativeMatrixLayoutRowMajor); - } else { - coopMatLoad(cache_a[cm_row], buf_a, (warp_r * WM + cm_row * TM) * TK / 2 + i * BM / 2, TK / 2, gl_CooperativeMatrixLayoutRowMajor); - } + coopMatLoad(cache_a[cm_row], buf_a, a_shmem_index(warp_r * WM + cm_row * TM, i / 2), a_shmem_stride(), gl_CooperativeMatrixLayoutRowMajor); } [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index ce3767f8991..b85e5ee13b4 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -1,91 +1,73 @@ +// k_pair is the K coordinate measured in FLOAT_TYPEV2 elements. +uint a_shmem_index(uint m, uint k_pair) { + if (APPLY_SLM_A_RESHAPE) { + const uint tile_width = TK / 2; + return (k_pair / tile_width) * BM * tile_width + + m * tile_width + + k_pair % tile_width; + } + return m * SHMEM_STRIDE + k_pair; +} + +uint a_shmem_stride() { + return APPLY_SLM_A_RESHAPE ? TK / 2 : SHMEM_STRIDE; +} + +void store_a(uint m, uint k_pair, FLOAT_TYPEV2 value) { + buf_a[a_shmem_index(m, k_pair)] = value; +} + void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uint idx_m, const uint block, const uint end_k) { #if defined(DATA_A_F32) || defined(DATA_A_F16) #if LOAD_VEC_A == 8 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; - } + const uint k_pair = row * LOAD_VEC_A / 2; FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); - buf_a[buf_idx ] = aa[0].xy; - buf_a[buf_idx + 1] = aa[0].zw; - buf_a[buf_idx + 2] = aa[1].xy; - buf_a[buf_idx + 3] = aa[1].zw; + store_a(col, k_pair, aa[0].xy); + store_a(col, k_pair + 1, aa[0].zw); + store_a(col, k_pair + 2, aa[1].xy); + store_a(col, k_pair + 3, aa[1].zw); return; } #elif LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; - } + const uint k_pair = row * LOAD_VEC_A / 2; FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); - buf_a[buf_idx ] = aa.xy; - buf_a[buf_idx + 1] = aa.zw; + store_a(col, k_pair, aa.xy); + store_a(col, k_pair + 1, aa.zw); return; } #endif const uint idx = pos_a + col * p.stride_a + row * 2; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row; - } else { - const uint buf_idx_outer = row / (TK / 2); - const uint buf_idx_inner = row % (TK / 2); - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner; - } if (idx_m < p.M && block + row * 2 + 1 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], - data_a_scalar[idx + 1]); + store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], + data_a_scalar[idx + 1])); } else if (idx_m < p.M && block + row * 2 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(data_a_scalar[idx], 0.0f); + store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], 0.0f)); } else { - buf_a[buf_idx] = FLOAT_TYPEV2(0.0f); + store_a(col, row, FLOAT_TYPEV2(0.0f)); } #elif defined(DATA_A_BF16) #if LOAD_VEC_A == 4 if (ALIGNED != 0) { const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; - } + const uint k_pair = row * LOAD_VEC_A / 2; FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx])); - buf_a[buf_idx ] = aa.xy; - buf_a[buf_idx + 1] = aa.zw; + store_a(col, k_pair, aa.xy); + store_a(col, k_pair + 1, aa.zw); return; } #endif const uint idx = pos_a + col * p.stride_a + row * 2; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row; - } else { - const uint outer = row / (TK / 2); - const uint inner = row % (TK / 2); - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner; - } if (idx_m < p.M && block + row * 2 + 1 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), - TO_FLOAT_TYPE(data_a_scalar[idx + 1])); + store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), + TO_FLOAT_TYPE(data_a_scalar[idx + 1]))); } else if (idx_m < p.M && block + row * 2 < end_k) { - buf_a[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f); + store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f)); } else { - buf_a[buf_idx] = FLOAT_TYPEV2(0.0f); + store_a(col, row, FLOAT_TYPEV2(0.0f)); } #elif defined(DATA_A_Q4_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -98,22 +80,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d; const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d; - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.zw); - } else { - const uint eff_row = row * LOAD_VEC_A / 4; - const uint outer = eff_row / TK; - const uint buf_idx_0 = outer * BM * TK + col * TK / 2 + eff_row; - const uint buf_idx_1 = outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; - buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.zw); - } + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); #elif defined(DATA_A_Q4_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -125,22 +96,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y; const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y; - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx + 1 ] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx + 8 ] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx + 9 ] = FLOAT_TYPEV2(v1.zw); - } else { - const uint eff_row = row * LOAD_VEC_A / 4; - const uint outer = eff_row / TK; - const uint buf_idx_0 = outer * BM * TK + col * TK / 2 + eff_row; - const uint buf_idx_1 = outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; - buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xy); - buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v0.zw); - buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v1.xy); - buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.zw); - } + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); #elif defined(DATA_A_Q5_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -154,17 +114,8 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint vui = uint(data_a_packed16[ib].qs[iqs]); const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d; - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + row; - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v.yw); - } else { - const uint buf_idx_outer = row / (TK); - const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; - const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; - buf_a[buf_idx_0] = FLOAT_TYPEV2(v.xz); - buf_a[buf_idx_1] = FLOAT_TYPEV2(v.yw); - } + store_a(col, row, FLOAT_TYPEV2(v.xz)); + store_a(col, row + 8, FLOAT_TYPEV2(v.yw)); #elif defined(DATA_A_Q5_1) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -182,32 +133,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y; const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y; - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; - buf_a[buf_idx ] = FLOAT_TYPEV2(v0.xz); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v1.xz); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(v0.yw); - buf_a[buf_idx + 9] = FLOAT_TYPEV2(v1.yw); - } else { - const uint eff_row = row * LOAD_VEC_A / 4; - const uint buf_idx_outer = eff_row / TK; - const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + eff_row; - const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; - buf_a[buf_idx_0 ] = FLOAT_TYPEV2(v0.xz); - buf_a[buf_idx_0 + 1] = FLOAT_TYPEV2(v1.xz); - buf_a[buf_idx_1 ] = FLOAT_TYPEV2(v0.yw); - buf_a[buf_idx_1 + 1] = FLOAT_TYPEV2(v1.yw); - } + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xz)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v1.xz)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v0.yw)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.yw)); #elif defined(DATA_A_Q8_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -217,18 +149,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const i8vec2 v1 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs + 1])).xy; const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d; - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_Q1_0) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 16; const uint iqs = idx & 0xfu; @@ -236,20 +161,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float d = float(data_a[ib].d); const uint bits = uint(data_a[ib].qs[iqs]); - buf_a[buf_idx ] = FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d); - buf_a[buf_idx + 1] = FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d); - buf_a[buf_idx + 2] = FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d); - buf_a[buf_idx + 3] = FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d)); + store_a(col, k_pair + 2, FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d)); + store_a(col, k_pair + 3, FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d)); #elif defined(DATA_A_Q2_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -264,18 +182,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4); - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_Q3_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -295,18 +206,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec2 qs = vec2(unpack8((uint(data_a_packed16[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy); const vec2 hm = vec2(unpack8(((uint(data_a_packed16[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy); - buf_a[buf_idx] = FLOAT_TYPEV2(dl * (qs.x - hm.x), - dl * (qs.y - hm.y)); + store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(dl * (qs.x - hm.x), + dl * (qs.y - hm.y))); #elif defined(DATA_A_Q4_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; - } const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -353,18 +256,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin float(bitfieldExtract(qs_word, base + 16, 4)), float(bitfieldExtract(qs_word, base + 24, 4))); - buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); + store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); #elif defined(DATA_A_Q5_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; - } const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -400,18 +296,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4; const vec4 q = vec4(unpack8(qs | qh)); - buf_a[buf_idx ] = FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m)); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m)); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); + store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); #elif defined(DATA_A_Q6_K) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint buf_idx_outer = (row * LOAD_VEC_A) / TK; - const uint buf_idx_inner = (row * LOAD_VEC_A) % TK; - buf_idx = buf_idx_outer * BM * TK / 2 + col * TK / 2 + buf_idx_inner / 2; - } const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -430,17 +319,9 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint qh = (uint(data_a_packed16[ib].qh[qhi]) >> qhshift) & 0x0303; const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale; - buf_a[buf_idx] = FLOAT_TYPEV2(q.x, q.y); + store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(q.x, q.y)); #elif defined(DATA_A_IQ1_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -453,20 +334,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA; const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]); + const uint k_pair = row * LOAD_VEC_A / 2; [[unroll]] for (int k = 0; k < 4; ++k) { - buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), - dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)); + store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), + dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); } #elif defined(DATA_A_IQ1_M) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; @@ -482,20 +356,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA; const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]); + const uint k_pair = row * LOAD_VEC_A / 2; [[unroll]] for (int k = 0; k < 4; ++k) { - buf_a[buf_idx + k] = FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), - dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta)); + store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), + dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); } #elif defined(DATA_A_IQ2_XXS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -516,24 +383,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 grid0 = vec4(unpack8(grid.x)); const vec4 grid1 = vec4(unpack8(grid.y)); - buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y); - buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w); - buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y); - buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ2_XS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -549,24 +409,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 grid0 = vec4(unpack8(grid.x)); const vec4 grid1 = vec4(unpack8(grid.y)); - buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y); - buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w); - buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y); - buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ2_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; // 0..31 @@ -584,24 +437,17 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const vec4 grid0 = vec4(unpack8(grid.x)); const vec4 grid1 = vec4(unpack8(grid.y)); - buf_a[buf_idx ] = db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y); - buf_a[buf_idx + 1] = db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w); - buf_a[buf_idx + 2] = db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y); - buf_a[buf_idx + 3] = db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ3_XXS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -619,20 +465,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint grid = iq3xxs_grid[qs]; const vec4 v = db * vec4(unpack8(grid)); - buf_a[buf_idx ] = FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, - (sign & 2) != 0 ? -v.y : v.y); - buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, - (sign & 8) != 0 ? -v.w : v.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, + (sign & 2) != 0 ? -v.y : v.y)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, + (sign & 8) != 0 ? -v.w : v.w)); #elif defined(DATA_A_IQ3_S) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -648,20 +487,13 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)]; const vec4 v = db * vec4(unpack8(grid)); - buf_a[buf_idx ] = FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, - (sign & 2) != 0 ? -v.y : v.y); - buf_a[buf_idx + 1] = FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, - (sign & 8) != 0 ? -v.w : v.w); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, + (sign & 2) != 0 ? -v.y : v.y)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, + (sign & 8) != 0 ? -v.w : v.w)); #elif defined(DATA_A_IQ4_XS) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - uint buf_idx; - if (!APPLY_SLM_A_RESHAPE) { - buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - } else { - const uint outer = (row * LOAD_VEC_A) / TK; - const uint inner = (row * LOAD_VEC_A) % TK; - buf_idx = outer * BM * TK / 2 + col * TK / 2 + inner / 2; - } const uint ib = idx / 64; // 4 values per idx const uint ib32 = (idx % 64) / 8; // 0..7 @@ -675,8 +507,9 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const float d = float(data_a[ib].d); const vec4 v = d * float(int(sl | (sh << 4)) - 32) * vec4(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y], kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]); - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xy); - buf_a[buf_idx + 1] = FLOAT_TYPEV2(v.zw); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_IQ4_NL) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -686,22 +519,11 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d); const uint vui = uint(data_a_packed16[ib].qs[iqs]); - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; - buf_a[buf_idx ] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], - kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]); - buf_a[buf_idx + 8] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], - kvalues_iq4nl[vui >> 12]); - } else { - const uint eff_row = row * LOAD_VEC_A / 4; - const uint outer = eff_row / TK; - const uint buf_idx_0 = outer * BM * TK + col * TK / 2 + eff_row; - const uint buf_idx_1 = outer * BM * TK + BM * TK / 2 + col * TK / 2 + eff_row; - buf_a[buf_idx_0] = d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], - kvalues_iq4nl[bitfieldExtract(vui, 8, 4)]); - buf_a[buf_idx_1] = d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], - kvalues_iq4nl[vui >> 12]); - } + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], + kvalues_iq4nl[bitfieldExtract(vui, 8, 4)])); + store_a(col, k_pair + 8, d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], + kvalues_iq4nl[vui >> 12])); #elif defined(DATA_A_MXFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; @@ -719,21 +541,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin // --------------------------------------------------------------------------- const uint MXFP4_LUT = 0xC8643210u; #define MXFP4_VAL(nibble, scale) (float(bitfieldExtract(MXFP4_LUT, int(((nibble) & 7u) << 2), 4)) * (((nibble) >= 8u) ? -(scale) : (scale))) - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; - buf_a[buf_idx ] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), - MXFP4_VAL(vui2 & 0xF, d)); - buf_a[buf_idx + 8] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), - MXFP4_VAL(vui2 >> 4, d)); - } else { - const uint buf_idx_outer = row / (TK); - const uint buf_idx_0 = buf_idx_outer * BM * TK + col * TK / 2 + row; - const uint buf_idx_1 = buf_idx_outer * BM * TK + BM * TK / 2 + col * TK / 2 + row; - buf_a[buf_idx_0] = FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), - MXFP4_VAL(vui2 & 0xF, d)); - buf_a[buf_idx_1] = FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), - MXFP4_VAL(vui2 >> 4, d)); - } + store_a(col, row, FLOAT_TYPEV2(MXFP4_VAL(vui & 0xF, d), + MXFP4_VAL(vui2 & 0xF, d))); + store_a(col, row + 8, FLOAT_TYPEV2(MXFP4_VAL(vui >> 4, d), + MXFP4_VAL(vui2 >> 4, d))); #elif defined(DATA_A_NVFP4) const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; const uint ib = idx / 16u; @@ -746,25 +557,10 @@ void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uin // lo and hi nibbles are 8 elements apart, which doesn't quite line up with // how the thread mapping and buf_idx calculation works for other types. const uint eff_row = (row & 3) + (row & ~3) * 2; - if (!APPLY_SLM_A_RESHAPE) { - const uint buf_idx = col * SHMEM_STRIDE + eff_row; - buf_a[buf_idx ] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d); - buf_a[buf_idx + 4] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d); - } else { - // Reshape [BM, BK/2] -> [BK/TK, BM, TK/2] (FLOAT_TYPEV2). Because the - // hi pair is only 8 K-elements (= 4 FLOAT_TYPEV2) away, it stays inside - // the same outer K-tile as the lo pair (TK = 16 > 8) — no need for the - // +BM*TK/2 split MXFP4 uses. - const uint buf_idx_outer = eff_row / (TK / 2); - const uint buf_idx_0 = buf_idx_outer * (BM * TK / 2) + col * (TK / 2) + (eff_row % (TK / 2)); - const uint buf_idx_1 = buf_idx_0 + 4; - buf_a[buf_idx_0] = FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d); - buf_a[buf_idx_1] = FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d); - } + store_a(col, eff_row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d)); + store_a(col, eff_row + 4, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d)); #endif }