diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index c0ab9f1c658..0df95908be3 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]; @@ -3841,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, @@ -3950,6 +3961,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 +3978,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_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 + } 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 @@ -4249,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; }; @@ -4349,20 +4372,46 @@ 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); \ + 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) \ @@ -4373,6 +4422,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, 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 ## 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) { \ + 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 +4498,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_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_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, ); + 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_XE1) { + 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 +4604,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 +5015,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 +7338,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 +7375,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 +7464,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 +7510,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 +7530,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 +8430,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 +8828,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 +9108,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 +9175,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 +9198,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 +9356,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 +9694,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 +9785,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 +10304,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); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 57c0410e455..bcfb10af0b7 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -119,10 +119,13 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit #endif #ifdef COOPMAT -#define SHMEM_STRIDE (BK / 2 + 4) +layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4; +layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false; #else -#define SHMEM_STRIDE (BK / 2 + 1) +const uint SHMEM_STRIDE_PAD = 1; +const bool APPLY_SLM_A_RESHAPE = false; #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]; @@ -235,6 +238,10 @@ void main() { // Workgroup has no work if (ic * BN >= _ne1) return; + + 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 #ifdef MUL_MAT_ID @@ -259,8 +266,6 @@ void main() { #endif #ifdef COOPMAT - coopmat cache_a; - coopmat cache_b; coopmat sums[cms_per_row * cms_per_col]; [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) { @@ -285,6 +290,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 +300,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 +310,30 @@ 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) { + 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 - coopMatLoad(cache_a, buf_a, (warp_r * WM + cm_row * TM) * SHMEM_STRIDE + i / 2, SHMEM_STRIDE, 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++) { - 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..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,60 +1,76 @@ +// 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 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; - const uint buf_idx = col * SHMEM_STRIDE + row; 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 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; - const uint buf_idx = col * SHMEM_STRIDE + row; 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -64,13 +80,13 @@ 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); + 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -80,13 +96,13 @@ 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); + 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -98,12 +114,10 @@ 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; - - buf_a[buf_idx ] = FLOAT_TYPEV2(v.xz); - buf_a[buf_idx + 8] = 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 4; const uint iqs = idx & 0x03; @@ -119,13 +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; - 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); + 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -135,11 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 16; const uint iqs = idx & 0xfu; @@ -147,13 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -168,11 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -192,47 +206,61 @@ 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; - - 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 float d = loadd.x * sc; - const float m = -loadd.y * mbyte; - - const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F)); + 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); + } - 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 float d = loadd.x * float(sc); + const float m = -loadd.y * float(mbyte); + + // 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))); + + 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = (idx % 64) * 2; // 0,2,4..126 @@ -243,35 +271,36 @@ 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); - - 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; + // 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 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 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 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; 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 128; // 2 values per idx const uint iqs = idx % 128; // 0..127 @@ -290,10 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -306,13 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; @@ -328,13 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -355,17 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib32 = (idx % 32) / 4; // 0..7 @@ -381,17 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 32; // 8 values per idx const uint ib8 = idx % 32; // 0..31 @@ -409,17 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -437,13 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint iqs = idx % 64; // 0..63 @@ -459,13 +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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 2; const uint ib = idx / 64; // 4 values per idx const uint ib32 = (idx % 64) / 8; // 0..7 @@ -479,11 +507,11 @@ 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = idx & 0x07; @@ -491,13 +519,13 @@ 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]); + 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; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_A / 4; const uint ib = idx / 8; const uint iqs = (idx & 0x07) * 2; @@ -505,17 +533,20 @@ 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))) + 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; - // 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; @@ -523,10 +554,13 @@ 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; + 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 }