From 8512c365f10dab613539d00495dc8a94378ad6f7 Mon Sep 17 00:00:00 2001 From: Piotr Wilkin Date: Fri, 21 Aug 2026 10:36:26 +0200 Subject: [PATCH] vulkan: optimize Krea2 kernels on RDNA3 Assisted-by: OpenAI Codex --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 1207 ++++++++++++++++- .../vulkan-shaders/flash_attn_cm1.comp | 777 +++++++++-- .../ggml-vulkan/vulkan-shaders/mul_add.comp | 37 + .../vulkan-shaders/mul_mmq_cm1.comp | 444 ++++++ .../vulkan-shaders/quantize_q8_1.comp | 33 +- .../ggml-vulkan/vulkan-shaders/rms_norm.comp | 31 +- .../vulkan-shaders/rms_norm_channel_last.comp | 50 + .../vulkan-shaders/rms_norm_modulate.comp | 93 ++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 14 + tests/test-backend-ops.cpp | 268 ++++ 10 files changed, 2815 insertions(+), 139 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/mul_add.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_channel_last.comp create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_modulate.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f6cbaecb7f5..7601261f486 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -1,7 +1,6 @@ #include "ggml-vulkan.h" #include #if defined(GGML_VULKAN_RUN_TESTS) || defined(GGML_VULKAN_CHECK_RESULTS) -#include #include "ggml-cpu.h" #endif @@ -49,7 +48,11 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #endif #include +#include #include +#include +#include +#include #include #include #include @@ -568,10 +571,15 @@ struct vk_fa_pipeline_state { uint32_t limit_occupancy_shmem; ggml_type k_type; ggml_type v_type; + // Selects the PV staging layout inside flash_attn_cm1.comp. Must take part in the + // map key: two pipelines differing only here are different shaders. + uint32_t transpose_pv; + // Selects the RADV wave32 64-query x 32-KV work-per-tile arm. + uint32_t work_per_tile; bool operator<(const vk_fa_pipeline_state &b) const { - return std::tie(HSK, HSV, Br, Bc, D_split, row_split, shmem_staging, path, workgroup_size, subgroup_size, aligned, f32acc, flags, limit_occupancy_shmem, k_type, v_type) < - std::tie(b.HSK, b.HSV, b.Br, b.Bc, b.D_split, b.row_split, b.shmem_staging, b.path, b.workgroup_size, b.subgroup_size, b.aligned, b.f32acc, b.flags, b.limit_occupancy_shmem, b.k_type, b.v_type); + return std::tie(HSK, HSV, Br, Bc, D_split, row_split, shmem_staging, path, workgroup_size, subgroup_size, aligned, f32acc, flags, limit_occupancy_shmem, k_type, v_type, transpose_pv, work_per_tile) < + std::tie(b.HSK, b.HSV, b.Br, b.Bc, b.D_split, b.row_split, b.shmem_staging, b.path, b.workgroup_size, b.subgroup_size, b.aligned, b.f32acc, b.flags, b.limit_occupancy_shmem, b.k_type, b.v_type, b.transpose_pv, b.work_per_tile); } }; @@ -889,6 +897,12 @@ struct vk_device_struct { bool mul_mat_id_m_int[GGML_TYPE_COUNT]; bool mul_mat_id_s_int[GGML_TYPE_COUNT]; + // Tile availability for the int8 cooperative-matrix mmq path (mul_mmq_cm1.comp). + // Only Q8_0 uses it, so this is a plain flag rather than a per-type array. + bool mul_mat_cm_int_l {}; + bool mul_mat_cm_int_m {}; + bool mul_mat_cm_int_s {}; + vk::DescriptorSetLayout dsl; vk_matmul_pipeline pipeline_matmul_f32 {}; @@ -911,6 +925,11 @@ struct vk_device_struct { vk_pipeline pipeline_matmul_split_k_reduce; vk_pipeline pipeline_quantize_q8_1_x4; + vk_pipeline pipeline_quantize_q8_1_x4_swiglu; + vk_pipeline pipeline_quantize_q8_1_x4_sigmoid_mul; + vk_pipeline pipeline_quantize_i8_convrot; + vk_pipeline pipeline_mul_mat_i8_tensorwise; + vk_pipeline pipeline_mul_mat_i8_tensorwise_cm1; vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; vk_pipeline pipeline_dequant_transpose[GGML_TYPE_COUNT]; // fused dequant+transpose for FA quant-KV @@ -971,10 +990,14 @@ struct vk_device_struct { vk_pipeline pipeline_group_norm_f32; vk_pipeline pipeline_rms_norm_f32; vk_pipeline pipeline_rms_norm_mul_f32; + vk_pipeline pipeline_rms_norm_channel_last_f32[4]; + vk_pipeline pipeline_rms_norm_modulate_f32; + vk_pipeline pipeline_mul_add_f32; vk_pipeline pipeline_rms_norm_partials_f32; vk_pipeline pipeline_rms_norm_mul_partials_f32; vk_pipeline pipeline_rms_norm_mul_rope_f32_f32; vk_pipeline pipeline_rms_norm_mul_rope_f32_f16; + vk_pipeline pipeline_rms_norm_mul_mrope_pack_f32_f32; vk_pipeline pipeline_rms_norm_back_f32; vk_pipeline pipeline_l2_norm_f32; @@ -1103,6 +1126,12 @@ struct vk_device_struct { ggml_backend_buffer_type buffer_type; bool disable_fusion; + bool disable_rms_norm_channel_last; + bool disable_rms_norm_modulate; + bool disable_rms_norm_mrope_pack; + bool disable_mul_add; + bool disable_swiglu_quant; + bool disable_sigmoid_quant; bool disable_host_visible_vidmem; bool allow_sysmem_fallback; bool disable_graph_optimize; @@ -1596,6 +1625,25 @@ struct vk_op_binary_push_constants { float param1; float param2; int32_t param3; }; +struct vk_op_rms_norm_modulate_push_constants { + uint32_t ne00; uint32_t ne01; uint32_t ne02; uint32_t ne03; + uint32_t nb01; uint32_t nb02; uint32_t nb03; + float eps; +}; + +struct vk_op_rms_norm_channel_last_push_constants { + uint32_t width; + uint32_t height; + uint32_t depth; + uint32_t channels; + float eps; +}; + +struct vk_op_mul_add_push_constants { + uint32_t ne0_vec4; + uint32_t nrows; +}; + // Distinct type with the same layout so concat can overload tensor offset initialization. struct vk_op_concat_push_constants : vk_op_binary_push_constants {}; static_assert(sizeof(vk_op_concat_push_constants) == sizeof(vk_op_binary_push_constants)); @@ -2006,6 +2054,24 @@ struct vk_quantize_q8_1_push_constants { uint32_t num_blocks; }; +enum vk_quantize_q8_1_fusion { + VK_QUANTIZE_Q8_1_NONE, + VK_QUANTIZE_Q8_1_SWIGLU, + VK_QUANTIZE_Q8_1_SIGMOID_MUL, +}; + +struct vk_quantize_i8_convrot_push_constants { + uint32_t k; + uint32_t rows; +}; + +struct vk_mul_mat_i8_tensorwise_push_constants { + uint32_t n; + uint32_t rows; + uint32_t k; + uint32_t has_bias; + uint32_t row_offset; +}; struct vk_op_flash_attn_split_k_reduce_push_constants { uint32_t D; uint32_t ne1; @@ -2272,6 +2338,15 @@ class vk_perf_logger { name += "(" + std::to_string(node->ne[0]) + "," + std::to_string(node->ne[1]) + "," + std::to_string(node->ne[2]) + "," + std::to_string(node->ne[3]) + ")"; return fusion_str + name; } + if (node->op == GGML_OP_CONT) { + std::string name = ggml_op_name(node->op); + name += " " + std::string(ggml_type_name(node->type)) + + "(" + std::to_string(node->ne[0]) + "," + + std::to_string(node->ne[1]) + "," + + std::to_string(node->ne[2]) + "," + + std::to_string(node->ne[3]) + ")"; + return fusion_str + name; + } if (node->op == GGML_OP_FLASH_ATTN_EXT) { const ggml_tensor * dst = node; const ggml_tensor * q = node->src[0]; @@ -2391,6 +2466,12 @@ struct ggml_backend_vk_context { int fused_ops_write_mask {}; topk_moe_mode fused_topk_moe_mode {}; bool fused_topk_moe_scale {}; + bool fused_rms_norm_channel_last {}; + bool fused_rms_norm_modulate {}; + bool fused_rms_norm_mrope_pack {}; + bool fused_mul_add {}; + bool fused_swiglu_quant {}; + bool fused_sigmoid_quant {}; // for GGML_VK_PERF_LOGGER std::unique_ptr perf_logger; @@ -3000,8 +3081,11 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin pipeline_shader_create_info.setPNext(&pipeline_shader_stage_required_subgroup_size_create_info); } + const bool capture_pipeline_stats = device->pipeline_executable_properties_support && + !vk_pipeline_stats_filter.empty() && + pipeline->name.find(vk_pipeline_stats_filter) != std::string::npos; vk::ComputePipelineCreateInfo compute_pipeline_create_info( - device->pipeline_executable_properties_support ? + capture_pipeline_stats ? vk::PipelineCreateFlagBits::eCaptureStatisticsKHR : vk::PipelineCreateFlags{}, pipeline_shader_create_info, @@ -3020,7 +3104,7 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin if (pipeline->is_64b_indexing) { pipelineFlags2CreateInfo.flags = vk::PipelineCreateFlagBits2::e64BitIndexingEXT; - if (device->pipeline_executable_properties_support) { + if (capture_pipeline_stats) { pipelineFlags2CreateInfo.flags |= vk::PipelineCreateFlagBits2::eCaptureStatisticsKHR; } pipelineFlags2CreateInfo.setPNext(compute_pipeline_create_info.pNext); @@ -3029,7 +3113,7 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin #endif try { - pipeline->pipeline = device->device.createComputePipeline(VK_NULL_HANDLE, compute_pipeline_create_info).value; + pipeline->pipeline = device->device.createComputePipeline(nullptr, compute_pipeline_create_info).value; } catch (const vk::SystemError& e) { std::cerr << "ggml_vulkan: Compute pipeline creation failed for " << pipeline->name << std::endl; std::cerr << "ggml_vulkan: " << e.what() << std::endl; @@ -3044,20 +3128,18 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin vk_instance.pfn_vkSetDebugUtilsObjectNameEXT(device->device, &static_cast(duoni)); } - if (device->pipeline_executable_properties_support) { + if (capture_pipeline_stats) { vk::PipelineExecutableInfoKHR executableInfo; executableInfo.pipeline = pipeline->pipeline; auto statistics = device->device.getPipelineExecutableStatisticsKHR(executableInfo); - bool print_stats = !vk_pipeline_stats_filter.empty() && - pipeline->name.find(vk_pipeline_stats_filter) != std::string::npos; - if (print_stats) { + if (capture_pipeline_stats) { std::cerr << "ggml_vulkan: pipeline stats for " << pipeline->name << ":" << std::endl; } for (auto & s : statistics) { - if (print_stats) { + if (capture_pipeline_stats) { std::cerr << "ggml_vulkan: " << s.name.data() << ": "; switch (s.format) { case vk::PipelineExecutableStatisticFormatKHR::eBool32: @@ -3082,7 +3164,6 @@ static void ggml_vk_create_pipeline_func(vk_device& device, vk_pipeline& pipelin } } } - { std::lock_guard guard(device->compile_mutex); device->all_pipelines.push_back(pipeline); @@ -3696,6 +3777,10 @@ struct vk_fa_tuning_params { bool shmem_staging; bool disable_subgroups; uint32_t limit_occupancy_shmem; + // flash_attn_cm1.comp specialization constant 16. Zero on every other FA path. + uint32_t transpose_pv; + // flash_attn_cm1.comp specialization constant 17. + uint32_t work_per_tile; void print() const { std::cerr << "path=" << path << " workgroup_size=" << workgroup_size << " subgroup_size=" << subgroup_size << @@ -3809,8 +3894,28 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device const uint32_t D_lsb = D ^ (D & (D-1)); // extract lowest set bit result.d_split = std::min(std::min(result.subgroup_size, 8u), D_lsb / 4); + // Whole-tile staging increases LDS pressure on AMD hardware, so retain the + // existing NVIDIA-only selection for the portable cooperative-matrix path. result.shmem_staging = (device->vendor_id == VK_VENDOR_ID_NVIDIA && hsk < 256 && hsv < 256) ? 1 : 0; + // Transposed P/V staging for the PV multiply (flash_attn_cm1.comp spec constant 16). + // Both PV coopMatLoads then run along the fragment's native wide-load axis, which + // LLPC lowers to ds_load_b128: ds_load 324 -> 104 per dispatch and -7.8% on the + // measured gfx1151 workload. ACO does not form those wide loads from the same GLSL, + // keeps the per-fragment scalar gather, and measured +3.5% on RADV - so this is + // gated on the LLPC stack rather than on the AMD vendor id, and RADV, NVIDIA and + // Intel stay on the default arm. GCN/RDNA1/RDNA2 are excluded as unmeasured. + // It is a specialization constant rather than a runtime branch on purpose: the two + // arms differ in LDS layout and loop-carried state, and keeping both live would + // hold two accumulator sets across the KV back-edge and cross the 192-VGPR cliff. + const bool llpc_amd = device->vendor_id == VK_VENDOR_ID_AMD && + (device->driver_id == vk::DriverId::eAmdProprietary || + device->driver_id == vk::DriverId::eAmdOpenSource); + result.transpose_pv = (llpc_amd && + device->architecture != AMD_GCN && + device->architecture != AMD_RDNA1 && + device->architecture != AMD_RDNA2) ? 1 : 0; + return result; } @@ -3842,7 +3947,7 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat2(const vk_device& device return result; } -static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc) { +static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_t hsk, uint32_t hsv, uint32_t n_rows, uint32_t n_kv, ggml_type k_type, ggml_type v_type, bool f32acc, bool allow_work_per_tile = false) { FaCodePath path = device->coopmat2 ? FA_COOPMAT2 : device->coopmat1_fa_support ? FA_COOPMAT1 : FA_SCALAR; @@ -3874,16 +3979,42 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ path = FA_SCALAR; } + vk_fa_tuning_params tuned; switch (path) { case FA_SCALAR: - return get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); + tuned = get_fa_tuning_params_scalar(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); + break; case FA_COOPMAT1: - return get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); + tuned = get_fa_tuning_params_coopmat1(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); + break; case FA_COOPMAT2: - return get_fa_tuning_params_coopmat2(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); + tuned = get_fa_tuning_params_coopmat2(device, hsk, hsv, n_rows, n_kv, k_type, v_type, f32acc); + break; default: throw std::runtime_error("unsupported FaCodePath"); } + + const bool work_per_tile_hardware = path == FA_COOPMAT1 && + device->vendor_id == VK_VENDOR_ID_AMD && + device->driver_id == vk::DriverId::eMesaRadv && + device->architecture == vk_device_architecture::AMD_RDNA3 && + device->subgroup_size_control && + device->subgroup_min_size <= 32 && device->subgroup_max_size >= 32; + const bool use_work_per_tile = allow_work_per_tile && work_per_tile_hardware && + n_rows == n_kv && n_rows >= 1024 && + hsk == 128 && hsv == 128; + if (use_work_per_tile) { + tuned.work_per_tile = 1; + tuned.block_rows = 64; + tuned.block_cols = 32; + tuned.row_split = 4; + tuned.subgroup_size = 32; + tuned.workgroup_size = 128; + tuned.shmem_staging = false; + tuned.transpose_pv = 0; + } + + return tuned; } static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool aligned, bool f32acc, @@ -3898,7 +4029,7 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size; - return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type}; + return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type, params.transpose_pv, params.work_per_tile}; } static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& state) { @@ -3923,6 +4054,8 @@ static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& s /*13 FaTypeV */ static_cast(state.v_type), /*14 FaBlockBytesK */ fa_block_bytes(state.k_type), /*15 FaBlockBytesV */ fa_block_bytes(state.v_type), + /*16 TRANSPOSE_PV */ state.transpose_pv, + /*17 WORK_PER_TILE */ state.work_per_tile, }; } @@ -4051,6 +4184,32 @@ static bool ggml_vk_matmul_int_shmem_support(const vk_device& device, const std: return supported; } +// Shmem usage for the int8 cooperative-matrix mmq shader (mul_mmq_cm1.comp). +// Mirrors the shared declarations in that shader; keep the two in sync. +static bool ggml_vk_matmul_coopmat_int_shmem_support(const vk_device& device, const std::vector& warptile) { + const uint32_t BM = warptile[1]; + const uint32_t BN = warptile[2]; + const uint32_t CM_M = warptile[7]; + const uint32_t CM_N = warptile[8]; + + // BK is fixed by the quant format; BK_STEP is specialization constant 11. + const uint32_t BK = 32; + const uint32_t BK_STEP = warptile.size() > 11 ? warptile[11] : 2; + const uint32_t qpitch = BK_STEP * (BK / 4) + 4; + + const uint32_t buf_qs = (BM + BN) * qpitch * 4u; // buf_a_qs + buf_b_qs + const uint32_t buf_d = (BM + BN) * BK_STEP * 4u; // buf_a_d + buf_b_d (float) + const uint32_t buf_probe = warptile[10] == 32 ? 0u : CM_M * CM_N * 4u; + + const uint32_t total_size = buf_qs + buf_d + buf_probe; + const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; + + VK_LOG_DEBUG("ggml_vk_matmul_coopmat_int_shmem_support(BM=" << BM << ", BN=" << BN + << ", total=" << total_size << ", supported=" << supported); + + return supported; +} + struct GpuPipelineConfig { // GPU architecture identifier. // Example: vk_device_architecture::AMD_GCN @@ -4185,7 +4344,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { l_warptile_mmq_k, m_warptile_mmq_k, s_warptile_mmq_k, l_warptile_mmqid, m_warptile_mmqid, s_warptile_mmqid, l_warptile_mmqid_int, m_warptile_mmqid_int, s_warptile_mmqid_int, - l_warptile_mmqid_int_k, m_warptile_mmqid_int_k, s_warptile_mmqid_int_k; + l_warptile_mmqid_int_k, m_warptile_mmqid_int_k, s_warptile_mmqid_int_k, + l_warptile_mmq_cm_int, m_warptile_mmq_cm_int, s_warptile_mmq_cm_int; std::array l_wg_denoms, m_wg_denoms, s_wg_denoms, l_mmq_wg_denoms, m_mmq_wg_denoms, s_mmq_wg_denoms, l_mmq_wg_denoms_k, m_mmq_wg_denoms_k, s_mmq_wg_denoms_k, @@ -4260,6 +4420,39 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { m_warptile_mmq = { 128, 64, 64, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; s_warptile_mmq = { subgroup_size_32, 32, 32, 32, s_warptile_wm, 32, 2, tm_s, tn_s, tk_s, subgroup_size_8 }; + // int8 cooperative-matrix MMQ (mul_mmq_cm1.comp). TM/TN/TK carry the coopmat + // shape instead of a register subtile, and the shader assumes one subgroup per + // (WM, WN) tile, so BLOCK_SIZE / WARP must equal (BM / WM) * (BN / WN). + { + const uint32_t cm_i_m = device->coopmat_int_support ? device->coopmat_int_m : 16; + const uint32_t cm_i_n = device->coopmat_int_support ? device->coopmat_int_n : 16; + const uint32_t cm_i_k = device->coopmat_int_support ? device->coopmat_int_k : 16; + // ACO can pair the f32 dequant epilogue as VOPD in wave32. Keeping the + // per-subgroup tile at 32x32 preserves the same eight accumulator tiles + // while adding subgroups to cover the workgroup tile. + const bool use_wave32 = device->architecture == vk_device_architecture::AMD_RDNA3 && + device->driver_id == vk::DriverId::eMesaRadv && + device->subgroup_size_control && + device->subgroup_min_size <= 32 && device->subgroup_max_size >= 32; + const uint32_t sg = use_wave32 ? 32 : device->subgroup_size; + const uint32_t wm = use_wave32 ? 32 : 64; + const uint32_t l_bk_step = use_wave32 ? 3 : 2; + const auto block_size = [sg, wm](uint32_t bm, uint32_t bn) { + return sg * (bm / wm) * (bn / 32); + }; + + // Each subgroup owns (WM / CM_M) * (WN / CM_N) f32 accumulator tiles; keeping + // that product at 8 holds the accumulators in registers. Wider workgroup + // tiles are reached by adding subgroups rather than growing the per-subgroup + // tile, which measured much better than a 4x4 tile per subgroup. + // Accumulator tiles per subgroup are BM*BN/(subgroups*256); holding that at + // 8 keeps them in registers. Larger per-subgroup tiles raise register + // pressure enough to reduce occupancy on current AMD implementations. + l_warptile_mmq_cm_int = { block_size(128, 160), 128, 160, 32, wm, 32, 1, cm_i_m, cm_i_n, cm_i_k, sg, l_bk_step }; + m_warptile_mmq_cm_int = { block_size(128, 64), 128, 64, 32, wm, 32, 1, cm_i_m, cm_i_n, cm_i_k, sg, 2 }; + s_warptile_mmq_cm_int = { block_size( 32, 32), 32, 32, 32, 32, 32, 1, cm_i_m, cm_i_n, cm_i_k, sg, 2 }; + } + // Integer MMQ has a smaller shared memory profile, but heavier register use l_warptile_mmq_int = { 128, 128, 128, 32, subgroup_size_8 * 2, 64, 2, 4, 4, 1, subgroup_size_8 }; m_warptile_mmq_int = { 128, 64, 64, 32, subgroup_size_8, 32, 2, 2, 2, 1, subgroup_size_8 }; @@ -4308,6 +4501,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { m_align = 64; s_align = 32; + // The int8 coopmat mmq shader hardcodes a 32-wide K block (the Q8_0/Q8_1 block + // size) and tiles it with whole cooperative matrices, so anything other than a + // 16x16x16 int shape would need different warptiles. Fall back rather than + // misdispatch on such a device. + if (device->coopmat_int_support && + device->coopmat_int_m == 16 && device->coopmat_int_n == 16 && device->coopmat_int_k == 16) { + device->mul_mat_cm_int_s = ggml_vk_matmul_coopmat_int_shmem_support(device, s_warptile_mmq_cm_int); + device->mul_mat_cm_int_m = device->mul_mat_cm_int_s && ggml_vk_matmul_coopmat_int_shmem_support(device, m_warptile_mmq_cm_int); + device->mul_mat_cm_int_l = device->mul_mat_cm_int_m && ggml_vk_matmul_coopmat_int_shmem_support(device, l_warptile_mmq_cm_int); + } + 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 @@ -4752,6 +4956,29 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); +#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) + // Q8_0 x Q8_1 through sint8 cooperative matrices. Q8_0 quants need no unpacking, + // so this avoids the f16 dequantization the CREATE_MM2 shaders above perform. + // Filling the q8_1 slot makes ggml_vk_get_mul_mat_mat_pipeline() pick it up. + if (device->integer_dot_product) { + // Workgroup tile sizes of the *_warptile_mmq_cm_int entries above. + const std::array l_cm_int_wg_denoms = { l_warptile_mmq_cm_int[1], l_warptile_mmq_cm_int[2], 1 }; + const std::array m_cm_int_wg_denoms = { m_warptile_mmq_cm_int[1], m_warptile_mmq_cm_int[2], 1 }; + const std::array s_cm_int_wg_denoms = { s_warptile_mmq_cm_int[1], s_warptile_mmq_cm_int[2], 1 }; +#define CREATE_MMQ_CM(SUFFIX, TILE, WARPTILE, WG_DENOMS) \ + if (device->mul_mat_cm_int_ ## TILE) { \ + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q8_0].f32acc->SUFFIX, \ + "matmul_q8_0_q8_1_cm1_" #SUFFIX, matmul_q8_0_q8_1_cm1_len, matmul_q8_0_q8_1_cm1_data, "main", 3, \ + sizeof(vk_mat_mat_push_constants), WG_DENOMS, WARPTILE, 1, false, true, WARPTILE[10]); \ + } + + CREATE_MMQ_CM(l, l, l_warptile_mmq_cm_int, l_cm_int_wg_denoms) + CREATE_MMQ_CM(m, m, m_warptile_mmq_cm_int, m_cm_int_wg_denoms) + CREATE_MMQ_CM(s, s, s_warptile_mmq_cm_int, s_cm_int_wg_denoms) +#undef CREATE_MMQ_CM + } +#endif + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); @@ -5480,8 +5707,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { if (device->subgroup_clustered && device->subgroup_require_full_support) { ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4, "quantize_q8_1_x4", quantize_q8_1_x4_subgroup_len, quantize_q8_1_x4_subgroup_data, "main", 2, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1, true, true); + ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4_swiglu, "quantize_q8_1_x4_swiglu", quantize_q8_1_x4_swiglu_subgroup_len, quantize_q8_1_x4_swiglu_subgroup_data, "main", 3, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1, true, true); + ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4_sigmoid_mul, "quantize_q8_1_x4_sigmoid_mul", quantize_q8_1_x4_sigmoid_mul_subgroup_len, quantize_q8_1_x4_sigmoid_mul_subgroup_data, "main", 3, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1, true, true); } else { ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4, "quantize_q8_1_x4", quantize_q8_1_x4_len, quantize_q8_1_x4_data, "main", 2, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4_swiglu, "quantize_q8_1_x4_swiglu", quantize_q8_1_x4_swiglu_len, quantize_q8_1_x4_swiglu_data, "main", 3, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1); + ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4_sigmoid_mul, "quantize_q8_1_x4_sigmoid_mul", quantize_q8_1_x4_sigmoid_mul_len, quantize_q8_1_x4_sigmoid_mul_data, "main", 3, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1); } for (uint32_t i = 0; i < p021_max_gqa_ratio; ++i) { @@ -5498,12 +5729,19 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_rms_norm_f32, "rms_norm_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_f32, "rms_norm_mul_f32", rms_norm_f32_len, rms_norm_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_channel_last_f32[0], "rms_norm_channel_last_f32_64", rms_norm_channel_last_f32_len, rms_norm_channel_last_f32_data, "main", 3, sizeof(vk_op_rms_norm_channel_last_push_constants), {1, 1, 1}, {64}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_channel_last_f32[1], "rms_norm_channel_last_f32_128", rms_norm_channel_last_f32_len, rms_norm_channel_last_f32_data, "main", 3, sizeof(vk_op_rms_norm_channel_last_push_constants), {1, 1, 1}, {128}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_channel_last_f32[2], "rms_norm_channel_last_f32_256", rms_norm_channel_last_f32_len, rms_norm_channel_last_f32_data, "main", 3, sizeof(vk_op_rms_norm_channel_last_push_constants), {1, 1, 1}, {256}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_channel_last_f32[3], "rms_norm_channel_last_f32_512", rms_norm_channel_last_f32_len, rms_norm_channel_last_f32_data, "main", 3, sizeof(vk_op_rms_norm_channel_last_push_constants), {1, 1, 1}, {512}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_modulate_f32, "rms_norm_modulate_f32", rms_norm_modulate_f32_len, rms_norm_modulate_f32_data, "main", 5, sizeof(vk_op_rms_norm_modulate_push_constants), {1, 1, 1}, {}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_mul_add_f32, "mul_add_f32", mul_add_f32_len, mul_add_f32_data, "main", 4, sizeof(vk_op_mul_add_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_partials_f32, "rms_norm_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 0}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_partials_f32, "rms_norm_mul_partials_f32", rms_norm_partials_f32_len, rms_norm_partials_f32_data, "main", 4, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {0, 1}, 1, true); if (sizeof(vk_op_rms_norm_mul_rope_push_constants) <= device->properties.limits.maxPushConstantsSize) { ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_rope_f32_f32, "rms_norm_mul_rope_f32_f32", rms_norm_mul_rope_f32_f32_len, rms_norm_mul_rope_f32_f32_data, "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1}, 1, true); ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_rope_f32_f16, "rms_norm_mul_rope_f32_f16", rms_norm_mul_rope_f32_f16_len, rms_norm_mul_rope_f32_f16_data, "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_rms_norm_mul_mrope_pack_f32_f32, "rms_norm_mul_mrope_pack_f32_f32", rms_norm_mul_rope_f32_f32_len, rms_norm_mul_rope_f32_f32_data, "main", 7, sizeof(vk_op_rms_norm_mul_rope_push_constants), {1, 1, 1}, {0, 1, 1}, 1, true); } ggml_vk_create_pipeline(device, device->pipeline_rms_norm_back_f32, "rms_norm_back_f32", rms_norm_back_f32_len, rms_norm_back_f32_data, "main", 3, sizeof(vk_op_push_constants), {1, 1, 1}, {}, 1); @@ -5970,7 +6208,15 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { s != CONV_SHAPE_128x128; #endif - const uint32_t conv2d_cm1_shmem_pad = 8; + // Doubling the reduction tile halves staging rounds on the RADV RDNA3 cm1 path. + const bool conv2d_cm1_tuned = conv2d_use_cm1 && + device->architecture == vk_device_architecture::AMD_RDNA3 && + device->driver_id == vk::DriverId::eMesaRadv; + if (conv2d_cm1_tuned) { + conv2d_BS.CRS = 32; + } + + const uint32_t conv2d_cm1_shmem_pad = conv2d_cm1_tuned ? 4 : 8; auto shmem_req = [&](uint32_t pad, bool csh_store, bool fp16_shmem) { const uint32_t elem_size = fp16_shmem ? (uint32_t)sizeof(uint16_t) : (uint32_t)sizeof(float); @@ -6382,8 +6628,11 @@ static vk_device ggml_vk_get_device(size_t idx) { if (GGML_VK_SUBALLOCATION_BLOCK_SIZE != nullptr) { device->suballocation_block_size = std::stoull(GGML_VK_SUBALLOCATION_BLOCK_SIZE); } else { - // Limit batching of allocations to 1GB by default to avoid fragmentation issues - device->suballocation_block_size = 1024*1024*1024; + // Smaller UMA chunks expose enough independent mapped buffers for parallel + // weight population; discrete GPUs retain the fragmentation-oriented 1 GiB default. + device->suballocation_block_size = device->properties.deviceType == vk::PhysicalDeviceType::eIntegratedGpu + ? 256ull * 1024ull * 1024ull + : 1024ull * 1024ull * 1024ull; } device->suballocation_block_size = std::min(device->suballocation_block_size, device->max_memory_allocation_size); @@ -7071,6 +7320,12 @@ static vk_device ggml_vk_get_device(size_t idx) { device->serialize_submissions = getenv("GGML_VK_SERIALIZE_SUBMISSIONS") != nullptr; device->disable_fusion = getenv("GGML_VK_DISABLE_FUSION") != nullptr; + device->disable_rms_norm_channel_last = getenv("GGML_VK_DISABLE_RMS_NORM_CHANNEL_LAST") != nullptr; + device->disable_rms_norm_modulate = getenv("GGML_VK_DISABLE_RMS_NORM_MODULATE") != nullptr; + device->disable_rms_norm_mrope_pack = getenv("GGML_VK_DISABLE_RMS_NORM_MROPE_PACK") != nullptr; + device->disable_mul_add = getenv("GGML_VK_DISABLE_MUL_ADD") != nullptr; + device->disable_swiglu_quant = getenv("GGML_VK_DISABLE_SWIGLU_QUANT") != nullptr; + device->disable_sigmoid_quant = getenv("GGML_VK_DISABLE_SIGMOID_QUANT") != nullptr; device->add_rms_fusion = !device->disable_fusion && device->subgroup_arithmetic && @@ -8766,9 +9021,15 @@ static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, // The q8_1 (integer dot) mmq path uses a different shader with its own // shared-memory layout, so use the int-specific availability flags. const bool is_q8_1 = (src1_type == GGML_TYPE_Q8_1); - const bool mm_l = is_q8_1 ? ctx->device->mul_mat_l_int[src0_type] : ctx->device->mul_mat_l[src0_type]; - const bool mm_m = is_q8_1 ? ctx->device->mul_mat_m_int[src0_type] : ctx->device->mul_mat_m[src0_type]; - const bool mm_s = is_q8_1 ? ctx->device->mul_mat_s_int[src0_type] : ctx->device->mul_mat_s[src0_type]; + + // The int8 coopmat mmq shader (mul_mmq_cm1.comp) owns the Q8_0 q8_1 slot when it + // exists. It has its own shared-memory budget and does not inherit the float + // path's per-driver large-tile heuristic, so it needs its own flags. + const bool is_cm_int = is_q8_1 && src0_type == GGML_TYPE_Q8_0 && ctx->device->mul_mat_cm_int_s; + + const bool mm_l = is_cm_int ? ctx->device->mul_mat_cm_int_l : (is_q8_1 ? ctx->device->mul_mat_l_int[src0_type] : ctx->device->mul_mat_l[src0_type]); + const bool mm_m = is_cm_int ? ctx->device->mul_mat_cm_int_m : (is_q8_1 ? ctx->device->mul_mat_m_int[src0_type] : ctx->device->mul_mat_m[src0_type]); + const bool mm_s = is_cm_int ? ctx->device->mul_mat_cm_int_s : (is_q8_1 ? ctx->device->mul_mat_s_int[src0_type] : ctx->device->mul_mat_s[src0_type]); if (ctx->device->coopmat2) { const uint32_t shader_core_count = ctx->device->shader_core_count; @@ -9146,6 +9407,46 @@ static void ggml_vk_quantize_q8_1(ggml_backend_vk_context * ctx, vk_context& sub ggml_vk_sync_buffers(ctx, subctx); } +static void ggml_vk_quantize_q8_1_swiglu(ggml_backend_vk_context * ctx, vk_context& subctx, + const vk_subbuffer & gate, const vk_subbuffer & up, + const vk_subbuffer & out, uint32_t ne) { + vk_pipeline pipeline = ctx->device->pipeline_quantize_q8_1_x4_swiglu; + + const uint32_t num_blocks = CEIL_DIV(ne, pipeline->wg_denoms[0]); + const uint64_t max_elements = std::min( + uint64_t{ctx->device->properties.limits.maxComputeWorkGroupCount[0]} * pipeline->wg_denoms[0], + std::numeric_limits::max()); + const uint32_t elements = std::min(ne, static_cast(max_elements)); + + const vk_quantize_q8_1_push_constants pc = { + ne, + num_blocks, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { gate, up, out }, pc, { elements, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); +} + +static void ggml_vk_quantize_q8_1_sigmoid_mul(ggml_backend_vk_context * ctx, vk_context& subctx, + const vk_subbuffer & value, const vk_subbuffer & gate, + const vk_subbuffer & out, uint32_t ne) { + vk_pipeline pipeline = ctx->device->pipeline_quantize_q8_1_x4_sigmoid_mul; + + const uint32_t num_blocks = CEIL_DIV(ne, pipeline->wg_denoms[0]); + const uint64_t max_elements = std::min( + uint64_t{ctx->device->properties.limits.maxComputeWorkGroupCount[0]} * pipeline->wg_denoms[0], + std::numeric_limits::max()); + const uint32_t elements = std::min(ne, static_cast(max_elements)); + + const vk_quantize_q8_1_push_constants pc = { + ne, + num_blocks, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { value, gate, out }, pc, { elements, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); +} + static vk_pipeline ggml_vk_get_64b_indexing_pipeline(ggml_backend_vk_context * ctx, vk_pipeline &pipeline) { GGML_UNUSED(ctx); #if defined(VK_EXT_shader_64bit_indexing) @@ -9160,7 +9461,15 @@ static vk_pipeline ggml_vk_get_64b_indexing_pipeline(ggml_backend_vk_context * c return pipeline; } -static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, bool disable_split_k) { +static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, + const ggml_tensor * src0, const ggml_tensor * src1, + ggml_tensor * dst, bool disable_split_k, + const ggml_tensor * quant_src0 = nullptr, + const ggml_tensor * quant_src1 = nullptr, + vk_quantize_q8_1_fusion quant_fusion = VK_QUANTIZE_Q8_1_NONE) { + GGML_ASSERT((quant_src0 == nullptr) == (quant_src1 == nullptr)); + GGML_ASSERT((quant_fusion == VK_QUANTIZE_Q8_1_NONE) == (quant_src0 == nullptr)); + const bool fused_quant = quant_fusion != VK_QUANTIZE_Q8_1_NONE; VK_LOG_DEBUG("ggml_vk_mul_mat_q_f16((" << src0 << ", name=" << src0->name << ", type=" << ggml_type_name(src0->type) << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << ggml_type_name(src1->type) << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << ggml_type_name(dst->type) << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3]; @@ -9282,8 +9591,19 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub GGML_ASSERT(!qy_needs_dequant || to_fp16_vk_1 != nullptr); // NOLINT if (quantize_y) { - to_q8_1 = ggml_vk_get_quantize_pipeline(ctx, GGML_TYPE_Q8_1); + switch (quant_fusion) { + case VK_QUANTIZE_Q8_1_SWIGLU: + to_q8_1 = ctx->device->pipeline_quantize_q8_1_x4_swiglu; + break; + case VK_QUANTIZE_Q8_1_SIGMOID_MUL: + to_q8_1 = ctx->device->pipeline_quantize_q8_1_x4_sigmoid_mul; + break; + case VK_QUANTIZE_Q8_1_NONE: + to_q8_1 = ggml_vk_get_quantize_pipeline(ctx, GGML_TYPE_Q8_1); + break; + } } + GGML_ASSERT(!fused_quant || quantize_y); { const uint64_t split_k_size = split_k > 1 ? d_sz * split_k : 0; @@ -9392,7 +9712,25 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub if (ctx->prealloc_y_need_sync) { ggml_vk_sync_buffers(ctx, subctx); } - ggml_vk_quantize_q8_1(ctx, subctx, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0), y_ne); + switch (quant_fusion) { + case VK_QUANTIZE_Q8_1_SWIGLU: + ggml_vk_quantize_q8_1_swiglu( + ctx, subctx, + ggml_vk_tensor_subbuffer(ctx, quant_src0), + ggml_vk_tensor_subbuffer(ctx, quant_src1), + ggml_vk_subbuffer(ctx, d_Y, 0), y_ne); + break; + case VK_QUANTIZE_Q8_1_SIGMOID_MUL: + ggml_vk_quantize_q8_1_sigmoid_mul( + ctx, subctx, + ggml_vk_tensor_subbuffer(ctx, quant_src0), + ggml_vk_tensor_subbuffer(ctx, quant_src1), + ggml_vk_subbuffer(ctx, d_Y, 0), y_ne); + break; + case VK_QUANTIZE_Q8_1_NONE: + ggml_vk_quantize_q8_1(ctx, subctx, ggml_vk_subbuffer(ctx, d_Qy, qy_buf_offset), ggml_vk_subbuffer(ctx, d_Y, 0), y_ne); + break; + } ctx->prealloc_y_last_pipeline_used = to_q8_1.get(); ctx->prealloc_y_last_tensor_used = src1; ctx->prealloc_y_last_decode_vector_staging = false; @@ -10005,6 +10343,24 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c ggml_tensor * dst = cgraph->nodes[node_idx]; ggml_tensor * src0 = dst->src[0]; ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * quant_src0 = nullptr; + const ggml_tensor * quant_src1 = nullptr; + vk_quantize_q8_1_fusion quant_fusion = VK_QUANTIZE_Q8_1_NONE; + if (ctx->fused_swiglu_quant) { + GGML_ASSERT(src1->op == GGML_OP_GLU && ggml_get_glu_op(src1) == GGML_GLU_OP_SWIGLU); + quant_src0 = src1->src[0]; + quant_src1 = src1->src[1]; + quant_fusion = VK_QUANTIZE_Q8_1_SWIGLU; + } else if (ctx->fused_sigmoid_quant) { + GGML_ASSERT(src1->op == GGML_OP_MUL); + const bool src0_is_sigmoid = src1->src[0]->op == GGML_OP_UNARY && + ggml_get_unary_op(src1->src[0]) == GGML_UNARY_OP_SIGMOID; + const ggml_tensor * sigmoid = src0_is_sigmoid ? src1->src[0] : src1->src[1]; + GGML_ASSERT(sigmoid->op == GGML_OP_UNARY && ggml_get_unary_op(sigmoid) == GGML_UNARY_OP_SIGMOID); + quant_src0 = src1->src[0] == sigmoid ? src1->src[1] : src1->src[0]; + quant_src1 = sigmoid->src[0]; + quant_fusion = VK_QUANTIZE_Q8_1_SIGMOID_MUL; + } VK_LOG_DEBUG("ggml_vk_mul_mat(" << src0 << ", " << src1 << ", " << dst << ")"); // Handle huge A matrix by splitting the M dimensions. This works well for convolution use cases @@ -10014,6 +10370,7 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c const size_t nbytes = ggml_nbytes(src0); const bool needs_split = dst->ne[2] == 1 && dst->ne[3] == 1 && nbytes > ctx->device->properties.limits.maxStorageBufferRange; if (needs_split) { + GGML_ASSERT(quant_fusion == VK_QUANTIZE_Q8_1_NONE); // Choose the number of rows that can fit (and divide by two, to allow for any additional offsets) const uint32_t M_split = ctx->device->properties.limits.maxStorageBufferRange / (2 * src0->nb[1]); uint32_t m_offset = 0; @@ -10061,7 +10418,7 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c (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 { - ggml_vk_mul_mat_q_f16(ctx, subctx, src0, src1, dst, false); + ggml_vk_mul_mat_q_f16(ctx, subctx, src0, src1, dst, false, quant_src0, quant_src1, quant_fusion); } } @@ -10848,6 +11205,14 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type; const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type; + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga. // For coopmat2 FA, we always use the small size (which is still pretty large for gqa). vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc); @@ -10863,7 +11228,14 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx workgroups_y /= gqa_ratio; } - tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc); + const bool allow_work_per_tile = mask == nullptr && sinks == nullptr && + max_bias == 0.0f && logit_softcap == 0.0f && + gqa_ratio == 1 && (qk_ratio == 1 || qk_ratio == 4) && + qk_ratio * nek2 == neq2 && qk_ratio * nev2 == neq2 && + neq3 == nek3 && neq3 == nev3 && + HSK == 128 && HSV == 128 && + k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16; + tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc, allow_work_per_tile); const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type)); uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type)); @@ -10898,14 +11270,6 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx aligned = false; } - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; - - memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); - memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); - if (logit_softcap != 0) { scale /= logit_softcap; } @@ -10958,6 +11322,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } } + if (fa_pipeline_state.work_per_tile != 0) { + split_k = 1; + split_kv = KV; + } + if (split_k > 1) { // Try to evenly split KV into split_k chunks, but it needs to be a multiple // of "align", so recompute split_k based on that. @@ -12536,6 +12905,31 @@ static void ggml_vk_multi_add(ggml_backend_vk_context * ctx, vk_context& subctx, }, pc, elements); } +static void ggml_vk_mul_add(ggml_backend_vk_context * ctx, vk_context& subctx, + const ggml_cgraph * cgraph, int node_idx) { + const ggml_tensor * mul = cgraph->nodes[node_idx]; + const ggml_tensor * add = cgraph->nodes[node_idx + 1]; + const ggml_tensor * residual = add->src[0] == mul ? add->src[1] : add->src[0]; + const bool src0_is_value = ggml_are_same_shape(mul->src[0], mul); + const ggml_tensor * value = src0_is_value ? mul->src[0] : mul->src[1]; + const ggml_tensor * gate = src0_is_value ? mul->src[1] : mul->src[0]; + + const vk_op_mul_add_push_constants pc { + (uint32_t)add->ne[0] / 4, + (uint32_t)ggml_nrows(add), + }; + + vk_pipeline pipeline = ctx->device->pipeline_mul_add_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, value), + ggml_vk_tensor_subbuffer(ctx, gate), + ggml_vk_tensor_subbuffer(ctx, add), + }, pc, { pc.ne0_vec4, pc.nrows, 1 }); +} + static void ggml_vk_add(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { const uint32_t src0_type_size = ggml_type_size(src0->type); const uint32_t src1_type_size = ggml_type_size(src1->type); @@ -13258,7 +13652,157 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor * return rope; } +static void ggml_vk_rms_norm_channel_last(ggml_backend_vk_context * ctx, vk_context& subctx, + const struct ggml_cgraph * cgraph, int node_idx) { + const ggml_tensor * input_cont = cgraph->nodes[node_idx]; + const ggml_tensor * input_perm = input_cont->src[0]; + const ggml_tensor * src = input_perm->src[0]; + const ggml_tensor * rms = cgraph->nodes[node_idx + 1]; + const ggml_tensor * norm = cgraph->nodes[node_idx + 2]->src[0]; + const ggml_tensor * dst = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; + + const vk_op_rms_norm_channel_last_push_constants pc { + (uint32_t)src->ne[0], + (uint32_t)src->ne[1], + (uint32_t)src->ne[2], + (uint32_t)src->ne[3], + ggml_get_op_params_f32(rms, 0), + }; + + const uint32_t pipeline_idx = pc.channels < 128 ? 0 : pc.channels < 256 ? 1 : pc.channels < 512 ? 2 : 3; + vk_pipeline pipeline = ctx->device->pipeline_rms_norm_channel_last_f32[pipeline_idx]; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src), + ggml_vk_tensor_subbuffer(ctx, norm), + ggml_vk_tensor_subbuffer(ctx, dst), + }, pc, { pc.width, pc.height, pc.depth }); +} + +static void ggml_vk_rms_norm_modulate(ggml_backend_vk_context * ctx, vk_context& subctx, + const struct ggml_cgraph * cgraph, int node_idx, float eps) { + const ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * norm_mul = nullptr; + const ggml_tensor * mod_scale_cont = nullptr; + const ggml_tensor * scale_op = nullptr; + const ggml_tensor * shift_cont = nullptr; + for (int i = node_idx + 1; i <= node_idx + ctx->num_additional_fused_ops; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (!norm_mul && node->op == GGML_OP_MUL) { + norm_mul = node; + } else if (node->op == GGML_OP_CONT && !scale_op) { + mod_scale_cont = node; + } else if (node->op == GGML_OP_SCALE) { + scale_op = node; + } else if (node->op == GGML_OP_CONT) { + shift_cont = node; + } + } + const ggml_tensor * add = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; + GGML_ASSERT(norm_mul && scale_op && add->op == GGML_OP_ADD && + ((mod_scale_cont == nullptr) == (shift_cont == nullptr))); + + const ggml_tensor * src = rms->src[0]; + const ggml_tensor * norm = norm_mul->src[1]; + const ggml_tensor * mod_scale = mod_scale_cont ? mod_scale_cont->src[0] : scale_op->src[0]; + const ggml_tensor * shift = shift_cont ? shift_cont->src[0] : add->src[1]; + + vk_op_rms_norm_modulate_push_constants pc { + (uint32_t)src->ne[0], (uint32_t)src->ne[1], (uint32_t)src->ne[2], (uint32_t)src->ne[3], + (uint32_t)(src->nb[1] / sizeof(float)), + (uint32_t)(src->nb[2] / sizeof(float)), + (uint32_t)(src->nb[3] / sizeof(float)), + eps, + }; + + vk_pipeline pipeline = ctx->device->pipeline_rms_norm_modulate_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + ggml_vk_tensor_subbuffer(ctx, src), + ggml_vk_tensor_subbuffer(ctx, norm), + ggml_vk_tensor_subbuffer(ctx, mod_scale), + ggml_vk_tensor_subbuffer(ctx, shift), + ggml_vk_tensor_subbuffer(ctx, add), + }, pc, { (uint32_t)src->ne[1], (uint32_t)src->ne[2], (uint32_t)src->ne[3] }); +} + +static void ggml_vk_rms_norm_mrope_pack(ggml_backend_vk_context * ctx, vk_context& subctx, + const struct ggml_cgraph * cgraph, int node_idx, float eps) { + const ggml_tensor * rms = cgraph->nodes[node_idx]; + const ggml_tensor * mul = nullptr; + const ggml_tensor * rope = nullptr; + for (int i = node_idx + 1; i <= node_idx + ctx->num_additional_fused_ops; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (!mul && node->op == GGML_OP_MUL) { + mul = node; + } else if (node->op == GGML_OP_ROPE) { + rope = node; + } + } + + GGML_ASSERT(mul && rope && mul->src[0] == rms); + const ggml_tensor * src = rms->src[0]; + const ggml_tensor * norm = mul->src[1]; + const ggml_tensor * dst = cgraph->nodes[node_idx + ctx->num_additional_fused_ops]; + + const uint32_t src_type_size = ggml_type_size(src->type); + const uint32_t norm_type_size = ggml_type_size(norm->type); + const uint32_t dst_type_size = ggml_type_size(dst->type); + + vk_op_rms_norm_mul_rope_push_constants pc; + pc.bin = { + (uint32_t)ggml_nelements(src), + (uint32_t)src->ne[0], (uint32_t)src->ne[1], (uint32_t)src->ne[2], (uint32_t)src->ne[3], + (uint32_t)(src->nb[0] / src_type_size), (uint32_t)(src->nb[1] / src_type_size), + (uint32_t)(src->nb[2] / src_type_size), (uint32_t)(src->nb[3] / src_type_size), + (uint32_t)norm->ne[0], (uint32_t)norm->ne[1], (uint32_t)norm->ne[2], (uint32_t)norm->ne[3], + (uint32_t)(norm->nb[0] / norm_type_size), (uint32_t)(norm->nb[1] / norm_type_size), + (uint32_t)(norm->nb[2] / norm_type_size), (uint32_t)(norm->nb[3] / norm_type_size), + (uint32_t)dst->ne[0], (uint32_t)dst->ne[1], (uint32_t)dst->ne[2], (uint32_t)dst->ne[3], + (uint32_t)(dst->nb[0] / dst_type_size), (uint32_t)(dst->nb[1] / dst_type_size), + (uint32_t)(dst->nb[2] / dst_type_size), (uint32_t)(dst->nb[3] / dst_type_size), + 0, + eps, 0.0f, 0, + }; + pc.rope = ggml_vk_make_rope_constants(rope, rope->src[0], true, false, 0); + + // The ROPE invocation retains [head, token, batch] coordinates, but writes the + // already-packed [head_dim, token, head*batch] layout consumed by flash attention. + pc.rope.nb11 = dst->nb[2] / dst_type_size; + pc.rope.nb12 = dst->nb[1] / dst_type_size; + pc.rope.nb13 = dst->nb[3] / dst_type_size; + pc.rope.d_offset = get_misalign_bytes(ctx, dst) / dst_type_size; + + vk_pipeline pipeline = ctx->device->pipeline_rms_norm_mul_mrope_pack_f32_f32; + + const vk_subbuffer src_buf = ggml_vk_tensor_subbuffer(ctx, src); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { + src_buf, + ggml_vk_tensor_subbuffer(ctx, norm), + src_buf, + ggml_vk_tensor_subbuffer(ctx, rope->src[1]), + ggml_vk_tensor_subbuffer(ctx, rope->src[2]), + dst_buf, + src_buf, + }, pc, { (uint32_t)src->ne[1], (uint32_t)src->ne[2], (uint32_t)src->ne[3] }); +} + static void ggml_vk_rms_norm(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, float * op_params) { + if (ctx->fused_rms_norm_mrope_pack) { + ggml_vk_rms_norm_mrope_pack(ctx, subctx, cgraph, node_idx, op_params[0]); + return; + } + + if (ctx->fused_rms_norm_modulate) { + ggml_vk_rms_norm_modulate(ctx, subctx, cgraph, node_idx, op_params[0]); + return; + } + ggml_tensor * dst; const ggml_tensor * src0; const ggml_tensor * src1; @@ -15446,7 +15990,9 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; case GGML_OP_MUL: - if (ctx->num_additional_fused_ops) { + if (ctx->fused_mul_add) { + ggml_vk_mul_add(ctx, compute_ctx, cgraph, node_idx); + } else if (ctx->num_additional_fused_ops) { ggml_vk_snake_dispatch_fused(ctx, compute_ctx, cgraph, node_idx); } else { ggml_vk_mul(ctx, compute_ctx, src0, src1, node); @@ -15528,7 +16074,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_OP_CPY: case GGML_OP_CONT: case GGML_OP_DUP: - ggml_vk_cpy(ctx, compute_ctx, src0, node); + if (ctx->fused_rms_norm_channel_last) { + ggml_vk_rms_norm_channel_last(ctx, compute_ctx, cgraph, node_idx); + } else { + ggml_vk_cpy(ctx, compute_ctx, src0, node); + } break; case GGML_OP_SET_ROWS: @@ -15559,6 +16109,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; case GGML_OP_UNARY: + if (ctx->fused_sigmoid_quant) { + ggml_vk_mul_mat(ctx, compute_ctx, cgraph, node_idx + 2); + break; + } if (ctx->fused_topk_moe_mode != TOPK_MOE_COUNT) { ggml_vk_topk_moe(ctx, compute_ctx, cgraph, node_idx); break; @@ -15596,6 +16150,10 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr } break; case GGML_OP_GLU: + if (ctx->fused_swiglu_quant) { + ggml_vk_mul_mat(ctx, compute_ctx, cgraph, node_idx + 1); + break; + } switch (ggml_get_glu_op(node)) { case GGML_GLU_OP_GEGLU: case GGML_GLU_OP_REGLU: @@ -16494,7 +17052,7 @@ static void ggml_backend_vk_synchronize(ggml_backend_t backend) { ggml_vk_graph_cleanup(ctx); } -static bool ggml_vk_is_empty(ggml_tensor * node) { +static bool ggml_vk_is_empty(const ggml_tensor * node) { return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; } @@ -16651,6 +17209,125 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g return true; } +static bool ggml_vk_can_fuse_mul_add(const ggml_backend_vk_context * ctx, + const struct ggml_cgraph * cgraph, int node_idx) { + if (ctx->device->disable_mul_add || + !ggml_can_fuse(cgraph, node_idx, { GGML_OP_MUL, GGML_OP_ADD })) { + return false; + } + + const ggml_tensor * mul = cgraph->nodes[node_idx]; + const ggml_tensor * add = cgraph->nodes[node_idx + 1]; + const ggml_tensor * residual = add->src[0] == mul ? add->src[1] : add->src[0]; + + const bool src0_is_value = ggml_are_same_shape(mul->src[0], mul); + const bool src1_is_value = ggml_are_same_shape(mul->src[1], mul); + if (src0_is_value == src1_is_value) { + return false; + } + + const ggml_tensor * value = src0_is_value ? mul->src[0] : mul->src[1]; + const ggml_tensor * gate = src0_is_value ? mul->src[1] : mul->src[0]; + + if (mul->type != GGML_TYPE_F32 || add->type != GGML_TYPE_F32 || + residual->type != GGML_TYPE_F32 || value->type != GGML_TYPE_F32 || gate->type != GGML_TYPE_F32 || + !ggml_are_same_shape(residual, add) || !ggml_are_same_shape(value, add) || + gate->ne[0] != add->ne[0] || ggml_nrows(gate) != 1 || + add->ne[0] % 4 != 0 || + !ggml_is_contiguous(residual) || !ggml_is_contiguous(value) || + !ggml_is_contiguous(gate) || !ggml_is_contiguous(add) || + ggml_nrows(add) > ctx->device->properties.limits.maxComputeWorkGroupCount[1]) { + return false; + } + + const ggml_tensor * tensors[] = { residual, value, gate, add }; + for (const ggml_tensor * tensor : tensors) { + if (get_misalign_bytes(ctx, tensor) != 0) { + return false; + } + } + + return true; +} + +static bool ggml_vk_can_fuse_swiglu_quant(ggml_backend_vk_context * ctx, + const struct ggml_cgraph * cgraph, int node_idx) { + if (ctx->device->disable_swiglu_quant || + !ggml_can_fuse_subgraph(cgraph, node_idx, + { GGML_OP_GLU, GGML_OP_MUL_MAT }, { node_idx + 1 })) { + return false; + } + + const ggml_tensor * glu = cgraph->nodes[node_idx]; + const ggml_tensor * mm = cgraph->nodes[node_idx + 1]; + const ggml_tensor * gate = glu->src[0]; + const ggml_tensor * up = glu->src[1]; + + if (ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || up == nullptr || mm->src[1] != glu || + glu->type != GGML_TYPE_F32 || gate->type != GGML_TYPE_F32 || up->type != GGML_TYPE_F32 || + mm->src[0]->type != GGML_TYPE_Q8_0 || mm->type != GGML_TYPE_F32 || + !ctx->device->integer_dot_product || mm->ne[1] <= mul_mat_vec_max_cols || + ggml_nbytes(mm->src[0]) > ctx->device->properties.limits.maxStorageBufferRange || + !ggml_are_same_shape(gate, up) || !ggml_are_same_shape(gate, glu) || + !ggml_is_contiguous(gate) || !ggml_is_contiguous(up) || !ggml_is_contiguous(glu) || + (glu->ne[0] * glu->ne[1]) % 4 != 0) { + return false; + } + + const ggml_tensor * tensors[] = { gate, up, mm }; + for (const ggml_tensor * tensor : tensors) { + if (get_misalign_bytes(ctx, tensor) != 0) { + return false; + } + } + + return ggml_vk_get_mul_mat_mat_pipeline( + ctx, mm->src[0]->type, GGML_TYPE_Q8_1, (ggml_prec)mm->op_params[0]) != nullptr; +} + +static bool ggml_vk_can_fuse_sigmoid_quant(ggml_backend_vk_context * ctx, + const struct ggml_cgraph * cgraph, int node_idx) { + if (ctx->device->disable_sigmoid_quant || + !ggml_can_fuse_subgraph(cgraph, node_idx, + { GGML_OP_UNARY, GGML_OP_MUL, GGML_OP_MUL_MAT }, { node_idx + 2 })) { + return false; + } + + const ggml_tensor * sigmoid = cgraph->nodes[node_idx]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + const ggml_tensor * mm = cgraph->nodes[node_idx + 2]; + + if (ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID || + (mul->src[0] != sigmoid && mul->src[1] != sigmoid) || mm->src[1] != mul) { + return false; + } + + const ggml_tensor * value = mul->src[0] == sigmoid ? mul->src[1] : mul->src[0]; + const ggml_tensor * gate = sigmoid->src[0]; + + if (sigmoid->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32 || + value->type != GGML_TYPE_F32 || gate->type != GGML_TYPE_F32 || + mm->src[0]->type != GGML_TYPE_Q8_0 || mm->type != GGML_TYPE_F32 || + !ctx->device->integer_dot_product || mm->ne[1] <= mul_mat_vec_max_cols || + ggml_nbytes(mm->src[0]) > ctx->device->properties.limits.maxStorageBufferRange || + !ggml_are_same_shape(value, gate) || !ggml_are_same_shape(value, mul) || + !ggml_is_contiguous(value) || !ggml_is_contiguous(gate) || + !ggml_is_contiguous(sigmoid) || !ggml_is_contiguous(mul) || + (mul->ne[0] * mul->ne[1]) % 4 != 0) { + return false; + } + + const ggml_tensor * tensors[] = { value, gate, mm }; + for (const ggml_tensor * tensor : tensors) { + if (get_misalign_bytes(ctx, tensor) != 0) { + return false; + } + } + + return ggml_vk_get_mul_mat_mat_pipeline( + ctx, mm->src[0]->type, GGML_TYPE_Q8_1, (ggml_prec)mm->op_params[0]) != nullptr; +} + // Match SSM_CONV + UNARY(SILU) or SSM_CONV + ADD + UNARY(SILU). num_extra is 1 or 2. static bool ggml_vk_can_fuse_ssm_conv(const ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx, int num_extra) { @@ -16973,6 +17650,405 @@ static bool ggml_vk_can_fuse_rms_norm_mul_rope(ggml_backend_vk_context * ctx, co return true; } +static uint32_t ggml_vk_can_fuse_rms_norm_mrope_pack(ggml_backend_vk_context * ctx, + const struct ggml_cgraph * cgraph, + int node_idx) { + if (ctx->device->disable_rms_norm_mrope_pack) { + return 0; + } + if (node_idx >= cgraph->n_nodes || cgraph->nodes[node_idx]->op != GGML_OP_RMS_NORM) { + return 0; + } + + int compute_nodes[6] {}; + uint32_t n_compute = 0; + for (int i = node_idx; i < cgraph->n_nodes && i <= node_idx + 12 && n_compute < 6; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (ggml_vk_is_empty(node)) { + continue; + } + compute_nodes[n_compute++] = i; + } + + static constexpr ggml_op required[] { + GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_CONT, GGML_OP_ROPE, GGML_OP_CONT, + }; + if (n_compute < 5) { + return 0; + } + for (uint32_t i = 0; i < 5; ++i) { + const ggml_tensor * node = cgraph->nodes[compute_nodes[i]]; + if (node->op != required[i] || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return 0; + } + } + auto const reject = [node_idx](const char * reason) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MROPE_PACK rejected at node " << node_idx << ": " << reason << "\n"; + } + return 0u; + }; + + auto const resolves_to = [](const ggml_tensor * tensor, const ggml_tensor * target) { + while (tensor) { + if (tensor == target) { + return true; + } + if (tensor->op != GGML_OP_VIEW && tensor->op != GGML_OP_RESHAPE && + tensor->op != GGML_OP_PERMUTE && tensor->op != GGML_OP_TRANSPOSE) { + return false; + } + tensor = tensor->src[0]; + } + return false; + }; + auto const perm_matches = [](const ggml_tensor * tensor, const int32_t expected[GGML_MAX_DIMS]) { + if (!tensor || tensor->op != GGML_OP_PERMUTE) { + return false; + } + for (uint32_t d = 0; d < GGML_MAX_DIMS; ++d) { + if (ggml_get_op_params_i32(tensor, d) != expected[d]) { + return false; + } + } + return true; + }; + + const ggml_tensor * rms = cgraph->nodes[compute_nodes[0]]; + const ggml_tensor * mul = cgraph->nodes[compute_nodes[1]]; + const ggml_tensor * deinterleave = cgraph->nodes[compute_nodes[2]]; + const ggml_tensor * rope = cgraph->nodes[compute_nodes[3]]; + const ggml_tensor * pack = cgraph->nodes[compute_nodes[4]]; + const ggml_tensor * src = rms->src[0]; + + if (!src || mul->src[0] != rms || !mul->src[1] || + src->type != GGML_TYPE_F32 || rms->type != GGML_TYPE_F32 || + mul->type != GGML_TYPE_F32 || mul->src[1]->type != GGML_TYPE_F32 || + deinterleave->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 || + pack->type != GGML_TYPE_F32 || + !ggml_are_same_shape(src, rms) || !ggml_are_same_shape(rms, mul) || + !ggml_is_contiguous(src) || !ggml_is_contiguous_rows(mul->src[1]) || + !ggml_is_contiguous(deinterleave) || !ggml_is_contiguous(rope) || !ggml_is_contiguous(pack) || + ggml_nelements(mul->src[1]) != src->ne[0] || src->ne[3] != 1 || + src->ne[0] == 0 || (src->ne[0] & 1) != 0 || src->ne[0] > 1024) { + return reject("type, shape, or base layout mismatch"); + } + + const int32_t deinterleave_perm_order[GGML_MAX_DIMS] = { 1, 0, 2, 3 }; + const ggml_tensor * deinterleave_perm = deinterleave->src[0]; + if (!perm_matches(deinterleave_perm, deinterleave_perm_order) || + !resolves_to(deinterleave_perm->src[0], mul) || + deinterleave->ne[0] != src->ne[0] / 2 || deinterleave->ne[1] != 2 || + deinterleave->ne[2] != src->ne[1] || deinterleave->ne[3] != src->ne[2]) { + return reject("NORMAL-to-NEOX deinterleave mismatch"); + } + + const ggml_tensor * rope_src = rope->src[0]; + const int32_t * rope_params = (const int32_t *)rope->op_params; + if (!rope_src || !resolves_to(rope_src, deinterleave) || + !ggml_are_same_shape(rope_src, src) || + rope_params[1] != src->ne[0] || rope_params[2] != GGML_ROPE_TYPE_MROPE || + !rope->src[1] || rope->src[1]->type != GGML_TYPE_I32 || + !rope->src[2] || rope->src[2]->type != GGML_TYPE_F32 || + rope->src[1]->ne[0] != 4 * src->ne[2] || rope->src[2]->ne[0] < src->ne[0] / 2) { + return reject("MROPE parameters or input chain mismatch"); + } + + const int32_t pack_perm_order[GGML_MAX_DIMS] = { 0, 2, 1, 3 }; + const ggml_tensor * pack_perm = pack->src[0]; + if (!perm_matches(pack_perm, pack_perm_order) || !resolves_to(pack_perm->src[0], rope) || + pack->ne[0] != src->ne[0] || pack->ne[1] != src->ne[2] || + pack->ne[2] != src->ne[1] || pack->ne[3] != 1) { + return reject("flash-attention pack mismatch"); + } + + const ggml_tensor * dst = pack; + const int end_idx = compute_nodes[4]; + + if (end_idx - node_idx >= 13 || + sizeof(vk_op_rms_norm_mul_rope_push_constants) > ctx->device->properties.limits.maxPushConstantsSize || + get_misalign_bytes(ctx, src) != 0 || get_misalign_bytes(ctx, mul->src[1]) != 0 || + get_misalign_bytes(ctx, rope->src[1]) != 0 || get_misalign_bytes(ctx, rope->src[2]) != 0 || + get_misalign_bytes(ctx, dst) % ggml_type_size(dst->type) != 0) { + return reject("push-constant or storage alignment mismatch"); + } + + for (int i = node_idx; i < end_idx; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { + return reject("intermediate is a graph output"); + } + int internal_uses = 0; + for (int j = node_idx; j <= end_idx; ++j) { + for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { + if (cgraph->nodes[j]->src[s] == node) { + internal_uses++; + } + } + } + if (internal_uses != ggml_node_get_use_count(cgraph, i)) { + return reject("intermediate has an external use"); + } + } + + return end_idx - node_idx; +} + +static bool ggml_vk_can_fuse_rms_norm_channel_last(ggml_backend_vk_context * ctx, + const struct ggml_cgraph * cgraph, + int node_idx) { + if (ctx->device->disable_rms_norm_channel_last) { + return false; + } + + static constexpr ggml_op ops[] { + GGML_OP_CONT, + GGML_OP_RMS_NORM, + GGML_OP_RESHAPE, + GGML_OP_MUL, + GGML_OP_PERMUTE, + GGML_OP_CONT, + }; + const int output = node_idx + 5; + if (!ggml_can_fuse_subgraph(cgraph, node_idx, 6, ops, &output, 1)) { + return false; + } + + const ggml_tensor * input_cont = cgraph->nodes[node_idx + 0]; + const ggml_tensor * rms = cgraph->nodes[node_idx + 1]; + const ggml_tensor * norm_view = cgraph->nodes[node_idx + 2]; + const ggml_tensor * mul = cgraph->nodes[node_idx + 3]; + const ggml_tensor * out_perm = cgraph->nodes[node_idx + 4]; + const ggml_tensor * dst = cgraph->nodes[node_idx + 5]; + const ggml_tensor * input_perm = input_cont->src[0]; + + if (!input_perm || input_perm->op != GGML_OP_PERMUTE || !input_perm->src[0] || + rms->src[0] != input_cont || norm_view->src[0] == nullptr || + !((mul->src[0] == rms && mul->src[1] == norm_view) || + (mul->src[1] == rms && mul->src[0] == norm_view)) || + out_perm->src[0] != mul || dst->src[0] != out_perm) { + return false; + } + + const int32_t expected_input_perm[4] = { 1, 2, 3, 0 }; + const int32_t expected_output_perm[4] = { 3, 0, 1, 2 }; + for (uint32_t d = 0; d < GGML_MAX_DIMS; ++d) { + if (ggml_get_op_params_i32(input_perm, d) != expected_input_perm[d] || + ggml_get_op_params_i32(out_perm, d) != expected_output_perm[d]) { + return false; + } + } + + const ggml_tensor * src = input_perm->src[0]; + const ggml_tensor * norm = norm_view->src[0]; + const ggml_tensor * typed[] = { src, input_perm, input_cont, rms, norm_view, norm, mul, out_perm, dst }; + for (const ggml_tensor * tensor : typed) { + if (tensor->type != GGML_TYPE_F32) { + return false; + } + } + + if (!ggml_is_contiguous(src) || !ggml_is_contiguous(norm) || !ggml_is_contiguous(dst) || + !ggml_are_same_shape(src, dst) || !ggml_are_same_shape(input_perm, input_cont) || + !ggml_are_same_shape(input_cont, rms) || !ggml_are_same_shape(rms, mul) || + !ggml_are_same_shape(src, out_perm) || src->ne[3] != ggml_nelements(norm) || + norm_view->ne[0] != src->ne[3] || ggml_nelements(norm_view) != src->ne[3] || + input_cont->ne[0] != src->ne[3] || input_cont->ne[1] != src->ne[0] || + input_cont->ne[2] != src->ne[1] || input_cont->ne[3] != src->ne[2]) { + return false; + } + + return get_misalign_bytes(ctx, src) == 0 && + get_misalign_bytes(ctx, norm) == 0 && + get_misalign_bytes(ctx, dst) == 0; +} + +static uint32_t ggml_vk_can_fuse_rms_norm_modulate(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, + int node_idx) { + if (ctx->device->disable_rms_norm_modulate) { + return 0; + } + + static constexpr ggml_op sequence_with_cont[] { + GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_CONT, GGML_OP_SCALE, + GGML_OP_MUL, GGML_OP_CONT, GGML_OP_ADD, + }; + static constexpr ggml_op sequence_direct[] { + GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_SCALE, GGML_OP_MUL, GGML_OP_ADD, + }; + + int compute_nodes[7] {}; + uint32_t n_compute = 0; + int end_idx = -1; + for (int i = node_idx; i < cgraph->n_nodes && i < node_idx + 13; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (ggml_vk_is_empty(node)) { + if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { + return 0; + } + continue; + } + if (n_compute == 7) { + return 0; + } + compute_nodes[n_compute++] = i; + if (node->op == GGML_OP_ADD) { + end_idx = i; + break; + } + } + + const ggml_op * expected = nullptr; + if (n_compute == 7) { + expected = sequence_with_cont; + } else if (n_compute == 5) { + expected = sequence_direct; + } else { + return 0; + } + + for (uint32_t i = 0; i < n_compute; ++i) { + const ggml_tensor * node = cgraph->nodes[compute_nodes[i]]; + if (node->op != expected[i] || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return 0; + } + if (i + 1 < n_compute) { + int subgraph_uses = 0; + for (int j = node_idx; j <= end_idx; ++j) { + for (uint32_t src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + if (cgraph->nodes[j]->src[src_idx] == node) { + subgraph_uses++; + } + } + } + if (subgraph_uses != ggml_node_get_use_count(cgraph, compute_nodes[i]) || + (node->flags & GGML_TENSOR_FLAG_OUTPUT)) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MODULATE rejected at node " << node_idx + << ": external use/output on " << compute_nodes[i] + << " (subgraph=" << subgraph_uses + << ", total=" << ggml_node_get_use_count(cgraph, compute_nodes[i]) << ")\n"; + } + return 0; + } + } + } + + auto const same_value_alias = [](const ggml_tensor * lhs, const ggml_tensor * rhs) { + if (lhs == rhs) { + return true; + } + const ggml_tensor * lhs_base = lhs->view_src ? lhs->view_src : lhs; + const ggml_tensor * rhs_base = rhs->view_src ? rhs->view_src : rhs; + if (lhs_base != rhs_base || lhs->view_offs != rhs->view_offs || lhs->type != rhs->type || + !ggml_are_same_shape(lhs, rhs)) { + return false; + } + for (uint32_t d = 0; d < GGML_MAX_DIMS; ++d) { + if (lhs->nb[d] != rhs->nb[d]) { + return false; + } + } + return true; + }; + auto const resolves_to = [&same_value_alias](const ggml_tensor * tensor, const ggml_tensor * target) { + while (tensor) { + if (same_value_alias(tensor, target)) { + return true; + } + if (tensor->op != GGML_OP_VIEW && tensor->op != GGML_OP_RESHAPE && + tensor->op != GGML_OP_PERMUTE && tensor->op != GGML_OP_TRANSPOSE) { + return false; + } + tensor = tensor->src[0]; + } + return false; + }; + + const bool with_cont = n_compute == 7; + const ggml_tensor * rms = cgraph->nodes[compute_nodes[0]]; + const ggml_tensor * norm_mul = cgraph->nodes[compute_nodes[1]]; + const ggml_tensor * mod_scale_cont = with_cont ? cgraph->nodes[compute_nodes[2]] : nullptr; + const ggml_tensor * scale_op = cgraph->nodes[compute_nodes[with_cont ? 3 : 2]]; + const ggml_tensor * mod_mul = cgraph->nodes[compute_nodes[with_cont ? 4 : 3]]; + const ggml_tensor * shift_cont = with_cont ? cgraph->nodes[compute_nodes[5]] : nullptr; + const ggml_tensor * add = cgraph->nodes[compute_nodes[with_cont ? 6 : 4]]; + + const bool edge_rms_to_norm_mul = resolves_to(norm_mul->src[0], rms); + const bool edge_scale_cont = !with_cont || resolves_to(scale_op->src[0], mod_scale_cont); + const bool edge_norm_to_mod = resolves_to(mod_mul->src[0], norm_mul); + const bool edge_scale_to_mod = resolves_to(mod_mul->src[1], scale_op); + const bool edge_mod_to_add = resolves_to(add->src[0], mod_mul); + const bool edge_shift_cont = !with_cont || resolves_to(add->src[1], shift_cont); + if (!edge_rms_to_norm_mul || !edge_scale_cont || !edge_norm_to_mod || + !edge_scale_to_mod || !edge_mod_to_add || !edge_shift_cont) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MODULATE rejected at node " << node_idx + << ": edge mismatch [rms-mul=" << edge_rms_to_norm_mul + << " scale-cont=" << edge_scale_cont + << " norm-mod=" << edge_norm_to_mod + << " scale-mod=" << edge_scale_to_mod + << " mod-add=" << edge_mod_to_add + << " shift-cont=" << edge_shift_cont << "]\n"; + } + return 0; + } + + const ggml_tensor * src = rms->src[0]; + const ggml_tensor * norm = norm_mul->src[1]; + const ggml_tensor * mod_scale = with_cont ? mod_scale_cont->src[0] : scale_op->src[0]; + const ggml_tensor * shift = with_cont ? shift_cont->src[0] : add->src[1]; + const ggml_tensor * dst = add; + const float * scale_params = reinterpret_cast(scale_op->op_params); + + if (scale_params[0] != 1.0f || scale_params[1] != 1.0f || + src->type != GGML_TYPE_F32 || norm->type != GGML_TYPE_F32 || + mod_scale->type != GGML_TYPE_F32 || shift->type != GGML_TYPE_F32 || + rms->type != GGML_TYPE_F32 || norm_mul->type != GGML_TYPE_F32 || + scale_op->type != GGML_TYPE_F32 || mod_mul->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MODULATE rejected at node " << node_idx << ": params/type mismatch\n"; + } + return 0; + } + + if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst) || !ggml_are_same_shape(src, dst)) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MODULATE rejected at node " << node_idx + << ": src/dst layout mismatch (src_contig=" << ggml_is_contiguous(src) + << ", dst_contig=" << ggml_is_contiguous(dst) << ")\n"; + } + return 0; + } + + const ggml_tensor * vectors[] = { norm, mod_scale, shift }; + for (const ggml_tensor * vector : vectors) { + if (!ggml_is_contiguous(vector) || vector->ne[0] != src->ne[0] || + ggml_nelements(vector) != src->ne[0]) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MODULATE rejected at node " << node_idx + << ": vector layout mismatch (ne0=" << vector->ne[0] + << ", elements=" << ggml_nelements(vector) + << ", contig=" << ggml_is_contiguous(vector) << ")\n"; + } + return 0; + } + } + + const ggml_tensor * tensors[] = { src, norm, mod_scale, shift, dst }; + for (const ggml_tensor * tensor : tensors) { + if (get_misalign_bytes(ctx, tensor) != 0) { + if (vk_enable_sync_logger) { + std::cerr << "RMS_NORM_MUL_MODULATE rejected at node " << node_idx + << ": storage-buffer misalignment=" << get_misalign_bytes(ctx, tensor) << "\n"; + } + return 0; + } + } + + return end_idx - node_idx + 1; +} + static uint32_t ggml_vk_fuse_multi_add(ggml_backend_vk_context * ctx, const struct ggml_cgraph * cgraph, int node_idx) { const ggml_tensor *first_node = cgraph->nodes[node_idx]; @@ -17197,6 +18273,12 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; + ctx->fused_rms_norm_channel_last = false; + ctx->fused_rms_norm_modulate = false; + ctx->fused_rms_norm_mrope_pack = false; + ctx->fused_mul_add = false; + ctx->fused_swiglu_quant = false; + ctx->fused_sigmoid_quant = false; const char *fusion_string {}; if (!ctx->device->disable_fusion) { uint32_t num_adds = ggml_vk_fuse_multi_add(ctx, cgraph, i); @@ -17204,6 +18286,25 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->num_additional_fused_ops = num_adds - 1; fusion_string = "MULTI_ADD"; std::fill_n(op_srcs_fused_elementwise, ctx->num_additional_fused_ops + 1, true); + } else if (ggml_vk_can_fuse_mul_add(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = 1; + ctx->fused_mul_add = true; + fusion_string = "MUL_ADD"; + std::fill_n(op_srcs_fused_elementwise, 2, true); + } else if (ggml_vk_can_fuse_swiglu_quant(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = 1; + ctx->fused_swiglu_quant = true; + fusion_string = "SWIGLU_QUANT_MUL_MAT"; + std::fill_n(op_srcs_fused_elementwise, 2, false); + } else if (ggml_vk_can_fuse_sigmoid_quant(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = 2; + ctx->fused_sigmoid_quant = true; + fusion_string = "SIGMOID_MUL_QUANT_MUL_MAT"; + // Both elementwise inputs are fully consumed into Q8_1 scratch before + // the matmul writes its result, so exact input/output aliasing is safe. + op_srcs_fused_elementwise[0] = true; + op_srcs_fused_elementwise[1] = true; + op_srcs_fused_elementwise[2] = false; } else if (ggml_vk_can_fuse(ctx, cgraph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD })) { ctx->num_additional_fused_ops = 2; fusion_string = "MUL_MAT_ADD_ADD"; @@ -17231,6 +18332,22 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg fusion_string = "MUL_MAT_ID_MUL"; op_srcs_fused_elementwise[0] = false; op_srcs_fused_elementwise[1] = true; + } else if (ggml_vk_can_fuse_rms_norm_channel_last(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = 5; + ctx->fused_rms_norm_channel_last = true; + fusion_string = "RMS_NORM_CHANNEL_LAST"; + op_srcs_fused_elementwise[0] = true; + std::fill_n(op_srcs_fused_elementwise + 1, 5, false); + } else if (uint32_t fused_offset = ggml_vk_can_fuse_rms_norm_mrope_pack(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = fused_offset; + ctx->fused_rms_norm_mrope_pack = true; + fusion_string = "RMS_NORM_MUL_MROPE_PACK"; + std::fill_n(op_srcs_fused_elementwise, fused_offset + 1, false); + } else if (uint32_t fused_ops = ggml_vk_can_fuse_rms_norm_modulate(ctx, cgraph, i)) { + ctx->num_additional_fused_ops = fused_ops - 1; + ctx->fused_rms_norm_modulate = true; + fusion_string = "RMS_NORM_MUL_MODULATE"; + std::fill_n(op_srcs_fused_elementwise, fused_ops, true); } else if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, { i + 4 }) && ggml_check_edges(cgraph, i, rms_norm_mul_rope_view_set_rows_edges) && ggml_vk_can_fuse_rms_norm_mul_rope(ctx, cgraph, i) && @@ -17402,6 +18519,13 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->fused_ops_write_mask = 1; ctx->fused_topk_moe_mode = TOPK_MOE_COUNT; ctx->fused_topk_moe_scale = false; + ctx->fused_rms_norm_channel_last = false; + ctx->fused_rms_norm_modulate = false; + ctx->fused_rms_norm_mrope_pack = false; + ctx->fused_mul_add = false; + ctx->fused_swiglu_quant = false; + ctx->fused_sigmoid_quant = false; + fusion_string = nullptr; } } @@ -17445,6 +18569,9 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg i += ctx->num_additional_fused_ops; ctx->num_additional_fused_ops = 0; ctx->fused_ops_write_mask = 0; + ctx->fused_rms_norm_channel_last = false; + ctx->fused_rms_norm_modulate = false; + ctx->fused_rms_norm_mrope_pack = false; } ctx->last_total_flops = total_flops; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index 057ed739aa8..01ec4ab0719 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -13,6 +13,9 @@ #extension GL_KHR_shader_subgroup_basic : enable #extension GL_KHR_shader_subgroup_arithmetic : enable #extension GL_KHR_shader_subgroup_vote : enable +#extension GL_KHR_shader_subgroup_clustered : require +#extension GL_KHR_shader_subgroup_ballot : require +#extension GL_KHR_shader_subgroup_shuffle : enable #extension GL_KHR_memory_scope_semantics : enable #extension GL_KHR_cooperative_matrix : enable @@ -26,10 +29,64 @@ const uint32_t MatBr = 16; const uint32_t MatBc = 16; +// Driver-selected PV layout. Ids 0..15 belong to flash_attn_base.glsl; this is +// the next free one and is set from the host (vk_fa_pipeline_state::transpose_pv). +// +// 0 (DEFAULT, every driver the host does not explicitly opt in): P is staged as +// [KV column][query row] and V as [KV row][HSV column], and the PV A-fragment +// is read ColumnMajor. +// 1: P and V are both staged transposed so that BOTH PV coopMatLoads run along +// the fragment's native wide-load axis. +// +// WHY THIS IS A SPECIALIZATION CONSTANT AND NOT ONE LAYOUT FOR EVERYONE: on +// gfx1151 LLPC (AMDVLK) turns the transposed layout into ds_load_b128 and gains +// ~4% on top of arm 0; ACO (RADV) does not form the wide loads, keeps the same +// per-fragment scalar gather, and measured ~3.5% SLOWER than arm 0. The two arms +// also carry different loop-carried state (thread-space Of vs. register-resident +// O coopmats), which is exactly why this must fold at pipeline-compile time - a +// runtime branch would keep both live and cost ~10 VGPR, crossing the 192-VGPR +// 4-wave cliff. Do not "simplify" this into a single path. +layout (constant_id = 16) const uint32_t TRANSPOSE_PV = 0; +// RADV wave32 schedule: four subgroups each own a 16-row query panel and reuse +// the same staged K/V tile. +layout (constant_id = 17) const uint32_t WORK_PER_TILE = 0; +const bool WPT_ENABLED = WORK_PER_TILE != 0; + +// Q is invariant across the whole KV loop, so its coopmat fragments can live in +// registers instead of being re-read from Qf on every (j, d). Each fragment is +// 16x16 float16 = 2 VGPRs on wave64, so the cost is 2 * HSK_pad/16 VGPRs: 16 at +// HSK=128, but 48 at HSK=384, which would cross the occupancy cliff. Hoist only +// up to this many fragments and leave the general path exactly as it was. +const uint32_t Q_HOIST_MAX_FRAGS = 8; + const uint32_t rows_per_thread = Br / row_split; const uint32_t cols_per_iter = gl_WorkGroupSize.x / row_split; const uint32_t cols_per_thread = Bc / cols_per_iter; +// Fast row max: partition the subgroup as (query row) x (column group) instead of +// one lane per KV column, so the cross-lane reduction is one clustered max over +// LANES_PER_ROW lanes instead of rows_per_thread full-subgroup reductions. +// Written as single flat divisions so no denominator can be zero for any tile +// the host can dispatch. +const uint32_t LANES_PER_ROW = SubGroupSize * row_split / Br; +const uint32_t COLS_PER_LANE = Bc * Br / (SubGroupSize * row_split); +// Bitwise "badness" word rather than a chain of &&: glslang emits chained +// boolean OpSpecConstantOp out of order and the module then fails SPIR-V +// validation with "ID has not been defined". Zero iff the fast path applies. +// The cols_per_iter term is the row_tid == gl_SubgroupID invariant that the +// existing subgroupMax/subgroupAdd reductions already depend on unstated; a +// configuration that violates it now takes the general path instead of +// computing garbage. +const uint32_t FAST_ROWMAX_BAD = + (cols_per_iter ^ SubGroupSize) + | (LANES_PER_ROW * rows_per_thread ^ SubGroupSize) + | (COLS_PER_LANE * LANES_PER_ROW ^ Bc) + | (COLS_PER_LANE % 4) + | (LANES_PER_ROW & (LANES_PER_ROW - 1)) // power of two + | (2 / (LANES_PER_ROW | 1)) // >= 2 + | (LANES_PER_ROW / 65); // <= 64 +#define FAST_ROWMAX (FAST_ROWMAX_BAD == 0) + layout (binding = 0) readonly buffer Q {float data_q[];}; layout (binding = 0) readonly buffer QV4 {vec4 data_qv4[];}; @@ -37,31 +94,322 @@ layout (binding = 1) readonly buffer K {FLOAT_TYPE data_k[];}; layout (binding = 1) readonly buffer KV4 {FLOAT_TYPEV4 data_kv4[];}; layout (binding = 2) readonly buffer V {FLOAT_TYPE data_v[];}; layout (binding = 2) readonly buffer VV4 {FLOAT_TYPEV4 data_vv4[];}; +layout (binding = 2) readonly volatile buffer VV4_VOLATILE {FLOAT_TYPEV4 data_vv4_volatile[];}; layout (binding = 3) readonly buffer M {float16_t data_m[];}; shared float tmpsh[row_split]; const uint32_t qstride = HSK_pad / 4 + 2; -shared FLOAT_TYPEV4 Qf[Br * qstride]; +const uint32_t QF_LEN = Br * qstride; +shared FLOAT_TYPEV4 Qf[WPT_ENABLED ? 1 : QF_LEN]; +// The two arms' P staging buffers. Each is sized by a specialization-constant +// expression so the unused one collapses to a single element and the arm that +// runs never pays for the other's LDS - the sum of both at full size would not +// fit under maxComputeSharedMemorySize. const uint psh_stride = Br / 4 + 2; -shared FLOAT_TYPEV4 Psh[Bc * psh_stride]; +shared FLOAT_TYPEV4 Psh[(TRANSPOSE_PV != 0 || WPT_ENABLED) ? 1 : (Bc * psh_stride)]; + +// P staged TRANSPOSED: [Br query rows][Bc KV cols], KV contiguous, so the PV +// A-fragment loads RowMajor (2 wide ds_load) instead of ColumnMajor (16 +// ds_load_u16). Row stride is a multiple of 8 f16 = 16 B so ds_load_b128 stays +// legal. The PV A-fragment is MatBr x 16, so PshT must have at least MatBr rows; +// the host pins Br == MatBr == 16 for the coopmat1 path (ggml-vulkan.cpp:3824) +// and the pre-existing pvsh indexing already assumed Br <= MatBc. +const uint pt_stride = Bc + 8; +shared FLOAT_TYPE PshT[(TRANSPOSE_PV != 0 || WPT_ENABLED) ? ((Br > MatBr ? Br : MatBr) * pt_stride) : 1]; + +// sfsh holds S^T as [query row][KV column]; the QK accumulator is stored +// ColumnMajor to put it in that order. The thread-space softmax reads a row of +// S with one lane per KV column, so consecutive lanes now land on consecutive +// dwords. The upstream [KV column][query row] order gave those same reads a +// 4*sfshstride dword lane stride, which lands on 4 banks for any padding. +// +// HARD RULE: sfshstride must be ODD in ACC_TYPEV4 units. The accumulator store +// addresses dword 4*sfshstride*(lane&15) + (lane>>4). For odd s, 4s mod 32 is +// one of {4,12,20,28}, all with gcd 4, so the (lane&15) term covers 8 banks and +// the (lane>>4) term splits them into 16 -> 2-way, matching upstream. s == 2 mod +// 4 gives 8 -> 4-way and s == 0 mod 4 gives 2 -> 16-way. Rounding the pad to a +// "tidy" even number silently gives back part of what the transpose buys. +// Bc/4 + 7 is also chosen so the 512 B LDS granule total is unchanged from +// upstream; Bc/4 + 1 is the minimum odd pad for the HSK > 128 path that has to +// fit in 48KB shmem. +const uint32_t sfshstride = (HSK <= 128) ? (Bc / 4 + 7) : (Bc / 4 + 1); +shared ACC_TYPEV4 sfsh[WPT_ENABLED ? 1 : (Br * sfshstride)]; + +// subgroupClusteredMax needs a literal cluster size; the caller's value is a +// specialization-constant expression, so select the literal here and let the +// driver fold the dead arms away. +float fa_clustered_max(float v, const uint32_t n) { + if (n == 64) { return subgroupClusteredMax(v, 64); } + if (n == 32) { return subgroupClusteredMax(v, 32); } + if (n == 16) { return subgroupClusteredMax(v, 16); } + if (n == 8) { return subgroupClusteredMax(v, 8); } + if (n == 4) { return subgroupClusteredMax(v, 4); } + if (n == 2) { return subgroupClusteredMax(v, 2); } + return v; +} -// Avoid padding for hsk==256 to make it fit in 48KB shmem. -const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; -shared ACC_TYPEV4 sfsh[Bc * sfshstride]; +// Same reason, for the lane id: OpGroupNonUniformBroadcast only accepts a +// dynamically-uniform Id from SPIR-V 1.5 on, and glslang rejects a non-literal +// outright below --target-env=vulkan1.2. Selecting the literal here keeps the +// module valid at vulkan1.1 and costs nothing - the selector is a +// specialization-constant expression and folds to a single v_readlane. +float fa_row_broadcast(float v, const uint32_t lane) { + switch (lane) { + case 0u: return subgroupBroadcast(v, 0u); case 1u: return subgroupBroadcast(v, 1u); case 2u: return subgroupBroadcast(v, 2u); case 3u: return subgroupBroadcast(v, 3u); + case 4u: return subgroupBroadcast(v, 4u); case 5u: return subgroupBroadcast(v, 5u); case 6u: return subgroupBroadcast(v, 6u); case 7u: return subgroupBroadcast(v, 7u); + case 8u: return subgroupBroadcast(v, 8u); case 9u: return subgroupBroadcast(v, 9u); case 10u: return subgroupBroadcast(v, 10u); case 11u: return subgroupBroadcast(v, 11u); + case 12u: return subgroupBroadcast(v, 12u); case 13u: return subgroupBroadcast(v, 13u); case 14u: return subgroupBroadcast(v, 14u); case 15u: return subgroupBroadcast(v, 15u); + case 16u: return subgroupBroadcast(v, 16u); case 17u: return subgroupBroadcast(v, 17u); case 18u: return subgroupBroadcast(v, 18u); case 19u: return subgroupBroadcast(v, 19u); + case 20u: return subgroupBroadcast(v, 20u); case 21u: return subgroupBroadcast(v, 21u); case 22u: return subgroupBroadcast(v, 22u); case 23u: return subgroupBroadcast(v, 23u); + case 24u: return subgroupBroadcast(v, 24u); case 25u: return subgroupBroadcast(v, 25u); case 26u: return subgroupBroadcast(v, 26u); case 27u: return subgroupBroadcast(v, 27u); + case 28u: return subgroupBroadcast(v, 28u); case 29u: return subgroupBroadcast(v, 29u); case 30u: return subgroupBroadcast(v, 30u); case 31u: return subgroupBroadcast(v, 31u); + case 32u: return subgroupBroadcast(v, 32u); case 33u: return subgroupBroadcast(v, 33u); case 34u: return subgroupBroadcast(v, 34u); case 35u: return subgroupBroadcast(v, 35u); + case 36u: return subgroupBroadcast(v, 36u); case 37u: return subgroupBroadcast(v, 37u); case 38u: return subgroupBroadcast(v, 38u); case 39u: return subgroupBroadcast(v, 39u); + case 40u: return subgroupBroadcast(v, 40u); case 41u: return subgroupBroadcast(v, 41u); case 42u: return subgroupBroadcast(v, 42u); case 43u: return subgroupBroadcast(v, 43u); + case 44u: return subgroupBroadcast(v, 44u); case 45u: return subgroupBroadcast(v, 45u); case 46u: return subgroupBroadcast(v, 46u); case 47u: return subgroupBroadcast(v, 47u); + case 48u: return subgroupBroadcast(v, 48u); case 49u: return subgroupBroadcast(v, 49u); case 50u: return subgroupBroadcast(v, 50u); case 51u: return subgroupBroadcast(v, 51u); + case 52u: return subgroupBroadcast(v, 52u); case 53u: return subgroupBroadcast(v, 53u); case 54u: return subgroupBroadcast(v, 54u); case 55u: return subgroupBroadcast(v, 55u); + case 56u: return subgroupBroadcast(v, 56u); case 57u: return subgroupBroadcast(v, 57u); case 58u: return subgroupBroadcast(v, 58u); case 59u: return subgroupBroadcast(v, 59u); + case 60u: return subgroupBroadcast(v, 60u); case 61u: return subgroupBroadcast(v, 61u); case 62u: return subgroupBroadcast(v, 62u); case 63u: return subgroupBroadcast(v, 63u); + } + return v; +} const uint32_t D_pad = HSK_pad > HSV_pad ? HSK_pad : HSV_pad; const uint32_t kvsh_stride = (SHMEM_STAGING != 0 ? D_pad : MatBr) / 4 + 2; -const uint v_cols = MatBc / 4 * row_split; // total cols, 4 vec4s per MatBc * number of subgroups +const uint v_cols = WPT_ENABLED ? (HSV_pad / 4) : (MatBc / 4 * row_split); const uint vsh_stride = v_cols; -shared FLOAT_TYPEV4 kvsh[(kvsh_stride >= vsh_stride) ? (Bc * kvsh_stride) : (Bc * vsh_stride)]; +const uint32_t kvsh_len = (kvsh_stride >= vsh_stride) ? (Bc * kvsh_stride) : (Bc * vsh_stride); +// The WPT path aliases the Q staging region with K/V after Q has been hoisted +// into cooperative-matrix fragments. The full K tile then fits without +// increasing LDS, allowing all head-dimension slices to share one barrier. +const uint32_t wpt_k_stride = HSK_pad / 4 + 2; +const uint32_t wpt_k_len = Bc * wpt_k_stride; +const uint32_t wpt_kv_len = (wpt_k_len >= kvsh_len) ? wpt_k_len : kvsh_len; +const uint32_t wpt_union_len = (QF_LEN >= wpt_kv_len) ? QF_LEN : wpt_kv_len; +shared FLOAT_TYPEV4 kvsh[(TRANSPOSE_PV != 0) ? 1 : (WPT_ENABLED ? wpt_union_len : kvsh_len)]; + +// Arm 1's K/V staging buffer is a SCALAR f16 array: the V role needs a transposed +// layout whose row stride (Bc + 8 f16) is not a multiple of 4 components, and +// every f16vec4 store into it was already being scalarised into 4 ds_store_b16. +// kvshT_stride is kvsh_stride expressed in f16 rather than f16vec4. +const uint32_t kvshT_stride = kvsh_stride * 4; +// V^T tile: vt_rows HSV columns x Bc KV rows, KV contiguous -> B loads ColumnMajor. +const uint vt_rows = MatBc * row_split; +const uint vt_cols4 = vt_rows / 4; +const uint vt_stride = Bc + 8; +shared FLOAT_TYPE kvshT[(TRANSPOSE_PV != 0) ? ((Bc * kvshT_stride >= vt_rows * vt_stride) ? (Bc * kvshT_stride) : (vt_rows * vt_stride)) : 1]; + +// Only the LDS-staged V path is permuted; the direct-from-global path keeps +// hsv column == fragment column. A macro rather than a global const bool: glslc +// lowers the short-circuit || into branches, and a global-scope definition of +// those does not dominate the uses inside the KV loop. +#define V_PERMUTED ((SHMEM_STAGING == 0) && (USE_DECODE_V || KV_bounds_check)) const uint32_t osh_stride = row_split * MatBr / 4; -shared O_TYPEV4 pvsh[MatBc * osh_stride]; +shared O_TYPEV4 pvsh[(TRANSPOSE_PV != 0 || WPT_ENABLED) ? 1 : (MatBc * osh_stride)]; + +// Arm 1's pvshT serves two roles at different times: the MatBr x MatBc broadcast +// of the online-softmax rescale factor during the KV loop, and the accumulator +// drain after it. MatBc * oshT_stride >= MatBr * MatBc for every row_split >= 1. +const uint32_t oshT_stride = row_split * MatBr; +shared O_TYPE pvshT[(TRANSPOSE_PV != 0) ? (MatBc * oshT_stride) : 1]; + +const uint32_t num_hsv_tiles = (HSV + MatBc * row_split - 1) / (MatBc * row_split); // round up shared ACC_TYPE slope[Br]; +void main_work_per_tile() { + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint qbase = gl_SubgroupID * MatBr; + const uint KCHUNKS = Bc / MatBc; + const uint OTILES = HSV_pad / MatBc; + const uint VLOADS = (Bc * HSV_pad / 4 + WorkGroupSize - 1) / WorkGroupSize; + + uint q_offset = gqa_iq1*p.nb01 + (iq2*p.nb02 + iq3*p.nb03) / 4; + [[unroll]] for (uint idx = 0; idx < Br * HSK / 4; idx += gl_WorkGroupSize.x) { + const uint d = (idx + tid) % (HSK / 4); + const uint r = (idx + tid) / (HSK / 4); + if (r < Br) { + FLOAT_TYPEV4 qv = FLOAT_TYPEV4(0.0); + if (i * Br + r < N) { + qv = FLOAT_TYPEV4(data_qv4[q_offset / 4 + (i * Br + r) * q_stride / 4 + d] * p.scale); + } + kvsh[r * qstride + d] = qv; + } + } + barrier(); + + coopmat QMats[HSK_pad / 16]; + [[unroll]] for (uint d = 0; d < HSK_pad / 16; ++d) { + coopMatLoad(QMats[d], kvsh, qbase * qstride + d * 4, qstride, + gl_CooperativeMatrixLayoutColumnMajor); + } + // K overwrites the aliased Q region immediately below. Ensure every + // subgroup has completed its final cooperative Q load first. + barrier(); + + coopmat OMats[OTILES]; + [[unroll]] for (uint t = 0; t < OTILES; ++t) { + OMats[t] = coopmat(0); + } + + const float NEG_FLT_MAX_OVER_2 = uintBitsToFloat(0xFEFFFFFF); + float Lf = 0.0; + float Mf = NEG_FLT_MAX_OVER_2; + + const uint k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; + const uint v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; + const uint krow = tid / (MatBr / 4); + const uint kcv = tid % (MatBr / 4); + + FLOAT_TYPEV4 KPrefetch[HSK_pad / 16]; + if (start_j < end_j) { + [[unroll]] for (uint d = 0; d < HSK_pad / 16; ++d) { + FLOAT_TYPEV4 kv = FLOAT_TYPEV4(0.0); + if ((!KV_bounds_check || start_j * Bc + krow < KV) && d * 16 + kcv * 4 < HSK) { + kv = data_kv4[k_offset / 4 + (start_j * Bc + krow) * k_stride / 4 + d * 4 + kcv]; + } + KPrefetch[d] = kv; + } + } + + [[dont_unroll]] + for (uint j = start_j; j < end_j; ++j) { + coopmat SMats[KCHUNKS]; + [[unroll]] for (uint bc = 0; bc < KCHUNKS; ++bc) { + SMats[bc] = coopmat(0); + } + + coopmat KMat; + [[unroll]] for (uint d = 0; d < HSK_pad / 16; ++d) { + kvsh[krow * wpt_k_stride + d * 4 + kcv] = KPrefetch[d]; + } + barrier(); + + [[unroll]] for (uint d = 0; d < HSK_pad / 16; ++d) { + [[unroll]] for (uint bc = 0; bc < KCHUNKS; ++bc) { + coopMatLoad(KMat, kvsh, bc * MatBc * wpt_k_stride + d * 4, + wpt_k_stride, gl_CooperativeMatrixLayoutRowMajor); + SMats[bc] = coopMatMulAdd(KMat, QMats[d], SMats[bc]); + } + } + + FLOAT_TYPEV4 VPrefetch[VLOADS]; + [[unroll]] for (uint vi = 0; vi < VLOADS; ++vi) { + const uint idx = tid + vi * WorkGroupSize; + FLOAT_TYPEV4 vv = FLOAT_TYPEV4(0.0); + if (idx < Bc * HSV_pad / 4) { + const uint row = idx / (HSV_pad / 4); + const uint d = idx % (HSV_pad / 4); + if ((!KV_bounds_check || j * Bc + row < KV) && d * 4 < HSV) { + const uint v_index = v_offset / 4 + (j * Bc + row) * v_stride / 4 + d; + vv = data_vv4_volatile[v_index]; + } + } + VPrefetch[vi] = vv; + } + + float m = NEG_FLT_MAX_OVER_2; + [[unroll]] for (uint bc = 0; bc < KCHUNKS; ++bc) { + [[unroll]] for (uint e = 0; e < 8; ++e) { + const uint col = bc * MatBc + 2 * e + (lane >> 4); + const float x = (KV_bounds_check && j * Bc + col >= KV) + ? NEG_FLT_MAX_OVER_2 : float(SMats[bc][e]); + m = max(m, x); + } + } + m = max(m, subgroupShuffleXor(m, 16)) + FATTN_KQ_MAX_OFFSET; + + const float old_m = Mf; + Mf = max(m, old_m); + const float eMf = exp(old_m - Mf); + const bool rescale_output = !subgroupAll(eMf == 1.0); + float psum = 0.0; + + if (rescale_output) { + [[unroll]] for (uint t = 0; t < OTILES; ++t) { + [[unroll]] for (uint e = 0; e < 8; ++e) { + const uint row = 2 * e + (lane >> 4); + OMats[t][e] *= O_TYPE(subgroupShuffle(eMf, row)); + } + } + } + + const uint qrow = lane & 15u; + [[unroll]] for (uint bc = 0; bc < KCHUNKS; ++bc) { + [[unroll]] for (uint e = 0; e < 8; ++e) { + const uint col = bc * MatBc + 2 * e + (lane >> 4); + FLOAT_TYPE pf = FLOAT_TYPE(0.0); + if (!KV_bounds_check || j * Bc + col < KV) { + pf = FLOAT_TYPE(exp(float(SMats[bc][e]) - Mf)); + } + psum += float(pf); + PshT[(qbase + qrow) * pt_stride + col] = pf; + } + } + Lf = Lf * eMf + psum + subgroupShuffleXor(psum, 16); + + barrier(); + + [[unroll]] for (uint vi = 0; vi < VLOADS; ++vi) { + const uint idx = tid + vi * WorkGroupSize; + if (idx < Bc * HSV_pad / 4) { + const uint row = idx / (HSV_pad / 4); + const uint d = idx % (HSV_pad / 4); + kvsh[row * vsh_stride + d] = VPrefetch[vi]; + } + } + + barrier(); + + if (j + 1 < end_j) { + [[unroll]] for (uint d = 0; d < HSK_pad / 16; ++d) { + FLOAT_TYPEV4 kv = FLOAT_TYPEV4(0.0); + if ((!KV_bounds_check || (j + 1) * Bc + krow < KV) && d * 16 + kcv * 4 < HSK) { + kv = data_kv4[k_offset / 4 + ((j + 1) * Bc + krow) * k_stride / 4 + d * 4 + kcv]; + } + KPrefetch[d] = kv; + } + } + + coopmat PMat; + coopmat VMat; + [[unroll]] for (uint bc = 0; bc < KCHUNKS; ++bc) { + coopMatLoad(PMat, PshT, qbase * pt_stride + bc * MatBc, pt_stride, + gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint t = 0; t < OTILES; ++t) { + coopMatLoad(VMat, kvsh, bc * MatBc * vsh_stride + t * (MatBc / 4), + vsh_stride, gl_CooperativeMatrixLayoutRowMajor); + OMats[t] = coopMatMulAdd(PMat, VMat, OMats[t]); + } + } + barrier(); + } + + const float invL = Lf == 0.0 ? 0.0 : 1.0 / Lf; + const uint out_base = (iq3 * p.ne2 * p.ne1 * HSV) / 4; + [[unroll]] for (uint t = 0; t < OTILES; ++t) { + [[unroll]] for (uint e = 0; e < 8; ++e) { + const uint row = 2 * e + (lane >> 4); + const float oe = float(OMats[t][e]) * subgroupShuffle(invL, row); + const uint src = lane & ~3u; + const vec4 ov = vec4(subgroupShuffle(oe, src), + subgroupShuffle(oe, src + 1), + subgroupShuffle(oe, src + 2), + subgroupShuffle(oe, src + 3)); + const uint global_row = i * Br + qbase + row; + if ((lane & 3u) == 0 && global_row < N) { + const uint d = t * (MatBc / 4) + (lane & 15u) / 4; + data_ov4[out_base + (iq2 * HSV + global_row * p.ne1 * HSV) / 4 + d] = D_TYPEV4(ov); + } + } + } +} + void main() { #ifdef NEEDS_INIT_IQ_SHMEM if (fa_type_needs_shmem(FaTypeK) || fa_type_needs_shmem(FaTypeV)) { @@ -71,6 +419,11 @@ void main() { init_indices(); + if (WPT_ENABLED) { + main_work_per_tile(); + return; + } + const uint32_t tid = gl_LocalInvocationIndex; const uint32_t threads_per_rowgroup = gl_WorkGroupSize.x / row_split; @@ -102,6 +455,9 @@ void main() { } barrier(); + // Arm 0 accumulates O here across the whole KV loop; arm 1 accumulates in + // OMat and only drains into Of once, after the loop. Either way the tail + // (sinks, 1/L, split_k, the stores) reads Of and is shared. O_TYPEV4 Of[rows_per_thread][d_per_thread]; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { [[unroll]] for (uint32_t d = 0; d < d_per_thread; ++d) { @@ -109,6 +465,23 @@ void main() { } } + coopmat OMat[num_hsv_tiles]; + if (TRANSPOSE_PV != 0) { + // Rows Br..MatBr-1 of the rescale tile are never republished inside the + // loop (no thread owns them when Br < MatBr), so give them a defined + // value once. + [[unroll]] for (uint32_t idx = 0; idx < MatBr * MatBc; idx += gl_WorkGroupSize.x) { + if (idx + tid < MatBr * MatBc) { + pvshT[idx + tid] = O_TYPE(1.0); + } + } + barrier(); + + [[unroll]] for (uint32_t t = 0; t < num_hsv_tiles; ++t) { + OMat[t] = coopmat(0); + } + } + float Lf[rows_per_thread], Mf[rows_per_thread]; // Use -FLT_MAX/2 rather than -inf to reduce the possibility of NaNs, e.g. when computing Mold-M. @@ -150,6 +523,19 @@ void main() { uint32_t mask_opt_bits = 0; f16vec4 mask_cache[Bc * Br / 4 / WorkGroupSize]; + // Qf is filled once above and never written again, so these fragments are + // constant for the whole KV loop. Q_HOIST_N is 1 when hoisting is off, which + // keeps QMats[] in range on every path. + const uint32_t Q_FRAGS = HSK_pad / 16; + const bool HOIST_Q = Q_FRAGS <= Q_HOIST_MAX_FRAGS; + const uint32_t Q_HOIST_N = HOIST_Q ? Q_FRAGS : 1; + coopmat QMats[Q_HOIST_N]; + if (HOIST_Q) { + [[unroll]] for (uint32_t d = 0; d < Q_HOIST_N; ++d) { + coopMatLoad(QMats[d], Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); + } + } + [[dont_unroll]] for (uint32_t j = start_j; j < end_j; ++j) { @@ -173,40 +559,24 @@ void main() { float max_mask = NEG_FLT_MAX_OVER_2; [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { - uint32_t c = (idx + tid) / (Br / 4); - uint32_t r = (idx + tid) % (Br / 4); + uint32_t rq = (idx + tid) / (Bc / 4); + uint32_t cv = (idx + tid) % (Bc / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { - if ((!KV_bounds_check || j * Bc + c < KV)) { - f16vec4 m; - if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { - m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 3) * m_stride + (j * Bc + c)]); - max_mask = max(max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])), float(m[3])); - } else if (i * Br + r * 4 + 2 < p.nem1) { - m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], - 0.0); - max_mask = max(max(max(max_mask, float(m[0])), float(m[1])), float(m[2])); - } else if (i * Br + r * 4 + 1 < p.nem1) { - m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], - data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], - 0.0, - 0.0); - max_mask = max(max(max_mask, float(m[0])), float(m[1])); - } else if (i * Br + r * 4 < p.nem1) { - m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], - 0.0, - 0.0, - 0.0); - max_mask = max(max_mask, float(m[0])); - } else { - m = f16vec4(0.0); + // One query row, four consecutive KV columns: the mask row is + // contiguous in KV, so this is a wider and better coalesced read + // than the upstream one-column-four-rows gather, and the nem1 + // ladder collapses to a single row test. + f16vec4 m = f16vec4(0.0); + if (!nem1_bounds_check || i * Br + rq < p.nem1) { + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + const uint32_t col = cv * 4 + e; + if (!KV_bounds_check || j * Bc + col < KV) { + m[e] = data_m[m_offset + (i * Br + rq) * m_stride + (j * Bc + col)]; + max_mask = max(max_mask, float(m[e])); + } } - mask_cache[idx / WorkGroupSize] = m; } + mask_cache[idx / WorkGroupSize] = m; } } // skip the block if the mask is entirely -inf @@ -245,7 +615,13 @@ void main() { } } - kvsh[c * kvsh_stride + d] = K_Tf; + if (TRANSPOSE_PV != 0) { + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + kvshT[c * kvshT_stride + d * 4 + e] = K_Tf[e]; + } + } else { + kvsh[c * kvsh_stride + d] = K_Tf; + } } } barrier(); @@ -257,6 +633,10 @@ void main() { coopmat SfMat = coopmat(0); coopmat KMat; coopmat QMat; + // Arm 1's PV operands. A = P (query rows x KV), B = V^T (KV x HSV). + // Arm 0 reuses KMat/QMat for PV instead; these fold away there. + coopmat PMat; + coopmat VMat; [[unroll]] for (uint32_t d = 0; d < HSK_pad / 16; ++d) { // If SHMEM_STAGING is set, a Bc * HSK_pad size tile of K is loaded to shmem @@ -288,39 +668,60 @@ void main() { } } - kvsh[row * kvsh_stride + col_vec] = K_Tf; + if (TRANSPOSE_PV != 0) { + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + kvshT[row * kvshT_stride + col_vec * 4 + e] = K_Tf[e]; + } + } else { + kvsh[row * kvsh_stride + col_vec] = K_Tf; + } } } barrier(); } if (stage_k) { - uint coord = (gl_SubgroupID * MatBc) * kvsh_stride; - coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); + if (TRANSPOSE_PV != 0) { + uint coord = (gl_SubgroupID * MatBc) * kvshT_stride; + coopMatLoad(KMat, kvshT, coord, kvshT_stride, gl_CooperativeMatrixLayoutRowMajor); + } else { + uint coord = (gl_SubgroupID * MatBc) * kvsh_stride; + coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); + } } else { const uint coord = k_offset / 4 + (j * Bc + gl_SubgroupID * MatBc) * k_stride / 4 + d * 16 / 4; coopMatLoad(KMat, data_kv4, coord, k_stride / 4, gl_CooperativeMatrixLayoutRowMajor); } } else { - uint coord = (gl_SubgroupID * MatBc) * kvsh_stride + d * 16 / 4; - coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); + if (TRANSPOSE_PV != 0) { + uint coord = (gl_SubgroupID * MatBc) * kvshT_stride + d * 16; + coopMatLoad(KMat, kvshT, coord, kvshT_stride, gl_CooperativeMatrixLayoutRowMajor); + } else { + uint coord = (gl_SubgroupID * MatBc) * kvsh_stride + d * 16 / 4; + coopMatLoad(KMat, kvsh, coord, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); + } } - coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); - - SfMat = coopMatMulAdd(KMat, QMat, SfMat); + // Feed the hoisted fragment straight into the multiply: assigning it to + // QMat first makes the operand live twice over and costs registers. + if (HOIST_Q) { + SfMat = coopMatMulAdd(KMat, QMats[d < Q_HOIST_N ? d : 0], SfMat); + } else { + coopMatLoad(QMat, Qf, d * 16 / 4, qstride, gl_CooperativeMatrixLayoutColumnMajor); + SfMat = coopMatMulAdd(KMat, QMat, SfMat); + } } - uint coord = gl_SubgroupID * MatBc * sfshstride; - coopMatStore(SfMat, sfsh, coord, sfshstride, gl_CooperativeMatrixLayoutRowMajor); + uint coord = gl_SubgroupID * MatBc / 4; + coopMatStore(SfMat, sfsh, coord, sfshstride, gl_CooperativeMatrixLayoutColumnMajor); barrier(); if (LOGIT_SOFTCAP) { [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { - uint32_t c = (idx + tid) / (Br / 4); - uint32_t r = (idx + tid) % (Br / 4); + uint32_t rq = (idx + tid) / (Bc / 4); + uint32_t cv = (idx + tid) % (Bc / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { - sfsh[c * sfshstride + r] = ACC_TYPEV4(p.logit_softcap * tanh(sfsh[c * sfshstride + r])); + sfsh[rq * sfshstride + cv] = ACC_TYPEV4(p.logit_softcap * tanh(sfsh[rq * sfshstride + cv])); } } barrier(); @@ -328,69 +729,137 @@ void main() { if (MASK_ENABLE && mask_opt_bits != MASK_OPT_ALL_ZERO) { [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { - uint32_t c = (idx + tid) / (Br / 4); - uint32_t r = (idx + tid) % (Br / 4); + uint32_t rq = (idx + tid) / (Bc / 4); + uint32_t cv = (idx + tid) % (Bc / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { - if (!KV_bounds_check || j * Bc + c < KV) { - // Mask nem1 bounds check is handled when loading masks - ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); - ACC_TYPEV4 slopes = ACC_TYPEV4(slope[r * 4], slope[r * 4 + 1], slope[r * 4 + 2], slope[r * 4 + 3]); - sfsh[c * sfshstride + r] += slopes * masks; - } + // Mask nem1 / KV bounds are handled when loading masks (out of + // range entries are 0, and adding 0 to an out-of-range column is + // harmless because the softmax masks those columns anyway). + ACC_TYPEV4 masks = ACC_TYPEV4(mask_cache[idx / WorkGroupSize]); + sfsh[rq * sfshstride + cv] += ACC_TYPE(slope[rq]) * masks; } } barrier(); } - float eMf[rows_per_thread]; - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - const uint r_vec = tile_row(r) / 4; - const uint r_comp = tile_row(r) % 4; - - float rowmaxf = NEG_FLT_MAX_OVER_2; - [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { - continue; + float rowmax[rows_per_thread]; + if (FAST_ROWMAX) { + // One lane owns COLS_PER_LANE consecutive KV columns of ONE query row, + // so each lane reduces its own columns with plain v_max and the only + // cross-lane step is a single clustered max over the LANES_PER_ROW lanes + // that share a row. Replaces rows_per_thread serialized full-subgroup + // reductions (each a 4-DPP + 2-permlane + readfirstlane chain, and each + // preceded by its own exec-masked LDS load) with one. + const uint32_t lane = gl_SubgroupInvocationID; + const uint32_t rr = lane / LANES_PER_ROW; + const uint32_t cg = lane % LANES_PER_ROW; + const uint32_t qrow = row_tid * rows_per_thread + rr; + + float m = NEG_FLT_MAX_OVER_2; + [[unroll]] for (uint32_t k = 0; k < COLS_PER_LANE / 4; ++k) { + const uint32_t cv = cg * (COLS_PER_LANE / 4) + k; + const ACC_TYPEV4 sv = sfsh[qrow * sfshstride + cv]; + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + // Branch-free bounds mask: out-of-KV columns hold exactly 0 + // (K is staged as zero there), which would beat a negative max. + const float x = (KV_bounds_check && j * Bc + cv * 4 + e >= KV) + ? NEG_FLT_MAX_OVER_2 : float(sv[e]); + m = max(m, x); } - rowmaxf = max(rowmaxf, float(sfsh[r_vec + (c * cols_per_iter + col_tid) * sfshstride][r_comp])); } - rowmaxf += FATTN_KQ_MAX_OFFSET; - float Moldf = Mf[r]; - - // Compute max across the row - rowmaxf = subgroupMax(rowmaxf); + m += FATTN_KQ_MAX_OFFSET; + m = fa_clustered_max(m, LANES_PER_ROW); + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + // The outer subgroupBroadcastFirst is not redundant: without it + // LLPC keeps rowmax in VGPRs and moves the whole M/eM update from + // SALU to VALU, costing 4 VGPRs and 12 VALU. + rowmax[r] = subgroupBroadcastFirst(fa_row_broadcast(m, r * LANES_PER_ROW)); + } + } else { + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + const uint row = tile_row(r); + float rowmaxf = NEG_FLT_MAX_OVER_2; + [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { + const uint col = c * cols_per_iter + col_tid; + if (KV_bounds_check && j * Bc + col >= KV) { + continue; + } + rowmaxf = max(rowmaxf, float(sfsh[row * sfshstride + col / 4][col % 4])); + } + rowmax[r] = subgroupMax(rowmaxf + FATTN_KQ_MAX_OFFSET); + } + } + float eMf[rows_per_thread]; + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + const float Moldf = Mf[r]; // M = max(rowmax, Mold) // P = e^(S - M) // eM = e^(Mold - M) - Mf[r] = max(rowmaxf, Moldf); + Mf[r] = max(rowmax[r], Moldf); eMf[r] = exp(Moldf - Mf[r]); - Lf[r] = eMf[r]*Lf[r]; } - [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { - const uint d_local = d0 / threads_per_rowgroup; - [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Of[r][d_local] = O_TYPE(eMf[r]) * Of[r][d_local]; + if (TRANSPOSE_PV != 0) { + // Each subgroup only computed eM for the rows_per_thread rows it owns, but every + // subgroup's O accumulator spans all Br rows, so eM has to cross subgroups. + // Broadcast it as a MatBr x MatBc tile with eM constant along the row, so the + // rescale below is a componentwise coopmat multiply and never needs to know + // which (row, col) a component holds - that mapping is implementation defined. + // Fenced by the barrier that already publishes the P store below. + [[unroll]] for (uint32_t c0 = 0; c0 < MatBc; c0 += threads_per_rowgroup) { + const uint c = c0 + col_tid; + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + if (c < MatBc && tile_row(r) < MatBr) { + pvshT[tile_row(r) * MatBc + c] = O_TYPE(eMf[r]); + } + } + } + } else { + [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { + const uint d_local = d0 / threads_per_rowgroup; + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + Of[r][d_local] = O_TYPE(eMf[r]) * Of[r][d_local]; + } } } - // Calculate and store Pf in Psh + // Calculate and store Pf. In arm 1 the read side is the transposed sfsh + // ([query row][KV column], one lane per KV column) and the write side is + // the transposed P; both transposes are along different axes of the same + // value, so the four sfsh scalar reads and the four PshT scalar writes + // pair off. [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { const uint col = c * cols_per_iter + col_tid; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { const uint row = tile_row(r); if (KV_bounds_check && j * Bc + col >= KV) { - Psh[col * psh_stride + row / 4] = FLOAT_TYPEV4(0.0f); + if (TRANSPOSE_PV != 0) { + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + PshT[(row + vec_idx) * pt_stride + col] = FLOAT_TYPE(0.0f); + } + } else { + Psh[col * psh_stride + row / 4] = FLOAT_TYPEV4(0.0f); + } } else { const vec4 mfvec = vec4(Mf[r], Mf[r + 1], Mf[r + 2], Mf[r + 3]); - const FLOAT_TYPEV4 Pf = FLOAT_TYPEV4(exp(vec4(sfsh[row / 4 + col * sfshstride]) - mfvec)); + const uint cv = col / 4, ce = col % 4; + const vec4 svec = vec4(float(sfsh[(row ) * sfshstride + cv][ce]), + float(sfsh[(row + 1) * sfshstride + cv][ce]), + float(sfsh[(row + 2) * sfshstride + cv][ce]), + float(sfsh[(row + 3) * sfshstride + cv][ce])); + const FLOAT_TYPEV4 Pf = FLOAT_TYPEV4(exp(svec - mfvec)); [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { Lf[r + vec_idx] += Pf[vec_idx]; + if (TRANSPOSE_PV != 0) { + PshT[(row + vec_idx) * pt_stride + col] = Pf[vec_idx]; + } + } + if (TRANSPOSE_PV == 0) { + Psh[col * psh_stride + row / 4] = Pf; } - Psh[col * psh_stride + row / 4] = Pf; } } } @@ -415,14 +884,112 @@ void main() { } } - kvsh[c * kvsh_stride + d] = V_Tf; + if (TRANSPOSE_PV != 0) { + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + kvshT[c * kvshT_stride + d * 4 + e] = V_Tf[e]; + } + } else { + kvsh[c * kvsh_stride + d] = V_Tf; + } } } } barrier(); - const uint num_hsv_tiles = (HSV + MatBc * row_split - 1) / (MatBc * row_split); // round up + if (TRANSPOSE_PV != 0) { + // Rescale O in place, so the accumulator never leaves registers. eMMat has the + // same type and layout as OMat, so their components line up element for element + // and no knowledge of the component -> (row, col) mapping is needed. + coopmat eMMat; + coopMatLoad(eMMat, pvshT, 0, MatBc, gl_CooperativeMatrixLayoutRowMajor); + [[unroll]] for (uint32_t t = 0; t < num_hsv_tiles; ++t) { + OMat[t] = OMat[t] * eMMat; + } + // Each subgroup handles MatBc HSV columns per hsv_tile + [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { + const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16; + + // Preload V tiles for [Bc, 16 * num subgroups] + const uint v_total = Bc * vt_cols4; + const uint v_loads_per_thread = v_total / gl_WorkGroupSize.x; + + if (SHMEM_STAGING == 0) { + const bool stage_v = USE_DECODE_V || KV_bounds_check; + if (stage_v) { + if (hsv_tile > 0) { + // WAR: the previous tile's V fragments have been consumed. + barrier(); + } + [[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) { + const uint idx = i * gl_WorkGroupSize.x + tid; + const uint row = idx / vt_cols4; // KV row inside the Bc tile + const uint cb = idx % vt_cols4; // f16vec4 of HSV inside the staged tile + + const uint v_row = j * Bc + row; + const uint v_col = hsv_tile * vt_rows + cb * 4; + + const uint coord = v_row * v_stride * BLOCK_SIZE_V + v_col; + const uint ib = coord / BLOCK_SIZE_V; + const uint iqs = coord % BLOCK_SIZE_V; + + // Guard the VALUE, not the store: duplicating the store for the + // out-of-bounds arm doubles the ds_store count on every tile. + FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0.0f); + if (!KV_bounds_check || (v_row < KV && v_col < HSV)) { +#if !defined(BFLOAT16) + if (USE_DECODE_V) { + V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); + } else +#endif + { + V_Tf = data_vv4[(v_offset + v_row * v_stride + v_col) / 4]; + } + } + // Transposing scatter. The HSV row index is permuted to + // (hsv % 4) * vt_cols4 + hsv / 4 so that 16 consecutive lanes, + // which hold 16 consecutive f16vec4 of HSV, land on 16 + // CONSECUTIVE V^T rows instead of rows 4 apart. Rows 4 apart at a + // 16 B-aligned stride collapse onto 2 LDS banks. + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + kvshT[(e * vt_cols4 + cb) * vt_stride + row] = V_Tf[e]; + } + } + } + } + barrier(); + + // Under V_PERMUTED a subgroup owns the HSV columns congruent to its id mod 4 + // spread over the whole tile, not the contiguous 16 at hsv_offset, so the + // contiguous-range guard would silence subgroups that still hold live + // columns whenever HSV is not a multiple of MatBc*row_split. The staging + // scatter zero-fills out-of-range columns, so running all of them is safe. + if (V_PERMUTED || hsv_offset < HSV_pad) { + [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { + // A = P, M = query row, K = KV col. RowMajor on PshT -> wide loads. + coopMatLoad(PMat, PshT, bc_chunk * MatBc, pt_stride, gl_CooperativeMatrixLayoutRowMajor); + + if (SHMEM_STAGING == 0) { + if (!USE_DECODE_V && !KV_bounds_check) { + // F16/BF16 values can be loaded directly from global memory + const uint v_tile_row = j * Bc + bc_chunk * MatBc; + const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; + coopMatLoad(VMat, data_vv4, v_tile_offset, v_stride / 4, gl_CooperativeMatrixLayoutRowMajor); + } else { + // B = V^T, K = KV row contiguous. ColumnMajor -> wide loads. + const uint v_tile_offset = gl_SubgroupID * MatBc * vt_stride + bc_chunk * MatBc; + coopMatLoad(VMat, kvshT, v_tile_offset, vt_stride, gl_CooperativeMatrixLayoutColumnMajor); + } + } else { + const uint v_tile_offset = bc_chunk * MatBc * kvshT_stride + (hsv_tile * row_split + gl_SubgroupID) * MatBc; + coopMatLoad(VMat, kvshT, v_tile_offset, kvshT_stride, gl_CooperativeMatrixLayoutRowMajor); + } + + OMat[hsv_tile] = coopMatMulAdd(PMat, VMat, OMat[hsv_tile]); + } + } + } + } else { // Each subgroup handles HSV/4 columns [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16; @@ -523,10 +1090,44 @@ void main() { } } } + } barrier(); } + if (TRANSPOSE_PV != 0) { + // Drain the O accumulators into the thread-space layout the stores below expect. + // Executed once, not once per KV iteration: this is what the pvsh round trip used + // to cost 49 times in arm 0. + [[unroll]] for (uint32_t t = 0; t < num_hsv_tiles; ++t) { + barrier(); + if (V_PERMUTED || (t * row_split + gl_SubgroupID) * MatBc < HSV_pad) { + coopMatStore(OMat[t], pvshT, gl_SubgroupID * MatBc, oshT_stride, gl_CooperativeMatrixLayoutRowMajor); + } + barrier(); + + const uint hsv_base = t * vt_rows; + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { + const uint row = tile_row(r); + [[unroll]] for (uint32_t d_local = 0; d_local < d_per_thread; ++d_local) { + const uint d = d_local * threads_per_rowgroup + col_tid; + const uint hsv_col = 4 * d; + + if (hsv_col >= hsv_base && hsv_col < hsv_base + vt_rows && hsv_col < HSV) { + const uint local_hsv = (hsv_col - hsv_base) / 4; + [[unroll]] for (uint32_t e = 0; e < 4; ++e) { + // V_PERMUTED undoes the staging permutation; otherwise the + // fragment column is the HSV column, as before. + const uint pv_col = V_PERMUTED ? (e * vt_cols4 + local_hsv) + : (4 * local_hsv + e); + Of[r][d_local][e] = pvshT[row * oshT_stride + pv_col]; + } + } + } + } + } + } + [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Lf[r] = subgroupAdd(Lf[r]); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_add.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_add.comp new file mode 100644 index 00000000000..d1388c3cf93 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_add.comp @@ -0,0 +1,37 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +layout (push_constant) uniform parameter { + uint ne0_vec4; + uint nrows; +} p; + +layout (binding = 0) readonly buffer Residual { vec4 data_residual[]; }; +layout (binding = 1) readonly buffer Value { vec4 data_value[]; }; +layout (binding = 2) readonly buffer Gate { vec4 data_gate[]; }; +layout (binding = 3) writeonly buffer D { vec4 data_d[]; }; + +const uint num_threads = 256; +const uint num_iter = 2; + +layout(local_size_x = num_threads, local_size_y = 1, local_size_z = 1) in; + +void main() { + const uint row = gl_WorkGroupID.y; + uint channel = gl_WorkGroupID.x * num_threads * num_iter + gl_LocalInvocationID.x; + + if (row >= p.nrows) { + return; + } + + [[unroll]] for (uint i = 0; i < num_iter; ++i) { + if (channel < p.ne0_vec4) { + const uint idx = row * p.ne0_vec4 + channel; + precise vec4 product = data_value[idx] * data_gate[channel]; + precise vec4 sum = data_residual[idx] + product; + data_d[idx] = sum; + } + channel += num_threads; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp new file mode 100644 index 00000000000..d6fff9bf81d --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_cm1.comp @@ -0,0 +1,444 @@ +#version 460 + +// Integer cooperative-matrix (WMMA) matmul: Q8_0 weights x Q8_1-quantized activations. +// +// The float coopmat path (mul_mm.comp) dequantizes the weight tile to f16 before +// staging it in shared memory, which costs one multiply plus a narrow LDS store per +// weight element and leaves the matrix cores mostly idle. Q8_0 values are already +// int8, so here the raw quants are staged as-is and fed to sint8 cooperative +// matrices; the per-32-value block scales are folded in once per K-block on the +// int32 accumulator instead of once per element on the operands. +// +// Element ownership inside an accumulator fragment is implementation defined, so it +// is discovered at runtime from a shared-memory probe (see cm_layout_probe below) +// rather than assumed. That keeps the per-block scaling portable across drivers. +// +// The K loop is software pipelined through registers. Global loads for the next +// group of K blocks are issued after the current group is published to LDS, so they +// can overlap the current LDS reads, cooperative multiplies, and scale epilogue. +// +// Two hazards constrain where the prefetch may sit: +// * it overwrites the registers the LDS stores consume, so it must follow them; +// * it must stay after the store barrier, because a barrier placed after the loads +// could be lowered to a full vmcnt(0) wait and would re-expose the latency. +// With one LDS buffer the two barriers per iteration (write-visible, then +// reads-complete) are irreducible; the prefetch changes what they overlap, not how +// many there are. The workgroup swizzle in main() additionally trades repeated A +// fetches for the smaller B operand by keeping a bounded group of A panels together. + +#extension GL_EXT_control_flow_attributes : enable +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require + +// Unconditional: the Q8_0/Q8_1 block scales are fp16 and are held in shared memory. +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +#include "types.glsl" + + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A_PACKED16 {block_q8_0_packed16 data_a_packed16[];}; +layout (binding = 1) readonly buffer B {block_q8_1_x4_packed128 data_b[];}; +layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; + +layout (push_constant) uniform parameter +{ + uint M; + uint N; + uint K; + uint stride_a; + uint stride_b; + uint stride_d; + + uint batch_stride_a; + uint batch_stride_b; + uint batch_stride_d; + + uint base_work_group_z; + uint num_batches; + uint k_split; + uint ne02; + uint ne12; + uint broadcast2; + uint broadcast3; +} p; + +layout (constant_id = 0) const uint BLOCK_SIZE = 512; +layout (constant_id = 1) const uint BM = 128; +layout (constant_id = 2) const uint BN = 128; +// constant_id 3 (BK) is unused: the K blocking is fixed by the Q8_0/Q8_1 block size. +layout (constant_id = 4) const uint WM = 64; +layout (constant_id = 5) const uint WN = 32; +// constant_id 6 (WMITER) is unused: coopmat tiles replace the register subtiling. +layout (constant_id = 7) const uint CM_M = 16; +layout (constant_id = 8) const uint CM_N = 16; +layout (constant_id = 9) const uint CM_K = 16; +layout (constant_id = 10) const uint WARP = 64; +// Number of Q8_0 blocks staged per publication barrier. Three is selected for +// the RADV wave32 large tile; other pipelines retain two. +layout (constant_id = 11) const uint BK_STEP = 2; + +#define BK 32 + +// Bytes of A the tile-order swizzle in main() tries to keep resident. Half of the 32 MB +// last-level cache, leaving the other half to the B tiles and the output stream that are +// flowing through it at the same time. +#define GROUP_A_BUDGET (16u * 1024u * 1024u) + +// Row pitch of the quant tiles, in 32-bit words. The +4 padding both keeps every +// row 16-byte aligned and spreads consecutive rows across all 32 LDS banks +// (BK_STEP*8 alone would alias rows two apart onto the same banks). +const uint QPITCH = BK_STEP * (BK / 4) + 4; + +const uint CM_ELEMS = (CM_M * CM_N) / WARP; +const uint K_SUB = BK / CM_K; +// Magic-bias dequant, folded into an FMA pair. ACO charges wave64 VALU by ISSUE PASSES, not +// instructions: v_cvt_f32_i32 is absent from is_dual_issue_capable() (aco_statistics.cpp:72) so +// it costs 2 passes, while v_fma_f32/v_mul_f32/v_sub_f32 cost 1. Wave32 reverses that tradeoff: +// native conversion and scale arithmetic can be paired as VOPD, while the extra bias state gets +// in the way. WARP is a specialization constant, so the unused path is removed at pipeline build. +// +// Seeding the WMMA C operand with the bit pattern of 1.5*2^23 makes the accumulator readable as +// a float directly: K_SUB=2 chains 32 int8 products so |acc| <= 32*128*128 = 2^19, and +// 1.5*2^23 +/- 2^19 stays inside [2^23, 2^24) where the f32 ULP is exactly 1, so +// intBitsToFloat(BIAS_BITS + n) == BIAS + n exactly for every reachable n. +// +// nbias_a = -BIAS*scale_a is exact because BIAS = 3*2^22 has a 2-bit significand and +// scale_a is an f16 widened to f32 (11 bits), so the product needs 13 bits of the 24 +// available. The first fma then evaluates +// (BIAS + n)*sa - BIAS*sa = n*sa exactly before its single rounding, which is the same value and +// the same number of roundings as the v_mul it replaces. +#define ACC_BIAS_BITS 0x4B400000 +#define ACC_BIAS_F 12582912.0f +const bool USE_MAGIC_BIAS = WARP != 32; + +const uint CMS_PER_ROW = WM / CM_M; +const uint CMS_PER_COL = WN / CM_N; + +const uint loadstride_a = BLOCK_SIZE / (BK / 4); +const uint loadstride_b = BLOCK_SIZE / (BK / 16); + +// Upper bound on the staging iterations an invocation can run: loadc is strictly +// below its stride, so no invocation can reach index BM/BN before this many steps. +const uint A_LOADS = (BM + loadstride_a - 1) / loadstride_a; +const uint B_LOADS = (BN + loadstride_b - 1) / loadstride_b; + +// The quant tiles are typed int32_t, not i8vec4. coopMatLoad addresses them +// identically - its offset and stride are counted in elements of the array type and +// both types are 4 bytes wide, so QPITCH and every load offset below are unchanged, +// and the int8 fragment elements are still read packed four per dword. What changes +// is the store side: LLPC lowers a v4char store to workgroup memory as four +// ds_store_b8, because the Vulkan memory model marks every shared access +// NonPrivatePointer and that blocks the merge back into one dword. +shared int32_t buf_a_qs[BM * QPITCH]; +shared int32_t buf_b_qs[BN * QPITCH]; +// Scales are widened to f32 on the way into LDS. Keeping them as float16_t costs an +// extra v_cvt_f32_f16 plus 16-bit register shuffling on every gather in the K loop, +// which measured far more than the extra 1 KB of shared memory is worth. +shared float buf_a_d[BM * BK_STEP]; +shared float buf_b_d[BN * BK_STEP]; +shared int32_t cm_layout_probe[CM_M * CM_N]; + +// Global -> registers. Nothing is combined, converted or selected here: every ALU op +// applied to a load result would pull the s_waitcnt for that load back in front of the +// compute and undo the pipelining. The two halves of an A quant word therefore stay in +// separate registers (the packed16 struct puts qs at byte offset 2, so a 32-bit load is +// not possible anyway) and the k-guard is applied to the *address* rather than to the +// loaded value, with the matching zero-fill deferred to the LDS store. +// +// Only the B side is k-guarded, matching the unpipelined kernel: A rows are read for +// every staged block, and this is never invoked for a block at or past end_k, so ks==0 +// is always in range and is a safe clamp target. +#define PREFETCH_BLOCK(blk) \ + [[unroll]] for (uint li = 0; li < A_LOADS; li++) { \ + const uint buf_ib = loadc_a + li * loadstride_a; \ + if (buf_ib < BM) { \ + const uint ib = pos_a_ib + buf_ib * p.stride_a / BK; \ + [[unroll]] for (uint ks = 0; ks < BK_STEP; ks++) { \ + pre_a_q0[li * BK_STEP + ks] = int32_t(data_a_packed16[ib + ks].qs[loadr_a * 2]);\ + pre_a_q1[li * BK_STEP + ks] = \ + int32_t(data_a_packed16[ib + ks].qs[loadr_a * 2 + 1]); \ + pre_a_d[li * BK_STEP + ks] = data_a_packed16[ib + ks].d; \ + } \ + } \ + } \ + [[unroll]] for (uint li = 0; li < B_LOADS; li++) { \ + const uint buf_ib = loadc_b + li * loadstride_b; \ + if (buf_ib < BN) { \ + const uint ib = pos_b_ib + buf_ib * p.stride_b / BK; \ + [[unroll]] for (uint ks = 0; ks < BK_STEP; ks++) { \ + const uint ib_k = ((blk) + ks * BK < end_k) ? (ib + ks) : ib; \ + const uint ib_outer = ib_k / 4; \ + const uint ib_inner = ib_k % 4; \ + pre_b_qs[li * BK_STEP + ks] = data_b[ib_outer].qs[ib_inner * 2 + loadr_b]; \ + pre_b_d[li * BK_STEP + ks] = data_b[ib_outer].ds[ib_inner].x; \ + } \ + } \ + } + +// Registers -> LDS for the block starting at blk. The k-guard is re-evaluated here +// because PREFETCH_BLOCK clamped the address instead of the value. +#define STORE_BLOCK_TO_LDS(blk) \ + [[unroll]] for (uint li = 0; li < A_LOADS; li++) { \ + const uint buf_ib = loadc_a + li * loadstride_a; \ + if (buf_ib < BM) { \ + [[unroll]] for (uint ks = 0; ks < BK_STEP; ks++) { \ + buf_a_qs[buf_ib * QPITCH + ks * (BK / 4) + loadr_a] = \ + pack32(i16vec2(int16_t(pre_a_q0[li * BK_STEP + ks]), \ + int16_t(pre_a_q1[li * BK_STEP + ks]))); \ + if (loadr_a == 0) { \ + buf_a_d[ks * BM + buf_ib] = float(pre_a_d[li * BK_STEP + ks]); \ + } \ + } \ + } \ + } \ + [[unroll]] for (uint li = 0; li < B_LOADS; li++) { \ + const uint buf_ib = loadc_b + li * loadstride_b; \ + if (buf_ib < BN) { \ + [[unroll]] for (uint ks = 0; ks < BK_STEP; ks++) { \ + const bool in_bounds = (blk) + ks * BK < end_k; \ + const ivec4 v = in_bounds ? pre_b_qs[li * BK_STEP + ks] : ivec4(0); \ + const uint base = buf_ib * QPITCH + ks * (BK / 4) + loadr_b * 4; \ + buf_b_qs[base ] = v.x; \ + buf_b_qs[base + 1] = v.y; \ + buf_b_qs[base + 2] = v.z; \ + buf_b_qs[base + 3] = v.w; \ + if (loadr_b == 0) { \ + buf_b_d[ks * BN + buf_ib] = \ + in_bounds ? float(pre_b_d[li * BK_STEP + ks]) : 0.0f; \ + } \ + } \ + } \ + } + +void main() { + const uint batch_idx = gl_WorkGroupID.z + p.base_work_group_z; + + const uint i13 = batch_idx / p.ne12; + const uint i12 = batch_idx % p.ne12; + + const uint i03 = i13 / p.broadcast3; + const uint i02 = i12 / p.broadcast2; + + const uint batch_idx_a = i03 * p.ne02 + i02; + + const uint blocks_m = (p.M + BM - 1) / BM; + const uint blocks_n = (p.N + BN - 1) / BN; + const uint ik = gl_WorkGroupID.x / blocks_m; + + // Workgroups are dispatched over (blocks_m, blocks_n) with x varying fastest, so the + // natural order walks every A row-panel once per N tile. One B column tile then stays + // resident in the last-level cache for the whole sweep while A is re-fetched from DRAM + // blocks_n times - and A is the larger operand, so the cache is being spent on the + // wrong one. For m=16384 n=3079 k=6144 that is 25 x 107 MB of A against 20 MB of B. + // + // Re-index the (row-panel, N-tile) grid so that group_m consecutive row-panels are + // held while the N tiles sweep past them. A is then fetched once per group instead of + // once per (panel, N tile); the price is re-fetching B once per group, which is the + // cheaper direction whenever A is the bigger matrix. group_m is sized so the resident + // panels stay well inside the 32 MB Infinity Cache. This only permutes which tile a + // workgroup computes; every workgroup still does exactly the same work as before. + const uint a_panel_bytes = BM * p.K + (BM * p.K) / 16; // int8 quants + fp16 block scales + const uint group_m = clamp(GROUP_A_BUDGET / max(a_panel_bytes, 1u), 1u, min(blocks_m, 32u)); + + // Standard group-major remap. gsize handles a final group shorter than group_m; it is + // what keeps the map a bijection onto the full grid, so no tile is dropped or doubled. + const uint tiles_per_group = group_m * blocks_n; + const uint lin = gl_WorkGroupID.y * blocks_m + (gl_WorkGroupID.x % blocks_m); + const uint group_id = lin / tiles_per_group; + const uint first_m = group_id * group_m; + const uint gsize = min(blocks_m - first_m, group_m); + const uint in_group = lin - group_id * tiles_per_group; + + const uint ir = first_m + in_group % gsize; + const uint ic = in_group / gsize; + + const uint warp_i = gl_SubgroupID; + const uint warp_r = warp_i % (BM / WM); + const uint warp_c = warp_i / (BM / WM); + + uint elem_row[CM_ELEMS]; + uint elem_col[CM_ELEMS]; + if (WARP == 32) { + [[unroll]] for (uint e = 0; e < CM_ELEMS; ++e) { + elem_row[e] = gl_SubgroupInvocationID / CM_N + 2 * e; + elem_col[e] = gl_SubgroupInvocationID % CM_N; + } + } else { + // Accumulator element ownership is implementation defined. Keep runtime + // discovery for every path except the RADV-only wave32 specialization. + for (uint i = gl_LocalInvocationID.x; i < CM_M * CM_N; i += BLOCK_SIZE) { + cm_layout_probe[i] = int32_t(i); + } + barrier(); + + coopmat probe; + coopMatLoad(probe, cm_layout_probe, 0, CM_N, gl_CooperativeMatrixLayoutRowMajor); + + [[unroll]] for (uint e = 0; e < CM_ELEMS; ++e) { + elem_row[e] = uint(probe[e]) / CM_N; + elem_col[e] = uint(probe[e]) % CM_N; + } + } + + const uint loadr_a = gl_LocalInvocationID.x % (BK / 4); + const uint loadc_a = gl_LocalInvocationID.x / (BK / 4); + const uint loadr_b = gl_LocalInvocationID.x % (BK / 16); + const uint loadc_b = gl_LocalInvocationID.x / (BK / 16); + + const uint start_k = ik * p.k_split; + const uint end_k = min(p.K, (ik + 1) * p.k_split); + + uint pos_a_ib = batch_idx_a * (p.batch_stride_a / BK) + (ir * BM * p.stride_a + start_k) / BK; + uint pos_b_ib = (batch_idx * p.batch_stride_b + ic * BN * p.stride_b + start_k) / BK; + + // 32-bit, not int16_t: half-width registers make the compiler pack two loads into + // the lo/hi halves of one VGPR, and that v_mov_b16 forces a partial s_waitcnt back + // in front of the WMMAs. Sign-extension folds into the load itself, so this costs + // registers but no instructions. + int32_t pre_a_q0[A_LOADS * BK_STEP]; + int32_t pre_a_q1[A_LOADS * BK_STEP]; + float16_t pre_a_d [A_LOADS * BK_STEP]; + ivec4 pre_b_qs[B_LOADS * BK_STEP]; + float16_t pre_b_d [B_LOADS * BK_STEP]; + + // Prologue: the first K block has no compute to overlap, so it is issued as early + // as possible and covered by the accumulator zeroing and the pre-loop barrier. + if (start_k < end_k) { + PREFETCH_BLOCK(start_k) + } + + float sums[CMS_PER_ROW * CMS_PER_COL * CM_ELEMS]; + [[unroll]] for (uint i = 0; i < CMS_PER_ROW * CMS_PER_COL * CM_ELEMS; i++) { + sums[i] = 0.0f; + } + + // Row/column of this subgroup's tile within the workgroup tile. + const uint a_row0 = warp_r * WM; + const uint b_col0 = warp_c * WN; + const bool active_col_tile = ic * BN + b_col0 < p.N; + + barrier(); + + for (uint block = start_k; block < end_k; block += BK * BK_STEP) { + STORE_BLOCK_TO_LDS(block) + + barrier(); + + pos_a_ib += BK_STEP; + pos_b_ib += BK_STEP; + + const uint next_block = block + BK * BK_STEP; + if (next_block < end_k) { + PREFETCH_BLOCK(next_block) + } + + if (active_col_tile) { + [[unroll]] for (uint ks = 0; ks < BK_STEP; ks++) { + coopmat cache_a[CMS_PER_ROW * K_SUB]; + coopmat cache_b[CMS_PER_COL * K_SUB]; + + [[unroll]] for (uint r = 0; r < CMS_PER_ROW; r++) { + [[unroll]] for (uint h = 0; h < K_SUB; h++) { + coopMatLoad(cache_a[r * K_SUB + h], buf_a_qs, + (a_row0 + r * CM_M) * QPITCH + ks * (BK / 4) + h * (CM_K / 4), + QPITCH, gl_CooperativeMatrixLayoutRowMajor); + } + } + [[unroll]] for (uint c = 0; c < CMS_PER_COL; c++) { + [[unroll]] for (uint h = 0; h < K_SUB; h++) { + coopMatLoad(cache_b[c * K_SUB + h], buf_b_qs, + (b_col0 + c * CM_N) * QPITCH + ks * (BK / 4) + h * (CM_K / 4), + QPITCH, gl_CooperativeMatrixLayoutColumnMajor); + } + } + + // Per-block scales, gathered once per tile row/column and reused across + // the whole CMS_PER_ROW x CMS_PER_COL grid. + float scale_a[CMS_PER_ROW * CM_ELEMS]; + float nbias_a[CMS_PER_ROW * CM_ELEMS]; + float scale_b[CMS_PER_COL * CM_ELEMS]; + [[unroll]] for (uint r = 0; r < CMS_PER_ROW; r++) { + [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { + scale_a[r * CM_ELEMS + e] = (buf_a_d[ks * BM + a_row0 + r * CM_M + elem_row[e]]); + if (USE_MAGIC_BIAS) { + nbias_a[r * CM_ELEMS + e] = -ACC_BIAS_F * scale_a[r * CM_ELEMS + e]; + } + } + } + [[unroll]] for (uint c = 0; c < CMS_PER_COL; c++) { + [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { + scale_b[c * CM_ELEMS + e] = (buf_b_d[ks * BN + b_col0 + c * CM_N + elem_col[e]]); + } + } + + // Issue every coopMatMulAdd for this K block before applying any scales. + // Scaling a tile immediately after its own WMMA makes the f32 math read the + // accumulator that was just written, which serializes the vector ALU against + // the matrix pipe (visible as a wall of s_delay_alu in the ISA). Splitting the + // two passes gives the scheduler independent work to hide the dependency + // behind. The int32 accumulators stay in registers: CMS_PER_ROW*CMS_PER_COL + // tiles * (CM_M*CM_N/WARP) elements is 32 VGPRs at the tuned tile sizes. + coopmat accs[CMS_PER_ROW * CMS_PER_COL]; + + [[unroll]] for (uint r = 0; r < CMS_PER_ROW; r++) { + [[unroll]] for (uint c = 0; c < CMS_PER_COL; c++) { + coopmat acc = + coopmat( + USE_MAGIC_BIAS ? ACC_BIAS_BITS : 0); + + [[unroll]] for (uint h = 0; h < K_SUB; h++) { + acc = coopMatMulAdd(cache_a[r * K_SUB + h], cache_b[c * K_SUB + h], acc); + } + + accs[r * CMS_PER_COL + c] = acc; + } + } + + [[unroll]] for (uint r = 0; r < CMS_PER_ROW; r++) { + [[unroll]] for (uint c = 0; c < CMS_PER_COL; c++) { + const uint tile = r * CMS_PER_COL + c; + [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { + if (USE_MAGIC_BIAS) { + const float t = fma(intBitsToFloat(int(accs[tile][e])), + scale_a[r * CM_ELEMS + e], + nbias_a[r * CM_ELEMS + e]); + sums[tile * CM_ELEMS + e] = fma(t, scale_b[c * CM_ELEMS + e], + sums[tile * CM_ELEMS + e]); + } else { + sums[tile * CM_ELEMS + e] += float(accs[tile][e]) * + scale_a[r * CM_ELEMS + e] * + scale_b[c * CM_ELEMS + e]; + } + } + } + } + } + } + + barrier(); + } + + const uint dr = ir * BM + a_row0; + const uint dc = ic * BN + b_col0; + const uint offsets = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches; + + [[unroll]] for (uint r = 0; r < CMS_PER_ROW; r++) { + [[unroll]] for (uint c = 0; c < CMS_PER_COL; c++) { + const uint sums_base = (r * CMS_PER_COL + c) * CM_ELEMS; + [[unroll]] for (uint e = 0; e < CM_ELEMS; e++) { + const uint row = dr + r * CM_M + elem_row[e]; + const uint col = dc + c * CM_N + elem_col[e]; + if (row < p.M && col < p.N) { + data_d[offsets + col * p.stride_d + row] = D_TYPE(sums[sums_base + e]); + } + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/quantize_q8_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/quantize_q8_1.comp index 7ea29a07e37..30329c9ceeb 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/quantize_q8_1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/quantize_q8_1.comp @@ -24,10 +24,16 @@ layout(constant_id = 0) const uint GROUP_SIZE = 32; layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; layout (binding = 0) readonly buffer A {vec4 data_a[];}; +#if defined(FUSED_SWIGLU) || defined(FUSED_SIGMOID_MUL) +layout (binding = 1) readonly buffer B {vec4 data_b[];}; +#define D_BINDING 2 +#else +#define D_BINDING 1 +#endif #ifndef QBLOCK_X4 -layout (binding = 1) writeonly buffer D {block_q8_1_packed32 data_b[];}; +layout (binding = D_BINDING) writeonly buffer D {block_q8_1_packed32 data_d[];}; #else -layout (binding = 1) writeonly buffer D {block_q8_1_x4 data_b[];}; +layout (binding = D_BINDING) writeonly buffer D {block_q8_1_x4 data_d[];}; #endif #ifndef USE_SUBGROUPS @@ -57,7 +63,20 @@ void quantize(const uint wgid) { const uint a_idx = ib * 8 + iqs; - vec4 vals = a_idx < p.ne / 4 ? data_a[a_idx] : vec4(0.0f); + vec4 vals = vec4(0.0f); + if (a_idx < p.ne / 4) { +#ifdef FUSED_SWIGLU + precise vec4 gate = data_a[a_idx]; + precise vec4 activated = gate / (vec4(1.0f) + exp(-gate)); + vals = activated * data_b[a_idx]; +#elif defined(FUSED_SIGMOID_MUL) + precise vec4 gate = data_b[a_idx]; + precise vec4 activated = vec4(1.0f) / (vec4(1.0f) + exp(-gate)); + vals = data_a[a_idx] * activated; +#else + vals = data_a[a_idx]; +#endif + } const vec4 abs_vals = abs(vals); // Find absolute max for each block @@ -82,9 +101,9 @@ void quantize(const uint wgid) { vals = round(vals * d_inv); #ifndef QBLOCK_X4 - data_b[ib].qs[iqs] = pack32(i8vec4(round(vals))); + data_d[ib].qs[iqs] = pack32(i8vec4(round(vals))); #else - data_b[ibx4_outer].qs[ibx4_inner * 8 + iqs] = pack32(i8vec4(round(vals))); + data_d[ibx4_outer].qs[ibx4_inner * 8 + iqs] = pack32(i8vec4(round(vals))); #endif #ifndef USE_SUBGROUPS @@ -111,9 +130,9 @@ void quantize(const uint wgid) { #endif #ifndef QBLOCK_X4 - data_b[ib].ds = f16vec2(vec2(d, sum * d)); + data_d[ib].ds = f16vec2(vec2(d, sum * d)); #else - data_b[ibx4_outer].ds[ibx4_inner] = f16vec2(vec2(d, sum * d)); + data_d[ibx4_outer].ds[ibx4_inner] = f16vec2(vec2(d, sum * d)); #endif } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp index 55b89f19a7a..66fbe7865bd 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm.comp @@ -33,6 +33,9 @@ layout (binding = 6) readonly buffer R_I {uvec2 rope_data_i[];}; // indices for #define BLOCK_SIZE 512 layout (constant_id = 1) const bool do_multiply = false; +#if RMS_NORM_ROPE_FUSION +layout (constant_id = 2) const bool mrope_pack = false; +#endif layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; @@ -91,14 +94,26 @@ void rms_norm(uint num_iters) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); + uint dst_col = col; +#if RMS_NORM_ROPE_FUSION + if (mrope_pack) { + dst_col = col / 2 + (col % 2) * (ncols / 2); + } +#endif + data_d[d_offset + dst_col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)])); } } else { [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); + uint dst_col = col; +#if RMS_NORM_ROPE_FUSION + if (mrope_pack) { + dst_col = col / 2 + (col % 2) * (ncols / 2); + } +#endif + data_d[d_offset + dst_col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col])); } } } else { @@ -106,14 +121,22 @@ void rms_norm(uint num_iters) { if (col >= ncols) { continue; } - data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); + uint dst_col = col; +#if RMS_NORM_ROPE_FUSION + if (mrope_pack) { + dst_col = col / 2 + (col % 2) * (ncols / 2); + } +#endif + data_d[d_offset + dst_col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); } } #if RMS_NORM_ROPE_FUSION barrier(); rope_params rp = p.rope; for (uint t = 2*tid; t < ncols; t += 2*BLOCK_SIZE) { - if (rp.rope_mode == GGML_ROPE_TYPE_NEOX) { + if (mrope_pack) { + rope_multi(t, row, channel, samp, rp); + } else if (rp.rope_mode == GGML_ROPE_TYPE_NEOX) { rope_neox(t, row, channel, samp, rp); } else if (rp.rope_mode == GGML_ROPE_TYPE_NORMAL) { rope_norm(t, row, channel, samp, rp); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_channel_last.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_channel_last.comp new file mode 100644 index 00000000000..4be172c41ad --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_channel_last.comp @@ -0,0 +1,50 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable + +layout(constant_id = 0) const uint BLOCK_SIZE = 512; + +layout(push_constant) uniform parameter { + uint width; + uint height; + uint depth; + uint channels; + float eps; +} p; + +layout(binding = 0) readonly buffer A { float data_a[]; }; +layout(binding = 1) readonly buffer NORM { float data_norm[]; }; +layout(binding = 2) writeonly buffer D { float data_d[]; }; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +shared float sumsh[BLOCK_SIZE]; + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint plane = p.width * p.height * p.depth; + const uint spatial = (gl_WorkGroupID.z * p.height + gl_WorkGroupID.y) * p.width + gl_WorkGroupID.x; + + precise float sum = 0.0f; + [[unroll]] for (uint channel = tid; channel < p.channels; channel += BLOCK_SIZE) { + const float value = data_a[spatial + channel * plane]; + precise float square = value * value; + sum += square; + } + + sumsh[tid] = sum; + barrier(); + [[unroll]] for (uint s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum += sumsh[tid + s]; + sumsh[tid] = sum; + } + barrier(); + } + + const float scale = inversesqrt(sumsh[0] / float(p.channels) + p.eps); + [[unroll]] for (uint channel = tid; channel < p.channels; channel += BLOCK_SIZE) { + const uint offset = spatial + channel * plane; + data_d[offset] = scale * data_a[offset] * data_norm[channel]; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_modulate.comp b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_modulate.comp new file mode 100644 index 00000000000..0dd4dfa28aa --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rms_norm_modulate.comp @@ -0,0 +1,93 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : enable + +#define BLOCK_SIZE 512 + +layout(push_constant) uniform parameter { + uint ne00; + uint ne01; + uint ne02; + uint ne03; + uint nb01; + uint nb02; + uint nb03; + float eps; +} p; + +layout(binding = 0) readonly buffer A { float data_a[]; }; +layout(binding = 1) readonly buffer NORM { float data_norm[]; }; +layout(binding = 2) readonly buffer SCALE { float data_scale[]; }; +layout(binding = 3) readonly buffer SHIFT { float data_shift[]; }; +layout(binding = 4) writeonly buffer D { float data_d[]; }; + +layout(local_size_x = BLOCK_SIZE, local_size_y = 1, local_size_z = 1) in; + +shared float sumsh[BLOCK_SIZE]; + +void rms_norm_modulate(uint num_iters) { + const uint row = gl_WorkGroupID.x; + const uint channel = gl_WorkGroupID.y; + const uint samp = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationID.x; + + const uint a_offset = samp*p.nb03 + channel*p.nb02 + row*p.nb01; + const uint d_offset = ((samp*p.ne02 + channel)*p.ne01 + row)*p.ne00; + + float sum = 0.0f; + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + float xi = 0.0f; + if (col < p.ne00) { + xi = data_a[a_offset + col]; + } + sum += xi * xi; + } + + sumsh[tid] = sum; + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum += sumsh[tid + s]; + sumsh[tid] = sum; + } + barrier(); + } + + const float inv_rms = inversesqrt(sumsh[0] / float(p.ne00) + p.eps); + + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= p.ne00) { + continue; + } + + precise float normalized = inv_rms * data_a[a_offset + col] * data_norm[col]; + precise float factor = data_scale[col] * 1.0f + 1.0f; + precise float modulated = normalized * factor; + data_d[d_offset + col] = modulated + data_shift[col]; + } +} + +void main() { + uint num_blocks = (p.ne00 + BLOCK_SIZE - 1) / BLOCK_SIZE; + if (num_blocks > 32) { + rms_norm_modulate(num_blocks); + } else if (num_blocks > 16) { + rms_norm_modulate(32); + } else if (num_blocks > 12) { + rms_norm_modulate(16); + } else if (num_blocks > 10) { + rms_norm_modulate(12); + } else if (num_blocks > 8) { + rms_norm_modulate(10); + } else if (num_blocks > 4) { + rms_norm_modulate(8); + } else if (num_blocks == 4) { + rms_norm_modulate(4); + } else if (num_blocks == 3) { + rms_norm_modulate(3); + } else if (num_blocks == 2) { + rms_norm_modulate(2); + } else if (num_blocks == 1) { + rms_norm_modulate(1); + } +} 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 caa0c889a45..58c531f32cc 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -626,6 +626,13 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) { string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc); } +#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + // Integer cooperative-matrix mmq. Only Q8_0 qualifies: its quants are already + // int8, so they can be fed to sint8 cooperative matrices without unpacking. + if (!f16acc && coopmat && !coopmat2 && !dot2 && matmul_id_type == MatMulIdType::NONE && tname == "q8_0") { + string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq_cm1.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc); + } +#endif #endif } } @@ -805,6 +812,9 @@ void process_shaders() { string_to_spv("norm_f32", "norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("group_norm_f32", "group_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); + string_to_spv("rms_norm_channel_last_f32", "rms_norm_channel_last.comp", {}); + string_to_spv("rms_norm_modulate_f32", "rms_norm_modulate.comp", {}); + string_to_spv("mul_add_f32", "mul_add.comp", {}); string_to_spv("rms_norm_partials_f32", "rms_norm_partials.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("rms_norm_mul_rope_f32_f32", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float"}, {"RMS_NORM_ROPE_FUSION", "1"}})); string_to_spv("rms_norm_mul_rope_f32_f16", "rms_norm.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"ROPE_D_TYPE", "float16_t"}, {"RMS_NORM_ROPE_FUSION", "1"}})); @@ -882,6 +892,10 @@ void process_shaders() { string_to_spv("quantize_q8_1_x4", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}}); string_to_spv("quantize_q8_1_x4_subgroup", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"USE_SUBGROUPS", "1"}}); + string_to_spv("quantize_q8_1_x4_swiglu", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SWIGLU", "1"}}); + string_to_spv("quantize_q8_1_x4_swiglu_subgroup", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SWIGLU", "1"}, {"USE_SUBGROUPS", "1"}}); + string_to_spv("quantize_q8_1_x4_sigmoid_mul", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SIGMOID_MUL", "1"}}); + string_to_spv("quantize_q8_1_x4_sigmoid_mul_subgroup", "quantize_q8_1.comp", {{"QBLOCK_X4", "1"}, {"FUSED_SIGMOID_MUL", "1"}, {"USE_SUBGROUPS", "1"}}); string_to_spv("mul_f32", "mul.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 8e3b273a1e4..b84d58c920e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3206,6 +3206,82 @@ struct test_bin_bcast : public test_case { } }; +// GGML_OP_MUL + GGML_OP_ADD with a channel-wise gate +struct test_mul_add : public test_case { + const std::array ne; + + explicit test_mul_add(std::array ne) : ne(ne) {} + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "MUL_ADD"; + } + + std::string vars() override { + return VARS_TO_STR1(ne); + } + + bool run_whole_graph() override { return true; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * residual = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_tensor * value = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_tensor * gate = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]); + ggml_set_param(residual); + ggml_set_param(value); + ggml_set_param(gate); + ggml_set_name(residual, "residual"); + ggml_set_name(value, "value"); + ggml_set_name(gate, "gate"); + + ggml_tensor * product = ggml_mul(ctx, value, gate); + ggml_tensor * out = ggml_add(ctx, residual, product); + ggml_set_name(out, "out"); + return out; + } +}; + +struct test_activation_quant_fusion : public test_case { + const bool swiglu; + const int64_t m; + const int64_t n; + const int64_t k; + + test_activation_quant_fusion(bool swiglu, int64_t m, int64_t n, int64_t k) + : swiglu(swiglu), m(m), n(n), k(k) {} + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return swiglu ? "SWIGLU_QUANT_MUL_MAT" : "SIGMOID_MUL_QUANT_MUL_MAT"; + } + + std::string vars() override { + return VARS_TO_STR4(swiglu, m, n, k); + } + + bool run_whole_graph() override { return true; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, k, m); + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, n); + ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, n); + ggml_set_param(a); + ggml_set_param(b); + ggml_set_name(weights, "weights"); + ggml_set_name(a, "a"); + ggml_set_name(b, "b"); + + ggml_tensor * activation = swiglu + ? ggml_swiglu_split(ctx, a, b) + : ggml_mul(ctx, a, ggml_sigmoid(ctx, b)); + ggml_tensor * out = ggml_mul_mat(ctx, weights, activation); + ggml_set_name(out, "out"); + return out; + } + + double max_nmse_err() override { return 5e-4; } +}; + // GGML_OP_ADD_ID struct test_add_id : public test_case { const ggml_type type_a; @@ -3610,6 +3686,184 @@ struct test_rms_norm_mul_add : public test_case { } }; +// GGML_OP_RMS_NORM + GGML_OP_MUL + GGML_OP_SCALE + GGML_OP_MUL + GGML_OP_ADD +struct test_rms_norm_mul_modulate : public test_case { + const std::array ne; + const float eps; + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "RMS_NORM_MUL_MODULATE"; + } + + bool run_whole_graph() override { return true; } + + std::string vars() override { + return VARS_TO_STR2(ne, eps); + } + + test_rms_norm_mul_modulate(std::array ne, float eps = 1e-6f) + : ne(ne), eps(eps) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_tensor * norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[0]); + const int64_t mod_stride = GGML_PAD(ne[0], 64); + ggml_tensor * modulation = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, mod_stride * 6); + ggml_tensor * mod_scale = ggml_view_1d(ctx, modulation, ne[0], 3 * mod_stride * sizeof(float)); + ggml_tensor * shift = ggml_view_1d(ctx, modulation, ne[0], 4 * mod_stride * sizeof(float)); + + ggml_set_param(a); + ggml_set_param(norm); + ggml_set_param(modulation); + + // Visit the norm weight before the fusion chain. Krea materializes two views from a + // packed six-vector modulation result before using them as scale and shift. + a = ggml_add(ctx, a, norm); + + ggml_tensor * out = ggml_rms_norm(ctx, a, eps); + out = ggml_mul_inplace(ctx, out, norm); + mod_scale = ggml_cont(ctx, mod_scale); + ggml_tensor * factor = ggml_scale_bias(ctx, mod_scale, 1.0f, 1.0f); + out = ggml_mul(ctx, out, factor); + shift = ggml_cont(ctx, shift); + out = ggml_add(ctx, out, shift); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + init_tensor_uniform(t, -1.0f, 1.0f); + } + } +}; + +// Krea2 projection normalization followed by NORMAL-to-NEOX deinterleave, +// multi-axis RoPE, and the layout pack consumed by flash attention. +struct test_rms_norm_mrope_pack : public test_case { + const int64_t d_head; + const int64_t n_head; + const int64_t n_token; + const float eps; + + test_rms_norm_mrope_pack(int64_t d_head, int64_t n_head, int64_t n_token, float eps = 1e-5f) + : d_head(d_head), n_head(n_head), n_token(n_token), eps(eps) {} + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "RMS_NORM_MROPE_PACK"; + } + + std::string vars() override { + return VARS_TO_STR4(d_head, n_head, n_token, eps); + } + + bool run_whole_graph() override { return true; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + GGML_ASSERT(d_head % 16 == 0); + + ggml_tensor * src = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_head, n_head, n_token, 1); + ggml_tensor * norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, d_head); + ggml_set_param(src); + ggml_set_param(norm); + + ggml_tensor * out = ggml_mul(ctx, ggml_rms_norm(ctx, src, eps), norm); + + out = ggml_reshape_4d(ctx, out, 2, d_head / 2, n_head, n_token); + out = ggml_cont(ctx, ggml_permute(ctx, out, 1, 0, 2, 3)); + out = ggml_reshape_4d(ctx, out, d_head, n_head, n_token, 1); + + ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 4 * n_token); + ggml_tensor * freq = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, d_head / 2); + ggml_set_name(pos, "pos"); + ggml_set_name(freq, "freq"); + + int sections[GGML_MROPE_SECTIONS] = { + (int)(d_head / 8), + (int)(3 * d_head / 16), + (int)(3 * d_head / 16), + 0, + }; + out = ggml_rope_multi(ctx, out, pos, freq, (int)d_head, sections, + GGML_ROPE_TYPE_MROPE, 0, 1000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + out = ggml_cont(ctx, ggml_permute(ctx, out, 0, 2, 1, 3)); + out = ggml_reshape_3d(ctx, out, d_head, n_token, n_head); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "pos") == 0) { + std::vector data(4 * n_token); + for (int64_t stream = 0; stream < 4; ++stream) { + for (int64_t token = 0; token < n_token; ++token) { + data[stream * n_token + token] = (int32_t)((stream + 1) * token); + } + } + ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(data[0])); + } else if (strcmp(t->name, "freq") == 0) { + init_tensor_uniform(t, 0.9f, 1.1f); + } else { + init_tensor_uniform(t, -1.0f, 1.0f); + } + } + } +}; + +// WAN VAE channel-last RMS normalization: transpose channels to ne[0], +// normalize and scale, then transpose back to the original image layout. +struct test_rms_norm_channel_last : public test_case { + const std::array ne; + const float eps; + + test_rms_norm_channel_last(std::array ne, float eps = 1e-12f) + : ne(ne), eps(eps) {} + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "RMS_NORM_CHANNEL_LAST"; + } + + std::string vars() override { + return VARS_TO_STR2(ne, eps); + } + + bool run_whole_graph() override { return true; } + bool use_weight_context() override { return true; } + + ggml_tensor * build_graph(ggml_context * ctx) override { + GGML_UNUSED(ctx); + GGML_ABORT("weight context required"); + } + + ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { + GGML_ASSERT(ctx_weights); + + ggml_tensor * src = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_tensor * norm_storage = ggml_new_tensor_4d(ctx_weights, GGML_TYPE_F32, 1, 1, 1, ne[3]); + ggml_set_param(src); + + ggml_tensor * norm = ggml_reshape_1d(ctx, norm_storage, ne[3]); + + ggml_tensor * out = ggml_cont(ctx, ggml_permute(ctx, src, 1, 2, 3, 0)); + out = ggml_rms_norm(ctx, out, eps); + out = ggml_mul(ctx, out, norm); + out = ggml_cont(ctx, ggml_permute(ctx, out, 3, 0, 1, 2)); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + init_tensor_uniform(t, -1.0f, 1.0f); + } + } +}; + // GGML_OP_ADD + GGML_OP_RMS_NORM (fused operation) struct test_add_rms_norm : public test_case { const ggml_type type; @@ -9003,6 +9257,11 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {10, 5, 4, 3}, {1, 2, 2, 2}, 7)); test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {2, 2, 2, 2}, 8)); test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {16, 5, 4, 3}, {1, 1, 1, 1}, 16)); + test_cases.emplace_back(new test_mul_add({64, 5, 4, 3})); + test_cases.emplace_back(new test_mul_add({1024, 17, 1, 1})); + test_cases.emplace_back(new test_mul_add({65, 5, 1, 1})); // fallback path: vec4 fusion is intentionally inapplicable + test_cases.emplace_back(new test_activation_quant_fusion(true, 128, 17, 256)); + test_cases.emplace_back(new test_activation_quant_fusion(false, 128, 17, 256)); test_cases.emplace_back(new test_scale()); test_cases.emplace_back(new test_scale(GGML_TYPE_F32, {10, 10, 10, 10}, 2.0f, 1.0f)); @@ -9052,6 +9311,15 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_add_rms_norm(GGML_TYPE_F32, {n, 1, 1, 1}, 1e-6f, false)); } + test_cases.emplace_back(new test_rms_norm_mul_modulate({64, 5, 4, 3})); + test_cases.emplace_back(new test_rms_norm_mul_modulate({1025, 5, 2, 1})); + test_cases.emplace_back(new test_rms_norm_mul_modulate({6144, 7, 1, 1})); + test_cases.emplace_back(new test_rms_norm_mrope_pack(128, 4, 17)); + test_cases.emplace_back(new test_rms_norm_mrope_pack(128, 48, 50)); + test_cases.emplace_back(new test_rms_norm_channel_last({7, 5, 3, 64})); + test_cases.emplace_back(new test_rms_norm_channel_last({7, 5, 2, 192})); + test_cases.emplace_back(new test_rms_norm_channel_last({5, 3, 2, 513})); + for (auto multi_add : {false, true}) { for (auto set_rows : {false, true}) { for (auto broadcast : {false, true}) {